Script to add marks per electrode per spike time per unit (doubly indexed column) to Units table

write_unit_marks.py

import datetime
import numpy as np
from pynwb import NWBFile, NWBHDF5IO, validate


nwbfile = NWBFile(
    session_description='session_description',
    identifier='identifier',
    session_start_time=datetime.datetime.now(datetime.timezone.utc),
)

device = nwbfile.create_device(
    name='device'
)
elec_group = nwbfile.create_electrode_group(
    name='electrodes',
    description='description',
    location='location',
    device=device
)
for i in range(128):
    nwbfile.add_electrode(
        x=np.nan,
        y=np.nan,
        z=np.nan,
        imp=np.nan,
        location='location',
        filtering='filtering',
        group=elec_group
    )

nwbfile.add_unit_column(
    name='marks',
    description=('The amplitude at each channel at the time of the spike. The first index represents the spike times, '
                 'i.e., there are the same number of index entries as there are spike times. The second index '
                 'represents electrodes, i.e., there are the same number of index entries as there are electrodes.'),
    index=2,  # <-- create a doubly ragged column (each unit can have a different number of spike times and electrodes)
)

nwbfile.add_unit(
    spike_times=[1.0, 2.0, 3.0],
    electrodes=[0, 2, 4, 6],  # indices into electrodes table (electrodes must be added to nwbfile first)
    marks=[[1.0, 1.2, 1.4, 1.6], [2.0, 2.2, 2.4, 2.6], [3.0, 3.2, 3.4, 3.6]]
)

# marks is indexed first by spike time, then by electrode --
# e.g., [1.0, 1.2, 1.4, 1.6] are the marks for the first spike time for the four electrodes.

nwbfile.add_unit(
    spike_times=[10.0, 20.0],
    electrodes=[1, 3, 5],  # indices into electrodes table (electrodes must be added to nwbfile first)
    marks=[[10.1, 10.3, 10.5], [20.1, 20.3, 20.5]]
)

# alternatively, you can index marks first by electrode, then by spike time. this is similar to how waveforms are
# stored in NWB. which method is best depends on your typical access patterns. whichever method you choose,
# you must be consistent and should document the method in the column description.
# nwbfile.add_unit(
#     spike_times=[1.0, 2.0, 3.0],
#     electrodes=[0, 2, 4, 6],  # indices into electrodes table
#     marks=[[1.0, 2.0, 3.0], [1.2, 2.2, 3.2], [1.4, 2.4, 3.4], [1.6, 2.6, 3.6]]
# )

filename = 'nwbfile.nwb'

with NWBHDF5IO(filename, 'w') as io:
    io.write(nwbfile)

with NWBHDF5IO(filename, 'r') as io:
    errors = validate(io)
    assert not errors
    nwbfile = io.read()
    print('all marks from unit 0:', nwbfile.units[0, 'marks'])
    print('all marks from unit 0, spike time 1:', nwbfile.units[0, 'marks'][1])
    print('mark from unit 0, spike time 1, electrode 2:', nwbfile.units[0, 'marks'][1][2])