multineuronchat 2025.11.10.dev0__py3-none-any.whl

multineuronchat/__init__.py

@@ -0,0 +1,28 @@
+from .MultiNeuronChatObject import MultiNeuronChatObject
+
+from .InteractionDB import (
+    InteractionDB,
+    InteractionDBRow
+)
+
+from .MultiNeuronChat import (
+    compute_communication_score_matrix,
+    compute_subject_specific_communication_score_matrix,
+    compute_avg_expression,
+    compute_subject_specific_avg_expression
+)
+
+from .normalize import (
+    cell_wise_log_normalization,
+    subject_wise_max_normalization
+)
+
+from .utils import (
+    filter_genes,
+    gene_filter_and_subject_wise_normalize_dataset
+)
+
+__all__ = ['MultiNeuronChatObject', 'InteractionDB', 'InteractionDBRow', 'compute_communication_score_matrix',
+           'compute_subject_specific_communication_score_matrix', 'compute_avg_expression',
+           'compute_subject_specific_avg_expression', 'cell_wise_log_normalization',
+           'subject_wise_max_normalization', 'filter_genes', 'gene_filter_and_subject_wise_normalize_dataset']

multineuronchat/loompy_utils.py

@@ -0,0 +1,66 @@
+import numpy as np
+
+import loompy
+from loompy import timestamp
+import h5py
+
+from typing import Dict, Tuple, Optional
+
+def create_empty_loom_file(
+        path_to_loom: str,
+        shape: Tuple[int, int],
+        row_attrs: Optional[Dict[str, np.ndarray]] = None,
+        col_attrs: Optional[Dict[str, np.ndarray]] = None,
+        dtype_to_use = np.float32
+) -> None:
+    """
+    Creates an empty loom file at the specified path with the given number of rows and columns.
+
+    Parameters
+    ----------
+    path_to_loom: str
+        Path where the loom file should be created.
+    shape: Tuple[int, int]
+        Shape of the matrix to be created in the loom file, specified as (n_rows, n_cols).
+    row_attrs: Optional[Dict[str, np.ndarray]]
+        Dictionary of row attributes where keys are attribute names and values are arrays of attribute values.
+    col_attrs: Optional[Dict[str, np.ndarray]]
+        Dictionary of column attributes where keys are attribute names and values are arrays of attribute values.
+    dtype_to_use: np.dtype
+        Data type to use for the matrix in the loom file (default: np.float32).
+
+    Returns
+    -------
+    None
+    """
+    n_rows, n_cols = shape
+
+    with h5py.File(path_to_loom, 'w') as f:
+        f.create_group('/attrs')  # v3.0.0
+
+        f.create_group('/layers')
+        f.create_group('/row_attrs')
+        f.create_group('/col_attrs')
+        f.create_group('/row_graphs')
+        f.create_group('/col_graphs')
+
+        f.create_dataset(
+            '/matrix',
+            shape=(n_rows, n_cols),
+            dtype=dtype_to_use,
+            chunks=True,
+            compression='gzip',
+            compression_opts=9,
+        )
+
+        f.flush()
+
+        with loompy.connect(path_to_loom, 'r+') as ds:
+            ds.attrs['CreationDate'] = timestamp()
+            ds.attrs["LOOM_SPEC_VERSION"] = loompy.loom_spec_version
+
+            for key, vals in row_attrs.items():
+                ds.ra[key] = vals
+
+            for key, vals in col_attrs.items():
+                ds.ca[key] = vals

multineuronchat/masks.py

@@ -0,0 +1,358 @@
+import numpy as np
+import xarray as xr
+
+from .MultiNeuronChatObject import MultiNeuronChatObject
+
+from typing import Optional, Union, Tuple
+
+
+def compute_wasserstein_mask(
+    mnc_object: MultiNeuronChatObject,
+    top_n: Optional[int] = None,
+    top_percentile: Optional[float] = None,
+    p: Optional[float] = 1.0,
+    normalized: Optional[bool] = False,
+    exclude_zero_distributions: Optional[bool] = True,
+    return_wasserstein_distances: Optional[bool] = False
+) -> Union[xr.DataArray, Tuple[xr.DataArray, xr.DataArray]]:
+    """
+    Compute a mask for the hypotheses with the top_n or top_percentile largest Wasserstein distances between the
+    distributions of the communication scores for each hypothesis. The mask is returned as a xarray DataArray with
+    dimensions 'source_cell_type', 'receiver_cell_type', and 'interaction_type'.
+
+    See also: compute_earth_mover_mask
+
+    :param mnc_object: A MultiNeuronChatObject containing the communication scores
+    :param top_n: The number of hypotheses with the largest Wasserstein distances that should be included in the mask.
+                  If None, top_percentile must be provided.
+    :param top_percentile: The percentile of the hypotheses with the largest Wasserstein distances that should be
+                           included in the mask. If None, top_n must be provided.
+    :param normalized: If True, the Wasserstein distances are normalized by the range of the communication scores for
+                       each hypothesis. This is done to ensure that the Wasserstein distances are comparable.
+                       Default: True
+    :param exclude_zero_distributions: If True, hypotheses where both distributions consist solely of zeros are excluded.
+    :param return_wasserstein_distances: If True, the computed Wasserstein distances are returned as well.
+    :return: A boolean xarray DataArray with dimensions 'source_cell_type', 'receiver_cell_type', and 'interaction_type'
+             or a tuple with the mask and the Wasserstein distances as xarray DataArrays
+    """
+    return compute_earth_mover_mask(
+        mnc_object=mnc_object,
+        top_n=top_n,
+        top_percentile=top_percentile,
+        p=p,
+        normalized=normalized,
+        exclude_zero_distributions=exclude_zero_distributions,
+        return_wasserstein_distances=return_wasserstein_distances
+    )
+
+
+def compute_earth_mover_mask(
+    mnc_object: MultiNeuronChatObject,
+    top_n: Optional[int] = None,
+    top_percentile: Optional[float] = None,
+    p: Optional[float] = 1.0,
+    normalized: Optional[bool] = False,
+    exclude_zero_distributions: Optional[bool] = True,
+    return_wasserstein_distances: Optional[bool] = False
+) -> Union[xr.DataArray, Tuple[xr.DataArray, xr.DataArray]]:
+    """
+    Compute a mask for the hypotheses with the top_n or top_percentile largest Wasserstein distances between the
+    distributions of the communication scores for each hypothesis. The mask is returned as a xarray DataArray with
+    dimensions 'source_cell_type', 'receiver_cell_type', and 'interaction_type'.
+
+    :param mnc_object: A MultiNeuronChatObject containing the communication scores
+    :param top_n: The number of hypotheses with the largest Wasserstein distances that should be included in the mask.
+                  If None, top_percentile must be provided.
+    :param top_percentile: The percentile of the hypotheses with the largest Wasserstein distances that should be
+                           included in the mask. If None, top_n must be provided.
+    :param normalized: If True, the Wasserstein distances are normalized by the range of the communication scores for
+                       each hypothesis. This is done to ensure that the Wasserstein distances are comparable.
+                       Default: True
+    :param exclude_zero_distributions: If True, hypotheses where both distributions consist solely of zeros are excluded.
+    :param return_wasserstein_distances: If True, the computed Wasserstein distances are returned as well.
+    :return: A boolean xarray DataArray with dimensions 'source_cell_type', 'receiver_cell_type', and 'interaction_type'
+             or a tuple with the mask and the Wasserstein distances as xarray DataArrays
+    """
+    if mnc_object is None:
+        raise ValueError('A MultiNeuronChatObject must be provided')
+
+    if mnc_object.communication_score_dict is None:
+        raise ValueError('The MultiNeuronChatObject must contain the computed communication scores!'
+                         'Please run compute_communication_scores() first.')
+
+    n_source_cell_types = len(mnc_object.source_cell_types)
+    n_receiver_cell_types = len(mnc_object.receiver_cell_types)
+    n_interaction_types = len(mnc_object.interaction_names)
+    n_hypotheses = n_source_cell_types * n_receiver_cell_types * n_interaction_types
+
+    if top_n is None and top_percentile is None:
+        raise ValueError('Either top_n or top_percentile must be provided')
+
+    if top_n is not None and top_percentile is not None:
+        raise ValueError('Only one of top_n or top_percentile must be provided')
+
+    if top_n is not None:
+        if top_n <= 0:
+            raise ValueError('top_n must be a positive integer')
+
+        # Check if top_n is bigger than the number of hypotheses
+        # (product of number of source cell types, receiver cell, and interaction types)
+        if top_n > n_hypotheses:
+            raise ValueError(f'top_n must be smaller than the number of hypotheses ({n_hypotheses})')
+
+    if top_percentile is not None:
+        if top_percentile <= 0 or top_percentile >= 100:
+            raise ValueError('top_percentile must be a float between 0 and 100')
+
+    if p <= 0:
+        raise ValueError('p must be a positive value!')
+
+    # Compute the Wasserstein distance between the distributions of the communication scores for each hypothesis
+    wasserstein_distances = np.ones((n_source_cell_types, n_receiver_cell_types, n_interaction_types)) * np.nan
+
+    for s, source in enumerate(mnc_object.source_cell_types):
+        for r, receiver in enumerate(mnc_object.receiver_cell_types):
+            for lt, interaction in enumerate(mnc_object.interaction_names):
+                ctrl_dist = mnc_object.communication_score_dict[mnc_object.condition_names[0]][:, s, r, lt]
+                scz_dist = mnc_object.communication_score_dict[mnc_object.condition_names[1]][:, s, r, lt]
+
+                ctrl_dist = ctrl_dist[~np.isnan(ctrl_dist)]
+                scz_dist = scz_dist[~np.isnan(scz_dist)]
+
+                if len(ctrl_dist) == 0 or len(scz_dist) == 0:
+                    continue
+                if exclude_zero_distributions and np.all(ctrl_dist == 0) and np.all(scz_dist == 0):
+                    continue
+
+                min_value: float = min(ctrl_dist.min(), scz_dist.min())
+                max_value: float = max(ctrl_dist.max(), scz_dist.max())
+                value_range: float = max_value - min_value
+
+                #wasserstein_distances[s, r, lt] = stats.wasserstein_distance(ctrl_dist, scz_dist)
+                wasserstein_distances[s, r, lt] = _cdf_distance(p, ctrl_dist, scz_dist)
+
+                if normalized and value_range > 0:
+                    wasserstein_distances[s, r, lt] /= value_range
+
+    mask: np.array = np.zeros((n_source_cell_types, n_receiver_cell_types, n_interaction_types), dtype=bool)
+
+    if top_n is not None:
+        # Get the top_n indices with the smallest Wasserstein distances (which can include nan values)
+        top_n_indices = np.unravel_index(np.argsort(-wasserstein_distances, axis=None)[:top_n], wasserstein_distances.shape)
+        mask = np.zeros_like(wasserstein_distances, dtype=bool)
+        mask[top_n_indices] = True
+    else:
+        percentile = np.nanpercentile(wasserstein_distances, top_percentile)
+        mask = wasserstein_distances >= percentile
+
+    mask_xr = xr.DataArray(
+        mask,
+        dims=('source', 'receiver', 'interaction'),
+        coords={
+            'source': mnc_object.source_cell_types,
+            'receiver': mnc_object.receiver_cell_types,
+            'interaction': mnc_object.interaction_names
+        }
+    )
+
+    wasserstein_distances_xr = xr.DataArray(
+        wasserstein_distances,
+        dims=('source', 'receiver', 'interaction'),
+        coords={
+            'source': mnc_object.source_cell_types,
+            'receiver': mnc_object.receiver_cell_types,
+            'interaction': mnc_object.interaction_names
+        }
+    )
+
+    if return_wasserstein_distances:
+        return mask_xr, wasserstein_distances_xr
+    return mask_xr
+
+
+def compute_threshold_communication_score_mask(
+        mnc_object: MultiNeuronChatObject,
+        n_samples: Optional[int] = 1,
+        threshold: Optional[float] = 0.0,
+) -> xr.DataArray:
+    """
+    Compute a mask for the hypotheses where n_samples from each condition must have a communication score that is
+    higher than the threshold. The mask is returned as a xarray DataArray with dimensions 'source_cell_type',
+    'receiver_cell_type', and 'interaction_type'.
+
+    :param mnc_object: A MultiNeuronChatObject containing the communication scores
+    :param n_samples: The number of samples from each condition that must have a communication score higher than the threshold.
+    :param threshold: The threshold that the communication score must exceed.
+    :return: A boolean xarray DataArray with dimensions 'source_cell_type', 'receiver_cell_type', and 'interaction_type'
+    """
+
+    if mnc_object is None:
+        raise ValueError('A MultiNeuronChatObject must be provided')
+
+    if mnc_object.communication_score_dict is None:
+        raise ValueError('The MultiNeuronChatObject must contain the computed communication scores!'
+                         'Please run compute_communication_scores() first.')
+
+    n_source_cell_types = len(mnc_object.source_cell_types)
+    n_receiver_cell_types = len(mnc_object.receiver_cell_types)
+    n_interaction_types = len(mnc_object.interaction_names)
+
+    mask: np.array = np.zeros((n_source_cell_types, n_receiver_cell_types, n_interaction_types), dtype=bool)
+
+    for s, source in enumerate(mnc_object.source_cell_types):
+        for r, receiver in enumerate(mnc_object.receiver_cell_types):
+            for lt, interaction in enumerate(mnc_object.interaction_names):
+                ctrl_dist = mnc_object.communication_score_dict[mnc_object.condition_names[0]][:, s, r, lt]
+                scz_dist = mnc_object.communication_score_dict[mnc_object.condition_names[1]][:, s, r, lt]
+
+                ctrl_mask = ctrl_dist > threshold
+                scz_mask = scz_dist > threshold
+
+                if np.sum(ctrl_mask) >= n_samples and np.sum(scz_mask) >= n_samples:
+                    mask[s, r, lt] = True
+
+    mask_xr = xr.DataArray(
+        mask,
+        dims=('source', 'receiver', 'interaction'),
+        coords={
+            'source': mnc_object.source_cell_types,
+            'receiver': mnc_object.receiver_cell_types,
+            'interaction': mnc_object.interaction_names
+        }
+    )
+
+    return mask_xr
+
+
+########################################################################################################################
+# Helper functions
+########################################################################################################################
+
+def _cdf_distance(p, u_values, v_values, u_weights=None, v_weights=None):
+    r"""
+    Compute, between two one-dimensional distributions :math:`u` and
+    :math:`v`, whose respective CDFs are :math:`U` and :math:`V`, the
+    statistical distance that is defined as:
+
+    .. math::
+
+        l_p(u, v) = \left( \int_{-\infty}^{+\infty} |U-V|^p \right)^{1/p}
+
+    p is a positive parameter; p = 1 gives the Wasserstein distance, p = 2
+    gives the energy distance.
+
+    Parameters
+    ----------
+    u_values, v_values : array_like
+        Values observed in the (empirical) distribution.
+    u_weights, v_weights : array_like, optional
+        Weight for each value. If unspecified, each value is assigned the same
+        weight.
+        `u_weights` (resp. `v_weights`) must have the same length as
+        `u_values` (resp. `v_values`). If the weight sum differs from 1, it
+        must still be positive and finite so that the weights can be normalized
+        to sum to 1.
+
+    Returns
+    -------
+    distance : float
+        The computed distance between the distributions.
+
+    Notes
+    -----
+    The input distributions can be empirical, therefore coming from samples
+    whose values are effectively inputs of the function, or they can be seen as
+    generalized functions, in which case they are weighted sums of Dirac delta
+    functions located at the specified values.
+
+    References
+    ----------
+    .. [1] Bellemare, Danihelka, Dabney, Mohamed, Lakshminarayanan, Hoyer,
+           Munos "The Cramer Distance as a Solution to Biased Wasserstein
+           Gradients" (2017). :arXiv:`1705.10743`.
+
+    """
+    u_values, u_weights = _validate_distribution(u_values, u_weights)
+    v_values, v_weights = _validate_distribution(v_values, v_weights)
+
+    u_sorter = np.argsort(u_values)
+    v_sorter = np.argsort(v_values)
+
+    all_values = np.concatenate((u_values, v_values))
+    all_values.sort(kind='mergesort')
+
+    # Compute the differences between pairs of successive values of u and v.
+    deltas = np.diff(all_values)
+
+    # Get the respective positions of the values of u and v among the values of
+    # both distributions.
+    u_cdf_indices = u_values[u_sorter].searchsorted(all_values[:-1], 'right')
+    v_cdf_indices = v_values[v_sorter].searchsorted(all_values[:-1], 'right')
+
+    # Calculate the CDFs of u and v using their weights, if specified.
+    if u_weights is None:
+        u_cdf = u_cdf_indices / u_values.size
+    else:
+        u_sorted_cumweights = np.concatenate(([0],
+                                              np.cumsum(u_weights[u_sorter])))
+        u_cdf = u_sorted_cumweights[u_cdf_indices] / u_sorted_cumweights[-1]
+
+    if v_weights is None:
+        v_cdf = v_cdf_indices / v_values.size
+    else:
+        v_sorted_cumweights = np.concatenate(([0],
+                                              np.cumsum(v_weights[v_sorter])))
+        v_cdf = v_sorted_cumweights[v_cdf_indices] / v_sorted_cumweights[-1]
+
+    # Compute the value of the integral based on the CDFs.
+    # If p = 1 or p = 2, we avoid using np.power, which introduces an overhead
+    # of about 15%.
+    if p == 1:
+        return np.sum(np.multiply(np.abs(u_cdf - v_cdf), deltas))
+    if p == 2:
+        return np.sqrt(np.sum(np.multiply(np.square(u_cdf - v_cdf), deltas)))
+    return np.power(np.sum(np.multiply(np.power(np.abs(u_cdf - v_cdf), p),
+                                       deltas)), 1/p)
+
+
+def _validate_distribution(values, weights):
+    """
+    Validate the values and weights from a distribution input of `cdf_distance`
+    and return them as ndarray objects.
+
+    Parameters
+    ----------
+    values : array_like
+        Values observed in the (empirical) distribution.
+    weights : array_like
+        Weight for each value.
+
+    Returns
+    -------
+    values : ndarray
+        Values as ndarray.
+    weights : ndarray
+        Weights as ndarray.
+
+    """
+    # Validate the value array.
+    values = np.asarray(values, dtype=float)
+    if len(values) == 0:
+        raise ValueError("Distribution can't be empty.")
+
+    # Validate the weight array, if specified.
+    if weights is not None:
+        weights = np.asarray(weights, dtype=float)
+        if len(weights) != len(values):
+            raise ValueError('Value and weight array-likes for the same '
+                             'empirical distribution must be of the same size.')
+        if np.any(weights < 0):
+            raise ValueError('All weights must be non-negative.')
+        if not 0 < np.sum(weights) < np.inf:
+            raise ValueError('Weight array-like sum must be positive and '
+                             'finite. Set as None for an equal distribution of '
+                             'weight.')
+
+        return values, weights
+
+    return values, None

multineuronchat/normalize.py

@@ -0,0 +1,177 @@
+import os
+
+import numpy as np
+
+import loompy
+from .loompy_utils import create_empty_loom_file
+
+from tqdm import tqdm
+
+from typing import Optional
+
+
+def cell_wise_log_normalization(
+        path_to_loom: str,
+        path_to_normalized_loom: Optional[str] = '',
+        chunk_size: Optional[int] = -1,
+        verbose: Optional[bool] = False
+) -> str:
+    """
+    Perform cell-wise log-normalization on the loom file at path_to_loom. The normalized loom file is saved at
+    path_to_normalized_loom. If path_to_normalized_loom is not provided, the normalized loom file is saved in the same
+    directory as the original loom file with the suffix '_CellWiseNormalized.loom'.
+
+    If the system you are running this on has limited memory, you can set the chunk_size to a lower value to reduce the
+    memory usage. If the chunk_size is set to zero or a negative number, the complete matrix is loaded into memory.
+
+    :param path_to_loom: Path to the loom file that should be normalized
+    :param path_to_normalized_loom: Path where the normalized loom file should be saved (default: same directory as the
+                                    original loom file with the suffix '_CellWiseNormalized.loom')
+    :param chunk_size: Number of columns that are loaded into memory at once (default: -1)
+    :param verbose: If True, the progress is printed to the console (default: False)
+    :return: Path to the cell-wise log-normalized loom file. If path_to_normalized_loom is provided, the same path is
+             returned.
+    """
+
+    if not os.path.isfile(path_to_loom):
+        raise FileNotFoundError(f"File {path_to_loom} not found")
+
+    if not path_to_loom.endswith('.loom'):
+        raise ValueError('The file must be a loom file')
+
+    # if path_to_normalized_loom is not provided, save the normalized loom file in the same directory as the original
+    # loom file with the suffix '_CellWiseNormalized.loom'
+    if path_to_normalized_loom == '' or path_to_normalized_loom is None:
+        path_to_normalized_loom: str = path_to_loom.replace(".loom", "_CellWiseNormalized.loom")
+
+    with loompy.connect(path_to_loom, mode='r') as src:
+        n_rows, n_cols = src.shape
+
+        # if the chunk_size is set to zero or a negative number, we load the complete matrix into memory
+        if chunk_size <= 0:
+            chunk_size = n_cols
+
+        # Create a new loom file with the same row and column attributes as the original loom file
+        create_empty_loom_file(
+            path_to_loom=path_to_normalized_loom,
+            shape=(n_rows, n_cols),
+            row_attrs=src.ra,
+            col_attrs=src.ca,
+            dtype_to_use=np.float32
+        )
+
+        # Connect to the new loom file in read/write mode and perform normalization
+        with loompy.connect(path_to_normalized_loom, mode='r+') as dst:
+            if verbose:
+                print('Cell-Wise Log-Normalization:')
+
+            for i in tqdm(range(0, n_cols, chunk_size), disable=(not verbose)):
+                end_i: int = min(i + chunk_size, n_cols)
+
+                data = src[:, i:end_i]
+
+                column_sum: np.array = np.sum(data, axis=0)
+                column_sum[column_sum == 0] = 1
+                log_norm: np.array = np.log1p(10_000 * (data / column_sum[None, :]))
+
+                dst[:, i:end_i] = log_norm
+
+    return path_to_normalized_loom
+
+
+def subject_wise_max_normalization(
+        path_to_loom: str,
+        subject_label_column: str,
+        path_to_normalized_loom: Optional[str] = '',
+        chunk_size: Optional[int] = -1
+) -> str:
+    """
+    Perform subject-wise max-normalization on the loom file at path_to_loom. The normalized loom file is saved at
+    path_to_normalized_loom. If path_to_normalized_loom is not provided, the normalized loom file is saved in the same
+    directory as the original loom file with the suffix '_MaxNormalized.loom'.
+
+    If the system you are running this on has limited memory, you can set the chunk_size to a lower value to reduce the
+    memory usage. If the chunk_size is set to zero or a negative number, the complete matrix is loaded into memory.
+
+    :param path_to_loom: Path to the loom file that should be normalized
+    :param subject_label_column: Column name in the column attributes that contains the subject IDs
+    :param path_to_normalized_loom: Path where the normalized loom file should be saved (default: same directory as the
+                                    original loom file with the suffix '_MaxNormalized.loom')
+    :param chunk_size: Number of columns that are loaded into memory at once (default: -1)
+    :return: Path to the subject-wise max-normalized loom file. If path_to_normalized_loom is provided, the same path is
+             returned.
+    """
+
+    if not os.path.isfile(path_to_loom):
+        raise FileNotFoundError(f"File {path_to_loom} not found")
+
+    if not path_to_loom.endswith('.loom'):
+        raise ValueError('The file must be a loom file')
+
+    # if path_to_normalized_loom is not provided, save the normalized loom file in the same directory as the original
+    # loom file with the suffix '_MaxNormalized.loom'
+    if path_to_normalized_loom == '' or path_to_normalized_loom is None:
+        path_to_normalized_loom: str = path_to_loom.replace(".loom", "_MaxNormalized.loom")
+
+    with loompy.connect(path_to_loom, mode='r') as src:
+        n_rows, n_cols = src.shape
+
+        if chunk_size is None or chunk_size < 0:
+            chunk_size = n_cols
+
+        # Pass 0: collect subjects and build index groups (no I/O on matrix)
+        subjects = np.asarray(src.ca[subject_label_column])
+        # Make deterministic order of subjects for reproducibility
+        uniq_subjects, inv = np.unique(subjects, return_inverse=True)
+        # inv[j] gives the index in uniq_subjects for column j
+
+        # Pass 1: compute per-column maxima, streaming in chunks:
+        col_max = np.empty(n_cols, dtype=np.float32)
+        for i in tqdm(range(0, n_cols, chunk_size), desc='Computing column maxima'):
+            end_i = min(i + chunk_size, n_cols)
+            X = src[:, i:end_i].astype(np.float32, copy=False)
+            col_max[i:end_i] = np.max(X, axis=0)
+            del X
+
+        # Reduce per-subject: max across the subject's columns
+        subj_max = np.full(uniq_subjects.shape[0], -np.inf, dtype=np.float32)
+
+        # vectorized max-by-group using bincount trick
+        # We can’t max with bincount directly, so do a simple loop over subjects on col_max indices
+        # but that loop is tiny (n_subjects), not n_rows*n_cols
+        for s_idx in range(uniq_subjects.shape[0]):
+            # columns belonging to subject s_idx
+            mask = (inv == s_idx)
+            if not np.any(mask):
+                subj_max[s_idx] = 1 # just to be sure, but should not happen
+            else:
+                m = col_max[mask].max()
+                subj_max[s_idx] = m if m > 0 else 1 # m <= 0 should not happen, but just in case
+
+        # Map each column to its subject's max value
+        col_scale = subj_max[inv]
+
+        # Create the new loom file with the same row and column attributes as the original loom file
+        create_empty_loom_file(
+            path_to_loom=path_to_normalized_loom,
+            shape=(n_rows, n_cols),
+            row_attrs=src.ra,
+            col_attrs=src.ca,
+            dtype_to_use=np.float32
+        )
+
+        # Pass 2: write normalized data in chunks (in-place operation and in float32)
+        with loompy.connect(path_to_normalized_loom, mode='r+') as dst:
+            for i in tqdm(range(0, n_cols, chunk_size), desc='Normalizing data'):
+                end_i = min(i + chunk_size, n_cols)
+                X = src[:, i:end_i].astype(np.float32, copy=False)
+
+                # Normalize by the subject's max value
+                X /= col_scale[i:end_i][None, :].astype(np.float32, copy=False)
+
+                # Write the normalized data to the new loom file
+                dst[:, i:end_i] = X
+
+                del X
+
+    return path_to_normalized_loom