Feat: Improve cerebellum atlas

atlas_direction_vectors/algorithms/layer_based_direction_vectors.py

@@ -15,6 +15,7 @@
 
 This script is used to compute the mouse isocortex and the mouse thalamus directions vectors.
 """
+
 from __future__ import annotations
 
 import enum

atlas_direction_vectors/algorithms/utils.py

@@ -1,11 +1,19 @@
 """Low-level tools for the computation of direction vectors"""
 
+import logging
+from typing import Union
+
 import numpy as np  # type: ignore
+import pyquaternion as pyq
 from atlas_commons.utils import FloatArray, NumericArray, normalize
+from joblib import Parallel, delayed
 from scipy.ndimage import gaussian_filter  # type: ignore
 from scipy.ndimage import generate_binary_structure  # type: ignore
+from scipy.optimize import minimize
 from scipy.signal import correlate  # type: ignore
 
+L = logging.getLogger(__name__)
+
 
 def compute_blur_gradient(
     scalar_field: FloatArray, gaussian_stddev=3.0, gaussian_radius=None
@@ -45,26 +53,49 @@ def compute_blur_gradient(
 
 
 def _quaternion_from_vectors(  # pylint: disable=invalid-name
-    s: NumericArray, t: NumericArray
+    s: NumericArray, t: NumericArray, align_to: Union[NumericArray, None] = None
 ) -> NumericArray:
     """
     Returns the quaternion (s cross t) + (s dot t + |s||t|).
 
     This quaternion q maps s to t, i.e., qsq^{-1} = t.
 
+    If align_to is specified (either `x`, or `z`), we will rotate it around its axis ([0, 1, 0])
+    so that it is maximally aligned with `align_to`.
+
     Args:
         s: numeric array of shape (3,) or (N, 3).
         t: numeric array of shape (N, 3) if s has two dimensions and its first dimension is N.
+        align_to: if not None, the returned quaternion is aligned with the given axis.
     Returns:
         Numeric array of shape (N, 4) where N is the first dimension of t.
         This data is interpreted as a 1D array of quaternions with size N. A quaternion is a 4D
         vector [w, x, y, z] where [x, y, z] is the imaginary part.
     """
     w = np.matmul(s, np.array(t).T) + np.linalg.norm(s, axis=-1) * np.linalg.norm(t, axis=-1)
-    return np.hstack([w[:, np.newaxis], np.cross(s, t)])
+    quaternions = np.hstack([w[:, np.newaxis], np.cross(s, t)])
+
+    if align_to is not None:
+        target_dir = {"x": np.array([1.0, 0.0, 0.0]), "z": np.array([0.0, 0.0, 1.0])}[align_to]
+
+        def rotate_q(_q):
+            q = pyq.Quaternion(_q)
+
+            def cost(a):
+                q_rot = pyq.Quaternion(axis=[0, 1, 0], angle=a)
+                return -(q * q_rot).rotate(target_dir).dot(target_dir)
+
+            angle = minimize(cost, 0.0).x
+            return (q * pyq.Quaternion(axis=[0, 1, 0], angle=angle)).q.tolist()
+
+        L.info("We are aligning %s quaternions to %s", len(quaternions), align_to)
+        with Parallel(n_jobs=-1, batch_size=10000, verbose=10) as parallel:
+            quaternions = np.array(parallel(delayed(rotate_q)(q) for q in quaternions))
+
+    return quaternions
 
 
-def vector_to_quaternion(vector_field: FloatArray) -> FloatArray:
+def vector_to_quaternion(vector_field: FloatArray, align_to: Union[str, None] = None) -> FloatArray:
     """
     Find quaternions which rotate [0.0, 1.0, 0.0] to each vector in `vector_field`.
 
@@ -75,6 +106,7 @@ def vector_to_quaternion(vector_field: FloatArray) -> FloatArray:
     Arguments:
         vector_field: field of floating point 3D unit vectors, i.e., a float array of shape
             (..., 3).
+        align_to: if not None, the returned quaternion is aligned with the given axis.
 
     Returns:
         numpy.ndarray of shape (..., 4) and of the same type as the input
@@ -88,7 +120,7 @@ def vector_to_quaternion(vector_field: FloatArray) -> FloatArray:
     quaternions = np.full(vector_field.shape[:-1] + (4,), np.nan, dtype=vector_field.dtype)
     non_nan_mask = ~np.isnan(np.linalg.norm(vector_field, axis=-1))
     quaternions[non_nan_mask] = _quaternion_from_vectors(
-        [0.0, 1.0, 0.0], vector_field[non_nan_mask]
+        [0.0, 1.0, 0.0], vector_field[non_nan_mask], align_to=align_to
     )
     return quaternions
 

atlas_direction_vectors/app/orientation_field.py

@@ -17,7 +17,9 @@
 convention used by the placement algorithm, see
 https://bbpteam.epfl.ch/documentation/projects/placement-algorithm/latest/index.html.
 """
+
 import logging
+from typing import Union
 
 import click
 import voxcell  # type: ignore
@@ -45,15 +47,22 @@
     "NaN vectors indicate out-of-domain voxels but also voxels for which an orientation could"
     " not be derived.",
 )
+@click.option(
+    "--align-to",
+    type=str,
+    required=False,
+    help=("Direction to align the direction vectors to."),
+)
 @log_args(L)
 def cmd(
     direction_vectors_path: str,
     output_path: str,
+    align_to: Union[str, None] = None,
 ) -> None:
     """Turn direction vectors into quaternions interpreted as 3D orientations."""
 
     direction_vectors = voxcell.VoxelData.load_nrrd(direction_vectors_path)
-    quaternions = vector_to_quaternion(direction_vectors.raw)
+    quaternions = vector_to_quaternion(direction_vectors.raw, align_to=align_to)
     voxcell.OrientationField(
         quaternions, direction_vectors.voxel_dimensions, direction_vectors.offset
     ).save_nrrd(output_path)

atlas_direction_vectors/cerebellum.py

@@ -4,11 +4,14 @@
 and high values where fiber tracts are outgoing. The direction vectors are given by the gradient
 of this scalar field.
 """
+
 import logging
+from functools import partial
 from typing import TYPE_CHECKING, Optional
 
 import numpy as np
 from atlas_commons.typing import FloatArray
+from joblib import Parallel, delayed
 
 from atlas_direction_vectors.algorithms.layer_based_direction_vectors import (
     compute_layered_region_direction_vectors,
@@ -49,10 +52,16 @@ def compute_direction_vectors(region_map: "RegionMap", annotation: "VoxelData")
             and parent_id not in subregion_ids
         ):
             subregion_ids.append(parent_id)
-            L.info("Computing direction vectors for region %s.", region_map.get(parent_id, "name"))
-            subregion_direction_vectors = cereb_subregion_direction_vectors(
-                parent_id, region_map, annotation
-            )
+    L.info(
+        "Computing direction vectors for regions %s.",
+        [region_map.get(parent_id, "name") for parent_id in subregion_ids],
+    )
+
+    f = partial(cereb_subregion_direction_vectors, region_map=region_map, annotation=annotation)
+    with Parallel(n_jobs=-1, verbose=10) as parallel:
+        for subregion_direction_vectors in parallel(
+            delayed(f)(parent_id) for parent_id in subregion_ids
+        ):
             # Assembles subregion direction vectors.
             subregion_mask = np.logical_not(np.isnan(subregion_direction_vectors))
             direction_vectors[subregion_mask] = subregion_direction_vectors[subregion_mask]
@@ -91,23 +100,19 @@ def cereb_subregion_direction_vectors(
         "layers": {
             "names": [
                 name + ", molecular layer",
-                name + ", Purkinje layer",
                 name + ", granular layer",
                 "cerebellum related fiber tracts",
             ],
-            "queries": [acronym + "mo", acronym + "pu", acronym + "gr", "cbf"],
+            "queries": [acronym + "mo", acronym + "gr", "cbf"],
             "attribute": "acronym",
             "with_descendants": True,
         },
     }
     region_to_weight = {
-        acronym + "mo": 1.0 if "mo" not in weights else weights["mo"],
-        acronym + "pu": 0.0 if "pu" not in weights else weights["pu"],
-        acronym + "gr": -1.0 if "gr" not in weights else weights["gr"],
-        "cbf": -5.0 if "cbf" not in weights else weights["cbf"],
-        "outside_of_brain": 3.0
-        if "outside_of_brain" not in weights
-        else weights["outside_of_brain"],
+        acronym + "mo": weights.get("mo", 1.0),
+        acronym + "gr": weights.get("gr", -1.0),
+        "cbf": weights.get("cbf", -5.0),
+        "outside_of_brain": weights.get("outside_of_brain", 3),
     }
 
     return compute_layered_region_direction_vectors(
@@ -117,5 +122,5 @@ def cereb_subregion_direction_vectors(
         region_to_weight=region_to_weight,
         shading_width=4,
         expansion_width=8,
-        has_hemispheres=False,
+        has_hemispheres=True,
     )

setup.py

@@ -25,6 +25,8 @@
         "numpy>=1.15.0",
         "scipy>=1.10.0",
         "voxcell>=3.0.0",
+        "joblib>=1.3",
+        "pyquaternion>=0.9.9",
     ],
     extras_require={
         "tests": [

tests/test_cerebellum.py

@@ -183,17 +183,16 @@ def _check_vectors_direction_dominance(
 def test_compute_cerebellum_direction_vectors(region_map, annotation):
     res = tested.compute_direction_vectors(region_map, annotation)
     _check_vectors_defined_in_regions(
-        res, region_map, annotation, ["FLgr", "FLpu", "FLmo"] + ["LINGgr", "LINGpu", "LINGmo"]
+        res, region_map, annotation, ["FLgr", "FLmo"] + ["LINGgr", "LINGmo"]
     )
 
     _check_vectors_direction_dominance(
-        res, region_map, annotation, ["FLgr", "FLpu", "FLmo"], [-1.0, 0.0, 0.0]
+        res, region_map, annotation, ["FLgr", "FLmo"], [-1.0, 0.0, 0.0]
     )
 
     _check_vectors_direction_dominance(
-        res, region_map, annotation, ["LINGgr", "LINGpu", "LINGmo"], [-1.0, 0.0, 0.0]
+        res, region_map, annotation, ["LINGgr", "LINGmo"], [-1.0, 0.0, 0.0]
     )
-
     expected_direction_vectors = VoxelData.load_nrrd(
         DATA_DIR / "cerebellum_shading_gradient_orientations.nrrd"
     )
@@ -205,4 +204,4 @@ def test_cereb_subreg_direction_vectors(region_map, annotation):
     res = tested.cereb_subregion_direction_vectors(
         region_map.find("FL", "acronym").pop(), region_map, annotation
     )
-    _check_vectors_defined_in_regions(res, region_map, annotation, ["FLgr", "FLpu", "FLmo"])
+    _check_vectors_defined_in_regions(res, region_map, annotation, ["FLgr", "FLmo"])