Merge pull request #4 from AllenNeuralDynamics/update_si

Update function following spikeinterface upgrade

pyproject.toml

@@ -22,7 +22,7 @@ dependencies = [
     'fpdf2',
     'scipy',
     'rich',
-    'spikeinterface[full]==0.100.7',
+    'spikeinterface[full]==0.101.2',
     'harp-python @ git+https://github.com/jsiegle/harp-python@decode-clock',
     'numpy==1.26.4'
 ]

src/aind_ephys_rig_qc/driftmap.json

@@ -0,0 +1,21 @@
+{
+    "detection": {
+        "peak_sign": "neg",
+        "detect_threshold": 5,
+        "exclude_sweep_ms": 0.1
+    },
+    "localization": {
+        "ms_before": 0.1,
+        "ms_after": 0.3,
+        "radius_um": 100.0
+    },
+    "cmr": {"reference": "global", "operator": "median"},
+    "highpass_filter": {"freq_min": 300.0, "margin_ms": 5.0},
+    "n_skip": 30,
+    "alpha": 0.15,
+    "vmin": -200,
+    "vmax": 0,
+    "cmap": "Greys_r",
+    "figsize": [10, 10],
+    "phase_shift": {"margin_ms": 100.0}
+}

src/aind_ephys_rig_qc/generate_report.py

@@ -31,6 +31,7 @@ def generate_qc_report(
     timestamp_alignment_method="local",
     original_timestamp_filename="original_timestamps.npy",
     num_chunks=3,
+    psd_chunk_size=150000,
     plot_drift_map=True,
     flip_NIDAQ=False,
 ):
@@ -59,6 +60,14 @@ def generate_qc_report(
         Whether to plot the drift map
 
     """
+    # Define log file path
+    outfile = os.path.join(directory, "ephys-rig-QC_output.txt")
+    with open(outfile, "a") as f:
+        f.write(datetime.now().strftime("%Y-%m-%d %H:%M:%S") + "\n")
+        f.write("Start processing..." + "\n")
+    # Redirect stdout to capture printed output
+    output_stream = io.StringIO()
+    sys.stdout = output_stream
 
     pdf = PdfReport("aind-ephys-rig-qc v" + package_version)
     pdf.add_page()
@@ -97,19 +106,19 @@ def generate_qc_report(
 
     print("Creating QC plots...")
     create_qc_plots(
-        pdf, directory, num_chunks=num_chunks, plot_drift_map=plot_drift_map
+        pdf,
+        directory,
+        num_chunks=num_chunks,
+        plot_drift_map=plot_drift_map,
+        psd_chunk_size=psd_chunk_size,
     )
 
     print("Saving QC report...")
     pdf.output(os.path.join(directory, report_name))
     print("Finished.")
-    output_content = output_stream.getvalue()
-
-    outfile = os.path.join(directory, "ephys-rig-QC_output.txt")
-
-    with open(outfile, "a") as output_file:
-        output_file.write(datetime.now().strftime("%Y-%m-%d %H:%M:%S") + "\n")
-        output_file.write(output_content)
+    with open(outfile, "a") as f:
+        f.write(datetime.now().strftime("%Y-%m-%d %H:%M:%S") + "\n")
+        f.write(output_stream.getvalue())
 
 
 def get_stream_info(directory):
@@ -214,7 +223,7 @@ def create_qc_plots(
     directory,
     num_chunks=3,
     raw_chunk_size=1000,
-    psd_chunk_size=10000,
+    psd_chunk_size=150000,
     plot_drift_map=True,
 ):
     """
@@ -320,8 +329,6 @@ def create_qc_plots(
 
 
 if __name__ == "__main__":
-
-
 
     if len(sys.argv) != 3:
         print("Two input arguments are required:")
@@ -335,23 +342,20 @@ def create_qc_plots(
             parameters = json.load(f)
         directory = sys.argv[1]
 
+        output_stream = io.StringIO()
+        sys.stdout = output_stream
+        print(datetime.now().strftime("%Y-%m-%d %H:%M:%S"))
         print("Running generate_report.py with parameters:")
         for param in parameters:
             print(f"  {param}: {parameters[param]}")
 
         if not os.path.exists(directory):
             raise ValueError(f"Data directory {directory} does not exist.")
 
-        output_stream = io.StringIO()
-        sys.stdout = output_stream
-        output_content = output_stream.getvalue()
-
         outfile = os.path.join(directory, "ephys-rig-QC_output.txt")
 
         with open(outfile, "a") as output_file:
-            output_file.write(
-                datetime.now().strftime("%Y-%m-%d %H:%M:%S") + "\n"
-            )
-            output_file.write(output_content)
+            print("Output written to: ", outfile)
+            output_file.write(output_stream.getvalue())
 
         generate_qc_report(directory, **parameters)

src/aind_ephys_rig_qc/parameters.json

@@ -1,7 +1,8 @@
 {
     "report_name" : "QC.pdf",
-    "timestamp_alignment_method" : "harp",
+    "timestamp_alignment_method" : "local",
     "original_timestamp_filename" : "original_timestamps.npy",
-    "plot_drift_map" : false,
-    "flip_NIDAQ" : true
+    "plot_drift_map" : true,
+    "flip_NIDAQ" : false,
+    "psd_chunk_size" : 150000
 }

src/aind_ephys_rig_qc/qc_figures.py

@@ -20,6 +20,8 @@
 from spikeinterface.sortingcomponents.peak_localization import (
     LocalizeCenterOfMass,
 )
+import json
+from pathlib import Path
 
 
 def plot_raw_data(
@@ -171,26 +173,14 @@ def plot_drift(directory, stream_name, block_index=0):
     """
 
     """set parameters for drift visualization"""
-    visualization_drift_params = {
-        "detection": {
-            "peak_sign": "neg",
-            "detect_threshold": 5,
-            "exclude_sweep_ms": 0.1,
-        },
-        "localization": {
-            "ms_before": 0.1,
-            "ms_after": 0.3,
-            "radius_um": 100.0,
-        },
-        "cmr": {"reference": "global", "operator": "median"},
-        "highpass_filter": {"freq_min": 300.0, "margin_ms": 5.0},
-        "n_skip": 30,
-        "alpha": 0.15,
-        "vmin": -200,
-        "vmax": 0,
-        "cmap": "Greys_r",
-        "figsize": [10, 10],
-    }
+    script_dir = Path(__file__).resolve().parent
+
+    # Path to the JSON file
+    visualization_drift_file = script_dir / "driftmap.json"
+    with open(visualization_drift_file, "r") as file:
+        visualization_drift_params = json.load(file)
+
+    print(visualization_drift_params)
 
     """ get blocks/experiments and streams info """
     si.set_global_job_kwargs(n_jobs=-1)
@@ -205,7 +195,7 @@ def plot_drift(directory, stream_name, block_index=0):
     n_skip = visualization_drift_params["n_skip"]
     alpha = visualization_drift_params["alpha"]
 
-    stream_names, _ = se.get_neo_streams("openephys", directory)
+    stream_names, _ = se.get_neo_streams("openephysbinary", directory)
     spike_stream = [
         curr_stream_name
         for curr_stream_name in stream_names
@@ -217,6 +207,11 @@ def plot_drift(directory, stream_name, block_index=0):
         stream_name=spike_stream,
         block_index=block_index,
     )
+    # phase shift
+    if visualization_drift_params["phase_shift"] is not None:
+        recording = spre.phase_shift(
+            recording, **visualization_drift_params["phase_shift"]
+        )
     # high-pass filter
     recording = spre.highpass_filter(
         recording, **visualization_drift_params["highpass_filter"]

src/aind_ephys_rig_qc/temporal_alignment.py

@@ -149,7 +149,7 @@ def search_harp_line(recording, directory, pdf=None):
         & (events.state == 1)
     ].line.unique()
 
-    stream_folder_names, _ = se.get_neo_streams("openephys", directory)
+    stream_folder_names, _ = se.get_neo_streams("openephysbinary", directory)
     stream_folder_names = [
         stream_folder_name.split("#")[-1]
         for stream_folder_name in stream_folder_names
@@ -292,7 +292,7 @@ def align_timestamps(  # noqa
     """
 
     session = Session(directory, mmap_timestamps=False)
-    stream_folder_names, _ = se.get_neo_streams("openephys", directory)
+    stream_folder_names, _ = se.get_neo_streams("openephysbinary", directory)
     stream_folder_names = [
         stream_folder_name.split("#")[-1]
         for stream_folder_name in stream_folder_names
@@ -340,8 +340,6 @@ def align_timestamps(  # noqa
             # detect discontinuities from sample numbers
             # and remove residual chunks to avoid misalignment
             sample_numbers = main_stream.sample_numbers
-            main_stream_start_sample = np.min(sample_numbers)
-            main_stream_start_sample = np.min(sample_numbers)
             sample_intervals = np.diff(sample_numbers)
             sample_intervals_cat, sample_intervals_counts = np.unique(
                 sample_intervals, return_counts=True
@@ -748,7 +746,7 @@ def align_timestamps_harp(
     """
 
     session = Session(directory, mmap_timestamps=False)
-    stream_folder_names, _ = se.get_neo_streams("openephys", directory)
+    stream_folder_names, _ = se.get_neo_streams("openephysbinary", directory)
     stream_folder_names = [
         stream_folder_name.split("#")[-1]
         for stream_folder_name in stream_folder_names