Add brainmapper analysis functionality

brainglobe_utils/brainmapper/analysis.py

@@ -0,0 +1,290 @@
+"""
+Parts based on https://github.com/SainsburyWellcomeCentre/cell_count_analysis
+by Charly Rousseau (https://github.com/crousseau).
+"""
+
+from dataclasses import dataclass
+from pathlib import Path
+from typing import List, Set, Union
+
+import numpy as np
+import pandas as pd
+from bg_atlasapi import BrainGlobeAtlas
+
+from brainglobe_utils.general.system import ensure_directory_exists
+from brainglobe_utils.pandas.misc import safe_pandas_concat, sanitise_df
+
+
+@dataclass
+class Point:
+    """
+    A class to represent a point in both raw and atlas coordinate spaces,
+    along with associated anatomical information.
+
+    Attributes
+    ----------
+    raw_coordinate : np.ndarray
+        A numpy array representing the raw coordinates of the point
+        in the original image space.
+    atlas_coordinate : np.ndarray
+        A numpy array representing the coordinates of the point in atlas space.
+    structure : str
+        The name of the atlas structure associated with the point.
+    structure_id : int
+        The numerical ID of the anatomical structure associated with the point.
+    hemisphere : str
+        The hemisphere ('left' or 'right') in which the point is located.
+
+    """
+
+    raw_coordinate: np.ndarray
+    atlas_coordinate: np.ndarray
+    structure: str
+    structure_id: int
+    hemisphere: str
+
+
+def calculate_densities(
+    counts: pd.DataFrame, volume_csv_path: Union[str, Path]
+) -> pd.DataFrame:
+    """
+    Use region volumes from registration to calculate cell densities.
+
+    Parameters
+    ----------
+    counts : pd.DataFrame
+        Dataframe with cell counts.
+    volume_csv_path : Union[str, Path]
+        Path of the CSV file containing the volumes of each brain region.
+
+    Returns
+    -------
+    pd.DataFrame
+        A dataframe containing the original cell counts and the calculated cell
+        densities per mm³ for each brain region.
+
+    """
+
+    volumes = pd.read_csv(volume_csv_path, sep=",", header=0, quotechar='"')
+    df = pd.merge(counts, volumes, on="structure_name", how="outer")
+    df = df.fillna(0)
+    df["left_cells_per_mm3"] = df.left_cell_count / df.left_volume_mm3
+    df["right_cells_per_mm3"] = df.right_cell_count / df.right_volume_mm3
+    return df
+
+
+def combine_df_hemispheres(df: pd.DataFrame) -> pd.DataFrame:
+    """
+    Combine left and right hemisphere data onto a single row
+
+    Parameters
+    ----------
+    df : pd.DataFrame
+        A pandas DataFrame with separate rows for left and
+        right hemisphere data.
+
+    Returns
+    -------
+    pd.DataFrame
+        A DataFrame with combined hemisphere data. Each row
+        represents the combined data of left and right hemispheres for
+        each brain region.
+
+    """
+    left = df[df["hemisphere"] == "left"]
+    right = df[df["hemisphere"] == "right"]
+    left = left.drop(["hemisphere"], axis=1)
+    right = right.drop(["hemisphere"], axis=1)
+    left.rename(columns={"cell_count": "left_cell_count"}, inplace=True)
+    right.rename(columns={"cell_count": "right_cell_count"}, inplace=True)
+    both = pd.merge(left, right, on="structure_name", how="outer")
+    both = both.fillna(0)
+    both["total_cells"] = both.left_cell_count + both.right_cell_count
+    both = both.sort_values("total_cells", ascending=False)
+    return both
+
+
+def create_all_cell_csv(
+    points: List[Point], output_filename: Union[str, Path]
+) -> None:
+    """
+    Create a CSV file with cell data from a list of Point objects.
+
+    This function takes a list of Point objects, each representing cell
+    coordinates and brain region and converts this into a pandas DataFrame.
+    The DataFrame is then saved to a CSV file at the specified filename.
+
+    Parameters
+    ----------
+    points : List[Point]
+        A list of Point objects, each containing cell data such as
+        raw and atlas coordinates,
+        structure name, and hemisphere information.
+    output_filename : Union[str, Path]
+        The filename (including path) where the CSV file will be saved.
+        Can be a string or a Path object.
+
+    Returns
+    -------
+    None
+    """
+
+    ensure_directory_exists(Path(output_filename).parent)
+    df = pd.DataFrame(
+        columns=(
+            "coordinate_raw_axis_0",
+            "coordinate_raw_axis_1",
+            "coordinate_raw_axis_2",
+            "coordinate_atlas_axis_0",
+            "coordinate_atlas_axis_1",
+            "coordinate_atlas_axis_2",
+            "structure_name",
+            "hemisphere",
+        )
+    )
+
+    temp_matrix = [[] for i in range(len(points))]
+    for i, point in enumerate(points):
+        temp_matrix[i].append(point.raw_coordinate[0])
+        temp_matrix[i].append(point.raw_coordinate[1])
+        temp_matrix[i].append(point.raw_coordinate[2])
+        temp_matrix[i].append(point.atlas_coordinate[0])
+        temp_matrix[i].append(point.atlas_coordinate[1])
+        temp_matrix[i].append(point.atlas_coordinate[2])
+        temp_matrix[i].append(point.structure)
+        temp_matrix[i].append(point.hemisphere)
+
+    df = pd.DataFrame(temp_matrix, columns=df.columns, index=None)
+    df.to_csv(output_filename, index=False)
+
+
+def count_points_per_brain_region(
+    points: List[Point],
+    structures_with_points: Set[str],
+    brainreg_volume_csv_path: Union[str, Path],
+    output_filename: Union[str, Path],
+) -> None:
+    """
+    Count the number of points per brain region.
+
+    Parameters
+    ----------
+    points : List[Point]
+        A list of Point objects.
+    structures_with_points : Set[str]
+        A set of strings representing the names of all atlas structures
+        represented
+    brainreg_volume_csv_path : Union[str, Path]
+        The path to the CSV file containing volume information from the
+        brainreg registration.
+    output_filename : Union[str, Path]
+        The path where the summary of points by atlas region will be saved.
+
+    Returns
+    -------
+    None
+    """
+
+    structures_with_points = list(structures_with_points)
+
+    point_numbers = pd.DataFrame(
+        columns=("structure_name", "hemisphere", "cell_count")
+    )
+    for structure in structures_with_points:
+        for hemisphere in ("left", "right"):
+            n_points = len(
+                [
+                    point
+                    for point in points
+                    if point.structure == structure
+                    and point.hemisphere == hemisphere
+                ]
+            )
+            if n_points:
+                point_numbers = safe_pandas_concat(
+                    point_numbers,
+                    pd.DataFrame(
+                        data=[[structure, hemisphere, n_points]],
+                        columns=[
+                            "structure_name",
+                            "hemisphere",
+                            "cell_count",
+                        ],
+                    ),
+                )
+    sorted_point_numbers = point_numbers.sort_values(
+        by=["cell_count"], ascending=False
+    )
+
+    combined_hemispheres = combine_df_hemispheres(sorted_point_numbers)
+    df = calculate_densities(combined_hemispheres, brainreg_volume_csv_path)
+    df = sanitise_df(df)
+
+    df.to_csv(output_filename, index=False)
+
+
+def summarise_points_by_atlas_region(
+    points_in_raw_data_space: np.ndarray,
+    points_in_atlas_space: np.ndarray,
+    atlas: BrainGlobeAtlas,
+    brainreg_volume_csv_path: Union[str, Path],
+    points_list_output_filename: Union[str, Path],
+    summary_filename: Union[str, Path],
+) -> None:
+    """
+    Summarise points data by atlas region.
+
+    This function takes points in both raw data space and atlas space,
+    and generates a summary of these points based on the BrainGlobe atlas
+    region they are found in.
+    The summary is saved to a CSV file.
+
+    Parameters
+    ----------
+    points_in_raw_data_space : np.ndarray
+        A numpy array representing points in the raw data space.
+    points_in_atlas_space : np.ndarray
+        A numpy array representing points in the atlas space.
+    atlas : BrainGlobeAtlas
+        The BrainGlobe atlas object used for the analysis
+    brainreg_volume_csv_path : Union[str, Path]
+        The path to the CSV file containing volume information from the
+        brainreg registration.
+    points_list_output_filename : Union[str, Path]
+        The path where the detailed points list will be saved.
+    summary_filename : Union[str, Path]
+        The path where the summary of points by atlas region will be saved.
+
+    Returns
+    -------
+    None
+    """
+
+    points = []
+    structures_with_points = set()
+    for idx, point in enumerate(points_in_atlas_space):
+        try:
+            structure_id = atlas.structure_from_coords(point)
+            structure = atlas.structures[structure_id]["name"]
+            hemisphere = atlas.hemisphere_from_coords(point, as_string=True)
+            points.append(
+                Point(
+                    points_in_raw_data_space[idx],
+                    point,
+                    structure,
+                    structure_id,
+                    hemisphere,
+                )
+            )
+            structures_with_points.add(structure)
+        except Exception:
+            continue
+
+    create_all_cell_csv(points, points_list_output_filename)
+
+    count_points_per_brain_region(
+        points,
+        structures_with_points,
+        brainreg_volume_csv_path,
+        summary_filename,
+    )