Simplifiy np deps (#174)

* except other errors for getting ecephys paths * fix logic for overwriting opto table, remove npc_sessions dependency * working on simplifying * get mtrain without np_session * factor out np_session * linting fixes * sort imports * remove deps from toml * refactor json_settings into job_settings and inherit from genericetl * have camstimephys inherit from genericetl * add model for camstimephys job settings * lint * get job_settings models to work in camstimephys class when passed as a dict * linting and cleanup * use genericetl.load in run_job * fix mekhlas reviews #174 * lint * finish mekhla feedback * fixes npc_mvr pin * removes hardcoded path --------- Co-authored-by: Mekhla Kapoor <[email protected]>
AllenNeuralDynamics · Oct 2, 2024 · 3503ebd · 3503ebd
1 parent aa8af6c
commit 3503ebd
Show file tree

Hide file tree

Showing 6 changed files with 209 additions and 168 deletions.
diff --git a/pyproject.toml b/pyproject.toml
@@ -66,11 +66,11 @@ mesoscope = [
 openephys = [
     "aind-metadata-mapper[schema]",
     "h5py >= 3.11.0",
-    "np_session >= 0.1.39",
     "npc_ephys >= 0.1.18 ; python_version >= '3.9'",
     "scipy >= 1.11.0",
     "pandas >= 2.2.2",
     "numpy >= 1.26.4",
+    "npc_mvr >= 0.1.6"
 ]
 
 dynamicrouting = [

diff --git a/src/aind_metadata_mapper/open_ephys/camstim_ephys_session.py b/src/aind_metadata_mapper/open_ephys/camstim_ephys_session.py
@@ -1,42 +1,45 @@
 """
-File containing CamstimEphysSession class
+File containing CamstimEphysSessionEtl class
 """
 
-import argparse
 import datetime
 import json
 import logging
 import re
 from pathlib import Path
+from typing import Union
 
-import np_session
 import npc_ephys
 import npc_mvr
-import npc_sessions
 import numpy as np
 import pandas as pd
 from aind_data_schema.components.coordinates import Coordinates3d
 from aind_data_schema.core.session import ManipulatorModule, Session, Stream
 from aind_data_schema_models.modalities import Modality
 
-import aind_metadata_mapper.open_ephys.utils.constants as constants
 import aind_metadata_mapper.open_ephys.utils.sync_utils as sync
 import aind_metadata_mapper.stimulus.camstim
+from aind_metadata_mapper.core import GenericEtl
+from aind_metadata_mapper.open_ephys.models import JobSettings
 
 logger = logging.getLogger(__name__)
 
 
-class CamstimEphysSession(aind_metadata_mapper.stimulus.camstim.Camstim):
+class CamstimEphysSessionEtl(
+    aind_metadata_mapper.stimulus.camstim.Camstim, GenericEtl
+):
     """
     An Ephys session, designed for OpenScope, employing neuropixel
     probes with visual and optogenetic stimulus from Camstim.
     """
 
     json_settings: dict = None
-    npexp_path: Path
+    session_path: Path
     recording_dir: Path
 
-    def __init__(self, session_id: str, json_settings: dict) -> None:
+    def __init__(
+        self, session_id: str, job_settings: Union[JobSettings, str, dict]
+    ) -> None:
         """
         Determine needed input filepaths from np-exp and lims, get session
         start and end times from sync file, write stim tables and extract
@@ -47,42 +50,42 @@ def __init__(self, session_id: str, json_settings: dict) -> None:
         different laser states for this experiment. Otherwise, the default is
         used from naming_utils.
         """
-        if json_settings.get("opto_conditions_map", None) is None:
-            self.opto_conditions_map = constants.DEFAULT_OPTO_CONDITIONS
+        if isinstance(job_settings, str):
+            job_settings_model = JobSettings.model_validate_json(job_settings)
+        elif isinstance(job_settings, dict):
+            job_settings_model = JobSettings(**job_settings)
         else:
-            self.opto_conditions_map = json_settings["opto_conditions_map"]
-        overwrite_tables = json_settings.get("overwrite_tables", False)
-
-        self.json_settings = json_settings
-        session_inst = np_session.Session(session_id)
-        self.mtrain = session_inst.mtrain
-        self.npexp_path = session_inst.npexp_path
-        self.folder = session_inst.folder
+            job_settings_model = job_settings
+        GenericEtl.__init__(self, job_settings=job_settings_model)
+
+        sessions_root = Path(self.job_settings.sessions_root)
+        self.folder = self.get_folder(session_id, sessions_root)
+        self.session_path = self.get_session_path(session_id, sessions_root)
         # sometimes data files are deleted on npexp so try files on lims
-        try:
-            self.recording_dir = npc_ephys.get_single_oebin_path(
-                session_inst.lims_path
-            ).parent
-        except FileNotFoundError:
-            self.recording_dir = npc_ephys.get_single_oebin_path(
-                session_inst.npexp_path
-            ).parent
+        # try:
+        #     self.recording_dir = npc_ephys.get_single_oebin_path(
+        #         session_inst.lims_path
+        #     ).parent
+        # except:
+        self.recording_dir = npc_ephys.get_single_oebin_path(
+            self.session_path
+        ).parent
 
         self.motor_locs_path = (
-            self.npexp_path / f"{self.folder}.motor-locs.csv"
+            self.session_path / f"{self.folder}.motor-locs.csv"
         )
-        self.pkl_path = self.npexp_path / f"{self.folder}.stim.pkl"
-        self.opto_pkl_path = self.npexp_path / f"{self.folder}.opto.pkl"
+        self.pkl_path = self.session_path / f"{self.folder}.stim.pkl"
+        self.opto_pkl_path = self.session_path / f"{self.folder}.opto.pkl"
         self.opto_table_path = (
-            self.npexp_path / f"{self.folder}_opto_epochs.csv"
+            self.session_path / f"{self.folder}_opto_epochs.csv"
         )
         self.stim_table_path = (
-            self.npexp_path / f"{self.folder}_stim_epochs.csv"
+            self.session_path / f"{self.folder}_stim_epochs.csv"
         )
-        self.sync_path = self.npexp_path / f"{self.folder}.sync"
+        self.sync_path = self.session_path / f"{self.folder}.sync"
 
         platform_path = next(
-            self.npexp_path.glob(f"{self.folder}_platform*.json")
+            self.session_path.glob(f"{self.folder}_platform*.json")
         )
         self.platform_json = json.loads(platform_path.read_text())
         self.project_name = self.platform_json["project"]
@@ -95,13 +98,17 @@ def __init__(self, session_id: str, json_settings: dict) -> None:
             f" session end: {self.session_end}"
         )
 
-        if not self.stim_table_path.exists() or overwrite_tables:
+        self.session_uuid = self.get_session_uuid()
+        self.mtrain_regimen = self.get_mtrain()
+
+        if not self.stim_table_path.exists() or (
+            self.job_settings.overwrite_tables
+        ):
             logger.debug("building stim table")
             self.build_stimulus_table()
-        if (
-            self.opto_pkl_path.exists()
-            and not self.opto_table_path.exists()
-            or overwrite_tables
+        if self.opto_pkl_path.exists() and (
+            not self.opto_table_path.exists() or
+            self.job_settings.overwrite_tables
         ):
             logger.debug("building opto table")
             self.build_optogenetics_table()
@@ -113,7 +120,17 @@ def __init__(self, session_id: str, json_settings: dict) -> None:
 
         self.available_probes = self.get_available_probes()
 
-    def generate_session_json(self) -> Session:
+    def run_job(self):
+        """Transforms all metadata for the session into relevant files"""
+        self._extract()
+        self._transform()
+        return self._load(self.session_json, self.session_path)
+
+    def _extract(self):
+        """TODO: refactor a lot of the __init__ code here"""
+        pass
+
+    def _transform(self) -> Session:
         """
         Creates the session schema json
         """
@@ -123,30 +140,19 @@ def generate_session_json(self) -> Session:
             ],
             session_start_time=self.session_start,
             session_end_time=self.session_end,
-            session_type=self.json_settings.get("session_type", ""),
-            iacuc_protocol=self.json_settings.get("iacuc_protocol", ""),
+            session_type=self.job_settings.session_type,
+            iacuc_protocol=self.job_settings.iacuc_protocol,
             rig_id=self.platform_json["rig_id"],
             subject_id=self.folder.split("_")[1],
             data_streams=self.data_streams(),
             stimulus_epochs=self.stim_epochs,
-            mouse_platform_name=self.json_settings.get(
-                "mouse_platform", "Mouse Platform"
-            ),
-            active_mouse_platform=self.json_settings.get(
-                "active_mouse_platform", False
-            ),
+            mouse_platform_name=self.job_settings.mouse_platform_name,
+            active_mouse_platform=self.job_settings.active_mouse_platform,
             reward_consumed_unit="milliliter",
             notes="",
         )
         return self.session_json
 
-    def write_session_json(self) -> None:
-        """
-        Writes the session json to a session.json file
-        """
-        self.session_json.write_standard_file(self.npexp_path)
-        logger.debug(f"File created at {str(self.npexp_path)}/session.json")
-
     @staticmethod
     def extract_probe_letter(probe_exp, s):
         """
@@ -214,7 +220,7 @@ def ephys_modules(self) -> list:
         """
         Return list of schema ephys modules for each available probe.
         """
-        newscale_coords = npc_sessions.get_newscale_coordinates(
+        newscale_coords = npc_ephys.get_newscale_coordinates(
             self.motor_locs_path
         )
 
@@ -289,7 +295,7 @@ def video_stream(self) -> Stream:
         Returns schema behavior videos stream for video timing
         """
         video_frame_times = npc_mvr.mvr.get_video_frame_times(
-            self.sync_path, self.npexp_path
+            self.sync_path, self.session_path
         )
 
         stream_first_time = min(
@@ -320,36 +326,12 @@ def data_streams(self) -> tuple[Stream, ...]:
         return tuple(data_streams)
 
 
-def parse_args() -> argparse.Namespace:
-    """
-    Parse Arguments
-    """
-    parser = argparse.ArgumentParser(
-        description="Generate a session.json file for an ephys session"
-    )
-    parser.add_argument(
-        "session_id",
-        help=(
-            "session ID (lims or np-exp foldername) or path to session"
-            "folder"
-        ),
-    )
-    parser.add_argument(
-        "json-settings",
-        help=(
-            'json containing at minimum the fields "session_type" and'
-            '"iacuc protocol"'
-        ),
-    )
-    return parser.parse_args()
-
-
 def main() -> None:
     """
     Run Main
     """
-    sessionETL = CamstimEphysSession(**vars(parse_args()))
-    sessionETL.generate_session_json()
+    sessionETL = CamstimEphysSessionEtl(**vars)
+    sessionETL.run_job()
 
 
 if __name__ == "__main__":

diff --git a/src/aind_metadata_mapper/open_ephys/models.py b/src/aind_metadata_mapper/open_ephys/models.py
@@ -0,0 +1,82 @@
+"""Module defining JobSettings for Mesoscope ETL"""
+
+from typing import Literal, Union
+from pathlib import Path
+from aind_metadata_mapper.core_models import BaseJobSettings
+
+DEFAULT_OPTO_CONDITIONS = {
+    "0": {
+        "duration": 0.01,
+        "name": "1Hz_10ms",
+        "condition": "10 ms pulse at 1 Hz",
+    },
+    "1": {
+        "duration": 0.002,
+        "name": "1Hz_2ms",
+        "condition": "2 ms pulse at 1 Hz",
+    },
+    "2": {
+        "duration": 1.0,
+        "name": "5Hz_2ms",
+        "condition": "2 ms pulses at 5 Hz",
+    },
+    "3": {
+        "duration": 1.0,
+        "name": "10Hz_2ms",
+        "condition": "2 ms pulses at 10 Hz",
+    },
+    "4": {
+        "duration": 1.0,
+        "name": "20Hz_2ms",
+        "condition": "2 ms pulses at 20 Hz",
+    },
+    "5": {
+        "duration": 1.0,
+        "name": "30Hz_2ms",
+        "condition": "2 ms pulses at 30 Hz",
+    },
+    "6": {
+        "duration": 1.0,
+        "name": "40Hz_2ms",
+        "condition": "2 ms pulses at 40 Hz",
+    },
+    "7": {
+        "duration": 1.0,
+        "name": "50Hz_2ms",
+        "condition": "2 ms pulses at 50 Hz",
+    },
+    "8": {
+        "duration": 1.0,
+        "name": "60Hz_2ms",
+        "condition": "2 ms pulses at 60 Hz",
+    },
+    "9": {
+        "duration": 1.0,
+        "name": "80Hz_2ms",
+        "condition": "2 ms pulses at 80 Hz",
+    },
+    "10": {
+        "duration": 1.0,
+        "name": "square_1s",
+        "condition": "1 second square pulse: continuously on for 1s",
+    },
+    "11": {"duration": 1.0, "name": "cosine_1s", "condition": "cosine pulse"},
+}
+
+
+class JobSettings(BaseJobSettings):
+    """Data to be entered by the user."""
+
+    job_settings_name: Literal["OpenEphys"] = "OpenEphys"
+    session_type: str
+    project_name: str
+    iacuc_protocol: str
+    description: str
+    sessions_root: Union[Path, str]
+    opto_conditions_map: dict = DEFAULT_OPTO_CONDITIONS
+    overwrite_tables: bool = False
+    mtrain_server: str
+    # TODO: use input_source and replace sessions_root, camstimephys.getfolder
+    input_source: str = "blah"
+    active_mouse_platform: bool = False
+    mouse_platform_name: str = "Mouse Platform"
diff --git a/src/aind_metadata_mapper/open_ephys/utils/behavior_utils.py b/src/aind_metadata_mapper/open_ephys/utils/behavior_utils.py
@@ -730,9 +730,9 @@ def fix_omitted_end_frame(stim_pres_table: pd.DataFrame) -> pd.DataFrame:
         stim_pres_table[stim_pres_table["omitted"]]["start_frame"]
         + median_stim_frame_duration
     )
-    stim_pres_table.loc[
-        stim_pres_table["omitted"], "end_frame"
-    ] = omitted_end_frames
+    stim_pres_table.loc[stim_pres_table["omitted"], "end_frame"] = (
+        omitted_end_frames
+    )
 
     stim_dtypes = stim_pres_table.dtypes.to_dict()
     stim_dtypes["start_frame"] = int
@@ -796,9 +796,9 @@ def compute_is_sham_change(
                 if np.array_equal(
                     active_images, stim_image_names[passive_block_mask].values
                 ):
-                    stim_df.loc[
-                        passive_block_mask, "is_sham_change"
-                    ] = stim_df[active_block_mask]["is_sham_change"].values
+                    stim_df.loc[passive_block_mask, "is_sham_change"] = (
+                        stim_df[active_block_mask]["is_sham_change"].values
+                    )
 
     return stim_df.sort_index()