From a7df4a5e5b16558e63f647ef9ccfe8c2902be8ba Mon Sep 17 00:00:00 2001 From: AlgoFoe Date: Thu, 13 Aug 2026 18:47:18 +0530 Subject: [PATCH 1/4] docs: improve numpy-style docstrings in atlas_specific modules --- .../allen_brain_atlas/gene_expression/api.py | 143 ++++++++++++++---- .../gene_expression/ge_utils.py | 107 ++++++++++--- .../allen_brain_atlas/streamlines.py | 124 +++++++++++---- 3 files changed, 287 insertions(+), 87 deletions(-) diff --git a/brainrender/atlas_specific/allen_brain_atlas/gene_expression/api.py b/brainrender/atlas_specific/allen_brain_atlas/gene_expression/api.py index a07481f2..d2f6d033 100644 --- a/brainrender/atlas_specific/allen_brain_atlas/gene_expression/api.py +++ b/brainrender/atlas_specific/allen_brain_atlas/gene_expression/api.py @@ -1,7 +1,10 @@ +"""API client for downloading Allen Brain Atlas gene expression data.""" + import os import sys from time import sleep +import numpy.typing as npt import pandas as pd import requests from loguru import logger @@ -17,6 +20,8 @@ class GeneExpressionAPI: + """Client for querying and downloading Allen Brain Atlas gene expression data.""" + voxel_size = 200 # um grid_size = [58, 41, 67] # number of voxels along each direction @@ -38,22 +43,39 @@ class GeneExpressionAPI: download_url = "http://api.brain-map.org/grid_data/download/EXP_ID?include=energy,intensity,density" gene_expression_cache = base_dir / "GeneExpressionCache" - gene_name = None + gene_name: str | None = None - def __init__(self): + def __init__(self) -> None: # Get metadata about all available genes - self.genes = None # when necessary gene data can be downloaded with self.get_all_genes + self.genes: pd.DataFrame | None = None # when necessary gene data can be downloaded with self.get_all_genes self.gene_expression_cache.mkdir(exist_ok=True) @fail_on_no_connection - def get_all_genes(self): + def get_all_genes(self) -> pd.DataFrame: """ - Download metadata about all the genes available in the Allen gene expression dataset + Download metadata about all genes in the Allen gene expression dataset. + + Returns + ------- + pd.DataFrame """ res = request(self.all_genes_url) return pd.DataFrame(res.json()["msg"]) - def get_gene_id_by_name(self, gene_name): + def get_gene_id_by_name(self, gene_name: str) -> int | None: + """ + Return the Allen gene ID for a given gene symbol. + + Parameters + ---------- + gene_name + Gene symbol. + + Returns + ------- + int or None + Gene ID, or None if the gene is not found. + """ self.gene_name = self.gene_name or gene_name if self.genes is None: self.genes = self.get_all_genes() @@ -71,7 +93,19 @@ def get_gene_id_by_name(self, gene_name): ] ) - def get_gene_symbol_by_id(self, gene_id): + def get_gene_symbol_by_id(self, gene_id: int | str) -> str: + """ + Return the gene symbol for a given Allen gene ID. + + Parameters + ---------- + gene_id + Allen gene ID. + + Returns + ------- + str + """ if self.genes is None: self.genes = self.get_all_genes() @@ -80,12 +114,19 @@ def get_gene_symbol_by_id(self, gene_id): ].gene_symbol.values[0] @fail_on_no_connection - def get_gene_experiments(self, gene): + def get_gene_experiments(self, gene: str) -> list[int] | None: """ - Given a gene_symbol it returns the list of ISH - experiments for this gene + Return ISH experiment IDs for a given gene symbol. + + Parameters + ---------- + gene + Gene symbol. - :param gene_symbol: str + Returns + ------- + list of int or None + List of experiment IDs, or None if no experiments are found. """ url = self.gene_experiments_url.replace("-GENE_SYMBOL-", gene) max_retries = 8 @@ -108,13 +149,15 @@ def get_gene_experiments(self, gene): return [d["id"] for d in data] @fail_on_no_connection - def download_gene_data(self, gene): + def download_gene_data(self, gene: str) -> None: """ - Downloads a gene's data from the Allen Institute - Gene Expression dataset and saves to cache. - See: http://help.brain-map.org/display/api/Downloading+3-D+Expression+Grid+Data + Download a gene's expression data from the Allen Institute and save to cache. + See http://help.brain-map.org/display/api/Downloading+3-D+Expression+Grid+Data. - :param gene: int, the gene_id for the gene being downloaded. + Parameters + ---------- + gene + Gene symbol to download data for. """ # Get the gene's experiment id exp_ids = self.get_gene_experiments(gene) @@ -130,9 +173,35 @@ def download_gene_data(self, gene): url, os.path.join(self.gene_expression_cache, f"{gene}-{eid}") ) - def get_gene_data(self, gene, exp_id, use_cache=True, metric="energy"): + def get_gene_data( + self, + gene: str, + exp_id: int | str, + use_cache: bool = True, + metric: str = "energy", + ) -> npt.NDArray: """ - Given a list of gene ids + Load gene expression data for a given gene and experiment. + + Parameters + ---------- + gene + Gene symbol. + exp_id + Experiment ID. + use_cache + If True, load from cache if available. Default True. + metric + Expression metric to load. Default ``"energy"``. + + Returns + ------- + numpy.ndarray + + Raises + ------ + ValueError + If data could not be cached after downloading. """ logger.debug(f"Getting gene data for gene: {gene} experiment {exp_id}") self.gene_name = self.gene_name or gene @@ -161,21 +230,31 @@ def get_gene_data(self, gene, exp_id, use_cache=True, metric="energy"): def griddata_to_volume( self, - griddata, - min_quantile=None, - min_value=None, - cmap="bwr", - ): + griddata: npt.NDArray, + min_quantile: float | None = None, + min_value: float | None = None, + cmap: str = "bwr", + ) -> Volume: """ - Takes a 3d numpy array with volumetric gene expression - and returns a vedo.Volume.isosurface actor. - The isosurface needs a lower bound threshold, this can be - either a user defined hard value (min_value) or the value - corresponding to some percentile of the gene expression data. - - :param griddata: np.ndarray, 3d array with gene expression data - :param min_quantile: float, percentile for threshold - :param min_value: float, value for threshold + Convert a 3D gene expression array to a Volume actor. + + The isosurface threshold can be set as a hard value or as a + percentile of the expression data. + + Parameters + ---------- + griddata + 3D array with gene expression data. + min_quantile + Percentile threshold for isosurface extraction. + min_value + Hard value threshold for isosurface extraction. + cmap + Colormap name. Default ``"bwr"``. + + Returns + ------- + Volume """ return Volume( griddata, diff --git a/brainrender/atlas_specific/allen_brain_atlas/gene_expression/ge_utils.py b/brainrender/atlas_specific/allen_brain_atlas/gene_expression/ge_utils.py index fa350f5c..8f0f3371 100644 --- a/brainrender/atlas_specific/allen_brain_atlas/gene_expression/ge_utils.py +++ b/brainrender/atlas_specific/allen_brain_atlas/gene_expression/ge_utils.py @@ -1,9 +1,13 @@ +"""Utilities for caching and loading Allen Brain Atlas gene expression data.""" + import io import os import sys import zipfile +from pathlib import Path import numpy as np +import numpy.typing as npt from brainrender._io import check_file_exists, request from brainrender._utils import get_subdirs, listdir @@ -11,15 +15,35 @@ # ----------------------------------- Cache ---------------------------------- # -def check_gene_cached(cache_folder, gene_id, exp_id): +def check_gene_cached( + cache_folder: str | Path, + gene_id: str, + exp_id: str | int, +) -> str | bool: """ - A gene is saved in a folder in cache_folder - with gene_id-exp_id as name. If the folder doesn't - exist the gene is not cached. - - :param cache_folder: str, path to general cache folder for all data - :param gene_id: str name of gene - :param exp_id: id of experiment + Check whether a gene experiment is already cached. + + A gene is cached in a subfolder of ``cache_folder`` named + ``{gene_id}-{exp_id}``. + + Parameters + ---------- + cache_folder + Path to the general cache folder. + gene_id + Gene name. + exp_id + Experiment ID. + + Returns + ------- + str or bool + Path to the cached folder if found, False if not cached. + + Raises + ------ + ValueError + If more than one matching folder is found. """ cache = [ sub @@ -34,13 +58,16 @@ def check_gene_cached(cache_folder, gene_id, exp_id): return cache[0] -def download_and_cache(url, cachedir): +def download_and_cache(url: str, cachedir: str | Path) -> None: """ - Given a url to download a gene's ISH experiment data, - this function download and unzips the data - - :param url: str, utl to download data - :param cachedir: str, path to folder where data will be downloaded + Download and unzip a gene's ISH experiment data to a cache directory. + + Parameters + ---------- + url + URL to download the data from. + cachedir + Path to the folder where data will be saved. """ # Get data req = request(url) @@ -54,9 +81,32 @@ def download_and_cache(url, cachedir): z.extractall(cachedir) -def load_cached_gene(cache, metric, grid_size): +def load_cached_gene( + cache: str | Path, + metric: str, + grid_size: tuple[int, int, int], +) -> npt.NDArray | None: """ - Loads a gene's data from cache + Load a gene's data from cache. + + Parameters + ---------- + cache + Path to the gene's cache folder. + metric + Metric name used to filter files (e.g. ``"energy"``). + grid_size + Shape to use when reshaping the raw data array. + + Returns + ------- + numpy.ndarray or None + Array of gene expression values, or None if no file is found. + + Raises + ------ + NotImplementedError + If more than one matching file is found. """ files = [ f for f in listdir(cache) if metric in f and not f.endswith(".mhd") @@ -71,15 +121,26 @@ def load_cached_gene(cache, metric, grid_size): # --------------------------------- Open .raw -------------------------------- # @check_file_exists -def read_raw(filepath, grid_size): +def read_raw( + filepath: str | Path, + grid_size: tuple[int, int, int], +) -> npt.NDArray: """ - reads a .raw file with gene expression data - downloaded from the Allen atlas and returns - a numpy array with the correct grid_size. - See as reference: - http://help.brain-map.org/display/mousebrain/API#API-Expression3DGridsz + Read a ``.raw`` gene expression file from the Allen Brain Atlas. + + See http://help.brain-map.org/display/mousebrain/API#API-Expression3DGridsz + for the file format reference. + + Parameters + ---------- + filepath + Path to the ``.raw`` file. + grid_size + Shape to use when reshaping the data array. - :param filepath: str or Path object + Returns + ------- + numpy.ndarray """ filepath = str(filepath) diff --git a/brainrender/atlas_specific/allen_brain_atlas/streamlines.py b/brainrender/atlas_specific/allen_brain_atlas/streamlines.py index 30c4d690..5ca85f80 100644 --- a/brainrender/atlas_specific/allen_brain_atlas/streamlines.py +++ b/brainrender/atlas_specific/allen_brain_atlas/streamlines.py @@ -1,3 +1,7 @@ +"""Streamline download and conversion utilities for the Allen Mouse Brain Atlas.""" + +from typing import Any + import pandas as pd import requests from loguru import logger @@ -36,18 +40,19 @@ 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 +_ml_extent_um_cache: float | None = None -def _get_ml_extent_um(): +def _get_ml_extent_um() -> float: """ - Derives the full medial-lateral extent of the Allen CCF atlas in microns + Derive 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. + Returns + ------- + float """ global _ml_extent_um_cache if _ml_extent_um_cache is None: @@ -56,10 +61,19 @@ def _get_ml_extent_um(): return _ml_extent_um_cache -def experiments_source_search(SOI): +def experiments_source_search(SOI: str) -> pd.DataFrame | None: """ - 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 + Return data about experiments whose injection was in the structure of interest. + + Parameters + ---------- + SOI + Acronym of the structure of interest to use as the search seed. + + Returns + ------- + pd.DataFrame or None + DataFrame of matching experiments, or None if AllenSDK is not installed. """ transgenic_id = 0 # id = 0 means use only wild type primary_structure_only = True @@ -81,15 +95,26 @@ def experiments_source_search(SOI): ) -def _get_injection_site_um(eid, ml_extent_um): +def _get_injection_site_um( + eid: int, + ml_extent_um: float, +) -> dict[str, float] | None: """ - 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 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 + Fetch injection site coordinates for an experiment from the Allen Brain Atlas API. + Coordinates are returned in Allen CCF µm space with the Z (ML) axis + flipped to match brainrender's hemisphere convention. + + Parameters + ---------- + eid + Experiment ID. + ml_extent_um + Full ML extent of the atlas in µm, used for the left-right flip. + + Returns + ------- + dict or None + Dict with ``x``, ``y``, ``z`` keys in µm, or None if not found. """ try: url = ( @@ -114,10 +139,13 @@ def _get_injection_site_um(eid, ml_extent_um): return None -def _skeleton_to_dataframe(skeleton, eid, ml_extent_um): +def _skeleton_to_dataframe( + skeleton: Any, + eid: int, + ml_extent_um: float, +) -> pd.DataFrame: """ - Converts a cloudvolume Skeleton object to the pd.DataFrame format - expected by brainrender's Streamlines actor. + Convert a cloudvolume Skeleton object to a DataFrame for the Streamlines actor. Vertices are in nanometers in Allen CCF space. We: 1. Convert nm -> um (divide by VOXEL_SIZE_NM) @@ -126,10 +154,19 @@ def _skeleton_to_dataframe(skeleton, eid, ml_extent_um): 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. - :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 + Parameters + ---------- + skeleton + cloudvolume Skeleton object. + eid + Experiment ID, used to fetch the real injection site coordinates. + ml_extent_um + Full ML extent of the atlas in µm, used for the left-right flip. + + Returns + ------- + pd.DataFrame + DataFrame with ``lines`` and ``injection_sites`` columns. """ components = skeleton.components() @@ -164,14 +201,25 @@ def _skeleton_to_dataframe(skeleton, eid, ml_extent_um): ) -def get_streamlines_data(eids, force_download=False): +def get_streamlines_data( + eids: list[int], + force_download: bool = False, +) -> list[pd.DataFrame]: """ - Given a list of experiment IDs, downloads streamline data from the + Given a list of experiment IDs, download 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 + via cloud-volume, and save them as JSON files. + + Parameters + ---------- + eids + Experiment IDs to download. + force_download + If True, re-download even if a cached file exists. Default False. + + Returns + ------- + list of pd.DataFrame """ if not cloudvolume_installed: print( @@ -210,7 +258,10 @@ def get_streamlines_data(eids, force_download=False): return data -def get_streamlines_for_region(region, force_download=False): +def get_streamlines_for_region( + region: str, + force_download: bool = False, +) -> list[pd.DataFrame] | None: """ 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 @@ -218,8 +269,17 @@ def get_streamlines_for_region(region, force_download=False): 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 + Parameters + ---------- + region + Acronym of the brain region to search for. + force_download + If True, re-download even if cached. Default False. + + Returns + ------- + list of pd.DataFrame or None + Streamlines data, or None if no experiments are found. """ logger.debug(f"Getting streamlines data for region: {region}") region_experiments = experiments_source_search(region) From 87f30a50e349d87a6abb17cdeddafa5d20ccf91b Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 13 Aug 2026 13:22:12 +0000 Subject: [PATCH 2/4] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../atlas_specific/allen_brain_atlas/gene_expression/api.py | 4 +++- brainrender/atlas_specific/allen_brain_atlas/streamlines.py | 2 +- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/brainrender/atlas_specific/allen_brain_atlas/gene_expression/api.py b/brainrender/atlas_specific/allen_brain_atlas/gene_expression/api.py index d2f6d033..0cca75a4 100644 --- a/brainrender/atlas_specific/allen_brain_atlas/gene_expression/api.py +++ b/brainrender/atlas_specific/allen_brain_atlas/gene_expression/api.py @@ -47,7 +47,9 @@ class GeneExpressionAPI: def __init__(self) -> None: # Get metadata about all available genes - self.genes: pd.DataFrame | None = None # when necessary gene data can be downloaded with self.get_all_genes + self.genes: pd.DataFrame | None = ( + None # when necessary gene data can be downloaded with self.get_all_genes + ) self.gene_expression_cache.mkdir(exist_ok=True) @fail_on_no_connection diff --git a/brainrender/atlas_specific/allen_brain_atlas/streamlines.py b/brainrender/atlas_specific/allen_brain_atlas/streamlines.py index 5ca85f80..b51fa9f4 100644 --- a/brainrender/atlas_specific/allen_brain_atlas/streamlines.py +++ b/brainrender/atlas_specific/allen_brain_atlas/streamlines.py @@ -162,7 +162,7 @@ def _skeleton_to_dataframe( Experiment ID, used to fetch the real injection site coordinates. ml_extent_um Full ML extent of the atlas in µm, used for the left-right flip. - + Returns ------- pd.DataFrame From 30c7c91b22567eef984e401cdcd3efdba016fc8e Mon Sep 17 00:00:00 2001 From: AlgoFoe Date: Mon, 17 Aug 2026 20:34:16 +0530 Subject: [PATCH 3/4] apply reviewed suggestions --- .../allen_brain_atlas/gene_expression/api.py | 6 ++++-- .../allen_brain_atlas/gene_expression/ge_utils.py | 5 +++-- .../atlas_specific/allen_brain_atlas/streamlines.py | 7 ++++--- 3 files changed, 11 insertions(+), 7 deletions(-) diff --git a/brainrender/atlas_specific/allen_brain_atlas/gene_expression/api.py b/brainrender/atlas_specific/allen_brain_atlas/gene_expression/api.py index d2f6d033..56c7ed0e 100644 --- a/brainrender/atlas_specific/allen_brain_atlas/gene_expression/api.py +++ b/brainrender/atlas_specific/allen_brain_atlas/gene_expression/api.py @@ -179,7 +179,7 @@ def get_gene_data( exp_id: int | str, use_cache: bool = True, metric: str = "energy", - ) -> npt.NDArray: + ) -> npt.NDArray | None: """ Load gene expression data for a given gene and experiment. @@ -196,7 +196,9 @@ def get_gene_data( Returns ------- - numpy.ndarray + numpy.ndarray or None + Gene expression data, or None if no data is available for the + requested metric. Raises ------ diff --git a/brainrender/atlas_specific/allen_brain_atlas/gene_expression/ge_utils.py b/brainrender/atlas_specific/allen_brain_atlas/gene_expression/ge_utils.py index 8f0f3371..618d5262 100644 --- a/brainrender/atlas_specific/allen_brain_atlas/gene_expression/ge_utils.py +++ b/brainrender/atlas_specific/allen_brain_atlas/gene_expression/ge_utils.py @@ -4,6 +4,7 @@ import os import sys import zipfile +from typing import Literal from pathlib import Path import numpy as np @@ -19,7 +20,7 @@ def check_gene_cached( cache_folder: str | Path, gene_id: str, exp_id: str | int, -) -> str | bool: +) -> str | Literal[False]: """ Check whether a gene experiment is already cached. @@ -37,7 +38,7 @@ def check_gene_cached( Returns ------- - str or bool + str or False Path to the cached folder if found, False if not cached. Raises diff --git a/brainrender/atlas_specific/allen_brain_atlas/streamlines.py b/brainrender/atlas_specific/allen_brain_atlas/streamlines.py index 5ca85f80..db563f05 100644 --- a/brainrender/atlas_specific/allen_brain_atlas/streamlines.py +++ b/brainrender/atlas_specific/allen_brain_atlas/streamlines.py @@ -46,13 +46,14 @@ def _get_ml_extent_um() -> float: """ Derive 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. + dynamically from the brainglobe atlas API. The computed extent is cached + for subsequent calls and used to flip the Z (ML) axis when converting + from Allen CCF space to brainrender's coordinate system. Returns ------- float + Full medial-lateral extent of the atlas in microns. """ global _ml_extent_um_cache if _ml_extent_um_cache is None: From 9ea12a4ef479a7413c99fe5470216596e6b47809 Mon Sep 17 00:00:00 2001 From: AlgoFoe Date: Mon, 17 Aug 2026 20:44:50 +0530 Subject: [PATCH 4/4] fix linting issues --- .../allen_brain_atlas/gene_expression/ge_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/brainrender/atlas_specific/allen_brain_atlas/gene_expression/ge_utils.py b/brainrender/atlas_specific/allen_brain_atlas/gene_expression/ge_utils.py index 618d5262..2a6ed417 100644 --- a/brainrender/atlas_specific/allen_brain_atlas/gene_expression/ge_utils.py +++ b/brainrender/atlas_specific/allen_brain_atlas/gene_expression/ge_utils.py @@ -4,8 +4,8 @@ import os import sys import zipfile -from typing import Literal from pathlib import Path +from typing import Literal import numpy as np import numpy.typing as npt