remove  atlas instantiation from histology module

Author: mayofaulkner

examples/one/histology/register_lasagna_tracks_alyx.py

@@ -24,7 +24,7 @@
 """
 # Author: Olivier, Gaelle
 from pathlib import Path
-
+from ibllib.atlas import AllenAtlas
 from one.api import ONE
 
 from ibllib.pipes import histology
@@ -42,11 +42,12 @@
 
 # ======== DO NOT EDIT BELOW ====
 one = ONE(base_url=ALYX_URL)
+ba = AllenAtlas()
 
 if EXAMPLE_OVERWRITE:
     # TODO Olivier : Function to download examples folder
     cachepath = Path(one.alyx.cache_dir)
     path_tracks = cachepath.joinpath('examples', 'histology', 'tracks_to_add')
 
-histology.register_track_files(path_tracks=path_tracks, one=one, overwrite=True)
-histology.detect_missing_histology_tracks(path_tracks=path_tracks, one=one)
+histology.register_track_files(path_tracks=path_tracks, one=one, overwrite=True, brain_atlas=ba)
+histology.detect_missing_histology_tracks(path_tracks=path_tracks, one=one, brain_atlas=ba)

ibllib/atlas/atlas.py

@@ -2,7 +2,7 @@
 import logging
 import matplotlib.pyplot as plt
 from pathlib import Path, PurePosixPath
-
+from functools import lru_cache
 import numpy as np
 import nrrd
 
@@ -710,6 +710,7 @@ def get_brain_entry(traj, brain_atlas):
         return Insertion._get_surface_intersection(traj, brain_atlas, surface='top')
 
 
+@lru_cache(maxsize=1)
 class AllenAtlas(BrainAtlas):
     """
     Instantiates an atlas.BrainAtlas corresponding to the Allen CCF at the given resolution

ibllib/pipes/histology.py

@@ -16,16 +16,15 @@
 
 _logger = logging.getLogger('ibllib')
 
-# origin Allen left, front, up
-brain_atlas = atlas.AllenAtlas(res_um=25)
 
-
-def load_track_csv(file_track):
+def load_track_csv(file_track, brain_atlas=None):
     """
     Loads a lasagna track and convert to IBL-ALlen coordinate framework
     :param file_track:
     :return: xyz
     """
+
+    brain_atlas = brain_atlas or atlas.AllenAtlas(25)
     # apmldv in the histology file is flipped along y direction
     file_track = Path(file_track)
     if file_track.stat().st_size == 0:
@@ -38,20 +37,21 @@ def load_track_csv(file_track):
     return xyz
 
 
-def get_picked_tracks(histology_path, glob_pattern="*_pts_transformed.csv"):
+def get_picked_tracks(histology_path, glob_pattern="*_pts_transformed.csv", brain_atlas=None):
     """
     This outputs reads in the Lasagna output and converts the picked tracks in the IBL coordinates
     :param histology_path: Path object: folder path containing tracks
     :return: xyz coordinates in
     """
+    brain_atlas = brain_atlas or atlas.AllenAtlas()
     xyzs = []
     histology_path = Path(histology_path)
     if histology_path.is_file():
         files_track = [histology_path]
     else:
         files_track = list(histology_path.rglob(glob_pattern))
     for file_track in files_track:
-        xyzs.append(load_track_csv(file_track))
+        xyzs.append(load_track_csv(file_track, brain_atlas=brain_atlas))
     return {'files': files_track, 'xyz': xyzs}
 
 
@@ -100,13 +100,14 @@ def get_micro_manipulator_data(subject, one=None, force_extract=False):
     return probes
 
 
-def plot2d_all(trajectories, tracks):
+def plot2d_all(trajectories, tracks, brain_atlas=None):
     """
     Plot all tracks on a single 2d slice
     :param trajectories: dictionary output of the Alyx REST query on trajectories
     :param tracks:
     :return:
     """
+    brain_atlas = brain_atlas or atlas.AllenAtlas(25)
     plt.figure()
     axs = brain_atlas.plot_sslice(brain_atlas.bc.i2x(190), cmap=plt.get_cmap('bone'))
     plt.figure()
@@ -121,14 +122,15 @@ def plot2d_all(trajectories, tracks):
         axs.plot(xyz[:, 1] * 1e3, xyz[:, 2] * 1e6, 'r')
 
 
-def plot3d_all(trajectories, tracks):
+def plot3d_all(trajectories, tracks, brain_atlas=None):
     """
     Plot all tracks on a single 2d slice
     :param trajectories: dictionary output of the Alyx REST query on trajectories
     :param tracks:
     :return:
     """
     from mayavi import mlab
+    brain_atlas = brain_atlas or atlas.AllenAtlas()
     src = mlab.pipeline.scalar_field(brain_atlas.label)
     mlab.pipeline.iso_surface(src, contours=[0.5, ], opacity=0.3)
 
@@ -168,7 +170,7 @@ def interpolate_along_track(xyz_track, depths):
     return xyz_channels
 
 
-def get_brain_regions(xyz, channels_positions=None, brain_atlas=brain_atlas):
+def get_brain_regions(xyz, channels_positions=None, brain_atlas=None):
     """
     :param xyz: numpy array of 3D coordinates corresponding to a picked track or a trajectory
     the deepest point is assumed to be the tip.
@@ -183,6 +185,7 @@ def get_brain_regions(xyz, channels_positions=None, brain_atlas=brain_atlas):
     Due to the blockiness, depths may not be unique along the track so it has to be prepared
     """
 
+    brain_atlas = brain_atlas or atlas.AllenAtlas(25)
     if channels_positions is None:
         geometry = trace_header(version=1)
         channels_positions = np.c_[geometry['x'], geometry['y']]
@@ -215,7 +218,7 @@ def get_brain_regions(xyz, channels_positions=None, brain_atlas=brain_atlas):
     return brain_regions, insertion
 
 
-def register_track(probe_id, picks=None, one=None, overwrite=False, channels=True):
+def register_track(probe_id, picks=None, one=None, overwrite=False, channels=True, brain_atlas=None):
     """
     Register the user picks to a probe in Alyx
     Here we update Alyx models on the database in 3 steps
@@ -225,6 +228,7 @@ def register_track(probe_id, picks=None, one=None, overwrite=False, channels=Tru
     3) Channel locations are set in the table
     """
     assert one
+    brain_atlas = brain_atlas or atlas.AllenAtlas()
     # 0) if it's an empty track, create a null trajectory and exit
     if picks is None or picks.size == 0:
         tdict = {'probe_insertion': probe_id,
@@ -241,7 +245,7 @@ def register_track(probe_id, picks=None, one=None, overwrite=False, channels=Tru
         insertion_histology = None
         # Here need to change the track qc to critical and also extended qc to zero
     else:
-        brain_locations, insertion_histology = get_brain_regions(picks)
+        brain_locations, insertion_histology = get_brain_regions(picks, brain_atlas=brain_atlas)
         # 1) update the alyx models, first put the picked points in the insertion json
         one.alyx.json_field_update(endpoint='insertions', uuid=probe_id, field_name='json',
                                    data={'xyz_picks': np.int32(picks * 1e6).tolist()})
@@ -275,14 +279,15 @@ def register_track(probe_id, picks=None, one=None, overwrite=False, channels=Tru
 
 
 def register_aligned_track(probe_id, xyz_channels, chn_coords=None, one=None, overwrite=False,
-                           channels=True):
+                           channels=True, brain_atlas=None):
     """
     Register ephys aligned trajectory and channel locations to Alyx
     Here we update Alyx models on the database in 2 steps
     1) The trajectory computed from the final electrode channel locations
     2) Channel locations are set to the trajectory
     """
     assert one
+    brain_atlas = brain_atlas or atlas.AllenAtlas(25)
     if chn_coords is None:
         geometry = trace_header(version=1)
         chn_coords = np.c_[geometry['x'], geometry['y']]
@@ -374,14 +379,14 @@ def _parse_filename(track_file):
     return search_filter
 
 
-def register_track_files(path_tracks, one=None, overwrite=False):
+def register_track_files(path_tracks, one=None, overwrite=False, brain_atlas=None):
     """
     :param path_tracks: path to directory containing tracks; also works with a single file name
     :param one:
     :return:
     """
+    brain_atlas = brain_atlas or atlas.AllenAtlas()
     glob_pattern = "*_probe*_pts*.csv"
-
     path_tracks = Path(path_tracks)
 
     if not path_tracks.is_dir():
@@ -415,22 +420,23 @@ def register_track_files(path_tracks, one=None, overwrite=False):
             raise ValueError("Multiple probes found.")
         probe_id = probe['id']
         try:
-            xyz_picks = load_track_csv(track_file)
-            register_track(probe_id, xyz_picks, one=one, overwrite=overwrite)
+            xyz_picks = load_track_csv(track_file, brain_atlas=brain_atlas)
+            register_track(probe_id, xyz_picks, one=one, overwrite=overwrite, brain_atlas=brain_atlas)
         except Exception as e:
             _logger.error(str(track_file))
             raise e
         _logger.info(f"{ind + 1}/{ntracks}, {str(track_file)}")
 
 
-def detect_missing_histology_tracks(path_tracks=None, one=None, subject=None):
+def detect_missing_histology_tracks(path_tracks=None, one=None, subject=None, brain_atlas=None):
     """
     Compares the number of probe insertions to the number of registered histology tracks to see if
     there is a discrepancy so that missing tracks can be properly logged in the database
     :param path_tracks: path to track files to be registered
     :param subject: subject nickname for which to detect missing tracks
     """
 
+    brain_atlas = brain_atlas or atlas.AllenAtlas()
     if path_tracks:
         glob_pattern = "*_probe*_pts*.csv"
 
@@ -479,7 +485,7 @@ def detect_missing_histology_tracks(path_tracks=None, one=None, subject=None):
                     probe_id = insertions[idx]['id']
                     print(insertions[idx]['session'])
                     print(probe_id)
-                    register_track(probe_id, one=one)
+                    register_track(probe_id, one=one, brain_atlas=brain_atlas)
                 else:
                     _logger.info('Histology track for this probe insertion will not be registered')
                     continue

ibllib/qc/alignment_qc.py

@@ -307,7 +307,8 @@ def upload_channels(self, alignment_key, upload_alyx, upload_flatiron):
             if alignment_key != self.align_keys_sorted[0]:
                 histology.register_aligned_track(self.eid, channels_mlapdv / 1e6,
                                                  chn_coords=chns, one=self.one,
-                                                 overwrite=True, channels=self.channels_flag)
+                                                 overwrite=True, channels=self.channels_flag,
+                                                 brain_atlas=self.brain_atlas)
 
                 ephys_traj = self.one.alyx.get(f'/trajectories?&probe_insertion={self.eid}'
                                                '&provenance=Ephys aligned histology track',
@@ -416,4 +417,4 @@ def get_aligned_channels(ins, chn_coords, one, ba=None, save_dir=None):
         np.save(f_name, channels_atlas_ids)
         out_files.append(f_name)
 
-    return out_files
+    return out_files

ibllib/tests/qc/test_alignment_qc.py

@@ -56,15 +56,15 @@ def setUpClass(cls) -> None:
 
     def test_tracing_exists(self):
         register_track(self.probe00_id, picks=self.xyz_picks, one=one, overwrite=True,
-                       channels=False)
+                       channels=False, brain_atlas=brain_atlas)
         insertion = one.alyx.get('/insertions/' + self.probe00_id, clobber=True)
 
         assert (insertion['json']['qc'] == 'NOT_SET')
         assert (insertion['json']['extended_qc']['tracing_exists'] == 1)
 
     def test_tracing_not_exists(self):
         register_track(self.probe01_id, picks=None, one=one, overwrite=True,
-                       channels=False)
+                       channels=False, brain_atlas=brain_atlas)
         insertion = one.alyx.get('/insertions/' + self.probe01_id, clobber=True)
         assert (insertion['json']['qc'] == 'CRITICAL')
         assert (insertion['json']['extended_qc']['tracing_exists'] == 0)

ibllib/tests/test_histology.py

@@ -19,7 +19,7 @@ def setUp(self) -> None:
     def test_histology_get_brain_regions(self):
         # first part of the test is to check on an actual track file
         for file_track in self.path_tracks.rglob("*_pts.csv"):
-            xyz = histology.load_track_csv(file_track)
+            xyz = histology.load_track_csv(file_track, brain_atlas=brain_atlas)
             channels, ins = histology.get_brain_regions(xyz=xyz, brain_atlas=brain_atlas)
         # also check that it works from an insertion
         channels, ins2 = histology.get_brain_regions(xyz=ins.xyz, brain_atlas=brain_atlas)
@@ -32,7 +32,7 @@ def test_histology_get_brain_regions(self):
     def test_histology_insertion_from_track(self):
 
         for file_track in self.path_tracks.rglob("*_pts.csv"):
-            xyz = histology.load_track_csv(file_track)
+            xyz = histology.load_track_csv(file_track, brain_atlas=brain_atlas)
             insertion = atlas.Insertion.from_track(xyz, brain_atlas=brain_atlas)
             # checks that the tip coordinate is not the deepest point but its projection
             self.assertFalse(np.all(np.isclose(insertion.tip, xyz[-1])))