yabplot 0.4.0__tar.gz → 0.5.0__tar.gz

{yabplot-0.4.0/yabplot.egg-info → yabplot-0.5.0}/PKG-INFO

@@ -1,12 +1,13 @@
 Metadata-Version: 2.4
 Name: yabplot
-Version: 0.4.0
+Version: 0.5.0
 Summary: yet another brain plot
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: cmocean>=4.0.3
 Requires-Dist: ipywidgets>=8.1.8
+Requires-Dist: matplotlib>=3.10.7
 Requires-Dist: nibabel>=5.3.2
 Requires-Dist: pandas>=2.3.3
 Requires-Dist: pooch>=1.8.2
@@ -31,17 +32,18 @@ Dynamic: license-file
 [![Tests](https://github.com/teanijarv/yabplot/actions/workflows/tests.yml/badge.svg)](https://github.com/teanijarv/yabplot/actions/workflows/tests.yml)
 [![DOI](https://zenodo.org/badge/DOI/10.5281/zenodo.18237144.svg)](https://doi.org/10.5281/zenodo.18237144)
 
-**yabplot** is a Python library for creating beautiful, publication-quality 3D brain visualizations. it supports plotting cortical regions (+ vertexwise), subcortical structures, and white matter bundles.
+**yabplot** is a Python library for creating publication-quality 3D brain visualizations. it provides a unified interface for cortical regions, subcortical structures, white matter bundles, connectomes, voxel-wise maps, and vertex-wise maps.
 
-the idea is simple. while there are already amazing visualization tools available, they often focus on specific domains—using one tool for white matter tracts and another for cortical surfaces inevitably leads to inconsistent styles. i wanted a unified, simple-to-use tool that enables me (and hopefully others) to perform most brain visualizations in a single place. recognizing that neuroscience evolves daily, i designed **yabplot** to be modular: it supports standard pre-packaged atlases out of the box, but easily accepts any custom parcellation or tractography dataset you might need.
+the idea is simple. while there are already amazing visualization tools available, they often focus on specific domains—using one tool for white matter tracts and another for cortical surfaces inevitably leads to inconsistent styles. i wanted a unified, simple-to-use tool that enables me (and hopefully others) to perform most brain visualizations in a single place. recognizing that neuroscience evolves daily, i designed **yabplot** to be modular: it supports standard pre-packaged atlases out of the box, but easily accepts any custom parcellation or tractography dataset you might need. moreover, it enables to plot volumetric data either voxel-wise or by projecting the data to cortical surface or white matter tracts.
 
 ## features
 
-* **pre-existing atlases:** access many commonly used atlases (schaefer2018, brainnetome, aparc, aseg, musus100, xtract, etc) on demand.
-* **vertexwise plotting:** project volume (.nii) to cortical surface and plot.
-* **simple to use:** plug-n-play functions for cortex, subcortex, and tracts with a unified API.
-* **custom atlases:** easily use your own parcellations, segmentations (.nii/.gii), or tractograms (.trk).
-* **flexible inputs:** accepts data as dictionaries (for partial mapping) or arrays (for strict mapping).
+* **unified plotting API:** plot cortical regions, vertex-wise maps, subcortical structures, voxel-wise maps, white-matter tracts, and connectomes with a consistent interface.
+* **pre-packaged resources:** access commonly used atlases and meshes on demand, including schaefer, brainnetome, aparc, aseg, musus100, and xtract atlases.
+* **flexible data mapping:** pass data as arrays for strict ordering or dictionaries for partial/name-based mapping.
+* **custom atlases:** build and use custom cortical parcellations, subcortical segmentations, and tractography datasets.
+* **volume projection:** project nifti images to cortical surfaces or tractograms for vertex-wise and tractometry visualizations, or plot the volumes voxel-wise.
+* **publication-oriented output:** generate static figures, saved images, and interactive 3D views.
 
 ## installation
 
@@ -55,7 +57,7 @@ pip install yabplot # to install
 pip install yabplot --upgrade # to update
 ```
 
-dependencies: python 3.11 with ipywidgets, nibabel, pandas, pooch, pyvista, scikit-image, trame, trame-vtk, trame-vuetify
+dependencies: python 3.11 with ipywidgets, nibabel, matplotlib, pandas, pooch, pyvista, scikit-image, trame, trame-vtk, trame-vuetify
 
 (Connectome Workbench (`wb_command`) is a requirement to create custom cortical atlases unless you plan to only use pre-loaded atlases; see more in docs)
 
@@ -77,42 +79,53 @@ print(yab.get_available_resources())
 # see the region names for a specific atlas
 print(yab.get_atlas_regions(atlas='aseg', category='subcortical'))
 
-# cortical surface regions
+# plotting cortical surface regions
 atlas = 'aparc'
-dmap1 = {'L_lateraloccipital': 0.265, 'L_postcentral': 0.086, ...}
-yab.plot_cortical(data=dmap1, atlas=atlas, vminmax=[-0.1, 0.3], 
-                  bmesh='midthickness', views=['left_lateral', 'left_medial'],
-                  figsize=(600, 300), cmap='viridis')
+dmap1 = {'L_lateraloccipital': 0.265, 'L_postcentral': 0.086}
+ax = yab.plot_cortical(data=dmap1, atlas=atlas, vminmax=[0.0, 0.3], cmap='viridis',
+    bmesh='midthickness', views=['left_lateral', 'left_medial'])
 
-# subcortical structures
+# plotting subcortical regions
 atlas = 'aseg'
 regs = yab.get_atlas_regions(atlas=atlas, category='subcortical')
 data = np.arange(1, len(regs)+1)
-yab.plot_subcortical(data=data, atlas=atlas, vminmax=[2, 14], 
-                     views=['left_lateral', 'superior', 'right_lateral'], 
-                     bmesh_alpha=0.1, figsize=(600, 300), cmap='viridis')
+ax = yab.plot_subcortical(data=data, atlas=atlas, vminmax=[2, 14], 
+    views=['left_lateral', 'superior', 'right_lateral'], 
+    bmesh_alpha=0.1, cmap='plasma')
 
-# vertex-wise surface
+# plotting white matter tracts
+atlas = 'xtract_medium'
+regs = yab.get_atlas_regions(atlas=atlas, category='tracts')
+data = {reg: np.sin(i) for i, reg in enumerate(regs)}
+ax = yab.plot_tracts(data=data, atlas=atlas, style='matte', cmap='coolwarm',
+    views=['left_lateral', 'anterior', 'superior'], bmesh='pial')
+
+# plotting connectome
+data = np.random.rand(400, 400)
+ax = yab.plot_connectome(matrix=data, atlas='schaefer400',
+    edge_cmap='dense', node_cmap='binary', edge_threshold='95%', 
+    views=['left_lateral', 'superior', 'posterior']
+)
+
+# volume projection to cortical surface and vertexwise plotting
 threshold = 4
 b_lh_path, b_rh_path = yab.data.get_surface_paths('midthickness', 'bmesh')
 lh_data, rh_data = yab.project_vol2surf('path/to/yourdata.nii.gz', bmesh='midthickness')
 lh_data = np.where(lh_data > threshold, lh_data, np.nan)
 rh_data = np.where(rh_data > threshold, rh_data, np.nan)
 lh_mesh, rh_mesh = yab.load_vertexwise_mesh(b_lh_path, b_rh_path, lh_data, rh_data)
-yab.plot_vertexwise(lh_mesh, rh_mesh, cmap='viridis', vminmax=[-4, 9], 
-                    views=['left_lateral', 'left_medial'], figsize=(600, 300))
+ax = yab.plot_vertexwise(lh_mesh, rh_mesh, cmap='viridis', vminmax=[-10, 10], 
+                    views=['left_lateral', 'left_medial'])
 
-# white matter bundles
-atlas = 'xtract_tiny'
-regs = yab.get_atlas_regions(atlas=atlas, category='tracts')
-data = {reg: np.sin(i) for i, reg in enumerate(regs)}
-yab.plot_tracts(data=data, atlas=atlas, style='matte',
-                views=['left_lateral', 'anterior', 'superior'], bmesh='pial',
-                bmesh_alpha=0.1, figsize=(1600, 800), cmap='viridis')
+# plotting volume voxel-wise 
+threshold = '99.5%'
+nii_path = 'path/to/yourdata.nii.gz'
+ax = yab.plot_voxelwise(nii_path, threshold='99%', cmap='Reds',
+    views=['left_lateral', 'superior', 'anterior'])
 
 ```
 
-![examples](docs/assets/examples.png)
+![examples](docs/assets/overview.png)
 
 ## acknowledgements
 

{yabplot-0.4.0 → yabplot-0.5.0}/README.md

@@ -7,17 +7,18 @@
 [![Tests](https://github.com/teanijarv/yabplot/actions/workflows/tests.yml/badge.svg)](https://github.com/teanijarv/yabplot/actions/workflows/tests.yml)
 [![DOI](https://zenodo.org/badge/DOI/10.5281/zenodo.18237144.svg)](https://doi.org/10.5281/zenodo.18237144)
 
-**yabplot** is a Python library for creating beautiful, publication-quality 3D brain visualizations. it supports plotting cortical regions (+ vertexwise), subcortical structures, and white matter bundles.
+**yabplot** is a Python library for creating publication-quality 3D brain visualizations. it provides a unified interface for cortical regions, subcortical structures, white matter bundles, connectomes, voxel-wise maps, and vertex-wise maps.
 
-the idea is simple. while there are already amazing visualization tools available, they often focus on specific domains—using one tool for white matter tracts and another for cortical surfaces inevitably leads to inconsistent styles. i wanted a unified, simple-to-use tool that enables me (and hopefully others) to perform most brain visualizations in a single place. recognizing that neuroscience evolves daily, i designed **yabplot** to be modular: it supports standard pre-packaged atlases out of the box, but easily accepts any custom parcellation or tractography dataset you might need.
+the idea is simple. while there are already amazing visualization tools available, they often focus on specific domains—using one tool for white matter tracts and another for cortical surfaces inevitably leads to inconsistent styles. i wanted a unified, simple-to-use tool that enables me (and hopefully others) to perform most brain visualizations in a single place. recognizing that neuroscience evolves daily, i designed **yabplot** to be modular: it supports standard pre-packaged atlases out of the box, but easily accepts any custom parcellation or tractography dataset you might need. moreover, it enables to plot volumetric data either voxel-wise or by projecting the data to cortical surface or white matter tracts.
 
 ## features
 
-* **pre-existing atlases:** access many commonly used atlases (schaefer2018, brainnetome, aparc, aseg, musus100, xtract, etc) on demand.
-* **vertexwise plotting:** project volume (.nii) to cortical surface and plot.
-* **simple to use:** plug-n-play functions for cortex, subcortex, and tracts with a unified API.
-* **custom atlases:** easily use your own parcellations, segmentations (.nii/.gii), or tractograms (.trk).
-* **flexible inputs:** accepts data as dictionaries (for partial mapping) or arrays (for strict mapping).
+* **unified plotting API:** plot cortical regions, vertex-wise maps, subcortical structures, voxel-wise maps, white-matter tracts, and connectomes with a consistent interface.
+* **pre-packaged resources:** access commonly used atlases and meshes on demand, including schaefer, brainnetome, aparc, aseg, musus100, and xtract atlases.
+* **flexible data mapping:** pass data as arrays for strict ordering or dictionaries for partial/name-based mapping.
+* **custom atlases:** build and use custom cortical parcellations, subcortical segmentations, and tractography datasets.
+* **volume projection:** project nifti images to cortical surfaces or tractograms for vertex-wise and tractometry visualizations, or plot the volumes voxel-wise.
+* **publication-oriented output:** generate static figures, saved images, and interactive 3D views.
 
 ## installation
 
@@ -31,7 +32,7 @@ pip install yabplot # to install
 pip install yabplot --upgrade # to update
 ```
 
-dependencies: python 3.11 with ipywidgets, nibabel, pandas, pooch, pyvista, scikit-image, trame, trame-vtk, trame-vuetify
+dependencies: python 3.11 with ipywidgets, nibabel, matplotlib, pandas, pooch, pyvista, scikit-image, trame, trame-vtk, trame-vuetify
 
 (Connectome Workbench (`wb_command`) is a requirement to create custom cortical atlases unless you plan to only use pre-loaded atlases; see more in docs)
 
@@ -53,42 +54,53 @@ print(yab.get_available_resources())
 # see the region names for a specific atlas
 print(yab.get_atlas_regions(atlas='aseg', category='subcortical'))
 
-# cortical surface regions
+# plotting cortical surface regions
 atlas = 'aparc'
-dmap1 = {'L_lateraloccipital': 0.265, 'L_postcentral': 0.086, ...}
-yab.plot_cortical(data=dmap1, atlas=atlas, vminmax=[-0.1, 0.3], 
-                  bmesh='midthickness', views=['left_lateral', 'left_medial'],
-                  figsize=(600, 300), cmap='viridis')
+dmap1 = {'L_lateraloccipital': 0.265, 'L_postcentral': 0.086}
+ax = yab.plot_cortical(data=dmap1, atlas=atlas, vminmax=[0.0, 0.3], cmap='viridis',
+    bmesh='midthickness', views=['left_lateral', 'left_medial'])
 
-# subcortical structures
+# plotting subcortical regions
 atlas = 'aseg'
 regs = yab.get_atlas_regions(atlas=atlas, category='subcortical')
 data = np.arange(1, len(regs)+1)
-yab.plot_subcortical(data=data, atlas=atlas, vminmax=[2, 14], 
-                     views=['left_lateral', 'superior', 'right_lateral'], 
-                     bmesh_alpha=0.1, figsize=(600, 300), cmap='viridis')
+ax = yab.plot_subcortical(data=data, atlas=atlas, vminmax=[2, 14], 
+    views=['left_lateral', 'superior', 'right_lateral'], 
+    bmesh_alpha=0.1, cmap='plasma')
 
-# vertex-wise surface
+# plotting white matter tracts
+atlas = 'xtract_medium'
+regs = yab.get_atlas_regions(atlas=atlas, category='tracts')
+data = {reg: np.sin(i) for i, reg in enumerate(regs)}
+ax = yab.plot_tracts(data=data, atlas=atlas, style='matte', cmap='coolwarm',
+    views=['left_lateral', 'anterior', 'superior'], bmesh='pial')
+
+# plotting connectome
+data = np.random.rand(400, 400)
+ax = yab.plot_connectome(matrix=data, atlas='schaefer400',
+    edge_cmap='dense', node_cmap='binary', edge_threshold='95%', 
+    views=['left_lateral', 'superior', 'posterior']
+)
+
+# volume projection to cortical surface and vertexwise plotting
 threshold = 4
 b_lh_path, b_rh_path = yab.data.get_surface_paths('midthickness', 'bmesh')
 lh_data, rh_data = yab.project_vol2surf('path/to/yourdata.nii.gz', bmesh='midthickness')
 lh_data = np.where(lh_data > threshold, lh_data, np.nan)
 rh_data = np.where(rh_data > threshold, rh_data, np.nan)
 lh_mesh, rh_mesh = yab.load_vertexwise_mesh(b_lh_path, b_rh_path, lh_data, rh_data)
-yab.plot_vertexwise(lh_mesh, rh_mesh, cmap='viridis', vminmax=[-4, 9], 
-                    views=['left_lateral', 'left_medial'], figsize=(600, 300))
+ax = yab.plot_vertexwise(lh_mesh, rh_mesh, cmap='viridis', vminmax=[-10, 10], 
+                    views=['left_lateral', 'left_medial'])
 
-# white matter bundles
-atlas = 'xtract_tiny'
-regs = yab.get_atlas_regions(atlas=atlas, category='tracts')
-data = {reg: np.sin(i) for i, reg in enumerate(regs)}
-yab.plot_tracts(data=data, atlas=atlas, style='matte',
-                views=['left_lateral', 'anterior', 'superior'], bmesh='pial',
-                bmesh_alpha=0.1, figsize=(1600, 800), cmap='viridis')
+# plotting volume voxel-wise 
+threshold = '99.5%'
+nii_path = 'path/to/yourdata.nii.gz'
+ax = yab.plot_voxelwise(nii_path, threshold='99%', cmap='Reds',
+    views=['left_lateral', 'superior', 'anterior'])
 
 ```
 
-![examples](docs/assets/examples.png)
+![examples](docs/assets/overview.png)
 
 ## acknowledgements
 

{yabplot-0.4.0 → yabplot-0.5.0}/pyproject.toml

@@ -1,12 +1,13 @@
 [project]
 name = "yabplot"
-version = "0.4.0"
+version = "0.5.0"
 description = "yet another brain plot"
 readme = "README.md"
 requires-python = ">=3.10"
 dependencies = [
     "cmocean>=4.0.3",
     "ipywidgets>=8.1.8",
+    "matplotlib>=3.10.7",
     "nibabel>=5.3.2",
     "pandas>=2.3.3",
     "pooch>=1.8.2",

{yabplot-0.4.0 → yabplot-0.5.0}/yabplot/__init__.py

@@ -2,7 +2,8 @@ from importlib.metadata import version, PackageNotFoundError
 
 from .plotting import (
     plot_cortical, plot_subcortical, plot_tracts, 
-    clear_tract_cache, plot_vertexwise
+    clear_cache, plot_vertexwise, plot_connectome,
+    plot_voxelwise
 )
 from .data import (
     get_available_resources, get_atlas_regions
@@ -11,7 +12,7 @@ from .atlas_builder import (
     build_cortical_atlas, build_subcortical_atlas
 )
 from .mesh import (
-    load_vertexwise_mesh, make_cortical_mesh
+    load_vertexwise_mesh, make_cortical_mesh, load_nii_as_mesh
 )
 from .projection import (
     project_vol2surf, project_vol2tract, project_vol2tract_atlas

{yabplot-0.4.0 → yabplot-0.5.0}/yabplot/atlas_builder.py

@@ -224,10 +224,10 @@ def qc_custom_cortical_atlas(atlas_dir, atlasname='atlas'):
             plot_file = os.path.join(qc_dir, f"{rid:03d}_{name}.png")
             
             try:
-                plot_cortical(
+                ax = plot_cortical(
                     data={name: 1},
                     custom_atlas_path=atlas_dir,
-                    cmap='binary',vminmax=[0, 1],
+                    cmap='binary', vminmax=[0, 1],
                     export_path=plot_file
                 )
             except Exception as e:
@@ -397,7 +397,7 @@ def qc_custom_subcortical_atlas(atlas_dir):
             plot_file = os.path.join(qc_dir, f"{region_name}.png")
             
             try:
-                plot_subcortical(
+                ax = plot_subcortical(
                     data={region_name: 1},
                     custom_atlas_path=atlas_dir,
                     cmap='binary', vminmax=[0, 1],

{yabplot-0.4.0 → yabplot-0.5.0}/yabplot/mesh.py

@@ -1,6 +1,12 @@
+import warnings
+
+import nibabel as nib
 import numpy as np
 import pyvista as pv
 import scipy.sparse as sp
+from scipy.ndimage import gaussian_filter
+from skimage import measure
+from concurrent.futures import ThreadPoolExecutor
 
 def load_bmesh(bmesh):
     """
@@ -104,6 +110,76 @@ def load_vertexwise_mesh(lh_mesh_path, rh_mesh_path, lh_data, rh_data, scalar_na
     rh = make_cortical_mesh(*load_gii(rh_mesh_path), rh_data, scalar_name)
     return lh, rh
 
+def load_nii_as_mesh(
+    nii_path,
+    threshold=0.5,
+    blur_sigma=1.5,
+    smooth_i=10,
+    smooth_f=0.1
+):
+    """
+    Build a surface mesh from a 3D NIfTI volume using marching cubes, with optional Gaussian blurring and mesh smoothing.
+
+    Parameters
+    ----------
+    nii_path : str
+        Absolute path to a NIfTI file representing a 3D volume. If 4D, only the first volume will be used.
+    threshold : float, optional
+        Threshold applied after optional blur. Voxels ``> threshold`` are kept.
+    blur_sigma : float, optional
+        Gaussian blur (voxel units) before thresholding.
+    smooth_i : int, optional
+        Number of PyVista smoothing iterations after surface extraction.
+    smooth_f : float, optional
+        Relaxation factor for mesh smoothing.
+
+    Returns
+    -------
+    mesh : pyvista.PolyData
+        The extracted and smoothed surface mesh ready for plotting.
+    """
+
+    img = nib.load(nii_path)
+    vol = img.get_fdata()
+
+    if vol.ndim > 3:
+        warnings.warn(
+            f"[WARNING] detected {vol.ndim}d nifti volume. using the first volume (index 0)."
+        )
+        vol = vol[..., 0]
+
+    vol = np.nan_to_num(vol, nan=0.0)
+
+    if blur_sigma and blur_sigma > 0:
+        vol = gaussian_filter(vol, sigma=float(blur_sigma))
+
+    mask = vol > float(threshold)
+
+    if not np.any(mask):
+        raise ValueError("Mask is empty after thresholding. Adjust threshold/blur_sigma.")
+
+    verts_vox, faces, _, _ = measure.marching_cubes(mask.astype(np.float32), level=0.5)
+    verts_world = nib.affines.apply_affine(img.affine, verts_vox)
+
+    faces_pv = np.hstack([
+        np.full((faces.shape[0], 1), 3, dtype=np.int64),
+        faces.astype(np.int64)
+    ]).ravel()
+    mesh = pv.PolyData(verts_world.astype(np.float32), faces_pv)
+
+    if smooth_i and smooth_i > 0:
+        mesh = mesh.smooth(n_iter=int(smooth_i), relaxation_factor=float(smooth_f))
+
+    if mesh.n_points == 0:
+        raise ValueError("Extracted mesh has no vertices. Check input mask and parameters.")
+
+    # fill topological holes in the extracted meshes
+    try:
+        mesh = mesh.fill_holes(1000)
+    except Exception as e:
+        warnings.warn(f"Mesh hole filling failed: {e}. Continuing with unfilled meshes.")
+
+    return mesh
 
 def map_values_to_surface(data, target_labels, lut_ids, dense_lut_names):
     """maps data to vertices."""
@@ -204,35 +280,60 @@ def apply_dilation(faces, data, iterations=4):
     return data_out
 
 
+def get_smooth_masks_vectorized(faces, n_v, r_masks, iterations=4):
+    adj = get_adj(faces, n_v)
+    deg = np.array(adj.sum(axis=1)).flatten()
+    deg[deg == 0] = 1.0
+    mask = r_masks.astype(np.float64)
+    for _ in range(iterations):
+        mask = (mask + (adj.dot(mask) / deg[:, None])) / 2.0
+    return mask
+
 def get_puzzle_pieces(v, f, raw_vals):
-    """carve out geometric pieces with slight overlap to prevent gaps."""
-    pieces = []
+    """
+    Creates sharp boundaries without gaps by calculating smooth probability fields
+    and interpolating them onto a highly subdivided continuous mesh.
+    """
     valid_mask = ~np.isnan(raw_vals) & (raw_vals != 0.0)
     u_vals = np.unique(raw_vals[valid_mask])
+    
+    # if no data, skip
+    if len(u_vals) == 0:
+        master = make_cortical_mesh(v, f, np.full(len(v), np.nan))
+        return pv.PolyData(), [master]
+    
+    n_v = len(v)
+    n_k = len(u_vals)
+    
+    # vectorized smoothing
+    r_masks = np.zeros((n_v, n_k + 1), dtype=np.float64)
+    r_masks[:, 0] = np.where(~valid_mask, 1.0, 0.0) # medial wall
+    for i, val in enumerate(u_vals):
+        r_masks[:, i+1] = np.where(raw_vals == val, 1.0, 0.0)
+    
+    s_masks = get_smooth_masks_vectorized(f, n_v, r_masks, iterations=4)
+    
     master = make_cortical_mesh(v, f, np.zeros_like(raw_vals))
-
-    for val in u_vals:
-        r_mask = np.where(raw_vals == val, 1.0, 0.0)
-        s_mask = get_smooth_mask(f, r_mask, iterations=4)
-        temp = master.copy()
-        temp['Slice_Mask'] = s_mask
-        # reduce search space
-        patch = temp.threshold(0.01, scalars='Slice_Mask')
-        if patch.n_points > 0:
-            # use 0.48 (slightly expanded) for pieces to seal cracks
-            piece = patch.clip_scalar(scalars='Slice_Mask', value=0.48, invert=False)
-            if piece.n_points > 0:
-                piece['Data'] = np.full(piece.n_points, val)
-                pieces.append(piece)
+    master['Masks'] = s_masks
+    
+    # subdivide to increase resolution (2 levels = 16x faces)
+    sub = master.subdivide(2, subfilter='linear')
+    
+    # assign labels to high-res vertices
+    interp_masks = sub['Masks']
+    best_class = np.argmax(interp_masks, axis=1)
+    new_data = np.full(sub.n_points, np.nan)
+    valid_idx = best_class > 0
+    
+    # map the argmax indices back to their original scalar values
+    new_data[valid_idx] = u_vals[best_class[valid_idx] - 1]
+    sub['Data'] = new_data
+    
+    # clean up multi-component array to save memory
+    del sub.point_data['Masks']
     
-    # slice base brain
-    all_mask = np.where(valid_mask, 1.0, 0.0)
-    s_all = get_smooth_mask(f, all_mask, iterations=4)
-    master['Slice_Mask'] = s_all
-    # use 0.52 (slightly contracted) for the hole to ensure colored pieces cover the edge
-    base_p = master.clip_scalar(scalars='Slice_Mask', value=0.52, invert=True)
-    if base_p.n_points > 0:
-        base_p['Data'] = np.full(base_p.n_points, np.nan)
+    base_p = pv.PolyData()
+    pieces = [sub]
     
     return base_p, pieces