Report snapshot probe: re-usable subclass for all snapshots linked to a pid

Author: oliche

ibllib/plots/figures.py

@@ -11,7 +11,7 @@
 import matplotlib.pyplot as plt
 
 from ibllib.dsp import voltage
-from ibllib.plots.snapshot import ReportSnapshot
+from ibllib.plots.snapshot import ReportSnapshotProbe
 from one.api import ONE
 import one.alf.io as alfio
 from one.alf.exceptions import ALFObjectNotFound
@@ -22,7 +22,7 @@
 logger = logging.getLogger('ibllib')
 
 
-class BadChannelsAp(ReportSnapshot):
+class SpikeSorting(ReportSnapshotProbe):
     """
     Plots raw electrophysiology AP band
     :param session_path: session path
@@ -34,20 +34,44 @@ class BadChannelsAp(ReportSnapshot):
         'output_files': []  # see setUp method for declaration of inputs
     }
 
-    def __init__(self, session_path, probe_id, **kwargs):
-        self.content_type = 'probeinsertion'
-        self.pid = probe_id
-        super(BadChannelsAp, self).__init__(session_path, probe_id, content_type=self.content_type, **kwargs)
+    def _run(self):
+        """runs for initiated PID, streams data, destripe and check bad channels"""
+        assert self.pid
+
+    def get_probe_signature(self):
+        input_signature = [('spikes.times.npy', f'alf/{self.pname}', True),
+                           ('spikes.amps.npy', f'alf/{self.pname}', True),
+                           ('spikes.depths.npy', f'alf/{self.pname}', True)]
+        output_signature = [('spike_sorting_raster.png', f'snapshot/{self.pname}', True)]
+        self.signature = {'input_files': input_signature, 'output_file': output_signature}
+
+
+class BadChannelsAp(ReportSnapshotProbe):
+    """
+    Plots raw electrophysiology AP band
+    task = BadChannelsAp(pid, one=one=one)
+    :param session_path: session path
+    :param probe_id: str, UUID of the probe insertion for which to create the plot
+    :param **kwargs: keyword arguments passed to tasks.Task
+    """
+    signature = {
+        'input_files': [],  # see setUp method for declaration of inputs
+        'output_files': []  # see setUp method for declaration of inputs
+    }
 
     @staticmethod
-    def spike_sorting_signature(pname=None):
+    def get_probe_signature(pname=None):
         pname = pname if pname is not None else "probe*"
         input_signature = [('*ap.meta', f'raw_ephys_data/{pname}', True),
                            ('*ap.ch', f'raw_ephys_data/{pname}', False),
                            ('*ap.cbin', f'raw_ephys_data/{pname}', False)]
-        output_signature = [('destripe.png', f'snapshot/{pname}', True),
-                            ('highpass.png', f'snapshot/{pname}', True)]
-        return input_signature, output_signature
+        output_signature = [('raw_ephys_bad_channels.png', f'snapshot/{pname}', True),
+                            ('raw_ephys_bad_channels_highpass.png', f'snapshot/{pname}', True),
+                            ('raw_ephys_bad_channels_highpass.png', f'snapshot/{pname}', True),
+                            ('raw_ephys_bad_channels_destripe.png', f'snapshot/{pname}', True),
+                            ('raw_ephys_bad_channels_difference.png', f'snapshot/{pname}', True),
+                            ]
+        return {'input_files': input_signature, 'output_files': output_signature}
 
     def _run(self):
         """runs for initiated PID, streams data, destripe and check bad channels"""

ibllib/plots/snapshot.py

@@ -2,6 +2,7 @@
 import requests
 import traceback
 import json
+import abc
 
 from one.api import ONE
 from ibllib.pipes import tasks
@@ -33,6 +34,25 @@ def register_images(self, widths=None, function=None):
         return snapshot.register_images(self.outputs, jsons=jsons, texts=texts, widths=widths)
 
 
+class ReportSnapshotProbe(ReportSnapshot):
+
+    def __init__(self, pid, one=None, **kwargs):
+        assert one
+        self.one = one
+        self.content_type = 'probeinsertion'
+        self.pid = pid
+        self.eid, self.pname = self.one.pid2eid(self.pid)
+        self.session_path = self.one.eid2path(self.eid)
+        self.signature = self.get_probe_signature()
+        super(ReportSnapshotProbe, self).__init__(self.session_path, object_id=pid, content_type=self.content_type, **kwargs)
+
+    @abc.abstractmethod
+    def get_probe_signature(pname=None):
+        # method that gets input and output signatures from the probe name. The format is a dictionary as follows:
+        # return {'input_files': input_signature, 'output_files': output_signature}
+        pass
+
+
 class Snapshot:
     """
     A class to register images in form of Notes, linked to an object on Alyx.