time_neo-nixio.py

import neo
import nixio as nix
import quantities as pq
import numpy as np
import time
import matplotlib.pyplot as plt
import cProfile


# n_segs = 150
# n_sts = 271
n_sts = 10

data = pq.Quantity(np.random.random(10000), "s")
t_stop = data[-1]+pq.Quantity(1, "s")


def time_neo_nixio_save_operation(n_segs):
    print("Generating {} segments with {} spiketrains each".format(
        n_segs, n_sts
    ))
    block = neo.Block()
    for seg_id in range(n_segs):
        seg = neo.Segment()
        block.segments.append(seg)

        for st_id in range(n_sts):
            st = neo.SpikeTrain(data, t_start=0*pq.s, t_stop=t_stop,
                                sampling_rate=1000*pq.Hz)

            seg.spiketrains.append(st)

    print('Storing data...', end="", flush=True)
    t0 = time.time()
    with neo.io.nixio.NixIO("neo_nixfile.nix", mode="ow") as nixfile:
        nixfile.write_block(block)
    t1 = time.time()
    wdura = t1 - t0
    print("Done ({} s)".format(wdura))

    print("Loading data...", end="", flush=True)
    t0 = time.time()
    with neo.io.nixio.NixIO("neo_nixfile.nix", mode="ro") as nixfile:
        nixfile.read_all_blocks()
    t1 = time.time()
    rdura = t1 - t0
    print("Done ({} s)".format(rdura))
    return wdura, rdura


def time_nix_save_operation(n_groups):
    print("Doing direct nix write of {} groups with {} MultiTags each".format(
        n_groups, n_sts
    ))
    datamag = data.magnitude
    print("Storing data...", end="", flush=True)
    t0 = time.time()
    nixfile = nix.File.open("nixfile.nix", mode=nix.FileMode.Overwrite,
                            backend="h5py")
    block = nixfile.create_block("blk", "tst")
    block.metadata = nixfile.create_section("blk", "blk.metadata")
    for gidx in range(n_groups):
        grp = block.create_group("grp-{}".format(gidx), "tst")
        grp.metadata = block.metadata.create_section(grp.name, "grp.metadata")
        for sidx in range(n_sts):
            da = block.create_data_array("da-{}-{}".format(gidx, sidx),
                                         "tst", data=datamag)
            mt = block.create_multi_tag("mt-{}-{}".format(gidx, sidx),
                                        "tst", positions=da)
            grp._add_data_array_obj(da)
            grp._add_multi_tag_obj(mt)
            mt.metadata = block.metadata.create_section(mt.name, "st.metadata")
            mt.metadata["t_start"] = 0
            mt.metadata["t_start.units"] = "s"
            mt.metadata["t_stop"] = 1
            mt.metadata["t_stop.units"] = "s"
    t1 = time.time()
    wdura = t1 - t0
    print("Done ({} s)".format(wdura))
    nixfile.close()

    print("Loading data...", end="", flush=True)
    t0 = time.time()
    nixfile = nix.File.open("nixfile.nix", mode=nix.FileMode.ReadOnly,
                            backend="h5py")
    for blk in nixfile.blocks:
        for grp in blk.groups:
            for da in grp.data_arrays:
                da[:]
                da.name
                da.type
            for mt in grp.multi_tags:
                mt.positions[:]
                mt.name
                mt.type
                [prop for prop in mt.metadata.props]
    t1 = time.time()
    rdura = t1 - t0
    print("Done ({} s)".format(rdura))
    nixfile.close()
    return wdura, rdura


def prof(func, *args, **kwargs):
    pr = cProfile.Profile()
    pr.enable()
    ret = func(*args, **kwargs)
    pr.disable()
    pr.print_stats(sort='tottime')
    return ret


if __name__ == "__main__":
    imgprefix = str(int(time.time()))
    wimgname = "{}-save.png".format(imgprefix)
    rimgname = "{}-load.png".format(imgprefix)
    seg_counts = [5, 10, 20, 40]
    neo_write_times = []
    neo_load_times = []
    for sc in seg_counts:
        write_time, load_time = time_neo_nixio_save_operation(sc)
        neo_write_times.append(write_time)
        neo_load_times.append(load_time)

    wfig = plt.figure("Write times")
    wax = wfig.add_subplot(1, 1, 1)
    wax.set_xlabel("N segments")
    wax.set_ylabel("Write time (s)")
    first_rate = neo_write_times[0]/seg_counts[0]
    linear = [first_rate * sc for sc in seg_counts]
    wax.plot(seg_counts, neo_write_times, label="neo-nixio")
    wax.plot(seg_counts, linear, "--k", label="linear neo-nixio")
    wfig.savefig(wimgname)

    rfig = plt.figure("Read times")
    rax = rfig.add_subplot(1, 1, 1)
    rax.set_xlabel("N segments")
    rax.set_ylabel("Read time (s)")
    first_rate = neo_load_times[0]/seg_counts[0]
    linear = [first_rate * sc for sc in seg_counts]
    rax.plot(seg_counts, neo_load_times, label="neo-nixio")
    rax.plot(seg_counts, linear, "--k", label="linear neo-nixio")
    rfig.savefig(rimgname)

    nix_write_times = []
    nix_load_times = []
    for sc in seg_counts:
        write_time, load_time = time_nix_save_operation(sc)
        nix_write_times.append(write_time)
        nix_load_times.append(load_time)

    first_rate = nix_write_times[0]/seg_counts[0]
    linear = [first_rate * sc for sc in seg_counts]
    wax.plot(seg_counts, nix_write_times, label="nix")
    wax.plot(seg_counts, linear, "-.k", label="linear nix")
    wax.legend(loc="best")
    wfig.savefig(wimgname)

    first_rate = nix_load_times[0]/seg_counts[0]
    linear = [first_rate * sc for sc in seg_counts]
    rax.plot(seg_counts, nix_load_times, label="nix")
    rax.plot(seg_counts, linear, "-.k", label="linear nix")
    rax.legend(loc="best")
    rfig.savefig(rimgname)

    print("Timings")
    print("  N         Neo          NIX     Neo/NIX")
    print("    Writing")
    for sc, neot, nixt in zip(seg_counts, neo_write_times, nix_write_times):
        print("{:03d}: {:10.2f} s {:10.2f} s   {:5.2f}".format(sc, neot, nixt,
                                                               neot/nixt))

    print("    Reading")
    for sc, neot, nixt in zip(seg_counts, neo_load_times, nix_load_times):
        print("{:03d}: {:10.2f} s {:10.2f} s   {:5.2f}".format(sc, neot, nixt,
                                                               neot/nixt))