refactor check_free_recall.py

Author: NxNiki

pyproject.toml

@@ -24,6 +24,7 @@ statsmodels = "^0.14.2"
 matplotlib = "^3.9.1"
 seaborn = "^0.13.2"
 pyarrow = "^17.0.0"
+scipy = "^1.14.1"
 
 [tool.poetry.dev-dependencies]
 black = "^23.0"

scripts/run_model.py

@@ -0,0 +1,119 @@
+import string
+from pathlib import Path
+
+import wandb
+
+from movie_decoding.main import pipeline
+from movie_decoding.utils.initializer import *
+
+# for patient in ['562', '563', '566', 'i728', '567', '572']:
+patient_list = ["i728", "572", "567", "566", "563", "562"]
+sd_list = [4, 4, 3.5, 4, 4, 3.5]
+# data_list = ['notch CAR4.5', 'notch CAR3.5', 'notch CAR4.5', 'notch CAR4', 'notch CAR3.5', 'notch CAR3.5']
+data_list = [
+    "notch CAR-quant-neg",
+    "notch CAR-quant-neg",
+    "notch CAR-quant-neg",
+    "notch CAR-quant-neg",
+    "notch CAR-quant-neg",
+    "notch CAR-quant-neg",
+]
+early_stop = [100, 100, 100, 50, 50, 75]
+for patient, sd, dd in zip(patient_list, sd_list, data_list):
+    print()
+    print("start: ", patient)
+    for data_type in ["clusterless"]:
+        for run in range(5, 6):
+            # root_path = os.path.dirname(os.path.abspath(__file__))
+            root_path = Path(__file__).parent.parent
+            # save the results
+            letters = string.ascii_lowercase
+            # suffix = ''.join(random.choice(letters) for i in range(3))
+            suffix = f"test53_optimalX_CARX_{run}"
+            if data_type == "clusterless":
+                use_clusterless = True
+                use_lfp = False
+                use_combined = False
+                model_architecture = "multi-vit"  #'multi-vit'
+            elif data_type == "lfp":
+                use_clusterless = False
+                use_lfp = True
+                use_combined = False
+                model_architecture = "multi-vit"
+            elif data_type == "combined":
+                use_clusterless = True
+                use_lfp = True
+                use_combined = True
+                model_architecture = "multi-crossvit"
+            else:
+                ValueError(f"undefined data_type: {data_type}")
+
+            args = initialize_configs(architecture=model_architecture)
+            args["seed"] = 42
+            args["device"] = "cuda:1"
+            args["patient"] = patient
+            args["use_spike"] = use_clusterless
+            args["use_lfp"] = use_lfp
+            args["use_combined"] = use_combined
+            args["use_spontaneous"] = False
+            if use_clusterless:
+                args["use_shuffle"] = True
+            elif use_lfp:
+                args["use_shuffle"] = False
+
+            args["use_bipolar"] = False
+            args["use_sleep"] = False
+            args["use_overlap"] = False
+            args["model_architecture"] = model_architecture
+
+            args["spike_data_mode"] = dd
+            args["spike_data_mode_inference"] = dd
+            args["spike_data_sd"] = [sd]
+            args["spike_data_sd_inference"] = sd
+            args["use_augment"] = False
+            args["use_long_input"] = False
+            args["use_shuffle_diagnostic"] = False
+            args["model_aggregate_type"] = "sum"
+
+            train_save_path = os.path.join(
+                root_path,
+                "results/8concepts/{}_{}_{}_{}/train".format(args["patient"], data_type, model_architecture, suffix),
+            )
+            valid_save_path = os.path.join(
+                root_path,
+                "results/8concepts/{}_{}_{}_{}/valid".format(args["patient"], data_type, model_architecture, suffix),
+            )
+            test_save_path = os.path.join(
+                root_path,
+                "results/8concepts/{}_{}_{}_{}/test".format(args["patient"], data_type, model_architecture, suffix),
+            )
+            memory_save_path = os.path.join(
+                root_path,
+                "results/8concepts/{}_{}_{}_{}/memory".format(args["patient"], data_type, model_architecture, suffix),
+            )
+            os.makedirs(train_save_path, exist_ok=True)
+            os.makedirs(valid_save_path, exist_ok=True)
+            os.makedirs(test_save_path, exist_ok=True)
+            os.makedirs(memory_save_path, exist_ok=True)
+            args["train_save_path"] = train_save_path
+            args["valid_save_path"] = valid_save_path
+            args["test_save_path"] = test_save_path
+            args["memory_save_path"] = memory_save_path
+
+            os.environ["WANDB_MODE"] = "offline"
+            # os.environ['WANDB_API_KEY'] = '5a6051ed615a193c44eb9f655b81703925460851'
+            wandb.login()
+            if use_lfp:
+                run_name = "LFP Concept level {} MultiEncoder".format(args["patient"])
+            else:
+                run_name = "Clusterless Concept level {} MultiEncoder".format(args["patient"])
+            wandb.init(project="24_Concepts", name=run_name, reinit=True, entity="24")
+
+            trainer = pipeline(args)
+
+            print("Start training")
+            # start_time = time.time()
+
+            trainer.train(args["epochs"], 1)
+    print("done: ", patient)
+    print()

src/movie_decoding/dataloader/patients.py

@@ -0,0 +1,71 @@
+from typing import Dict, List, Optional
+
+import numpy as np
+from pydantic import BaseModel, Field
+
+
+# Define the Event model
+class Event(BaseModel):
+    values: List[int] = Field(..., description="List of integer values for the event")
+
+
+# Define the Experiment model
+class Experiment(BaseModel):
+    events: Dict[str, Event] = Field(default_factory=dict, description="Dictionary of events within the experiment")
+    neural_data: Optional[np.ndarray] = Field(None, description="Neural recording data as a NumPy array")
+
+    def __getitem__(self, event_name: str) -> Event:
+        if event_name not in self.events:
+            self.events[event_name] = Event(values=[])
+        return self.events[event_name]
+
+    def __setitem__(self, event_name: str, values: List[int]):
+        self.events[event_name] = Event(values=values)
+
+
+# Define the Patient model
+class Patient(BaseModel):
+    experiments: Dict[str, Experiment] = Field(
+        default_factory=dict, description="Dictionary of experiments for the patient"
+    )
+
+    def __getitem__(self, experiment_name: str) -> Experiment:
+        if experiment_name not in self.experiments:
+            self.experiments[experiment_name] = Experiment()
+        return self.experiments[experiment_name]
+
+    def __setitem__(self, experiment_name: str, event_data: Dict[str, List[int]]):
+        experiment = self.experiments.get(experiment_name, Experiment())
+        for event_name, values in event_data.items():
+            experiment[event_name] = values
+        self.experiments[experiment_name] = experiment
+
+
+# Define the overall PatientsData model
+class PatientsData(BaseModel):
+    patients: Dict[str, Patient] = Field(default_factory=dict, description="Dictionary of patients")
+
+    def __getitem__(self, patient_id: str) -> Patient:
+        if patient_id not in self.patients:
+            self.patients[patient_id] = Patient()
+        return self.patients[patient_id]
+
+    def __setitem__(self, patient_id: str, experiment_data: Dict[str, Dict[str, List[int]]]):
+        patient = self.patients.get(patient_id, Patient())
+        for experiment_name, event_data in experiment_data.items():
+            patient[experiment_name] = event_data
+        self.patients[patient_id] = patient
+
+
+# Example usage within the module (can be removed or commented out for production use)
+if __name__ == "__main__":
+    patients_data = PatientsData()
+
+    # Using direct assignment for events
+    patients_data["567"]["free_recall1"]["LA"] = [1234, 23456]
+
+    # Adding neural data to the experiment
+    patients_data["567"]["free_recall1"].neural_data = np.array([1.0, 2.0, 3.0])
+
+    print(patients_data["567"]["free_recall1"]["LA"].values)
+    print(patients_data["567"]["free_recall1"].neural_data)