spacelearn 0.1.0__py3-none-any.whl

spacelearn/loss/_alignment_loss.py

@@ -0,0 +1,59 @@
+"""
+Copyright (c) 2026 Tobias Karusseit
+This source code is licensed under the MIT license found in the
+LICENSE file in the root directory of this source tree.
+"""
+
+import torch
+import torch.nn.functional as F
+
+def alignment_loss(
+    Z: torch.Tensor,
+    W: dict[str, torch.Tensor],
+    A: dict[str, torch.Tensor],
+    Y: dict[str, torch.Tensor],
+) -> torch.Tensor:
+    """
+    Physics reconstruction loss: penalizes how well each subspace's
+    decoded projection matches its pooled physical target.
+
+    #### USAGE:
+        loss = alignment_loss(Z, W, A, Y)
+        loss.backward()
+
+    #### HOW IT WORKS:
+        For each quantity q in W:
+            S_q   = Z_c @ W_q.T        # (B, k) latent coordinates in subspace q
+            y_hat = S_q @ A_q          # (B, n^2) decoded prediction
+            loss += MSE(y_hat, Y_q_centered)
+        The per-quantity losses are averaged across all quantities in `W`.
+
+    #### ARGS:
+        - Z: (B, D) latent vectors.
+        - W: {name: (k, D)} subspace bases.
+        - A: {name: (k, n^2)} linear decoders.
+        - Y: {name: (B, n^2) or (B, n, n) ...} pooled physical targets.
+             Targets with more than 2 dims are flattened to (B, n^2).
+
+    #### RETURNS:
+        - A scalar torch.Tensor loss. Perfect reconstruction -> 0.
+          Averaged across quantities in `W`.
+    """
+    Z_c = Z - Z.mean(0, keepdim=True)
+    names = list(W.keys())
+
+    loss = torch.tensor(0.0, device=Z.device)
+    for name in names:
+        W_q = W[name].to(Z.device)
+        A_q = A[name].to(Z.device)
+        S_q = Z_c @ W_q.T  # (B, k)
+
+        target = Y[name].to(Z.device)
+        if target.dim() > 2:
+            target = target.flatten(1)
+        target_c = target - target.mean(0, keepdim=True)
+
+        y_hat = S_q @ A_q
+        loss = loss + F.mse_loss(y_hat, target_c)
+
+    return loss / max(len(names), 1)

spacelearn/loss/_disentanglement_loss.py

@@ -0,0 +1,111 @@
+"""
+Copyright (c) 2026 Tobias Karusseit
+This source code is licensed under the MIT license found in the
+LICENSE file in the root directory of this source tree.
+"""
+
+import torch
+
+def disentanglement_loss_original(
+    Z: torch.Tensor, 
+    W: dict[str, torch.Tensor]
+) -> torch.Tensor:
+    """
+    [LEGACY] Penalizes correlation between latent vectors projected into
+    different subspaces, using mean squared correlation per pair.
+
+    We recommend `disentanglement_loss` instead — see that function's
+    docstring for the difference.
+
+    #### USAGE:
+        loss = disentanglement_loss_original(Z, W)
+        loss.backward()
+
+    #### HOW IT WORKS:
+        For every pair of subspaces (i, j):
+            S_i, S_j = Z_c @ W_i.T, Z_c @ W_j.T   # project into each subspace
+            normalize S_i, S_j to unit std (no centering)
+            corr = S_i_n.T @ S_j_n / B            # (k_i, k_j) correlation matrix
+            total += mean(corr ** 2)              # squared correlation
+        The sum is averaged over the number of subspace pairs.
+
+    #### ARGS:
+        - Z: (B, D) latent vectors.
+        - W: {name: (D, k)} subspace bases.
+
+    #### RETURNS:
+        - A scalar torch.Tensor loss. Perfect disentanglement -> 0,
+          fully correlated -> 1.
+    """
+    total = torch.tensor(0.0, device=Z.device)
+    Z_c = Z - Z.mean(0, keepdim=True)
+    names = list(W.keys())
+    for i in range(len(names)):
+        for j in range(i + 1, len(names)):
+            Si   = Z_c @ W[names[i]].T # the latent vector in subspace i
+            Sj   = Z_c @ W[names[j]].T # the latent vector in subspace j
+            Si_n = Si / (Si.std(dim=0, keepdim=True) + 1e-8) # normed to unit length (for scale invariance in the comming step)
+            Sj_n = Sj / (Sj.std(dim=0, keepdim=True) + 1e-8) # normed to unit length
+            corr = Si_n.T @ Sj_n / Si.shape[0]       # correlation
+            total = total + corr.pow(2).mean()       # squared correlation to focus on large correlations
+    n_pairs  = len(names) * (len(names) - 1) / 2     # number of pairs for normalization
+    print(f"[disentanglement] n_pairs={n_pairs}")
+    total = total / max(n_pairs, 1)                  # normalized by number of pairs
+    return total
+
+def disentanglement_loss(
+    Z: torch.Tensor,
+    W: dict[str, torch.Tensor]
+) -> torch.Tensor:
+    """
+    Penalizes correlation between latent vectors projected into different
+    subspaces, using the normalized trace of the squared correlation
+    matrix per pair.
+
+    #### USAGE:
+        loss = disentanglement_loss(Z, W)
+        loss.backward()
+
+    #### HOW IT WORKS:
+        For every pair of subspaces (i, j):
+            S_i, S_j = Z_c @ W_i.T, Z_c @ W_j.T   # project into each subspace
+            standardize S_i, S_j (center and divide by std)
+            Corr = S_i_std.T @ S_j_std / B        # (k_i, k_j) correlation matrix
+            total += trace(Corr.T @ Corr) / k_i   # normalized squared correlation
+        The sum is averaged over the number of subspace pairs.
+
+    #### ARGS:
+        - Z: (B, D) latent vectors.
+        - W: {name: (D, k)} subspace bases.
+
+    #### RETURNS:
+        - A scalar torch.Tensor loss. Perfect disentanglement -> 0,
+          fully correlated -> 1.
+    """
+    Z_c = Z - Z.mean(0, keepdim=True)
+    total = torch.tensor(0.0, device=Z.device)
+    names = list(W.keys())
+    pairs = 0
+    
+    for i in range(len(names)):
+        for j in range(i + 1, len(names)):
+            # Project
+            Si = Z_c @ W[names[i]].T
+            Sj = Z_c @ W[names[j]].T
+            
+            # Standardize (Center and divide by std)
+            # This is the correct way to compute correlation
+            Si_std = (Si - Si.mean(0, keepdim=True)) / (Si.std(0, keepdim=True) + 1e-8)
+            Sj_std = (Sj - Sj.mean(0, keepdim=True)) / (Sj.std(0, keepdim=True) + 1e-8)
+            
+            # Correlation matrix (ki, kj)
+            # Dividing by Z.shape[0] gives the covariance of standardized variables
+            Corr = Si_std.T @ Sj_std / Z.shape[0]
+            
+            # Use the Trace of the squared correlation matrix
+            # For identical subspaces, trace(Corr^2) = k
+            # Dividing by k gives exactly 1.0
+            total += torch.trace(Corr.T @ Corr) / Si.shape[1]
+            pairs += 1
+            
+    return total / max(pairs, 1)

spacelearn/loss/_isotropy_loss.py

@@ -0,0 +1,56 @@
+"""
+Copyright (c) 2026 Tobias Karusseit
+This source code is licensed under the MIT license found in the
+LICENSE file in the root directory of this source tree.
+"""
+
+import torch
+from ..util import singular_values
+
+def isotropy_loss(
+    Z: torch.Tensor, 
+    W: dict[str, torch.Tensor], 
+    n_dir: int = 32
+) -> torch.Tensor:
+    """
+    Enforces isotropy in the latent residual space by penalizing variance
+    among its leading singular values.
+
+    #### USAGE:
+        loss = isotropy_loss(Z, W, n_dir=32)
+        loss.backward()
+
+    #### HOW IT WORKS:
+        Z_c     = Z - Z.mean(0)                     # center latents
+        W_all   = concat(W.values())                # (sum_k, D) stacked bases
+        Z_resid = Z_c - Z_c @ W_all.T @ W_all        # residual after projecting
+                                                      # out all known subspaces
+                                                      # (assumes subspaces are
+                                                      # mutually orthogonal)
+        S       = singular_values(Z_resid)            # singular values of residual
+        S_top   = S[:n_dir]                           # leading n_dir values
+        S_norm  = S_top / S_top.mean()                # normalize to mean 1
+        loss    = mean((S_norm - 1) ** 2)             # penalize deviation
+                                                      # from uniform spread
+
+    #### ARGS:
+        - Z: (B, D) latent vectors.
+        - W: {name: (D, k)} subspace bases already accounted for (e.g. by
+             alignment/disentanglement losses). These are projected out of
+             `Z` before computing isotropy on what remains.
+        - n_dir: Number of leading singular values of the residual to
+                 consider. Defaults to 32.
+
+    #### RETURNS:
+        - A scalar torch.Tensor loss. Perfectly isotropic residual
+          directions (uniform singular value spectrum) -> 0.
+    """
+    z_c = Z - Z.mean(0, keepdim=True) # centered for variance calc
+    names = list(W.keys())
+    W_all   = torch.cat([W[n] for n in names], dim=0).to(Z.device)
+    z_resid = z_c - z_c @ W_all.T @ W_all # residual subspace (approximation only, assumes the subspaces are orthogonal)
+    S_sv = singular_values(z_resid)
+    S_top   = S_sv[:n_dir]
+    S_norm  = S_top / (S_top.mean() + 1e-8)
+    loss_rank = ((S_norm - 1.0) ** 2).mean() # punish the model if the svd vals are not even
+    return loss_rank

spacelearn/loss/_stability_loss.py

@@ -0,0 +1,44 @@
+"""
+Copyright (c) 2026 Tobias Karusseit
+This source code is licensed under the MIT license found in the
+LICENSE file in the root directory of this source tree.
+"""
+
+import torch
+
+def stability_loss(
+    W: dict[str, torch.Tensor], 
+    W_prev: dict[str, torch.Tensor]
+) -> torch.Tensor:
+    """
+    Punishes large or rapid changes in subspace geometry between training
+    steps, encouraging the model to evolve its latent subspace structure
+    gradually rather than jumping between geometries.
+
+    #### USAGE:
+        loss = stability_loss(W, W_prev)
+        loss.backward()
+
+    #### HOW IT WORKS:
+        If W_prev is provided:
+            for each quantity name in W:
+                loss += sum((W[name] - W_prev[name]) ** 2)   # squared drift
+            loss /= number of quantities
+        Otherwise, the loss is 0 (e.g. on the first step, when there is no
+        previous subspace to compare against).
+
+    #### ARGS:
+        - W: {name: (D, k)} current subspace bases.
+        - W_prev: {name: (D, k)} subspace bases from the previous step, or
+                  None if there is no previous step to compare against.
+
+    #### RETURNS:
+        - A scalar torch.Tensor loss. No drift -> 0; larger or more
+          rapid changes in subspace geometry -> larger values.
+    """
+    loss_stab = torch.tensor(0.0, device=W[next(iter(W.keys()))].device)
+    if W_prev is not None:
+        for name in W.keys():
+            loss_stab = loss_stab + (W[name] - W_prev[name]).pow(2).sum()
+        loss_stab = loss_stab / len(W.keys())
+    return loss_stab

spacelearn/optim/__init__.py

@@ -0,0 +1,19 @@
+"""
+Copyright (c) 2026 Tobias Karusseit
+This source code is licensed under the MIT license found in the
+LICENSE file in the root directory of this source tree.
+"""
+
+from ._optimize_targets import optimal_bins
+from ._optimize_latent import minimal_latent_dim
+from ._fit_decoder import fit_decoders
+from ._subspace_change import subspace_change
+from ._optimize_subspaces import k_per_q
+
+__all__ = [
+    "optimal_bins",
+    "minimal_latent_dim",
+    "fit_decoders",
+    "subspace_change",
+    "k_per_q",
+]

spacelearn/optim/_fit_decoder.py

@@ -0,0 +1,56 @@
+"""
+Copyright (c) 2026 Tobias Karusseit
+This source code is licensed under the MIT license found in the
+LICENSE file in the root directory of this source tree.
+"""
+
+import torch
+
+def fit_decoders(W: dict, Z: torch.Tensor, pooled_global: dict[str, torch.Tensor]) -> dict:
+    """
+    Learns a linear mapping (decoder) from each latent subspace's
+    coordinates to its corresponding physical target space, via
+    least-squares regression.
+
+    #### USAGE:
+        A = fit_decoders(W, Z, pooled_global)
+        # A can then be passed directly to alignment_loss as the
+        # 'A' (decoder) argument
+
+    #### HOW IT WORKS:
+        Z_c = Z - Z.mean(0)                          # center latents
+        for each quantity name in W:
+            S_q   = Z_c @ W_q.T                       # project latents
+                                                        # into subspace q
+            Y_q   = pooled_global[name]                # physical target,
+                                                        # flattened if needed
+            Y_q_c = Y_q - Y_q.mean(0)                  # center target
+            A_q   = lstsq(S_q, Y_q_c).solution         # solve for A_q
+                                                        # minimizing
+                                                        # ||Y_q_c - S_q @ A_q||_F
+        Each A_q is detached from the autograd graph before being
+        returned, since it is fit via least squares rather than learned
+        through backpropagation.
+
+    #### ARGS:
+        - W: {name: (k, D)} subspace basis matrices.
+        - Z: (B, D) latent vectors.
+        - pooled_global: {name: (B, n^2) or (B, n, n) ...} pooled
+          physical targets. Targets with more than 2 dims are flattened
+          to (B, n^2) before fitting.
+
+    #### RETURNS:
+        - A dict {name: (k, n^2)} of optimal linear decoder matrices, one
+          per quantity in W, each detached from the autograd graph.
+    """
+    Z_c = Z - Z.mean(0, keepdim=True) # center
+    A   = {}
+    for name, W_q in W.items():
+        S_q = Z_c @ W_q.T          # projection of Z into the subspace
+        y   = pooled_global[name]  # pooled target
+        if y.dim() > 2:            # flatten
+            y = y.flatten(1)
+        y_c = y - y.mean(0, keepdim=True) # center
+        # solve (finds A, such that (y_c - S_q @ A) is approx 0)
+        A[name] = torch.linalg.lstsq(S_q, y_c).solution.detach()
+    return A

spacelearn/optim/_optimize_latent.py

@@ -0,0 +1,92 @@
+"""
+Copyright (c) 2026 Tobias Karusseit
+This source code is licensed under the MIT license found in the
+LICENSE file in the root directory of this source tree.
+"""
+
+import math
+
+
+def _align_latent_dim(dim: int, num_heads: int) -> int:
+    """Round dim up to the next multiple of num_heads (unchanged if already divisible)."""
+    if dim < 1:
+        raise ValueError(f"dim must be >= 1, got {dim}")
+    if num_heads < 1:
+        raise ValueError(f"num_heads must be >= 1, got {num_heads}")
+    if num_heads == 1:
+        return dim
+    remainder = dim % num_heads
+    if remainder == 0:
+        return dim
+    return dim + (num_heads - remainder)
+
+
+def minimal_latent_dim(
+    k: int | dict[str, int],
+    *,
+    Q: int | None = None,
+    free_frac: float = 0.30,
+    num_heads: int | None = None,
+) -> int:
+    """
+    Computes the minimum latent dimension needed to fit a set of physics
+    subspaces plus a fraction of free (unconstrained) capacity, optionally
+    aligned to a multiple of num_heads.
+
+    #### USAGE:
+        # Example 1: shared k across all quantities
+        D = minimal_latent_dim(k=4, Q=6, free_frac=0.30)
+
+        # Example 2: per-quantity k
+        D = minimal_latent_dim(k={"scaled": 4, "grammian": 6}, free_frac=0.30)
+
+        # Example 3: align result to attention head count
+        D = minimal_latent_dim(k=4, Q=6, num_heads=8)
+
+    #### HOW IT WORKS:
+        physics_dims = Q * k                  # if k is a shared scalar
+        physics_dims = sum(k.values())        # if k is per-quantity
+        D_min = ceil(physics_dims / (1 - free_frac))
+        if num_heads is set:
+            D_min = next multiple of num_heads >= D_min
+
+    #### ARGS:
+        - k: Subspace dimensionality. Either a single int shared across all
+             `Q` quantities, or a dict mapping each quantity name to its
+             own k.
+        - Q: Number of quantities. Required (and only used) when `k` is a
+             scalar; ignored when `k` is a dict, since the count is
+             inferred from its keys.
+        - free_frac: Fraction of the latent dimension to reserve as free,
+                     unconstrained capacity beyond the physics subspaces.
+                     Must be in [0, 1). Defaults to 0.30.
+        - num_heads: If set, the result is rounded up to the next multiple
+                     of num_heads (e.g. to match an attention head count).
+                     Defaults to None (no alignment).
+
+    #### RETURNS:
+        - The minimum latent dimension D_min satisfying the physics
+          subspace requirement, the free-capacity fraction, and (if
+          given) the num_heads alignment constraint.
+    """
+    if not 0.0 <= free_frac < 1.0:
+        raise ValueError(f"free_frac must be in [0, 1), got {free_frac}")
+
+    if isinstance(k, dict):
+        physics_dims = sum(k.values())
+        detail = ", ".join(f"{name}={k_q}" for name, k_q in sorted(k.items()))
+        # print(f"[latent-sizing] k per quantity: {detail}  (sum={physics_dims})")
+    else:
+        if Q is None:
+            raise ValueError("Q is required when k is a scalar")
+        physics_dims = Q * k
+        # print(f"[latent-sizing] Q={Q}, k={k}")
+
+    D_min = math.ceil(physics_dims / (1.0 - free_frac))
+    print(f"[latent-sizing] free={free_frac * 100:.0f}%  ->  D_min={D_min}")
+    if num_heads is not None:
+        aligned = _align_latent_dim(D_min, num_heads)
+        # if aligned != D_min:
+            # print(f"[latent-sizing] aligned D {D_min} -> {aligned} (num_heads={num_heads})")
+        D_min = aligned
+    return D_min

spacelearn/optim/_optimize_subspaces.py

@@ -0,0 +1,109 @@
+"""
+Copyright (c) 2026 Tobias Karusseit
+This source code is licensed under the MIT license found in the
+LICENSE file in the root directory of this source tree.
+"""
+
+import torch
+from ..util import singular_values
+
+def _k_from_target(target: torch.Tensor, threshold: float) -> int:
+    if target.dim() == 1:
+        target = target.unsqueeze(1)
+    elif target.dim() > 2:
+        target = target.flatten(1)
+    y_c = target - target.mean(0, keepdim=True)
+    S = singular_values(y_c)
+    var = S.pow(2)
+    spec = var / (var.sum() + 1e-8)
+    cumvar = spec.cumsum(0)
+    return int((cumvar < threshold).sum().item() + 1)
+
+
+def _resolve_threshold(
+    name: str,
+    threshold: float | dict[str, float],
+    default: float,
+) -> float:
+    if isinstance(threshold, dict):
+        return float(threshold.get(name, default))
+    return float(threshold)
+
+
+def k_per_q(
+    pooled_global: dict[str, torch.Tensor],
+    threshold: float | dict[str, float] = 0.90,
+    *,
+    default_threshold: float | None = None,
+    per_quantity: bool = True,
+) -> int | dict[str, int]:
+    """
+    Estimates the minimal subspace dimension k needed per quantity to
+    retain a target fraction of variance, via PCA-style spectral analysis
+    of each pooled target.
+
+    #### USAGE:
+        # Example 1: shared threshold, one k per quantity
+        k = k_per_q(pooled_global, threshold=0.90)
+        # k = {"scaled": 3, "grammian": 5}
+
+        # Example 2: per-quantity thresholds, with a fallback for missing names
+        k = k_per_q(
+            pooled_global,
+            threshold={"scaled": 0.95},
+            default_threshold=0.90,
+        )
+
+        # Example 3: collapse to a single global k (the max across quantities)
+        k = k_per_q(pooled_global, threshold=0.90, per_quantity=False)
+        # k = 5
+
+    #### HOW IT WORKS:
+        For each quantity name, target in pooled_global:
+            target_c = target - target.mean(0)       # center (flattened
+                                                       # if > 2D)
+            S        = singular_values(target_c)      # spectrum of target
+            spec     = S^2 / sum(S^2)                  # normalized variance
+                                                        # per component
+            k_q      = number of leading components needed for
+                       cumulative variance to reach `threshold`
+        If per_quantity is True, returns {name: k_q} for all quantities.
+        Otherwise, returns max(k_q for all quantities).
+
+    #### ARGS:
+        - pooled_global: {name: (B, n^2) or (B, n, n) ...} pooled physical
+          targets, one per quantity. Targets with more than 2 dims are
+          flattened to (B, n^2); 1D targets are treated as (B, 1).
+        - threshold: Either a single float shared across all quantities,
+          or a {name: threshold} dict for per-quantity thresholds. Each
+          threshold is the target fraction of cumulative variance to
+          retain (e.g. 0.90 -> keep enough components for 90% variance).
+          Defaults to 0.90.
+        - default_threshold: Fallback threshold used for any quantity name
+          missing from a `threshold` dict. If None, falls back to the
+          scalar `threshold` value (or 0.90 if `threshold` is itself a
+          dict with no scalar to fall back to). Defaults to None
+          (keyword-only).
+        - per_quantity: If True, returns a dict of k per quantity. If
+          False, returns the single largest k found across all
+          quantities. Defaults to True (keyword-only).
+
+    #### RETURNS:
+        - If per_quantity is True: a dict {name: int} mapping each
+          quantity to its estimated subspace dimension k.
+        - If per_quantity is False: a single int, the maximum k across
+          all quantities.
+    """
+    fallback = default_threshold if default_threshold is not None else (
+        float(threshold) if not isinstance(threshold, dict) else 0.90
+    )
+    
+    k_map = {
+        name: _k_from_target(target, _resolve_threshold(name, threshold, fallback))
+        for name, target in pooled_global.items()
+    }
+    
+    if per_quantity:
+        return k_map
+        
+    return max(k_map.values())

spacelearn/optim/_optimize_targets.py

@@ -0,0 +1,119 @@
+"""
+Copyright (c) 2026 Tobias Karusseit
+This source code is licensed under the MIT license found in the
+LICENSE file in the root directory of this source tree.
+"""
+
+import math
+import torch
+from ..util import singular_values
+
+def _resolve_threshold(
+    name: str,
+    threshold: float | dict[str, float],
+    default: float,
+) -> float:
+    if isinstance(threshold, dict):
+        return float(threshold.get(name, default))
+    return float(threshold)
+
+
+def optimal_bins(
+    pooled_targets_coarse: dict[str, torch.Tensor],
+    max_bins: int = 16,
+    variance_threshold: float | dict[str, float] = 0.95,
+    *,
+    default_threshold: float | None = None,
+    per_quantity: bool = False,
+) -> int | dict[str, int]:
+    """
+    Finds the smallest n such that an n x n spatial pooling grid captures
+    a target fraction of the spatial variance in each physical target,
+    via PCA-style spectral analysis on a coarse pooling of that target.
+
+    Inspired by Nyquist sampling: this finds the resolution at which
+    adding more bins gives diminishing returns on spatial information
+    captured, rather than picking a pooling resolution by hand.
+
+    #### USAGE:
+        # Example 1: shared threshold, single global n
+        n = optimal_bins(pooled_targets_coarse, max_bins=16, variance_threshold=0.95)
+
+        # Example 2: per-quantity thresholds and per-quantity result
+        n = optimal_bins(
+            pooled_targets_coarse,
+            variance_threshold={"scaled": 0.99},
+            default_threshold=0.95,
+            per_quantity=True,
+        )
+        # n = {"scaled": 6, "grammian": 4}
+
+    #### HOW IT WORKS:
+        For each quantity name, Y in pooled_targets_coarse:
+            eta  = resolved variance_threshold for this quantity
+            Y_c  = Y - Y.mean(0)                       # center (flattened
+                                                        # if > 2D)
+            S    = singular_values(Y_c)                 # spectrum of Y
+            cumvar = cumsum(S^2) / sum(S^2)              # cumulative
+                                                          # variance explained
+            n_modes  = number of leading modes needed for cumvar >= eta
+            n_bins_q = ceil(sqrt(n_modes))               # modes -> grid side
+                                                          # length, clamped
+                                                          # to [2, max_bins]
+        If per_quantity is True, returns {name: n_bins_q} for all quantities.
+        Otherwise, returns max(n_bins_q for all quantities), so every
+        quantity is adequately sampled by a single shared resolution.
+
+    #### ARGS:
+        - pooled_targets_coarse: {name: (N, n_coarse^2)} pooled physical
+          targets at some coarse resolution, used to estimate the spatial
+          spectrum. Targets with more than 2 dims are flattened to
+          (N, n_coarse^2).
+        - max_bins: Maximum grid side length n to consider. Defaults to
+          16 (i.e. up to a 16x16 = 256 bin grid).
+        - variance_threshold: Either a single float shared across all
+          quantities, or a {name: threshold} dict for per-quantity
+          thresholds. Each threshold is the target fraction of spatial
+          variance to capture. Defaults to 0.95.
+        - default_threshold: Fallback threshold for any quantity name
+          missing from a `variance_threshold` dict. If None, falls back
+          to the scalar `variance_threshold` value (or 0.95 if
+          `variance_threshold` is itself a dict with no scalar to fall
+          back to). Defaults to None (keyword-only).
+        - per_quantity: If True, returns a dict of n_optimal per
+          quantity. If False, returns the single largest n_optimal found
+          across all quantities. Defaults to False (keyword-only).
+
+    #### RETURNS:
+        - If per_quantity is True: a dict {name: int} mapping each
+          quantity to its optimal grid side length n.
+        - If per_quantity is False: a single int, the largest n across
+          all quantities, ensuring every quantity is adequately sampled.
+    """
+    fallback = default_threshold if default_threshold is not None else (
+        float(variance_threshold) if not isinstance(variance_threshold, dict) else 0.95
+    )
+    n_optimal_per_quantity = {}
+
+    for name, Y in pooled_targets_coarse.items():
+        eta = _resolve_threshold(name, variance_threshold, fallback)
+        if Y.dim() > 2:
+            Y = Y.flatten(1)
+        Y_c = Y - Y.mean(0, keepdim=True)
+
+        S = singular_values(Y_c)
+        var = S.pow(2)
+        cumvar = var.cumsum(0) / (var.sum() + 1e-8)
+
+        n_modes = int((cumvar < eta).sum().item() + 1)
+        n_bins_q = math.ceil(math.sqrt(n_modes))
+        n_bins_q = max(2, min(n_bins_q, max_bins))
+
+        n_optimal_per_quantity[name] = n_bins_q
+        print(f"  [{name}] {n_modes} modes (η={eta:.2f}) -> n_bins={n_bins_q}")
+
+    if per_quantity:
+        return n_optimal_per_quantity
+
+    # take maximum so every quantity is adequately sampled
+    return max(n_optimal_per_quantity.values())