File size: 5,899 Bytes
6a0dd5a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 | """LitCoder-style fsaverage5 surface plotting helpers."""
from __future__ import annotations
import os
import warnings
from io import BytesIO
from pathlib import Path
from typing import Any
import numpy as np
DEFAULT_SCRATCH_ROOT = Path(os.environ.get("BRAINENCODING_SCRATCH_ROOT", "/storage/scratch1/8/whuang409"))
DEFAULT_FSAVERAGE5 = Path("/storage/project/r-aivanova7-0/shared/.env/freesurfer/subjects/fsaverage5")
N_HEMI = 10242
N_VERTICES = 20484
os.environ.setdefault("MPLCONFIGDIR", str(DEFAULT_SCRATCH_ROOT / "matplotlib"))
import matplotlib # noqa: E402
matplotlib.use("Agg")
import matplotlib.pyplot as plt # noqa: E402
from matplotlib.colors import Normalize # noqa: E402
def save_litcoder_style_map(
values: np.ndarray,
output_path: Path,
*,
title: str,
fsaverage5: Path = DEFAULT_FSAVERAGE5,
significant_mask: np.ndarray | None = None,
positive: bool = False,
mask_zero: bool = True,
vmax: float | None = None,
cmap: str | Any | None = None,
dpi: int = 150,
) -> None:
fig = plot_litcoder_style_map(
values,
title=title,
fsaverage5=fsaverage5,
significant_mask=significant_mask,
positive=positive,
mask_zero=mask_zero,
vmax=vmax,
cmap=cmap,
)
output_path.parent.mkdir(parents=True, exist_ok=True)
fig.savefig(output_path, format="png", bbox_inches="tight", dpi=dpi)
plt.close(fig)
def litcoder_style_png_bytes(
values: np.ndarray,
*,
title: str,
fsaverage5: Path = DEFAULT_FSAVERAGE5,
significant_mask: np.ndarray | None = None,
positive: bool = False,
mask_zero: bool = True,
vmax: float | None = None,
cmap: str | Any | None = None,
dpi: int = 130,
) -> bytes:
fig = plot_litcoder_style_map(
values,
title=title,
fsaverage5=fsaverage5,
significant_mask=significant_mask,
positive=positive,
mask_zero=mask_zero,
vmax=vmax,
cmap=cmap,
)
buffer = BytesIO()
fig.savefig(buffer, format="png", bbox_inches="tight", dpi=dpi)
plt.close(fig)
return buffer.getvalue()
def plot_litcoder_style_map(
values: np.ndarray,
*,
title: str,
fsaverage5: Path = DEFAULT_FSAVERAGE5,
significant_mask: np.ndarray | None = None,
positive: bool = False,
mask_zero: bool = True,
vmax: float | None = None,
cmap: str | Any | None = None,
) -> plt.Figure:
"""Plot a 20484-vector in the LitCoder four-panel fsaverage5 layout."""
from nilearn import plotting
from nilearn.plotting.cm import cold_hot
display_values = prepare_values(
values,
significant_mask=significant_mask,
positive=positive,
mask_zero=mask_zero,
)
vmax = value_limit(display_values, vmax=vmax, positive=positive)
vmin = 0.0 if positive else -vmax
cmap = cmap or ("viridis" if positive else cold_hot)
norm = Normalize(vmin=vmin, vmax=vmax)
fsaverage5 = Path(fsaverage5)
meshes = {
"lh": fsaverage5 / "surf" / "lh.inflated",
"rh": fsaverage5 / "surf" / "rh.inflated",
}
for path in meshes.values():
if not path.exists():
raise FileNotFoundError(f"Missing fsaverage5 surface mesh: {path}")
fig = plt.figure(figsize=(15, 10))
views = [
("lh", "left", "lateral", 231, "Left Lateral"),
("lh", "left", "medial", 232, "Left Medial"),
("rh", "right", "lateral", 234, "Right Lateral"),
("rh", "right", "medial", 235, "Right Medial"),
]
for hemi_key, hemi_name, view, subplot, pane_title in views:
hemi_values = display_values[:N_HEMI] if hemi_key == "lh" else display_values[N_HEMI:]
ax = fig.add_subplot(subplot, projection="3d")
with warnings.catch_warnings():
warnings.filterwarnings("ignore", message="All-NaN slice encountered", category=RuntimeWarning)
plotting.plot_surf_stat_map(
str(meshes[hemi_key]),
hemi_values,
hemi=hemi_name,
view=view,
colorbar=False,
axes=ax,
cmap=cmap,
vmin=vmin,
vmax=vmax,
title=pane_title,
)
sm = plt.cm.ScalarMappable(norm=norm, cmap=cmap)
sm.set_array([])
cax = fig.add_axes([0.92, 0.15, 0.02, 0.7])
fig.colorbar(sm, cax=cax)
fig.suptitle(title, fontsize=16)
with warnings.catch_warnings():
warnings.filterwarnings("ignore", message="This figure includes Axes that are not compatible with tight_layout.*")
fig.tight_layout(rect=[0.03, 0.03, 0.9, 0.97])
return fig
def prepare_values(
values: np.ndarray,
*,
significant_mask: np.ndarray | None,
positive: bool,
mask_zero: bool,
) -> np.ndarray:
out = np.asarray(values, dtype=np.float32).reshape(-1).copy()
if out.shape != (N_VERTICES,):
raise ValueError(f"Expected {(N_VERTICES,)} values, got {out.shape}")
shown = np.isfinite(out)
if significant_mask is not None:
mask = np.asarray(significant_mask, dtype=bool).reshape(-1)
if mask.shape != (N_VERTICES,):
raise ValueError(f"Expected {(N_VERTICES,)} mask values, got {mask.shape}")
shown &= mask
if mask_zero:
shown &= (out > 0.0) if positive else (np.abs(out) > 0.0)
out[~shown] = np.nan
return out
def value_limit(values: np.ndarray, *, vmax: float | None, positive: bool) -> float:
if vmax is not None and np.isfinite(vmax) and vmax > 0:
return float(vmax)
finite = values[np.isfinite(values)]
if finite.size == 0:
return 1.0
finite = finite[finite > 0.0] if positive else np.abs(finite)
if finite.size == 0:
return 1.0
limit = float(np.nanmax(finite))
return limit if np.isfinite(limit) and limit > 0 else 1.0
|