Add mock units table (#1875)

Co-authored-by: Ben Dichter <[email protected]>

CHANGELOG.md

@@ -9,6 +9,7 @@
   - Support `stimulus_template` as optional predefined column in `IntracellularStimuliTable`. @stephprince [#1815](https://github.com/NeurodataWithoutBorders/pynwb/pull/1815)
   - Support `NWBDataInterface` and `DynamicTable` in `NWBFile.stimulus`. @rly [#1842](https://github.com/NeurodataWithoutBorders/pynwb/pull/1842)
 - Added support for python 3.12 and upgraded dependency versions. This also includes infrastructure updates for developers. @mavaylon1 [#1853](https://github.com/NeurodataWithoutBorders/pynwb/pull/1853)
+- Added `mock_Units` for generating Units tables.  @h-mayorquin [#1875](https://github.com/NeurodataWithoutBorders/pynwb/pull/1875)
 
 ### Bug fixes
 - Fix bug with reading file with linked `TimeSeriesReferenceVectorData` @rly [#1865](https://github.com/NeurodataWithoutBorders/pynwb/pull/1865)

src/pynwb/testing/mock/ecephys.py

@@ -9,6 +9,7 @@
 from ...ecephys import ElectricalSeries, ElectrodeGroup, SpikeEventSeries
 from .device import mock_Device
 from .utils import name_generator
+from ...misc import Units
 
 
 def mock_ElectrodeGroup(
@@ -119,3 +120,32 @@ def mock_SpikeEventSeries(
         nwbfile.add_acquisition(spike_event_series)
 
     return spike_event_series
+
+
+def mock_Units(
+    num_units: int = 10,
+    max_spikes_per_unit: int = 10,
+    seed: int = 0,
+    nwbfile: Optional[NWBFile] = None,
+) -> Units:
+
+    units_table = Units()
+    units_table.add_column(name="unit_name", description="a readable identifier for the unit")
+
+    rng = np.random.default_rng(seed=seed)
+
+    times = rng.random(size=(num_units, max_spikes_per_unit)).cumsum(axis=1)
+    spikes_per_unit = rng.integers(1, max_spikes_per_unit, size=num_units)
+
+    spike_times = []
+    for unit_index in range(num_units):
+
+        # Not all units have the same number of spikes
+        spike_times = times[unit_index, : spikes_per_unit[unit_index]]
+        unit_name = f"unit_{unit_index}"
+        units_table.add_unit(spike_times=spike_times, unit_name=unit_name)
+
+    if nwbfile is not None:
+        nwbfile.units = units_table
+
+    return units_table

tests/unit/test_mock.py

@@ -35,6 +35,7 @@
     mock_ElectrodeTable,
     mock_ElectricalSeries,
     mock_SpikeEventSeries,
+    mock_Units,
 )
 
 from pynwb.testing.mock.icephys import (
@@ -82,6 +83,7 @@
     mock_IntracellularElectrode,
     mock_CurrentClampStimulusSeries,
     mock_IntracellularRecordingsTable,
+    mock_Units,
 ]
 
 
@@ -119,3 +121,4 @@ def test_name_generator():
 
     assert name_generator("TimeSeries") == "TimeSeries"
     assert name_generator("TimeSeries") == "TimeSeries2"
+