StochasticForceInference 2.0.0__py3-none-any.whl

SFI/trajectory/reserved_extras.py

@@ -0,0 +1,129 @@
+"""Framework-owned extras and the single resolver that materialises them.
+
+A force or diffusion expression reads two kinds of data through its ``extras``
+mapping: **user** values attached to the trajectory (drive protocols,
+per-particle properties, geometry) and **reserved** values supplied by the
+framework. The reserved keys are defined once here, in a small registry, and
+assembled together with the user values by :func:`resolve_extras` — the single
+entry point used by simulation, inference, and diagnostics.
+
+Reserved keys:
+
+``time``
+    Absolute time at each resolved frame — lets time-dependent bases (e.g.
+    :func:`~SFI.bases.time_fourier`) read the clock.
+``duration``
+    Total trajectory span.
+``dataset_index``
+    Dense index of the dataset within its collection, for pooled
+    multi-experiment models (:func:`~SFI.bases.per_dataset_scalar`,
+    :func:`~SFI.bases.dataset_indicator`).
+``particle_index``
+    Per-particle integer ids, gathered per edge by interaction dispatchers.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Callable, Dict, Mapping, Optional
+
+import jax.numpy as jnp
+
+
+@dataclass(frozen=True)
+class ExtrasContext:
+    """Everything a reserved-key resolver needs for a set of frames.
+
+    Assembled by whoever drives the evaluation (the trajectory producer, the
+    diagnostics residual builder, or the simulator) and passed to
+    :func:`resolve_extras`.
+    """
+
+    n_particles: int
+    dataset_index: int
+    frame_times: Any
+    duration: Any
+
+
+@dataclass(frozen=True)
+class ReservedKey:
+    """A framework-owned extras key and how it is materialised."""
+
+    name: str
+    resolve: Callable[[ExtrasContext], Any]
+
+
+_REGISTRY: Dict[str, ReservedKey] = {}
+
+
+def register(key: ReservedKey) -> None:
+    """Add a reserved key to the registry."""
+    _REGISTRY[key.name] = key
+
+
+register(ReservedKey("time", lambda ctx: ctx.frame_times))
+register(ReservedKey("duration", lambda ctx: ctx.duration))
+register(ReservedKey("dataset_index", lambda ctx: jnp.asarray(ctx.dataset_index, dtype=jnp.int32)))
+register(ReservedKey("particle_index", lambda ctx: jnp.arange(int(ctx.n_particles), dtype=jnp.int32)))
+
+
+#: The set of reserved key names; user extras may not use these.
+RESERVED_NAMES = frozenset(_REGISTRY)
+
+
+def is_reserved(name: str) -> bool:
+    """True when ``name`` is a framework-owned reserved key."""
+    return name in _REGISTRY
+
+
+def resolve_reserved(ctx: ExtrasContext) -> Dict[str, Any]:
+    """Materialise every reserved key for ``ctx``."""
+    return {name: key.resolve(ctx) for name, key in _REGISTRY.items()}
+
+
+def resolve_extras(user_extras: Mapping[str, Any], ctx: ExtrasContext) -> Dict[str, Any]:
+    """Full per-frame extras: user values plus the resolved reserved keys.
+
+    Reserved names are framework-owned; a user entry colliding with one is
+    rejected so the meaning of a reserved key is never ambiguous.
+    """
+    out = dict(user_extras)
+    clash = RESERVED_NAMES.intersection(out)
+    if clash:
+        raise ValueError(f"extras keys {sorted(clash)} are reserved; rename the user entries.")
+    out.update(resolve_reserved(ctx))
+    return out
+
+
+def slice_frame_extras(
+    extras_global: Optional[Mapping[str, Any]],
+    extras_local: Optional[Mapping[str, Any]],
+    *,
+    frame_idx: Any,
+    context: Optional[str] = None,
+) -> Dict[str, Any]:
+    """Materialise user extras at ``frame_idx``.
+
+    * :class:`~SFI.trajectory.dataset.TimeSeriesExtra` → sliced ``value.data[frame_idx]``;
+    * :class:`~SFI.trajectory.dataset.FunctionExtra` → its callable, forwarded;
+    * plain callable → invoked as ``value(frame_idx, context=context)``;
+    * anything else → forwarded unchanged.
+
+    ``extras_local`` overrides ``extras_global`` on key conflicts.
+    """
+    from SFI.trajectory.dataset import FunctionExtra, TimeSeriesExtra
+
+    def _materialise(value: Any) -> Any:
+        if isinstance(value, FunctionExtra):
+            return value.func
+        if isinstance(value, TimeSeriesExtra):
+            return jnp.asarray(value.data)[frame_idx]
+        if callable(value):
+            return value(frame_idx, context=context)
+        return value
+
+    out: Dict[str, Any] = {}
+    for source in (extras_global or {}, extras_local or {}):
+        for key, value in source.items():
+            out[key] = _materialise(value)
+    return out

SFI/utils/__init__.py

@@ -0,0 +1,17 @@
+"""
+SFI.utils — Mathematical, formatting, and plotting utilities.
+"""
+
+from .formatting import model_summary, print_model_comparison
+from .maths import as_default_float, default_float_dtype, fd_velocity, solve_or_pinv, sqrtm_psd, stable_pinv
+
+__all__ = [
+    "stable_pinv",
+    "sqrtm_psd",
+    "solve_or_pinv",
+    "fd_velocity",
+    "model_summary",
+    "print_model_comparison",
+    "default_float_dtype",
+    "as_default_float",
+]

SFI/utils/formatting.py

@@ -0,0 +1,308 @@
+# SFI/utils/formatting.py
+"""Pretty-printing utilities for inferred models."""
+
+from __future__ import annotations
+
+from typing import Optional, Sequence
+
+import numpy as np
+
+# ANSI escape codes for terminal highlighting
+_BOLD = "\033[1m"
+_DIM = "\033[2m"
+_RESET = "\033[0m"
+
+
+def model_summary(
+    labels: Sequence[str],
+    coefficients: np.ndarray,
+    *,
+    stderr: Optional[np.ndarray] = None,
+    support: Optional[np.ndarray] = None,
+    coeffs_true: Optional[np.ndarray] = None,
+    support_true: Optional[np.ndarray] = None,
+    title: str = "Coefficient Table",
+    max_rows: int = 60,
+    significance_thresholds: tuple = (2.0, 10.0, 100.0),
+    auto_labels: bool = False,
+) -> str:
+    """
+    Build a human-readable coefficient table with SNR and significance.
+
+    Only active (support) coefficients are shown in the table body.  Zeroed
+    basis functions are listed separately below, unless labels are
+    auto-generated.
+
+    Parameters
+    ----------
+    labels : sequence of str
+        One label per basis function.
+    coefficients : 1-D array
+        Coefficient vector (length must match *labels* or *support*).
+    stderr : 1-D array or None
+        Standard errors (same length as *coefficients*).  These reflect
+        sampling error only; discretization (finite-time-step) bias is
+        not included.
+    support : 1-D int array or None
+        Indices into *labels* that *coefficients* correspond to.
+        If ``None``, full support is assumed (all basis functions
+        have non-zero coefficients).
+    title : str
+        Section header printed above the table.
+    max_rows : int
+        If the table exceeds this many rows, truncate the middle.
+    significance_thresholds : tuple of float
+        Three SNR thresholds (in multiples of stderr) for the ``*``, ``**``,
+        and ``***`` significance levels.  Defaults ``(2.0, 10.0, 100.0)``.
+    auto_labels : bool
+        If ``True``, labels were auto-generated (e.g. ``b0``, ``b1``, …)
+        and the list of zeroed functions is suppressed.
+
+    Returns
+    -------
+    str
+        Ready-to-print multi-line table.  Terms are marked *, **, or ***
+        according to *significance_thresholds*; *** terms are **bold**,
+        * and ** terms are normal weight, non-significant active terms are dimmed.
+    """
+    coefficients = np.asarray(coefficients).ravel()
+    n_labels = len(labels)
+
+    # Expand coefficients to full length when support is given
+    if support is not None:
+        support = np.asarray(support).ravel().astype(int)
+        if len(coefficients) == len(support):
+            # coefficients is sparse; expand to full length
+            full = np.zeros(n_labels)
+            full[support] = coefficients
+            coefficients = full
+        elif len(coefficients) != n_labels:
+            raise ValueError(
+                f"len(coefficients)={len(coefficients)} matches neither "
+                f"len(support)={len(support)} nor n_labels={n_labels}"
+            )
+        active = set(int(i) for i in support)
+    else:
+        active = None  # all active
+
+    has_stderr = stderr is not None
+    if has_stderr:
+        stderr = np.asarray(stderr).ravel()
+        if stderr.shape[0] != n_labels:
+            # Might be sparse; try to expand
+            if support is not None and stderr.shape[0] == len(support):
+                se_full = np.zeros(n_labels)
+                se_full[support] = stderr
+                stderr = se_full
+            else:
+                raise ValueError(
+                    f"len(stderr)={stderr.shape[0]} matches neither "
+                    f"n_labels={n_labels} nor len(support)={len(support) if support is not None else 'N/A'}"
+                )
+
+    # Expand ground-truth coefficients to full length (paired "True" column)
+    has_true = coeffs_true is not None
+    true_full = None
+    if has_true:
+        ct = np.asarray(coeffs_true).ravel()
+        true_full = np.zeros(n_labels)
+        if support_true is not None:
+            st = np.asarray(support_true).ravel().astype(int)
+            true_full[st] = ct
+        elif ct.shape[0] == n_labels:
+            true_full = ct
+        elif support is not None and ct.shape[0] == len(support):
+            true_full[support] = ct
+        else:
+            true_full[: ct.shape[0]] = ct
+
+    # ── Build rows ──
+    # (index, label, coefficient, stderr_or_None, is_active, snr_or_None, true_or_None)
+    rows: list = []
+    for i, lab in enumerate(labels):
+        c = float(coefficients[i])
+        se = float(stderr[i]) if (stderr is not None and i < len(stderr)) else None
+        is_active = (active is None) or (i in active)
+        snr = abs(c / se) if (se is not None and se > 0) else None
+        tv = float(true_full[i]) if has_true else None
+        rows.append((i, lab, c, se, is_active, snr, tv))
+
+    # ── Column widths ──
+    idx_w = max(3, len(str(n_labels - 1)))
+    lab_w = max(5, max(len(r[1]) for r in rows))
+    coeff_fmt = "{:>12.5e}"
+    se_fmt = "{:>12.5e}"
+    snr_fmt = "{:>6.1f}"
+
+    # Determine whether SNR column should be shown
+    has_snr = has_stderr and any(r[5] is not None for r in rows)
+
+    # ── Header ──
+    sep_w = idx_w + lab_w + 30
+    if has_stderr:
+        sep_w += 15
+    if has_snr:
+        sep_w += 9  # "  SNR   "
+    if has_true:
+        sep_w += 15
+    sep = "─" * sep_w
+
+    lines: list[str] = []
+    lines.append(f"  {title}")
+    lines.append(f"  {sep}")
+
+    hdr = f"  {'#':<{idx_w}}  {'Label':<{lab_w}}  {'Coefficient':>12}"
+    if has_stderr:
+        hdr += f"  {'Std.Err':>12}"
+    if has_snr:
+        hdr += f"  {'SNR':>6}"
+    if has_true:
+        hdr += f"  {'True':>12}"
+    hdr += "  Sig"
+    lines.append(hdr)
+    lines.append(f"  {sep}")
+
+    # ── Rows (with optional truncation) — only show active (support) entries ──
+    thr1, thr2, thr3 = significance_thresholds
+
+    def _sig_marker(snr):
+        """Return (marker_str, is_bold) for a given SNR (or None)."""
+        if snr is None:
+            return "·", False
+        if snr >= thr3:
+            return "***", True
+        if snr >= thr2:
+            return "**", False
+        if snr >= thr1:
+            return "*", False
+        return "·", False
+
+    def _fmt_row(r):
+        i, lab, c, se, act, snr, tv = r
+        marker, bold = _sig_marker(snr if act else None)
+
+        # Start with index and label
+        s = f"  {i:<{idx_w}}  {lab:<{lab_w}}  {coeff_fmt.format(c)}"
+        if has_stderr:
+            s += f"  {se_fmt.format(se)}" if se is not None else f"  {'---':>12}"
+        if has_snr:
+            s += f"  {snr_fmt.format(snr)}" if snr is not None else f"  {'---':>6}"
+        if has_true:
+            s += f"  {coeff_fmt.format(tv)}" if tv is not None else f"  {'---':>12}"
+
+        # Significance marker (padded to 3 chars for alignment)
+        s += f"  {marker:<3}"
+
+        # Apply ANSI formatting: bold for ***, dim for non-significant active
+        if bold:
+            s = _BOLD + s + _RESET
+        elif marker == "·" and has_snr:
+            s = _DIM + s + _RESET
+
+        return s
+
+    active_rows = [r for r in rows if r[4]]
+    zero_rows = [r for r in rows if not r[4]]
+
+    if len(active_rows) <= max_rows:
+        for r in active_rows:
+            lines.append(_fmt_row(r))
+    else:
+        half = (max_rows - 1) // 2
+        for r in active_rows[:half]:
+            lines.append(_fmt_row(r))
+        lines.append(f"  {'...':^{sep_w}}")
+        for r in active_rows[-half:]:
+            lines.append(_fmt_row(r))
+
+    lines.append(f"  {sep}")
+
+    n_active = len(active_rows)
+    summary_line = f"  {n_active}/{n_labels} basis functions in support"
+    if has_snr:
+        n1 = sum(1 for r in active_rows if r[5] is not None and r[5] >= thr1)
+        n2 = sum(1 for r in active_rows if r[5] is not None and r[5] >= thr2)
+        n3 = sum(1 for r in active_rows if r[5] is not None and r[5] >= thr3)
+        summary_line += f", sig: {n1}* / {n2}** / {n3}*** (|SNR| ≥ {thr1:.4g} / {thr2:.4g} / {thr3:.4g})"
+    lines.append(summary_line)
+    if has_stderr:
+        lines.append("  (Std.err. reflects sampling error only; discretization bias is not included.)")
+
+    # ── Zeroed basis functions ──
+    if zero_rows and not auto_labels:
+        zero_labels = ", ".join(r[1] for r in zero_rows)
+        lines.append(f"  Zeroed ({len(zero_rows)}): {zero_labels}")
+
+    return "\n".join(lines)
+
+
+def print_model_comparison(
+    inferences,
+    labels,
+    *,
+    extra_cols=None,
+    metrics=None,
+    title: str = "Model Comparison",
+) -> str:
+    """Build a multi-model comparison table from several inference objects.
+
+    Parameters
+    ----------
+    inferences :
+        Sequence of fitted inference objects.
+    labels :
+        One name per inference (row labels).
+    metrics :
+        Attribute names to read from each inference (default
+        ``["n_params", "NMSE_force", "force_predicted_MSE"]``).  The special
+        ``"n_params"`` counts ``force_coefficients_full``.
+    extra_cols :
+        Optional ``{column_name: {label: value}}`` of caller-supplied cells.
+    title :
+        Header line.
+
+    Returns
+    -------
+    str
+        Ready-to-print table.
+    """
+    if metrics is None:
+        metrics = ["n_params", "NMSE_force", "force_predicted_MSE"]
+    extra_cols = extra_cols or {}
+    cols = ["Model"] + list(metrics) + list(extra_cols.keys())
+
+    def _metric(inf, name):
+        if name == "n_params":
+            cf = getattr(inf, "force_coefficients_full", None)
+            return int(np.asarray(cf).size) if cf is not None else None
+        v = getattr(inf, name, None)
+        return float(v) if v is not None else None
+
+    rows = []
+    for inf, lab in zip(inferences, labels):
+        row = {"Model": lab}
+        for name in metrics:
+            row[name] = _metric(inf, name)
+        for col, vals in extra_cols.items():
+            row[col] = vals.get(lab) if isinstance(vals, dict) else None
+        rows.append(row)
+
+    def _fmt(v):
+        if v is None:
+            return "—"
+        if isinstance(v, bool):
+            return str(v)
+        if isinstance(v, int):
+            return str(v)
+        if isinstance(v, float):
+            return f"{v:.4g}"
+        return str(v)
+
+    widths = {c: max(len(str(c)), *(len(_fmt(r.get(c))) for r in rows)) for c in cols}
+    lines = [f"  {title}"]
+    header = "  " + "  ".join(f"{str(c):>{widths[c]}}" for c in cols)
+    lines.append(header)
+    lines.append("  " + "─" * (len(header) - 2))
+    for r in rows:
+        lines.append("  " + "  ".join(f"{_fmt(r.get(c)):>{widths[c]}}" for c in cols))
+    return "\n".join(lines)

SFI/utils/maths.py

@@ -0,0 +1,185 @@
+import jax
+import jax.numpy as jnp
+import jax.scipy.linalg as jsp_la
+from jax import dtypes as _jdtypes
+from jax import jit
+
+# ---------------------------------------------------------------------
+#  Mathematical utilities
+# ---------------------------------------------------------------------
+
+
+def default_float_dtype():
+    """Return JAX's currently-active default float dtype.
+
+    Returns ``float64`` when ``JAX_ENABLE_X64`` is set, ``float32`` otherwise.
+    Use this everywhere user-supplied arrays enter the simulation pipeline
+    so that all internal arithmetic shares a single dtype and ``lax.scan``
+    carry-in / carry-out types always agree.
+    """
+    return _jdtypes.canonicalize_dtype(jnp.float64)
+
+
+def as_default_float(x):
+    """Return ``x`` cast to the JAX default float dtype.
+
+    Accepts any array-like; passes through ``None`` unchanged so callers
+    can use it on optional inputs (e.g. ``v0``).
+    """
+    if x is None:
+        return None
+    return jnp.asarray(x, dtype=default_float_dtype())
+
+
+def fd_velocity(x, dt, *, scheme: str = "central"):
+    """Finite-difference velocity along the leading (time) axis.
+
+    Reconstructs ``v(t) ≈ dx/dt`` from a position array using a
+    finite-difference stencil — the same secant-velocity convention the
+    underdamped inference engine uses internally.  The output keeps the
+    same leading time dimension as ``x`` (boundary frames fall back to a
+    one-sided stencil).
+
+    Parameters
+    ----------
+    x : array_like, shape ``(T, ...)``
+        Positions sampled in time along axis 0 (e.g. ``(T, d)`` or
+        ``(T, N, d)``).
+    dt : float or array_like, shape ``(T - 1,)``
+        Time step.  Scalar for uniform sampling, or per-interval spacings.
+    scheme : {"central", "forward", "backward"}
+        Interior stencil.  ``"central"`` is second-order accurate and the
+        default; ``"forward"`` / ``"backward"`` are first-order.
+
+    Returns
+    -------
+    v : jax.Array, shape ``(T, ...)``
+        Velocity estimate, same shape as ``x``.
+    """
+    x = jnp.asarray(x)
+    if x.shape[0] < 2:
+        raise ValueError("fd_velocity needs at least 2 time points.")
+    dt = jnp.asarray(dt)
+    n_extra = x.ndim - 1
+
+    def _bcast(a):
+        # reshape a per-step (T-1,) array to broadcast over the trailing axes
+        return a.reshape(a.shape + (1,) * n_extra)
+
+    # forward differences between consecutive frames, shape (T-1, ...)
+    fwd = (x[1:] - x[:-1]) / (dt if dt.ndim == 0 else _bcast(dt))
+
+    if scheme == "forward":
+        return jnp.concatenate([fwd, fwd[-1:]], axis=0)
+    if scheme == "backward":
+        return jnp.concatenate([fwd[:1], fwd], axis=0)
+    if scheme == "central":
+        denom = (2.0 * dt) if dt.ndim == 0 else _bcast(dt[:-1] + dt[1:])
+        central = (x[2:] - x[:-2]) / denom
+        return jnp.concatenate([fwd[:1], central, fwd[-1:]], axis=0)
+    raise ValueError(f"Unknown scheme {scheme!r}; expected central/forward/backward.")
+
+
+def stable_pinv(G):
+    """Numerically-stable pseudo-inverse of a Gram matrix.
+
+    Normalizes G by its diagonal before inversion so that all diagonal
+    entries of the rescaled matrix are 1.  This avoids ill-conditioning
+    when basis functions have very different scales.
+
+    Algorithm: let d_i = sqrt(G_{ii}) (clamped to 1 when zero).  Compute
+    pinv(G / outer(d, d)) and rescale by outer(1/d, 1/d) to recover the
+    pseudo-inverse of the original G.
+    """
+    # Normalize the matrix before computing the pseudo-inverse, for numerical stability
+    G_norm = jnp.sqrt(jnp.diag(G))  # Extract diagonal elements as normalization factors
+    # Avoid division by zero: if G_norm[i] is zero, keep it zero in scaling
+    safe_G_norm = jnp.where(G_norm > 0, G_norm, 1.0)  # Replace 0 with 1 to avoid NaN in division
+    return jnp.linalg.pinv(G / jnp.outer(safe_G_norm, safe_G_norm)) * jnp.outer(1 / safe_G_norm, 1 / safe_G_norm)
+
+
+@jit
+def sqrtm_psd(A: jax.Array, eps: float = 0.0) -> jax.Array:
+    """
+    Symmetric PSD matrix square root, applied to the last two (matrix) axes.
+    - Symmetrizes input for stability.
+    - Clips negative eigenvalues to 0 (PSD) then takes sqrt.
+    - Re-symmetrizes the result to kill small asymmetries from numerics.
+    """
+    A = 0.5 * (A + jnp.swapaxes(A, -1, -2))
+    w, V = jnp.linalg.eigh(A)
+    w = jnp.clip(w, min=0.0)  # PSD safeguard
+    if eps:
+        w = w + eps  # optional jitter if you like
+    S = (V * jnp.sqrt(w)[..., None, :]) @ jnp.swapaxes(V, -1, -2)
+    return 0.5 * (S + jnp.swapaxes(S, -1, -2))
+
+
+def solve_or_pinv(A: jax.Array, b: jax.Array, tol: float = 1e-15) -> jax.Array:
+    """
+    Solve A ⋅ x = b for x, with a fallback to the Moore–Penrose pseudo-inverse
+    if A is singular or not square.  To improve numerical stability, we first
+    normalize A by its diagonal: A_norm = D^{-1} A D^{-1}, b_norm = D^{-1} b,
+    solve A_norm ⋅ x_norm = b_norm, and then recover x = D^{-1} x_norm.
+
+    This ensures that the diagonal entries of A_norm are 1 (assuming A has
+    positive diagonal), which often makes the linear solve or pseudo-inverse
+    more robust when A has widely varying scales on its diagonal.
+
+    Parameters
+    ----------
+    A : jax.Array, shape (k, k)
+        The matrix to solve against.  We assume that A has nonnegative diagonal
+        entries; if any diagonal entry is zero, we clip it to a small floor
+        to avoid division by zero.
+    b : jax.Array, shape (k,)
+        The right-hand side vector.
+    tol : float, default=1e-15
+        The tolerance for the pseudo-inverse.  If A_norm is effectively singular,
+        we compute x_norm = pinv(A_norm, rcond=tol) @ b_norm.
+
+    Returns
+    -------
+    x : jax.Array, shape (k,)
+        The solution vector to A ⋅ x = b, computed as follows:
+          1) d_i = sqrt(max(A_{ii}, tol))
+             (we floor each diagonal entry to tol > 0 to avoid zero divides)
+          2) A_norm = D_inv @ A @ D_inv  where D_inv = diag(1 / d_i)
+             b_norm = b / d
+          3) Solve A_norm ⋅ x_norm = b_norm:
+             - if A_norm is non-singular, use a direct solver
+             - otherwise, fall back to x_norm = pinv(A_norm) @ b_norm
+          4) Recover x = x_norm / d
+    """
+    # 1) Extract and “floor” the diagonal of A to avoid zeros or negatives.
+    #    If A_{ii} is very small or negative (due to numerical noise), we floor it.
+    diag_A = jnp.diag(A)  # shape (k,)
+    # Clip to tol to avoid sqrt of zero or negative
+    diag_clipped = jnp.clip(diag_A, min=tol)  # shape (k,)
+    d = jnp.sqrt(diag_clipped)  # shape (k,)
+
+    # 2) Build the inverse scaling matrix D_inv = diag(1 / d_i)
+    #    We do this by dividing each row of A by d_i and each column by d_j.
+    #    More precisely: A_norm[i,j] = A[i,j] / (d[i] * d[j]).
+    D_inv = 1.0 / d  # shape (k,)
+    # Use broadcasting to normalize A: first divide each row by d,
+    # then each column by d.
+    A_norm = (A * D_inv[:, None]) * D_inv[None, :]
+
+    # 3) Normalize the RHS vector b as well: b_norm = b / d.
+    b_norm = b / d  # shape (k,)
+
+    # 4) Attempt a direct solve of A_norm ⋅ x_norm = b_norm.
+    #    If A_norm is non-singular, this will succeed.
+    try:
+        # We do not assume positive-definite here, so use 'gen' unless we detect
+        # symmetry and positive definiteness.  For simplicity, assume 'gen'.
+        x_norm = jsp_la.solve(A_norm, b_norm, assume_a="gen")
+    except Exception:  # broad catch: JAX/XLA doesn't expose a stable LinAlgError type
+        # 4a) Fallback: if A_norm is singular or not square, use pseudo-inverse.
+        x_norm = jnp.linalg.pinv(A_norm, rtol=tol) @ b_norm
+
+    # 5) Scale back to obtain the final solution: x_i = x_norm[i] / d[i].
+    x = x_norm / d  # shape (k,)
+
+    return x