Atlas correlation matrix

from pathlib import Path

import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy as np
import xarray as xr

import confusius as cf

# Adapt background color to the current Matplotlib style.
bg_color = mpl.colors.to_hex(mpl.rcParams["figure.facecolor"])

xr.set_options(display_expand_data=False)

template = cf.datasets.fetch_template_pepe_mariani_2026()

bids_root = cf.datasets.fetch_nunez_elizalde_2022(
    subjects="CR022", sessions="20201007", tasks="spontaneous", acqs="slice02"
)

data_path = (
    Path(bids_root)
    / "sub-CR022"
    / "ses-20201007"
    / "fusi"
    / "sub-CR022_ses-20201007_task-spontaneous_acq-slice02_pwd.nii.gz"
)
# The recording's timepoints are not perfectly uniformly spaced so we resample to a
# uniform grid before any time-domain processing (filtering below requires it).
data = cf.timing.resample_to_uniform_time(cf.load(data_path))
moving = data.mean(dim="time").fusi.scale.db().compute()

data

napari_affine = np.array(
    [
        [1.0, 0.0, 0.0, 5.594638656430411],
        [0.0, 1.0, 0.0, -2.50293925701927],
        [0.0, 0.0, 1.0, 5.6650243788545875],
        [0.0, 0.0, 0.0, 1.0],
    ]
)
initialization = np.linalg.inv(napari_affine)

# Crop the template to a thin band around the recording's expected location to improve
# registration speed and visualization.
target_z = napari_affine[0, 3] + float(moving.z.values[0])
fixed = template.sel(z=slice(target_z - 1.0, target_z + 1.0)).fusi.scale.db()

initialized = cf.registration.resample_like(
    moving, fixed, initialization, default_value=float(moving.min())
)
_ = cf.plotting.plot_composite(
    fixed,
    initialized,
    slice_coords=[target_z],
    normalize_strategy="per_slice",
    bg_color=bg_color,
)

registered, affine, _ = cf.registration.register_volume(
    moving=moving,
    fixed=fixed,
    transform_type="affine",
    metric="correlation",
    convergence_window_size=50,
    number_of_iterations=500,
    learning_rate=1,
    initialization=initialization,
    show_progress=True,
)

fig, axes = plt.subplots(1, 2, figsize=(10, 4))
fig.patch.set_facecolor(bg_color)
for ax, moving_view, title in [
    (axes[0], initialized, "Manual initialization"),
    (axes[1], registered, "Affine registration refinement"),
]:
    cf.plotting.plot_composite(
        fixed,
        moving_view,
        axes=ax,
        slice_coords=[target_z],
        normalize_strategy="per_slice",
        bg_color=bg_color,
    )
    ax.set_title(title)

_ = fig.suptitle("Template (red) / recording (cyan)")

physical_to_sform = template.attrs["affines"]["physical_to_sform"]
subject_to_atlas = physical_to_sform @ np.linalg.inv(affine)

atlas = cf.atlas.Atlas.from_brainglobe("allen_mouse_100um")
atlas_native = atlas.resample_like(moving, subject_to_atlas)

plotter = cf.plotting.plot_volume(
    moving, slice_mode="z", cmap="gray", show_colorbar=False, bg_color=bg_color
)
_ = plotter.add_contours(atlas_native.annotation)

groups = {
    "cortex": ["RSPv", "MOp", "SSp-bfd"],
    "hippocampus": ["DG", "CA3", "CA2"],
    "thalamus": ["PO", "VPM", "RT"],
    "hypothalamus": ["PH", "LHA", "ZI"],
}
# Allen CCF ids for each area's parent division, used below to color the group strips
# with the atlas's own official colors instead of arbitrary ones.
division_ids = {"cortex": 315, "hippocampus": 1089, "thalamus": 549, "hypothalamus": 1097}
group_colors = {
    area: "#{:02x}{:02x}{:02x}".format(*atlas.lookup.loc[division_id, "rgb_triplet"])
    for area, division_id in division_ids.items()
}
region_acronyms = [acronym for acronyms in groups.values() for acronym in acronyms]

region_order = []
group_labels = []
for area, acronyms in groups.items():
    region_order += [f"{acronym}_L" for acronym in reversed(acronyms)]
    region_order += [f"{acronym}_R" for acronym in acronyms]
    group_labels += [area] * (2 * len(acronyms))

sides = ["left"] * len(region_acronyms) + ["right"] * len(region_acronyms)
mask_names = [f"{acronym}_L" for acronym in region_acronyms] + [
    f"{acronym}_R" for acronym in region_acronyms
]
masks = atlas_native.get_masks(region_acronyms * 2, sides=sides).assign_coords(
    mask=mask_names
)
# get_masks reuses the same region id for both hemispheres of a region, which
# extract_with_labels would reject as a duplicate id across stacked-mask layers. Give
# every layer a unique nonzero id instead; extract_with_labels names each output
# region from the mask coordinate above, not from these ids.
layer_ids = xr.DataArray(np.arange(1, masks.sizes["mask"] + 1), dims="mask")
masks = xr.where(masks != 0, layer_ids, 0)

signals = cf.extract.extract_with_labels(data, masks, reduction="mean")
# extract_with_labels does not guarantee any particular region order, so reindex
# explicitly into the left-right sweep computed above.
signals = signals.sel(region=region_order)

signals

white_matter = atlas_native.get_masks("fiber tracts").isel(mask=0)
acompcor = cf.signal.compute_compcor_confounds(
    data, noise_mask=white_matter, n_components=1
)
signals = cf.signal.clean(signals, low_cutoff=0.01, confounds=acompcor)

connectivity = cf.connectivity.ConnectivityMatrix(kind="correlation").fit_transform(
    [signals]
)[0]

fig, ax = cf.plotting.plot_matrix(
    connectivity,
    labels=region_order,
    groups=group_labels,
    group_colors=group_colors,
    vmax=0.8,
    cbar_label="correlation",
    title="Region correlation matrix",
    bg_color=bg_color,
)