Implement headstage sensor data separation with units and scaling

Co-authored-by: edeno <[email protected]>

Author: Copilot

sample_metadata_with_sensor_units.yml

@@ -0,0 +1,84 @@
+experimenter_name:
+  - lastname, firstname
+lab: Loren Frank Lab
+institution: UCSF
+experiment_description: Test Conversion with Headstage Sensor Units
+session_description: test yaml insertion with separate sensor TimeSeries
+session_id: "12345"
+keywords:
+  - testing
+  - headstage_sensors
+subject:
+  description: Long-Evans Rat
+  genotype: Obese Prone CD Rat
+  sex: M
+  species: Rattus pyctoris
+  subject_id: "54321"
+  date_of_birth: 2000-01-01T00:00:00.000Z
+  weight: 100
+data_acq_device:
+  - name: SpikeGadgets
+    system: SpikeGadgets
+    amplifier: Intan
+    adc_circuit: Intan
+cameras:
+  - id: 0
+    meters_per_pixel: 0.001
+    manufacturer: Allied Vision
+    model: model1
+    lens: lens1
+    camera_name: test camera 1
+tasks:
+  - task_name: Sleep
+    task_description: sleeping
+    task_environment: sleep box
+    camera_id:
+      - 0
+    task_epochs:
+      - 1
+associated_files: []
+associated_video_files: []
+
+# Updated units specification with individual sensor units
+units:
+  analog: "unspecified"  # Default for uncategorized analog channels
+  behavioral_events: "unspecified"  # Default for DIO events
+
+# Optional: specify custom units for specific sensor types
+sensor_units:
+  accelerometer: "g"  # Will override default
+  gyroscope: "d/s"    # Will override default
+  magnetometer: "unspecified"
+  analog_input: "V"   # For ECU analog inputs
+
+times_period_multiplier: 1
+raw_data_to_volts: 1.95e-07
+default_header_file_path: /stelmo/sam/test_data
+
+device:
+  name:
+    - device1
+
+# Enhanced behavioral_events with individual unit specifications
+behavioral_events:
+  - description: Din1
+    name: Light_1
+    comments: Indicator for reward delivery
+    unit: "unspecified"  # Digital events are dimensionless
+  - description: Din2
+    name: Light_2
+    unit: "unspecified"
+  - description: Dout2
+    name: Poke_1
+    unit: "unspecified"
+  - description: ECU_Ain1
+    name: Analog_Input_1
+    comments: Custom analog input channel
+    unit: "V"  # Voltage
+  - description: ECU_Ain2  
+    name: Analog_Input_2
+    comments: Another analog input
+    unit: "mV"  # Millivolts
+
+electrode_groups: []
+ntrode_electrode_group_channel_map: []

src/trodes_to_nwb/convert.py

@@ -309,6 +309,7 @@ def _create_nwb(
         rec_filepaths,
         timestamps=rec_dci_timestamps,
         behavior_only=behavior_only,
+        metadata=metadata,
     )
     logger.info("ADDING SAMPLE COUNTS")
     add_sample_count(nwb_file, rec_dci)

src/trodes_to_nwb/convert_analog.py

@@ -1,94 +1,230 @@
 """Module for handling the conversion of ECU analog and headstage sensor data streams from Trodes .rec files to NWB format."""
 
+import re
 from xml.etree import ElementTree
 
 import numpy as np
 import pynwb
 from hdmf.backends.hdf5 import H5DataIO
-from pynwb import NWBFile
+from pynwb import NWBFile, TimeSeries
 
 from trodes_to_nwb import convert_rec_header
 from trodes_to_nwb.convert_ephys import RecFileDataChunkIterator
 
 DEFAULT_CHUNK_TIME_DIM = 16384
 DEFAULT_CHUNK_MAX_CHANNEL_DIM = 32
 
+# Sensor type definitions with scaling factors and units
+SENSOR_TYPE_CONFIG = {
+    'accelerometer': {
+        'pattern': r'Headstage_Accel[XYZ]',
+        'scaling_factor': 0.000061,  # Convert to g units
+        'unit': 'g',
+        'description': 'Headstage accelerometer data'
+    },
+    'gyroscope': {
+        'pattern': r'Headstage_Gyro[XYZ]',
+        'scaling_factor': 0.061,  # Convert to degrees/second
+        'unit': 'd/s', 
+        'description': 'Headstage gyroscope data'
+    },
+    'magnetometer': {
+        'pattern': r'Headstage_Mag[XYZ]',
+        'scaling_factor': 1.0,  # No scaling specified in issue
+        'unit': 'unspecified',
+        'description': 'Headstage magnetometer data'
+    },
+    'analog_input': {
+        'pattern': r'(ECU_Ain\d+|Controller_Ain\d+)',
+        'scaling_factor': 1.0,
+        'unit': 'unspecified',
+        'description': 'Analog input channel'
+    }
+}
+
+
+def _categorize_sensor_channels(channel_names: list[str]) -> dict[str, list[str]]:
+    """Categorize sensor channels by type based on naming patterns.
+    
+    Parameters
+    ----------
+    channel_names : list[str]
+        List of channel names to categorize
+        
+    Returns
+    -------
+    dict[str, list[str]]
+        Dictionary mapping sensor types to lists of channel names
+    """
+    categorized = {}
+    
+    for sensor_type, config in SENSOR_TYPE_CONFIG.items():
+        pattern = config['pattern']
+        matching_channels = [name for name in channel_names if re.match(pattern, name)]
+        if matching_channels:
+            categorized[sensor_type] = matching_channels
+    
+    # Handle uncategorized channels
+    categorized_flat = [name for channels in categorized.values() for name in channels]
+    uncategorized = [name for name in channel_names if name not in categorized_flat]
+    if uncategorized:
+        categorized['other'] = uncategorized
+    
+    return categorized
+
+
+def _create_sensor_timeseries(
+    sensor_type: str,
+    channel_names: list[str], 
+    data: np.ndarray,
+    timestamps: np.ndarray,
+    metadata: dict = None
+) -> TimeSeries:
+    """Create a TimeSeries object for a specific sensor type.
+    
+    Parameters
+    ----------
+    sensor_type : str
+        Type of sensor (accelerometer, gyroscope, etc.)
+    channel_names : list[str]
+        Names of channels for this sensor type
+    data : np.ndarray
+        Raw sensor data 
+    timestamps : np.ndarray
+        Timestamps for the data
+    metadata : dict, optional
+        Metadata dictionary for custom units/scaling
+        
+    Returns
+    -------
+    TimeSeries
+        Configured TimeSeries object for the sensor type
+    """
+    config = SENSOR_TYPE_CONFIG.get(sensor_type, {
+        'scaling_factor': 1.0,
+        'unit': 'unspecified', 
+        'description': f'{sensor_type} data'
+    })
+    
+    # Apply scaling factor
+    scaled_data = data * config['scaling_factor']
+    
+    # Create description with channel names
+    description = f"{config['description']}: {', '.join(channel_names)}"
+    
+    # Use custom units from metadata if available
+    unit = config['unit']
+    if metadata and 'sensor_units' in metadata and sensor_type in metadata['sensor_units']:
+        unit = metadata['sensor_units'][sensor_type]
+    
+    return TimeSeries(
+        name=sensor_type,
+        description=description,
+        data=scaled_data,
+        unit=unit,
+        timestamps=timestamps,
+    )
+
 
 def add_analog_data(
     nwbfile: NWBFile,
     rec_file_path: list[str],
     timestamps: np.ndarray = None,
     behavior_only: bool = False,
+    metadata: dict = None,
     **kwargs,
 ) -> None:
-    """Adds analog streams to the nwb file.
+    """Adds analog streams to the nwb file as separate TimeSeries objects for each sensor type.
 
     Parameters
     ----------
     nwbfile : NWBFile
         nwb file being assembled
-    recfile : list[str]
+    rec_file_path : list[str]
         ordered list of file paths to all recfiles with session's data
+    timestamps : np.ndarray, optional
+        timestamps for the data
+    behavior_only : bool, optional
+        if True, only include behavioral data
+    metadata : dict, optional
+        metadata dictionary for custom units and scaling
     """
-    # TODO: ADD HEADSTAGE DATA
-
-    # get the ids of the analog channels from the first rec file header
+    # Get the ids of the analog channels from the first rec file header
     root = convert_rec_header.read_header(rec_file_path[0])
     hconf = root.find("HardwareConfiguration")
     ecu_conf = None
     for conf in hconf:
         if conf.attrib["name"] == "ECU":
             ecu_conf = conf
             break
-    analog_channel_ids = []
+    
+    # Get ECU analog channel IDs
+    ecu_analog_channel_ids = []
     for channel in ecu_conf:
         if channel.attrib["dataType"] == "analog":
-            analog_channel_ids.append(channel.attrib["id"])
+            ecu_analog_channel_ids.append(channel.attrib["id"])
 
-    # make the data chunk iterator
-    # TODO use the stream name instead of the stream index to be more robust
+    # Make the data chunk iterator for ECU analog data
     rec_dci = RecFileDataChunkIterator(
         rec_file_path,
-        nwb_hw_channel_order=analog_channel_ids,
+        nwb_hw_channel_order=ecu_analog_channel_ids,
         stream_id="ECU_analog",
         is_analog=True,
         timestamps=timestamps,
         behavior_only=behavior_only,
     )
 
-    # add headstage channel IDs to the list of analog channel IDs
-    analog_channel_ids.extend(rec_dci.neo_io[0].multiplexed_channel_xml.keys())
-
-    # (16384, 32) chunks of dtype int16 (2 bytes) is 1 MB, which is recommended
-    # by studies by the NWB team.
-    # could also add compression here. zstd/blosc-zstd are recommended by the NWB team, but
-    # they require the hdf5plugin library to be installed. gzip is available by default.
-    data_data_io = H5DataIO(
-        rec_dci,
-        chunks=(
-            DEFAULT_CHUNK_TIME_DIM,
-            min(len(analog_channel_ids), DEFAULT_CHUNK_MAX_CHANNEL_DIM),
-        ),
-    )
+    # Get headstage sensor channel IDs from multiplexed channels
+    headstage_channel_ids = list(rec_dci.neo_io[0].multiplexed_channel_xml.keys())
+    all_analog_channel_ids = ecu_analog_channel_ids + headstage_channel_ids
 
-    # make the objects to add to the nwb file
-    nwbfile.create_processing_module(
-        name="analog", description="Contains all analog data"
-    )
-    analog_events = pynwb.behavior.BehavioralEvents(name="analog")
-    analog_events.add_timeseries(
-        pynwb.TimeSeries(
-            name="analog",
-            description=__merge_row_description(
-                analog_channel_ids
-            ),  # NOTE: matches rec_to_nwb system
-            data=data_data_io,
-            timestamps=rec_dci.timestamps,
-            unit="-1",
-        )
-    )
-    # add it to the nwb file
-    nwbfile.processing["analog"].add(analog_events)
+    # Process ECU analog channels if any exist
+    if ecu_analog_channel_ids:
+        # Get ECU analog data (without headstage data)
+        ecu_data = rec_dci._get_data((slice(None), slice(0, len(ecu_analog_channel_ids))))
+        
+        # Categorize ECU analog channels
+        ecu_categorized = _categorize_sensor_channels(ecu_analog_channel_ids)
+        
+        # Create TimeSeries for each ECU sensor type
+        for sensor_type, channel_names in ecu_categorized.items():
+            channel_indices = [ecu_analog_channel_ids.index(name) for name in channel_names]
+            sensor_data = ecu_data[:, channel_indices]
+            
+            timeseries = _create_sensor_timeseries(
+                sensor_type=f"ecu_{sensor_type}",
+                channel_names=channel_names,
+                data=sensor_data,
+                timestamps=rec_dci.timestamps,
+                metadata=metadata
+            )
+            
+            # Add to acquisition
+            nwbfile.add_acquisition(timeseries)
+
+    # Process headstage sensor channels if any exist  
+    if headstage_channel_ids:
+        # Get headstage sensor data
+        headstage_data = rec_dci.neo_io[0].get_analogsignal_multiplexed(headstage_channel_ids)
+        
+        # Categorize headstage channels by sensor type
+        headstage_categorized = _categorize_sensor_channels(headstage_channel_ids)
+        
+        # Create separate TimeSeries for each sensor type
+        for sensor_type, channel_names in headstage_categorized.items():
+            channel_indices = [headstage_channel_ids.index(name) for name in channel_names]
+            sensor_data = headstage_data[:, channel_indices]
+            
+            timeseries = _create_sensor_timeseries(
+                sensor_type=sensor_type,
+                channel_names=channel_names,
+                data=sensor_data,
+                timestamps=rec_dci.timestamps,
+                metadata=metadata
+            )
+            
+            # Add to acquisition  
+            nwbfile.add_acquisition(timeseries)
 
 
 def __merge_row_description(row_ids: list[str]) -> str:

src/trodes_to_nwb/convert_dios.py

@@ -29,11 +29,20 @@ def _get_channel_name_map(metadata: dict) -> dict[str, str]:
             raise ValueError(
                 f"Duplicate channel name {dio_event['description']} in metadata YAML"
             )
+        
+        # Get unit for this specific event type, or use default
+        unit = "unspecified"  # Default unit for digital events
+        if "unit" in dio_event:
+            unit = dio_event["unit"]
+        elif "units" in metadata and "behavioral_events" in metadata["units"]:
+            unit = metadata["units"]["behavioral_events"]
+            
         channel_name_map[dio_event["description"]] = {
             "name": dio_event["name"],
             "comments": (
                 dio_event["comments"] if "comments" in dio_event else "no comments"
             ),
+            "unit": unit
         }
     return channel_name_map
 
@@ -101,7 +110,7 @@ def add_dios(nwbfile: NWBFile, recfile: list[str], metadata: dict) -> None:
             comments=channel_name_map[channel_name]["comments"],
             description=channel_name,
             data=state_changes,
-            unit="-1",  # TODO change to "N/A",
+            unit=channel_name_map[channel_name]["unit"],
             timestamps=timestamps,  # TODO adjust timestamps
         )
         beh_events.add_timeseries(ts)