update convert session script

src/constantinople_lab_to_nwb/schierek_embargo_2024/schierek_embargo_2024_convert_session.py

@@ -2,7 +2,7 @@
 
 import os
 from pathlib import Path
-from typing import Union
+from typing import Union, Optional
 
 from dateutil import tz
 from neuroconv.datainterfaces import OpenEphysRecordingInterface
@@ -16,7 +16,10 @@ def session_to_nwb(
     openephys_recording_folder_path: Union[str, Path],
     spike_sorting_folder_path: Union[str, Path],
     processed_spike_sorting_file_path: Union[str, Path],
+    raw_behavior_file_path: Union[str, Path],
     nwbfile_path: Union[str, Path],
+    column_name_mapping: Optional[dict] = None,
+    column_descriptions: Optional[dict] = None,
     stub_test: bool = False,
     overwrite: bool = False,
 ):
@@ -76,14 +79,35 @@ def session_to_nwb(
             dict(
                 ProcessedSorting=dict(
                     file_path=processed_spike_sorting_file_path, sampling_frequency=sampling_frequency
-                )
+                ),
+                ProcessedBehavior=dict(
+                    file_path=processed_spike_sorting_file_path,
+                ),
             )
         )
         conversion_options.update(dict(ProcessedSorting=dict(write_as="processing", stub_test=False)))
+        conversion_options.update(
+            dict(
+                ProcessedBehavior=dict(column_name_mapping=column_name_mapping, column_descriptions=column_descriptions)
+            )
+        )
 
     # Add Behavior
-    # source_data.update(dict(Behavior=dict()))
-    # conversion_options.update(dict(Behavior=dict()))
+    source_data.update(dict(RawBehavior=dict(file_path=raw_behavior_file_path)))
+    # Exclude some task arguments from the trials table that are the same for all trials
+    task_arguments_to_exclude = [
+        "BlockLengthTest",
+        "BlockLengthAd",
+        "TrialsStage2",
+        "TrialsStage3",
+        "TrialsStage4",
+        "TrialsStage5",
+        "TrialsStage6",
+        "TrialsStage8",
+        "CTrial",
+        "HiITI",
+    ]
+    conversion_options.update(dict(RawBehavior=dict(task_arguments_to_exclude=task_arguments_to_exclude)))
 
     recording_folder_name = recording_folder_path.stem
     subject_id, session_id = recording_folder_name.split("_", maxsplit=1)
@@ -108,13 +132,19 @@ def session_to_nwb(
     metadata["NWBFile"].update(
         session_start_time=session_start_time.replace(tzinfo=tzinfo),
         session_id=session_id,
+        # TODO: add protocol name for behavior task
     )
 
     # Update default metadata with the editable in the corresponding yaml file
     editable_metadata_path = Path(__file__).parent / "schierek_embargo_2024_metadata.yaml"
     editable_metadata = load_dict_from_file(editable_metadata_path)
     metadata = dict_deep_update(metadata, editable_metadata)
 
+    # Update behavior metadata
+    behavior_metadata_path = Path(__file__).parent / "metadata" / "schierek_embargo_2024_behavior_metadata.yaml"
+    behavior_metadata = load_dict_from_file(behavior_metadata_path)
+    metadata = dict_deep_update(metadata, behavior_metadata)
+
     metadata["Subject"].update(subject_id=subject_id)
 
     # Run conversion
@@ -134,6 +164,49 @@ def session_to_nwb(
         openephys_recording_folder_path / "Record Node 117/experiment1/recording1/continuous/Neuropix-PXI-119.ProbeA-AP"
     )
     processed_sorting_file_path = Path("/Volumes/T9/Constantinople/Ephys Data/J076_2023-12-12.mat")
+    bpod_file_path = Path("/Volumes/T9/Constantinople/raw_Bpod/J076/DataFiles/J076_RWTautowait2_20231212_145250.mat")
+
+    # The column name mapping is used to rename the columns in the processed data to more descriptive column names. (optional)
+    column_name_mapping = dict(
+        hits="is_rewarded",
+        vios="is_violation",
+        optout="is_opt_out",
+        wait_time="wait_time",
+        wait_thresh="wait_time_threshold",
+        wait_for_cpoke="wait_for_center_poke",
+        zwait_for_cpoke="z_scored_wait_for_center_poke",
+        RewardedSide="rewarded_port",
+        Cled="center_poke_times",
+        Lled="left_poke_times",
+        Rled="right_poke_times",
+        l_opt="left_opt_out_times",
+        r_opt="right_opt_out_times",
+        ReactionTime="reaction_time",
+        slrt="short_latency_reaction_time",
+        iti="inter_trial_interval",
+    )
+    # The column descriptions are used to add descriptions to the columns in the processed data. (optional)
+    column_descriptions = dict(
+        hits="Whether the subject received reward for each trial.",
+        vios="Whether the subject violated the trial by not maintaining center poke for the time required by 'nose_in_center'.",
+        optout="Whether the subject opted out for each trial.",
+        wait_time="The wait time for the subject for for each trial in seconds, after removing outliers."
+        " For hit trials (when reward was delivered) the wait time is equal to the reward delay."
+        " For opt-out trials, the wait time is equal to the time waited from trial start to opting out.",
+        wait_for_cpoke="The time between side port poke and center poke in seconds, includes the time when the subject is consuming the reward.",
+        zwait_for_cpoke="The z-scored wait_for_cpoke using all trials.",
+        RewardedSide="The rewarded port (Left or Right) for each trial.",
+        Cled="The time of center port LED on/off for each trial (2 x ntrials).",
+        Lled="The time of left port LED on/off for each trial (2 x ntrials).",
+        Rled="The time of right port LED on/off for each trial (2 x ntrials).",
+        l_opt="The time of left port entered/exited for each trial (2 x ntrials).",
+        r_opt="The time of right port entered/exited for each trial (2 x ntrials).",
+        ReactionTime="The reaction time in seconds.",
+        slrt="The short-latency reaction time in seconds.",
+        iti="The time to initiate trial in seconds (the time between the end of the consummatory period and the time to initiate the next trial).",
+        wait_thresh="The threshold in seconds to remove wait-times (mean + 1*std of all cumulative wait-times).",
+    )
+
     nwbfile_path = Path("/Volumes/T9/Constantinople/nwbfiles/J076_2023-12-12_14-52-04.nwb")
     if not nwbfile_path.parent.exists():
         os.makedirs(nwbfile_path.parent, exist_ok=True)
@@ -145,6 +218,9 @@ def session_to_nwb(
         openephys_recording_folder_path=openephys_recording_folder_path,
         spike_sorting_folder_path=phy_sorting_folder_path,
         processed_spike_sorting_file_path=processed_sorting_file_path,
+        raw_behavior_file_path=bpod_file_path,
+        column_name_mapping=column_name_mapping,
+        column_descriptions=column_descriptions,
         nwbfile_path=nwbfile_path,
         stub_test=stub_test,
         overwrite=overwrite,