fix: replace dead streamlines URL with Allen mesoscale connectivity via cloud-volume

brainrender/atlas_specific/allen_brain_atlas/streamlines.py

@@ -1,4 +1,5 @@
 import pandas as pd
+import requests
 from loguru import logger
 from myterial import orange
 from rich import print
@@ -14,28 +15,58 @@
 except ModuleNotFoundError:  # pragma: no cover
     allen_sdk_installed = False  # pragma: no cover
 
+try:
+    import cloudvolume
+
+    cloudvolume_installed = True
+except ModuleNotFoundError:  # pragma: no cover
+    cloudvolume_installed = False  # pragma: no cover
+
+from brainglobe_atlasapi import BrainGlobeAtlas
 
 from brainrender import base_dir
-from brainrender._io import request
 from brainrender._utils import listify
 
 streamlines_folder = base_dir / "streamlines"
 streamlines_folder.mkdir(exist_ok=True)
 
+ALLEN_MESOSCALE_URL = (
+    "precomputed://gs://allen_neuroglancer_ccf/allen_mesoscale"
+)
+ALLEN_API_URL = "https://api.brain-map.org/api/v2/data/query.json"
+VOXEL_SIZE_NM = 1000  # skeleton vertices are in nanometers
+
+_ml_extent_um_cache = None
+
+
+def _get_ml_extent_um():
+    """
+    Derives the full medial-lateral extent of the Allen CCF atlas in microns
+    dynamically from the brainglobe atlas API. Used to flip the Z (ML) axis
+    when converting from Allen CCF space to brainrender's coordinate system,
+    where left and right hemispheres are mirrored relative to the Allen CCF.
+
+    Result is cached after the first call to avoid reinstantiating the atlas
+    on every experiment download.
+    """
+    global _ml_extent_um_cache
+    if _ml_extent_um_cache is None:
+        atlas = BrainGlobeAtlas("allen_mouse_25um", check_latest=False)
+        _ml_extent_um_cache = float(atlas.shape[2] * atlas.resolution[2])
+    return _ml_extent_um_cache
+
 
 def experiments_source_search(SOI):
     """
     Returns data about experiments whose injection was in the SOI, structure of interest
     :param SOI: str, structure of interest. Acronym of structure to use as seed for the search
-    :param source:  (Default value = True)
     """
-
     transgenic_id = 0  # id = 0 means use only wild type
     primary_structure_only = True
 
     if not allen_sdk_installed:
         print(
-            f"[{orange}]Streamlines cannot be download because the AllenSDK package is not installed. "
+            f"[{orange}]Streamlines cannot be downloaded because the AllenSDK package is not installed. "
             "Please install `allensdk` with `pip install allensdk`"
         )
         return None
@@ -50,57 +81,147 @@ def experiments_source_search(SOI):
     )
 
 
-def get_streamlines_data(eids, force_download=False):
+def _get_injection_site_um(eid, ml_extent_um):
     """
-    Given a list of expeirmental IDs, it downloads the streamline data
-    from the https://neuroinformatics.nl cache and saves them as
-    json files.
+    Fetches the injection site coordinates for an experiment from the Allen
+    Brain Atlas API. Coordinates are in Allen CCF um space with the Z (ML)
+    axis flipped to match brainrender's hemisphere convention.
 
-    :param eids: list of integers with experiments IDs
+    :param eid: int, experiment ID
+    :param ml_extent_um: float, full ML extent of the atlas in um for LR flip
+    :return: dict with x, y, z keys or None if not found
     """
-    data = []
-    for eid in track(eids, total=len(eids), description="downloading"):
-        url = "https://neuroinformatics.nl/HBP/allen-connectivity-viewer/json/streamlines_{}.json.gz".format(
-            eid
+    try:
+        url = (
+            f"{ALLEN_API_URL}?q=model::ProjectionStructureUnionize,"
+            f"rma::criteria,section_data_set[id$eq{eid}],"
+            f"rma::criteria,[is_injection$eqtrue],"
+            f"rma::options[num_rows$eq1][order$eq'projection_volume desc']"
+        )
+        response = requests.get(url, timeout=10)
+        data = response.json()
+        if data["success"] and data["num_rows"] > 0:
+            voxel = data["msg"][0]
+            return {
+                "x": float(voxel["max_voxel_x"]),
+                "y": float(voxel["max_voxel_y"]),
+                "z": float(ml_extent_um - voxel["max_voxel_z"]),
+            }
+    except Exception as e:
+        logger.warning(
+            f"Could not fetch injection site for experiment {eid}: {e}"
         )
+    return None
 
-        jsonpath = streamlines_folder / f"{eid}.json"
 
-        if not jsonpath.exists() or force_download:
-            response = request(url)
+def _skeleton_to_dataframe(skeleton, eid, ml_extent_um):
+    """
+    Converts a cloudvolume Skeleton object to the pd.DataFrame format
+    expected by brainrender's Streamlines actor.
 
-            # Write the response content as a temporary compressed file
-            temp_path = streamlines_folder / "temp.gz"
-            with open(str(temp_path), "wb") as temp:
-                temp.write(response.content)
+    Vertices are in nanometers in Allen CCF space. We:
+    1. Convert nm -> um (divide by VOXEL_SIZE_NM)
+    2. Flip Z (ML) axis to match brainrender's hemisphere convention
 
-            # Open in pandas and delete temp
-            url_data = pd.read_json(
-                str(temp_path), lines=True, compression="gzip"
-            )
-            temp_path.unlink()
+    X (AP) and Y (DV) are passed through as-is because brainrender's
+    brain mesh uses the same orientation as the Allen CCF for those axes.
 
-            # save json
-            url_data.to_json(str(jsonpath))
+    :param skeleton: cloudvolume Skeleton object
+    :param eid: int, experiment ID used to fetch real injection coordinates
+    :param ml_extent_um: float, full ML extent of the atlas in um for LR flip
+    :return: pd.DataFrame with 'lines' and 'injection_sites' columns
+    """
+    components = skeleton.components()
+
+    lines = []
+    for component in components:
+        verts_um = component.vertices / VOXEL_SIZE_NM
+        points = [
+            {
+                "x": float(v[0]),
+                "y": float(v[1]),
+                "z": float(ml_extent_um - v[2]),
+            }
+            for v in verts_um
+        ]
+        lines.append(points)
+
+    injection_site = _get_injection_site_um(eid, ml_extent_um)
+    if injection_site is None:
+        logger.warning(
+            f"Falling back to centroid for injection site of experiment {eid}"
+        )
+        all_verts_um = skeleton.vertices / VOXEL_SIZE_NM
+        centroid = all_verts_um.mean(axis=0)
+        injection_site = {
+            "x": float(centroid[0]),
+            "y": float(centroid[1]),
+            "z": float(ml_extent_um - centroid[2]),
+        }
 
-            # append to lists and return
-            data.append(url_data)
+    return pd.DataFrame(
+        {"lines": [lines], "injection_sites": [[injection_site]]}
+    )
+
+
+def get_streamlines_data(eids, force_download=False):
+    """
+    Given a list of experiment IDs, downloads streamline data from the
+    Allen mesoscale connectivity dataset hosted on Google Cloud Storage
+    via cloud-volume, and saves them as JSON files.
+
+    :param eids: list of integers with experiment IDs
+    :param force_download: bool, if True re-download even if cached
+    """
+    if not cloudvolume_installed:
+        print(
+            f"[{orange}]Streamlines cannot be downloaded because the cloud-volume package is not installed. "
+            "Please install it with `pip install cloud-volume`"
+        )
+        return []
+
+    ml_extent_um = _get_ml_extent_um()
+
+    cv = cloudvolume.CloudVolume(
+        ALLEN_MESOSCALE_URL,
+        use_https=True,
+        progress=False,
+    )
+
+    data = []
+    for eid in track(eids, total=len(eids), description="downloading"):
+        jsonpath = streamlines_folder / f"{eid}.json"
+
+        if not jsonpath.exists() or force_download:
+            try:
+                skeleton = cv.skeleton.get(int(eid))
+            except Exception as e:
+                logger.warning(
+                    f"Could not fetch streamlines for experiment {eid}: {e}"
+                )
+                continue
+
+            df = _skeleton_to_dataframe(skeleton, int(eid), ml_extent_um)
+            df.to_json(str(jsonpath))
+            data.append(df)
         else:
             data.append(pd.read_json(str(jsonpath)))
+
     return data
 
 
 def get_streamlines_for_region(region, force_download=False):
     """
-    Using the Allen Mouse Connectivity data and corresponding API, this function finds experiments whose injections
-    were targeted to the region of interest and downloads the corresponding streamlines data. By default, experiments
-    are selected for only WT mice and only when the region was the primary injection target.
-
-    :param region: str with region to use for research
-
+    Using the Allen Mouse Connectivity data and corresponding API, this function finds experiments
+    whose injections were targeted to the region of interest and downloads the corresponding
+    streamlines data from the Allen mesoscale connectivity dataset on Google Cloud Storage.
+    By default, experiments are selected for only WT mice and only when the region was
+    the primary injection target.
+
+    :param region: str with region to use for search
+    :param force_download: bool, if True re-download even if cached
     """
     logger.debug(f"Getting streamlines data for region: {region}")
-    # Get experiments whose injections were targeted to the region
     region_experiments = experiments_source_search(region)
     if region_experiments is None or region_experiments.empty:
         logger.debug("No experiments found from allen data")