enforce same size spatial node arrays

Author: samuelbray32

src/trodes_to_nwb/convert.py

@@ -109,6 +109,7 @@ def create_nwbs(
     output_dir: str = "/stelmo/nwb/raw",
     video_directory: str = "",
     convert_video: bool = False,
+    fs_gui_dir: str = "",
     n_workers: int = 1,
     query_expression: str | None = None,
     disable_ptp: bool = False,
@@ -170,6 +171,8 @@ def pass_func(args):
                     output_dir,
                     video_directory,
                     convert_video,
+                    fs_gui_dir,
+                    disable_ptp,
                     behavior_only=behavior_only,
                 )
                 return True
@@ -198,6 +201,7 @@ def pass_func(args):
                 output_dir,
                 video_directory,
                 convert_video,
+                fs_gui_dir,
                 disable_ptp,
                 behavior_only=behavior_only,
             )
@@ -211,6 +215,7 @@ def _create_nwb(
     output_dir: str = "/stelmo/nwb/raw",
     video_directory: str = "",
     convert_video: bool = False,
+    fs_gui_dir: str = "",
     disable_ptp: bool = False,
     behavior_only: bool = False,
 ):
@@ -278,7 +283,6 @@ def _create_nwb(
         nwb_file, metadata, session_df, video_directory, convert_video
     )
     add_optogenetics(nwb_file, metadata, device_metadata)
-    add_optogenetic_epochs(nwb_file, metadata)
 
     if not behavior_only:
         add_electrode_groups(
@@ -312,6 +316,7 @@ def _create_nwb(
         session_df=session_df,
         neo_io=rec_dci.neo_io,
     )
+    add_optogenetic_epochs(nwb_file, metadata, fs_gui_dir)
     logger.info("ADDING POSITION")
     # add position
     if disable_ptp:

src/trodes_to_nwb/convert_optogenetics.py

@@ -321,14 +321,46 @@ def add_optogenetic_epochs(
     )
 
     # loop through each fsgui script, which can apply to multiple epochs
+    rows = []
     for fs_gui_metadata in opto_epochs_metadata:
         new_rows = compile_opto_entries(
             fs_gui_metadata=fs_gui_metadata,
             nwbfile=nwbfile,
             file_dir=file_dir,
         )
-        for row in new_rows:
-            opto_epochs_table.add_row(**row)
+        rows.extend(new_rows)
+
+    # ensure spatial geometry nodes have same shape in each row
+    max_shape = [0, 0, 0]
+    for row in rows:
+        nodes = row.get("spatial_filter_region_node_coordinates_in_pixels", None)
+        if nodes is not None:
+            max_shape = [
+                max(max_shape[0], nodes.shape[0]),
+                max(max_shape[1], nodes.shape[1]),
+                max(max_shape[2], nodes.shape[2]),
+            ]
+    if max_shape != [0, 0, 0]:
+        for row in rows:
+            nodes = row.get("spatial_filter_region_node_coordinates_in_pixels", None)
+            if nodes is None:
+                row["spatial_filter_region_node_coordinates_in_pixels"] = (
+                    np.ones(max_shape) * np.nan
+                )
+            else:
+                nodes = np.pad(
+                    nodes,
+                    (
+                        (0, max_shape[0] - nodes.shape[0]),
+                        (0, max_shape[1] - nodes.shape[1]),
+                        (0, max_shape[2] - nodes.shape[2]),
+                    ),
+                    mode="empty",
+                )
+                row["spatial_filter_region_node_coordinates_in_pixels"] = nodes
+
+    for row in rows:
+        opto_epochs_table.add_row(**row)
 
     nwbfile.add_time_intervals(opto_epochs_table)