nwbio.py

"""
NWBIO
=====

IO class for reading data from a Neurodata Without Borders (NWB) dataset

Documentation : https://www.nwb.org/
Depends on: h5py, nwb, dateutil
Supported: Read, Write
Python API -  https://pynwb.readthedocs.io
Sample datasets from CRCNS - https://crcns.org/NWB
Sample datasets from Allen Institute
- http://alleninstitute.github.io/AllenSDK/cell_types.html#neurodata-without-borders
"""

from __future__ import absolute_import, division

import json
import logging
import os
from collections import defaultdict
from itertools import chain
from json.decoder import JSONDecodeError

import numpy as np
import quantities as pq

from neo.core.baseneo import _check_annotations
from neo.core import Segment, SpikeTrain, Epoch, Event, AnalogSignal, IrregularlySampledSignal, Block, ImageSequence
from neo.io.baseio import BaseIO
from neo.io.proxyobjects import (
    AnalogSignalProxy as BaseAnalogSignalProxy,
    EventProxy as BaseEventProxy,
    EpochProxy as BaseEpochProxy,
    SpikeTrainProxy as BaseSpikeTrainProxy,
)


logger = logging.getLogger("Neo")

GLOBAL_ANNOTATIONS = (
    "session_start_time",
    "identifier",
    "timestamps_reference_time",
    "experimenter",
    "experiment_description",
    "session_id",
    "institution",
    "keywords",
    "notes",
    "pharmacology",
    "protocol",
    "related_publications",
    "slices",
    "source_script",
    "source_script_file_name",
    "data_collection",
    "surgery",
    "virus",
    "stimulus_notes",
    "lab",
    "session_description",
    "rec_datetime",
)

POSSIBLE_JSON_FIELDS = ("source_script", "description")

prefix_map = {1e9: "giga", 1e6: "mega", 1e3: "kilo", 1: "", 1e-3: "milli", 1e-6: "micro", 1e-9: "nano", 1e-12: "pico"}


def try_json_field(content):
    """
    Try to interpret a string as JSON data.

    If successful, return the JSON data (dict or list)
    If unsuccessful, return the original string
    """
    try:
        return json.loads(content)
    except JSONDecodeError:
        return content


def get_class(module, name):
    """
    Given a module path and a class name, return the class object
    """
    import pynwb

    module_path = module.split(".")
    if len(module_path) != 2:
        raise ValueError(
            f"`module_path` must be 2, not {module_path}"
        )  # todo: handle the general case where this isn't 2
    return getattr(getattr(pynwb, module_path[1]), name)


def statistics(block):  # todo: move this to be a property of Block
    """
    Return simple statistics about a Neo Block.
    """
    stats = {
        "SpikeTrain": {"count": 0},
        "AnalogSignal": {"count": 0},
        "IrregularlySampledSignal": {"count": 0},
        "Epoch": {"count": 0},
        "Event": {"count": 0},
    }
    for segment in block.segments:
        stats["SpikeTrain"]["count"] += len(segment.spiketrains)
        stats["AnalogSignal"]["count"] += len(segment.analogsignals)
        stats["IrregularlySampledSignal"]["count"] += len(segment.irregularlysampledsignals)
        stats["Epoch"]["count"] += len(segment.epochs)
        stats["Event"]["count"] += len(segment.events)
    return stats


def get_units_conversion(signal, timeseries_class):
    """
    Given a quantity array and a TimeSeries subclass, return
    the conversion factor and the expected units
    """
    # it would be nice if the expected units was an attribute of the PyNWB class
    if "CurrentClamp" in timeseries_class.__name__:
        expected_units = pq.volt
    elif "VoltageClamp" in timeseries_class.__name__:
        expected_units = pq.ampere
    else:
        # todo: warn that we don't handle this subclass yet
        expected_units = signal.units
    return float((signal.units / expected_units).simplified.magnitude), expected_units


def time_in_seconds(t):
    return float(t.rescale("second"))


def _decompose_unit(unit):
    """
    Given a quantities unit object, return a base unit name and a conversion factor.

    Example:

    >>> _decompose_unit(pq.mV)
    ('volt', 0.001)
    """
    if not isinstance(unit, pq.quantity.Quantity):
        raise TypeError(f"`unit` must be of type pq.quantity.Quantity and not type {type(unit)}")
    if unit.magnitude != 1:
        raise ValueError(f"The magnitude of the `unit` must be 1 not {unit.magnitude}")
    conversion = 1.0

    def _decompose(unit):
        dim = unit.dimensionality
        if len(dim) != 1:
            raise NotImplementedError("Compound units not yet supported")  # e.g. volt-metre
        uq, n = list(dim.items())[0]
        if n != 1:
            raise NotImplementedError("Compound units not yet supported")  # e.g. volt^2
        uq_def = uq.definition
        return float(uq_def.magnitude), uq_def

    conv, unit2 = _decompose(unit)
    while conv != 1:
        conversion *= conv
        unit = unit2
        conv, unit2 = _decompose(unit)
    return list(unit.dimensionality.keys())[0].name, conversion


def _recompose_unit(base_unit_name, conversion):
    """
    Given a base unit name and a conversion factor, return a quantities unit object

    Example:

    >>> _recompose_unit("ampere", 1e-9)
    UnitCurrent('nanoampere', 0.001 * uA, 'nA')

    """
    unit_name = None
    for cf in prefix_map:
        # conversion may have a different float precision to the keys in
        # prefix_map, so we can't just use `prefix_map[conversion]`
        if abs(conversion - cf) / cf < 1e-6:
            unit_name = prefix_map[cf] + base_unit_name
    if unit_name is None:
        raise ValueError(f"Can't handle this conversion factor: {conversion}")

    if unit_name[-1] == "s":  # strip trailing 's', e.g. "volts" --> "volt"
        unit_name = unit_name[:-1]
    try:
        return getattr(pq, unit_name)
    except AttributeError:
        logger.warning(f"Can't handle unit '{unit_name}'. Returning dimensionless")
        return pq.dimensionless


def nwb_obj_to_dict(obj):
    if not hasattr(obj, "fields"):
        raise TypeError("Does not seem to be an NWB object")
    result = {}
    for key, value in obj.fields.items():
        if hasattr(value, "fields"):
            result[key] = nwb_obj_to_dict(value)
        else:
            result[key] = value
    return result


class NWBIO(BaseIO):
    """
    Class for "reading" experimental data from a .nwb file, and "writing" a .nwb file from Neo
    """

    supported_objects = [
        Block,
        Segment,
        AnalogSignal,
        IrregularlySampledSignal,
        SpikeTrain,
        Epoch,
        Event,
        ImageSequence,
    ]
    readable_objects = supported_objects
    writeable_objects = supported_objects

    has_header = False
    support_lazy = True

    name = "NeoNWB IO"
    description = "This IO reads/writes experimental data from/to an .nwb dataset"
    extensions = ["nwb"]
    mode = "one-file"

    is_readable = True
    is_writable = True
    is_streameable = False

    def __init__(self, filename, mode="r", **annotations):
        """
        Arguments:
            filename : the filename
        """
        import pynwb

        BaseIO.__init__(self, filename=filename)
        self.filename = filename
        self.blocks_written = 0
        self.nwb_file_mode = mode
        self._blocks = {}
        self.annotations = annotations
        self._io_nwb = None

    def read_all_blocks(self, lazy=False, **kwargs):
        """
        Load all blocks in the file.
        """
        import hdmf
        import pynwb

        if self.nwb_file_mode not in ("r",):
            raise ValueError("`mode` at init needs to be set to 'r' to read files")
        self._io_nwb = pynwb.NWBHDF5IO(
            self.filename, mode=self.nwb_file_mode, load_namespaces=True
        )  # Open a file with NWBHDF5IO
        try:
            self._file = self._io_nwb.read()
        except ValueError:
            print("Error: Unable to read this version of NWB file.")
            print("Please convert to a later NWB format.")
            raise

        self.global_block_metadata = {}
        for annotation_name in GLOBAL_ANNOTATIONS:
            value = getattr(self._file, annotation_name, None)
            if value is not None:
                if annotation_name in POSSIBLE_JSON_FIELDS:
                    value = try_json_field(value)
                elif isinstance(value, hdmf.utils.StrDataset):
                    value = list(value)
                # placing this check here for easier debugging, but it's redundant so we should remove it
                # once we're handling all possible annotation types
                _check_annotations(value)
                self.global_block_metadata[annotation_name] = value
        if "session_description" in self.global_block_metadata:
            self.global_block_metadata["description"] = self.global_block_metadata["session_description"]
        self.global_block_metadata["file_origin"] = self.filename
        if "session_start_time" in self.global_block_metadata:
            self.global_block_metadata["rec_datetime"] = self.global_block_metadata["session_start_time"]
        if "file_create_date" in self.global_block_metadata:
            self.global_block_metadata["file_datetime"] = self.global_block_metadata["rec_datetime"]

        self._blocks = {}
        self._read_acquisition_group(lazy=lazy)
        self._read_stimulus_group(lazy)
        self._read_units(lazy=lazy)
        self._read_epochs_group(lazy)

        return list(self._blocks.values())

    def read_block(self, lazy=False, block_index=0, **kargs):
        """
        Load the first block in the file.
        """
        return self.read_all_blocks(lazy=lazy)[block_index]

    def _get_segment(self, block_name, segment_name):
        # If we've already created a Block with the given name return it,
        #   otherwise create it now and store it in self._blocks.
        # If we've already created a Segment in the given block, return it,
        #   otherwise create it now and return it.
        if block_name in self._blocks:
            block = self._blocks[block_name]
        else:
            block = Block(name=block_name, **self.global_block_metadata)
            self._blocks[block_name] = block
        segment = None
        for seg in block.segments:
            if segment_name == seg.name:
                segment = seg
                break
        if segment is None:
            segment = Segment(name=segment_name)
            block.segments.append(segment)
        return segment

    def _read_epochs_group(self, lazy):
        if self._file.epochs is not None:
            try:
                # NWB files created by Neo store the segment, block and epoch names as extra
                # columns
                segment_names = self._file.epochs.segment[:]
                block_names = self._file.epochs.block[:]
                epoch_names = self._file.epochs._name[:]
            except AttributeError:
                epoch_names = None

            if epoch_names is not None:
                unique_epoch_names = np.unique(epoch_names)
                for epoch_name in unique_epoch_names:
                    (index,) = np.where((epoch_names == epoch_name))
                    epoch = EpochProxy(self._file.epochs, epoch_name, index)
                    if not lazy:
                        epoch = epoch.load()
                    segment_name = np.unique(segment_names[index])
                    block_name = np.unique(block_names[index])
                    if segment_name.size != block_name.size == 1:
                        raise ValueError("the `segment_name` and the `block_name` should be the same")
                    segment = self._get_segment(block_name[0], segment_name[0])
                    segment.epochs.append(epoch)
            else:
                epoch = EpochProxy(self._file.epochs)
                if not lazy:
                    epoch = epoch.load()
                segment = self._get_segment("default", "default")
                segment.epochs.append(epoch)

    def _read_timeseries_group(self, group_name, lazy):
        import pynwb

        group = getattr(self._file, group_name)
        for timeseries in group.values():
            try:
                # NWB files created by Neo store the segment and block names in the comments field
                hierarchy = json.loads(timeseries.comments)
            except JSONDecodeError:
                # For NWB files created with other applications, we put everything in a single
                # segment in a single block
                # todo: investigate whether there is a reliable way to create multiple segments,
                #       e.g. using Trial information
                block_name = "default"
                segment_name = "default"
            else:
                block_name = hierarchy["block"]
                segment_name = hierarchy["segment"]
            segment = self._get_segment(block_name, segment_name)
            if isinstance(timeseries, pynwb.misc.AnnotationSeries):
                event = EventProxy(timeseries, group_name)
                if not lazy:
                    event = event.load()
                segment.events.append(event)
            elif timeseries.rate:  # AnalogSignal
                signal = AnalogSignalProxy(timeseries, group_name)
                if not lazy:
                    signal = signal.load()
                segment.analogsignals.append(signal)
            else:  # IrregularlySampledSignal
                signal = AnalogSignalProxy(timeseries, group_name)
                if not lazy:
                    signal = signal.load()
                segment.irregularlysampledsignals.append(signal)

    def _read_units(self, lazy):
        if self._file.units:
            for id in range(len(self._file.units)):
                try:
                    # NWB files created by Neo store the segment and block names as extra columns
                    segment_name = self._file.units.segment[id]
                    block_name = self._file.units.block[id]
                except AttributeError:
                    # For NWB files created with other applications, we put everything in a single
                    # segment in a single block
                    segment_name = "default"
                    block_name = "default"
                segment = self._get_segment(block_name, segment_name)
                spiketrain = SpikeTrainProxy(self._file.units, id)
                if not lazy:
                    spiketrain = spiketrain.load()
                segment.spiketrains.append(spiketrain)

    def _read_acquisition_group(self, lazy):
        self._read_timeseries_group("acquisition", lazy)

    def _read_stimulus_group(self, lazy):
        self._read_timeseries_group("stimulus", lazy)

    def _build_global_annotations(self, blocks):
        annotations = defaultdict(set)
        for annotation_name in GLOBAL_ANNOTATIONS:
            if annotation_name in self.annotations:
                annotations[annotation_name] = self.annotations[annotation_name]
            else:
                for block in blocks:
                    if annotation_name in block.annotations:
                        try:
                            annotations[annotation_name].add(block.annotations[annotation_name])
                        except TypeError:
                            if annotation_name in POSSIBLE_JSON_FIELDS:
                                encoded = json.dumps(block.annotations[annotation_name])
                                annotations[annotation_name].add(encoded)
                            else:
                                raise
                if annotation_name in annotations:
                    if len(annotations[annotation_name]) > 1:
                        raise NotImplementedError(f"We don't yet support multiple values for {annotation_name}")
                    # take single value from set
                    (annotations[annotation_name],) = annotations[annotation_name]

        if "identifier" not in annotations:
            annotations["identifier"] = str(self.filename)
        if "session_description" not in annotations:
            annotations["session_description"] = blocks[0].description or str(self.filename)
            # need to use str() here because self.filename may be a pathlib path object
            # todo: concatenate descriptions of multiple blocks if different
        if annotations.get("session_start_time", None) is None:
            if "rec_datetime" in annotations:
                annotations["session_start_time"] = annotations["rec_datetime"]
            else:
                raise Exception("Writing to NWB requires an annotation 'session_start_time'")
        return annotations

    def write_all_blocks(self, blocks, validate=True, **kwargs):
        """
        Write list of blocks to the file
        """
        import pynwb

        global_annotations = self._build_global_annotations(blocks)
        self._nwbfile = pynwb.NWBFile(**global_annotations)

        if sum(statistics(block)["SpikeTrain"]["count"] for block in blocks) > 0:
            self._nwbfile.add_unit_column("_name", "the name attribute of the SpikeTrain")
            # nwbfile.add_unit_column('_description',
            # 'the description attribute of the SpikeTrain')
            self._nwbfile.add_unit_column("segment", "the name of the Neo Segment to which the SpikeTrain belongs")
            self._nwbfile.add_unit_column("block", "the name of the Neo Block to which the SpikeTrain belongs")

        if sum(statistics(block)["Epoch"]["count"] for block in blocks) > 0:
            self._nwbfile.add_epoch_column("_name", "the name attribute of the Epoch")
            # nwbfile.add_epoch_column('_description', 'the description attribute of the Epoch')
            self._nwbfile.add_epoch_column("segment", "the name of the Neo Segment to which the Epoch belongs")
            self._nwbfile.add_epoch_column("block", "the name of the Neo Block to which the Epoch belongs")

        for i, block in enumerate(blocks):
            self._write_block(block)

        if self.nwb_file_mode not in ("w",):
            raise ValueError("mode must be 'w' in order to write files")  # possibly expand to 'a'ppend later
        if self.nwb_file_mode == "w" and os.path.exists(self.filename):
            os.remove(self.filename)
        io_nwb = pynwb.NWBHDF5IO(self.filename, mode=self.nwb_file_mode)
        io_nwb.write(self._nwbfile)
        io_nwb.close()

        if validate:
            self.validate_file()

    def validate_file(self):
        import pynwb

        with pynwb.NWBHDF5IO(self.filename, "r") as io_validate:
            errors = pynwb.validate(io_validate, namespace="core")
            if errors:
                raise Exception(f"Errors found when validating {self.filename}")

    def write_block(self, block, **kwargs):
        """
        Write a single Block to the file
            :param block: Block to be written
        """
        return self.write_all_blocks([block], **kwargs)

    def _write_block(self, block):
        """
        Write a Block to the file
            :param block: Block to be written
        """
        electrodes = self._write_electrodes(self._nwbfile, block)
        if not block.name:
            block.name = f"block{self.blocks_written}"
        for i, segment in enumerate(block.segments):
            if segment.block is not block:
                raise TypeError(f"segment.block must be block it is {segment.block}")
            if not segment.name:
                segment.name = f"{block.name} : segment{i}"
            self._write_segment(self._nwbfile, segment, electrodes)
        self.blocks_written += 1

    def _write_electrodes(self, nwbfile, block):
        # this handles only icephys_electrode for now
        electrodes = {}
        devices = {}
        for segment in block.segments:
            for signal in chain(segment.analogsignals, segment.irregularlysampledsignals):
                if "nwb_electrode" in signal.annotations:
                    elec_meta = signal.annotations["nwb_electrode"].copy()
                    if elec_meta["name"] not in electrodes:
                        # todo: check for consistency if the name is already there
                        if elec_meta["device"]["name"] in devices:
                            device = devices[elec_meta["device"]["name"]]
                        else:
                            device = self._nwbfile.create_device(**elec_meta["device"])
                            devices[elec_meta["device"]["name"]] = device
                        elec_meta.pop("device")
                        electrodes[elec_meta["name"]] = self._nwbfile.create_icephys_electrode(
                            device=device, **elec_meta
                        )
        return electrodes

    def _write_segment(self, nwbfile, segment, electrodes):
        # maybe use NWB trials to store Segment metadata?
        for i, signal in enumerate(chain(segment.analogsignals, segment.irregularlysampledsignals)):
            if signal.segment is not segment:
                raise TypeError(f"signal.segment must be segment and is {signal.segment}")
            if hasattr(signal, "name"):
                signal.name = f"{segment.name} {signal.name} {i}"
                logging.warning(f"Warning signal name exists. New name: {signal.name}")
            else:
                signal.name = f"{segment.name} : analogsignal{signal.name} {i}"
            self._write_signal(self._nwbfile, signal, electrodes)

        for i, train in enumerate(segment.spiketrains):
            if train.segment is not segment:
                raise TypeError(f"train.segment must be segment and is {train.segment}")
            if not train.name:
                train.name = f"{segment.name} : spiketrain{i}"
            self._write_spiketrain(self._nwbfile, train)

        for i, event in enumerate(segment.events):
            if event.segment is not segment:
                raise TypeError(f"event.segment mst be segment and is {event.segment}")
            if hasattr(event, "name"):
                event.name = f"{segment.name}  {event.name} {i}"
                logging.warning(f"Warning event name exists. New name: {event.name}")
            else:
                event.name = f"{segment.name} : event{event.name} {i}"
            self._write_event(self._nwbfile, event)

        for i, epoch in enumerate(segment.epochs):
            if not epoch.name:
                epoch.name = f"{segment.name} : epoch{i}"
            self._write_epoch(self._nwbfile, epoch)

    def _write_signal(self, nwbfile, signal, electrodes):
        import pynwb

        hierarchy = {"block": signal.segment.block.name, "segment": signal.segment.name}
        if "nwb_neurodata_type" in signal.annotations:
            timeseries_class = get_class(*signal.annotations["nwb_neurodata_type"])
        else:
            timeseries_class = pynwb.TimeSeries  # default
        additional_metadata = {name[4:]: value for name, value in signal.annotations.items() if name.startswith("nwb:")}
        if "nwb_electrode" in signal.annotations:
            electrode_name = signal.annotations["nwb_electrode"]["name"]
            additional_metadata["electrode"] = electrodes[electrode_name]
        if timeseries_class != pynwb.TimeSeries:
            conversion, units = get_units_conversion(signal, timeseries_class)
            additional_metadata["conversion"] = conversion
        else:
            units = signal.units
        if hasattr(signal, "proxy_for") and signal.proxy_for in [AnalogSignal, IrregularlySampledSignal]:
            signal = signal.load()
        if issubclass(timeseries_class, pynwb.icephys.PatchClampSeries):
            if signal.shape[1] != 1:
                raise ValueError(
                    "To store patch clamp data in NWB, please ensure that each AnalogSignal"
                    f"contains only one channel. The current signal has {signal.shape[1]} channels."
                )
            # see https://github.com/NeurodataWithoutBorders/pynwb/issues/1300
            data = signal.ravel()  # convert to 1D
        else:
            data = signal
        if isinstance(signal, AnalogSignal):
            sampling_rate = signal.sampling_rate.rescale("Hz")
            tS = timeseries_class(
                name=signal.name,
                starting_time=time_in_seconds(signal.t_start),
                data=data,
                unit=units.dimensionality.string,
                rate=float(sampling_rate),
                comments=json.dumps(hierarchy),
                **additional_metadata,
            )
            # todo: try to add array_annotations via "control" attribute
        elif isinstance(signal, IrregularlySampledSignal):
            tS = timeseries_class(
                name=signal.name,
                data=data,
                unit=units.dimensionality.string,
                timestamps=signal.times.rescale("second").magnitude,
                comments=json.dumps(hierarchy),
                **additional_metadata,
            )
        else:
            raise TypeError(
                f"signal has type { signal.__class__.__name__}, should be AnalogSignal or IrregularlySampledSignal"
            )
        nwb_group = signal.annotations.get("nwb_group", "acquisition")
        add_method_map = {"acquisition": self._nwbfile.add_acquisition, "stimulus": self._nwbfile.add_stimulus}
        if nwb_group in add_method_map:
            add_time_series = add_method_map[nwb_group]
        else:
            raise NotImplementedError(f"NWB group '{nwb_group}' not yet supported")
        add_time_series(tS)
        return tS

    def _write_spiketrain(self, nwbfile, spiketrain):
        segment = spiketrain.segment
        if hasattr(spiketrain, "proxy_for") and spiketrain.proxy_for is SpikeTrain:
            spiketrain = spiketrain.load()
        self._nwbfile.add_unit(
            spike_times=spiketrain.rescale("s").magnitude,
            obs_intervals=[[float(spiketrain.t_start.rescale("s")), float(spiketrain.t_stop.rescale("s"))]],
            _name=spiketrain.name,
            # _description=spiketrain.description,
            segment=segment.name,
            block=segment.block.name,
        )
        # todo: handle annotations (using add_unit_column()?)
        # todo: handle Neo Units
        # todo: handle spike waveforms, if any (see SpikeEventSeries)
        return self._nwbfile.units

    def _write_event(self, nwbfile, event):
        import pynwb

        segment = event.segment
        if hasattr(event, "proxy_for") and event.proxy_for == Event:
            event = event.load()
        hierarchy = {"block": segment.block.name, "segment": segment.name}
        tS_evt = pynwb.misc.AnnotationSeries(
            name=event.name,
            data=event.labels,
            timestamps=event.times.rescale("second").magnitude,
            description=event.description or "",
            comments=json.dumps(hierarchy),
        )
        self._nwbfile.add_acquisition(tS_evt)
        return tS_evt

    def _write_epoch(self, nwbfile, epoch):
        segment = epoch.segment
        if hasattr(epoch, "proxy_for") and epoch.proxy_for == Epoch:
            epoch = epoch.load()
        for t_start, duration, label in zip(
            epoch.rescale("s").magnitude, epoch.durations.rescale("s").magnitude, epoch.labels
        ):
            self._nwbfile.add_epoch(
                t_start,
                t_start + duration,
                [label],
                [],
                _name=epoch.name,
                segment=segment.name,
                block=segment.block.name,
            )
        return self._nwbfile.epochs

    def close(self):
        if self._io_nwb:
            self._io_nwb.close()


class AnalogSignalProxy(BaseAnalogSignalProxy):
    common_metadata_fields = (
        # fields that are the same for all TimeSeries subclasses
        "comments",
        "description",
        "unit",
        "starting_time",
        "timestamps",
        "rate",
        "data",
        "starting_time_unit",
        "timestamps_unit",
        "electrode",
        "stream_id",
    )

    def __init__(self, timeseries, nwb_group):
        import pynwb

        self._timeseries = timeseries
        self.units = timeseries.unit
        if timeseries.conversion:
            self.units = _recompose_unit(timeseries.unit, timeseries.conversion)
        if timeseries.starting_time is not None:
            self.t_start = timeseries.starting_time * pq.s
        else:
            self.t_start = timeseries.timestamps[0] * pq.s
        if timeseries.rate:
            self.sampling_rate = timeseries.rate * pq.Hz
        else:
            self.sampling_rate = None
        self.name = timeseries.name
        self.annotations = {"nwb_group": nwb_group}
        self.description = try_json_field(timeseries.description)
        if isinstance(self.description, dict):
            self.annotations["notes"] = self.description
            if "name" in self.annotations:
                self.annotations.pop("name")
            self.description = None
        self.shape = self._timeseries.data.shape
        if len(self.shape) == 1:
            self.shape = (self.shape[0], 1)
        metadata_fields = list(timeseries.__nwbfields__)
        for field_name in self.__class__.common_metadata_fields:  # already handled
            try:
                metadata_fields.remove(field_name)
            except ValueError:
                pass
        for field_name in metadata_fields:
            value = getattr(timeseries, field_name)
            if hasattr(value, "fields"):
                value = nwb_obj_to_dict(value)
            if value is not None:
                self.annotations[f"nwb:{field_name}"] = value
            _check_annotations(value)  # tmp for easier debugging
        self.annotations["nwb_neurodata_type"] = (timeseries.__class__.__module__, timeseries.__class__.__name__)
        if hasattr(timeseries, "electrode"):
            # todo: once the Group class is available, we could add electrode metadata
            #       to a Group containing all signals that share that electrode
            #       This would reduce the amount of redundancy (repeated metadata in every signal)
            electrode_metadata = {"device": {}}
            metadata_fields = list(timeseries.electrode.__class__.__nwbfields__) + ["name"]
            metadata_fields.remove("device")  # needs special handling
            for field_name in metadata_fields:
                value = getattr(timeseries.electrode, field_name)
                if value is not None:
                    electrode_metadata[field_name] = value
            for field_name in timeseries.electrode.device.__class__.__nwbfields__:
                value = getattr(timeseries.electrode.device, field_name)
                if value is not None:
                    electrode_metadata["device"][field_name] = value
            self.annotations["nwb_electrode"] = electrode_metadata

    def load(self, time_slice=None, strict_slicing=True):
        """
        Load AnalogSignalProxy args:
            :param time_slice: None or tuple of the time slice expressed with quantities.
                            None is the entire signal.
            :param strict_slicing: True by default.
                Control if an error is raised or not when one of the time_slice members
                (t_start or t_stop) is outside the real time range of the segment.
        """
        i_start, i_stop, sig_t_start = None, None, self.t_start
        if time_slice:
            if self.sampling_rate is None:
                i_start, i_stop = np.searchsorted(self._timeseries.timestamps, time_slice)
            else:
                i_start, i_stop, sig_t_start = self._time_slice_indices(time_slice, strict_slicing=strict_slicing)
        signal = self._timeseries.data[i_start:i_stop]
        if self.sampling_rate is None:
            return IrregularlySampledSignal(
                self._timeseries.timestamps[i_start:i_stop] * pq.s,
                signal,
                units=self.units,
                t_start=sig_t_start,
                sampling_rate=self.sampling_rate,
                name=self.name,
                description=self.description,
                array_annotations=None,
                **self.annotations,
            )  # todo: timeseries.control / control_description

        else:
            return AnalogSignal(
                signal,
                units=self.units,
                t_start=sig_t_start,
                sampling_rate=self.sampling_rate,
                name=self.name,
                description=self.description,
                array_annotations=None,
                **self.annotations,
            )  # todo: timeseries.control / control_description


class EventProxy(BaseEventProxy):

    def __init__(self, timeseries, nwb_group):
        self._timeseries = timeseries
        self.name = timeseries.name
        self.annotations = {"nwb_group": nwb_group}
        self.description = try_json_field(timeseries.description)
        if isinstance(self.description, dict):
            self.annotations.update(self.description)
            self.description = None
        self.shape = self._timeseries.data.shape

    def load(self, time_slice=None, strict_slicing=True):
        """
        Load EventProxy args:
            :param time_slice: None or tuple of the time slice expressed with quantities.
                            None is the entire signal.
            :param strict_slicing: True by default.
                Control if an error is raised or not when one of the time_slice members
                (t_start or t_stop) is outside the real time range of the segment.
        """
        if time_slice:
            raise NotImplementedError("todo")
        else:
            times = self._timeseries.timestamps[:]
            labels = self._timeseries.data[:]
        return Event(times * pq.s, labels=labels, name=self.name, description=self.description, **self.annotations)


class EpochProxy(BaseEpochProxy):

    def __init__(self, time_intervals, epoch_name=None, index=None):
        """
        :param time_intervals: An epochs table,
            which is a specific TimeIntervals table that stores info about long periods
        :param epoch_name: (str)
            Name of the epoch object
        :param index: (np.array, slice)
            Slice object or array of bool values masking time_intervals to be used. In case of
            an array it has to have the same shape as `time_intervals`.
        """
        self._time_intervals = time_intervals
        if index is not None:
            self._index = index
            self.shape = (index.sum(),)
        else:
            self._index = slice(None)
            self.shape = (len(time_intervals),)
        self.name = epoch_name

    def load(self, time_slice=None, strict_slicing=True):
        """
        Load EpochProxy args:
            :param time_slice: None or tuple of the time slice expressed with quantities.
                            None is all of the intervals.
            :param strict_slicing: True by default.
                Control if an error is raised or not when one of the time_slice members
                (t_start or t_stop) is outside the real time range of the segment.
        """
        if time_slice:
            raise NotImplementedError("todo")
        else:
            start_times = self._time_intervals.start_time[self._index]
            stop_times = self._time_intervals.stop_time[self._index]
            durations = stop_times - start_times
            labels = self._time_intervals.tags[self._index]

        return Epoch(times=start_times * pq.s, durations=durations * pq.s, labels=labels, name=self.name)


class SpikeTrainProxy(BaseSpikeTrainProxy):

    def __init__(self, units_table, id):
        """
        :param units_table: A Units table
        (see https://pynwb.readthedocs.io/en/stable/pynwb.misc.html#pynwb.misc.Units)
        :param id: the cell/unit ID (integer)
        """
        self._units_table = units_table
        self.id = id
        self.units = pq.s
        try:
            obs_intervals = units_table.get_unit_obs_intervals(id)
        except KeyError:
            logger.warn("Unable to retrieve obs_intervals")
            t_start, t_stop = None, None
        else:
            if len(obs_intervals) == 0:
                t_start, t_stop = None, None
            elif len(obs_intervals) == 1:
                t_start, t_stop = obs_intervals[0]
                t_start = t_start * pq.s
                t_stop = t_stop * pq.s
            else:
                raise NotImplementedError("Can't yet handle multiple observation intervals")
        self.t_start = t_start
        self.t_stop = t_stop
        self.annotations = {"nwb_group": "acquisition"}
        try:
            # NWB files created by Neo store the name as an extra column
            self.name = units_table._name[id]
        except AttributeError:
            self.name = None
        self.shape = None  # no way to get this without reading the data

    def load(self, time_slice=None, strict_slicing=True):
        """
        Load SpikeTrainProxy args:
            :param time_slice: None or tuple of the time slice expressed with quantities.
                            None is the entire spike train.
            :param strict_slicing: True by default.
                Control if an error is raised or not when one of the time_slice members
                (t_start or t_stop) is outside the real time range of the segment.
        """
        interval = None
        if time_slice:
            interval = (float(t) for t in time_slice)  # convert from quantities
        spike_times = self._units_table.get_unit_spike_times(self.id, in_interval=interval)
        return SpikeTrain(
            spike_times * self.units,
            self.t_stop,
            units=self.units,
            # sampling_rate=array(1.) * Hz,
            t_start=self.t_start,
            # waveforms=None,
            # left_sweep=None,
            name=self.name,
            # file_origin=None,
            # description=None,
            # array_annotations=None,
            **self.annotations,
        )