Merge pull request #1 from AllenNeuralDynamics/addpsd

Added drift map, harp alignment and modified psd

.github/workflows/test_and_lint.yml

@@ -10,7 +10,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: [ '3.8', '3.9', '3.10' ]
+        python-version: [ '3.9', '3.10', '3.11', '3.12' ]
     steps:
     - uses: actions/checkout@v3
     - name: Set up Python ${{ matrix.python-version }}

.gitignore

@@ -137,3 +137,14 @@ dmypy.json
 
 # MacOs
 **/.DS_Store
+
+# Vscode
+./vscode
+
+# generated files
+**/*localsync_timestamps.npy
+**/*original_timestamps.npy
+**/*qc.pdf
+**/*temporal_alignment.png
+**/*harp_line_search.png
+**/*ephys-rig-QC_output.txt

pyproject.toml

@@ -17,12 +17,13 @@ readme = "README.md"
 dynamic = ["version"]
 
 dependencies = [
-    'open-ephys-python-tools',
+    'open-ephys-python-tools>=0.1.9',
     'matplotlib',
     'fpdf2',
     'scipy',
     'rich',
-    'harp-python'
+    'spikeinterface[full]',
+    'harp-python @ git+https://github.com/jsiegle/harp-python@decode-clock'
 ]
 
 [project.optional-dependencies]
@@ -64,15 +65,15 @@ exclude = '''
 '''
 
 [tool.coverage.run]
-omit = ["*__init__*"]
+omit = ["*__init__*", "*temporal_alignment*"]
 source = ["aind_ephys_rig_qc", "tests"]
 
 [tool.coverage.report]
 exclude_lines = [
     "if __name__ == .__main__.:",
     "from",
     "import",
-    "pragma: no cover"
+    "pragma: no cover",
 ]
 fail_under = 100
 

src/aind_ephys_rig_qc/generate_report.py

@@ -2,25 +2,36 @@
 Generates a PDF report from an Open Ephys data directory
 """
 
+import io
 import json
 import os
 import sys
+from datetime import datetime
 
 import numpy as np
 import pandas as pd
 from open_ephys.analysis import Session
 
 from aind_ephys_rig_qc import __version__ as package_version
 from aind_ephys_rig_qc.pdf_utils import PdfReport
-from aind_ephys_rig_qc.qc_figures import plot_power_spectrum, plot_raw_data
-from aind_ephys_rig_qc.temporal_alignment import align_timestamps
+from aind_ephys_rig_qc.qc_figures import (
+    plot_drift,
+    plot_power_spectrum,
+    plot_raw_data,
+)
+from aind_ephys_rig_qc.temporal_alignment import (
+    align_timestamps,
+    align_timestamps_harp,
+)
 
 
 def generate_qc_report(
     directory,
     report_name="QC.pdf",
     timestamp_alignment_method="local",
     original_timestamp_filename="original_timestamps.npy",
+    num_chunks=3,
+    plot_drift_map=True,
 ):
     """
     Generates a PDF report from an Open Ephys data directory
@@ -40,44 +51,68 @@ def generate_qc_report(
         Option 3: 'none' (don't align timestamps)
     original_timestamp_filename : str
         The name of the file for archiving the original timestamps
+    num_chunks : int
+        The number of chunks to split the data into for plotting raw data
+        and PSD
+    plot_drift_map : bool
+        Whether to plot the drift map
 
     """
 
+    output_stream = io.StringIO()
+    sys.stdout = output_stream
+
     pdf = PdfReport("aind-ephys-rig-qc v" + package_version)
     pdf.add_page()
     pdf.set_font("Helvetica", "B", size=12)
     pdf.set_y(30)
-    pdf.write(h=12, text=f"Overview of recordings in {directory}")
+    pdf.write(h=12, text="Overview of recordings in:")
+    pdf.set_y(40)
+    pdf.set_font("Helvetica", size=10)
+    pdf.write(h=8, text=f"{directory}")
 
     pdf.set_font("Helvetica", "", size=10)
-    pdf.set_y(45)
-    pdf.embed_table(get_stream_info(directory), width=pdf.epw)
+    pdf.set_y(60)
+    stream_info = get_stream_info(directory)
+    pdf.embed_table(stream_info, width=pdf.epw)
 
     if (
         timestamp_alignment_method == "local"
         or timestamp_alignment_method == "harp"
     ):
         # perform local alignment first in either case
+        print("Aligning timestamps to local clock...")
         align_timestamps(
             directory,
-            align_timestamps_to=timestamp_alignment_method,
             original_timestamp_filename=original_timestamp_filename,
             pdf=pdf,
         )
 
         if timestamp_alignment_method == "harp":
             # optionally align to Harp timestamps
-            align_timestamps(
-                directory,
-                align_timestamps_to=timestamp_alignment_method,
-                original_timestamp_filename=original_timestamp_filename,
-                pdf=pdf,
+            print("Aligning timestamps to Harp clock...")
+            align_timestamps_harp(
+                directory, pdf=pdf,
             )
 
-    create_qc_plots(pdf, directory)
+    print("Creating QC plots...")
+    create_qc_plots(
+        pdf, directory, num_chunks=num_chunks, plot_drift_map=plot_drift_map
+    )
 
+    print("Saving QC report...")
     pdf.output(os.path.join(directory, report_name))
 
+    output_content = output_stream.getvalue()
+
+    outfile = os.path.join(directory, "ephys-rig-QC_output.txt")
+
+    print("Finished.")
+
+    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)
+
 
 def get_stream_info(directory):
     """
@@ -107,18 +142,15 @@ def get_stream_info(directory):
     }
 
     for recordnode in session.recordnodes:
-
         current_record_node = os.path.basename(recordnode.directory).split(
             "Record Node "
         )[1]
 
         for recording in recordnode.recordings:
-
             current_experiment_index = recording.experiment_index
             current_recording_index = recording.recording_index
 
             for stream in recording.continuous:
-
                 sample_rate = stream.metadata["sample_rate"]
                 data_shape = stream.samples.shape
                 channel_count = data_shape[1]
@@ -179,33 +211,41 @@ def get_event_info(events, stream_name):
     return pd.DataFrame(data=event_info)
 
 
-def create_qc_plots(pdf, directory):
+def create_qc_plots(
+    pdf,
+    directory,
+    num_chunks=3,
+    raw_chunk_size=1000,
+    psd_chunk_size=10000,
+    plot_drift_map=True,
+):
     """
     Create QC plots for an Open Ephys data directory
     """
 
     session = Session(directory)
 
     for recordnode in session.recordnodes:
-
         current_record_node = os.path.basename(recordnode.directory).split(
             "Record Node "
         )[1]
 
-        for recording in recordnode.recordings:
-
+        for block_index, recording in enumerate(recordnode.recordings):
             current_experiment_index = recording.experiment_index
             current_recording_index = recording.recording_index
 
             events = recording.events
 
             for stream in recording.continuous:
-
                 duration = (
                     stream.samples.shape[0] / stream.metadata["sample_rate"]
                 )
+                start_frames = np.linspace(
+                    0, stream.samples.shape[0], num_chunks + 1, endpoint=False
+                )[1:]
 
                 stream_name = stream.metadata["stream_name"]
+                sample_rate = stream.metadata["sample_rate"]
 
                 pdf.add_page()
                 pdf.set_font("Helvetica", "B", size=12)
@@ -235,9 +275,7 @@ def create_qc_plots(pdf, directory):
                 pdf.write(h=10, text=f"Duration: {duration} s")
                 pdf.set_y(65)
                 pdf.write(
-                    h=10,
-                    text=f"Sample Rate: "
-                    f"{stream.metadata['sample_rate']} Hz",
+                    h=10, text=f"Sample Rate: " f"{sample_rate} Hz",
                 )
                 pdf.set_y(70)
                 pdf.write(h=10, text=f"Channels: {stream.samples.shape[1]}")
@@ -254,37 +292,49 @@ def create_qc_plots(pdf, directory):
                 pdf.set_y(120)
                 pdf.embed_figure(
                     plot_raw_data(
-                        stream.samples,
-                        stream.metadata["sample_rate"],
-                        stream_name,
+                        data=stream.samples,
+                        start_frames=start_frames,
+                        sample_rate=sample_rate,
+                        stream_name=stream_name,
+                        chunk_size=raw_chunk_size,
                     )
                 )
 
                 pdf.set_y(200)
                 pdf.embed_figure(
                     plot_power_spectrum(
-                        stream.samples,
-                        stream.metadata["sample_rate"],
-                        stream_name,
+                        data=stream.samples,
+                        start_frames=start_frames,
+                        sample_rate=sample_rate,
+                        stream_name=stream_name,
+                        chunk_size=psd_chunk_size,
                     )
                 )
 
+                if plot_drift_map:
+                    print("Plotting drift map for stream: ", stream_name)
+                    if "Probe" in stream_name and "LFP" not in stream_name:
+                        pdf.set_y(200)
+                        pdf.embed_figure(
+                            plot_drift(directory, stream_name, block_index)
+                        )
 
-if __name__ == "__main__":
 
+if __name__ == "__main__":
     if len(sys.argv) != 3:
         print("Two input arguments are required:")
         print(" 1. A data directory")
         print(" 2. A JSON parameters file")
     else:
-        with open(
-            sys.argv[2],
-            "r",
-        ) as f:
+        with open(sys.argv[2], "r",) as f:
             parameters = json.load(f)
 
         directory = sys.argv[1]
 
+        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.")
 

src/aind_ephys_rig_qc/parameters.json

@@ -1,5 +1,6 @@
 {
     "report_name" : "QC.pdf",
-    "timestamp_alignment_method" : "local",
-    "original_timestamp_filename" : "original_timestamps.npy"
+    "timestamp_alignment_method" : "harp",
+    "original_timestamp_filename" : "original_timestamps.npy",
+    "plot_drift_map" : false
 }

src/aind_ephys_rig_qc/pdf_utils.py

@@ -56,7 +56,7 @@ def set_matplotlib_defaults(self):
         if platform.system() == "Linux":
             plt.rcParams["font.sans-serif"] = ["Nimbus Sans"]
         else:
-            plt.rcParams["font.sans-serif"] = ["Arial"]
+            plt.rcParams["font.sans-serif"] = ["Arial"]  # pragma: no cover
 
     def embed_figure(self, fig, width=190):
         """
@@ -87,7 +87,7 @@ def embed_table(self, df, width=190):
             The width of the image in the PDF
         """
 
-        DF = df.map(str)  # convert all elements to string
+        DF = df.astype(str)  # convert all elements to string
         DATA = [
             list(DF)
         ] + DF.values.tolist()  # Combine columns and rows in one list