FastLSQ 0.1.5__tar.gz → 0.2.1__tar.gz

{fastlsq-0.1.5 → fastlsq-0.2.1}/CHANGELOG.md

@@ -2,29 +2,90 @@
 
 All notable changes to FastLSQ will be documented in this file.
 
-## [Unreleased]
+## [0.2.1] - 2026-06-02
 
 ### Added
 
-- **Vector-valued solutions (`u: ℝᵈ → ℝᵏ`)**: first-class support for coupled
-  systems and decoupled multi-output PDEs. Problems opt in with
-  `self.n_outputs = k`; `solver.beta` has shape `(N, k)` and
-  `solver.predict(x)` returns `(M, k)`. Scalar problems are the `k=1` case
-  and remain bit-for-bit identical (Helmholtz 2D / Poisson 5D / Bratu 2D
-  regressions verified).
-  - New module `fastlsq.block`: `block_concat` assembles a 2-D nested list
-    of `(M_i, N_j)` tensors into a block matrix (with `None` for zero
-    blocks); `pack_beta` / `unpack_beta` convert between `(N, k)` and the
-    block-stacked `(N*k, 1)` solve representation.
-  - `predict_with_grad` / `predict_with_laplacian` einsum fixed to
-    `"idh,hk->idk"` so the output dim no longer collapses for `k>1`. Squeezes
-    back to `(M, d)` when `k=1` for backward compatibility.
-  - `ElasticWave2D` refactored: gains `n_outputs = 2`, `build()` uses
-    `block_concat`, and now exposes `exact_grad` with shape `(M, d, k)`.
-  - `LearnableFastLSQ` accepts `n_outputs`; legacy subclasses that store a
-    flat `(Nk, 1)` `beta` (e.g. `ElasticLearnable` in the elastic wave
-    example) keep working under the default `n_outputs=1`.
-- **Learnable operator coefficients**: `Op` now accepts `nn.Parameter` (and tensors) as coefficients in scalar multiplication. Use `k = nn.Parameter(...)` with `Op.laplacian(d=2) + k**2 * Op.identity(d=2)` and optimise via AdamW; gradients flow through the prebuilt linear solve. See `examples/learnable_helmholtz.py`.
+- **Device abstraction** (`fastlsq/device.py`): `resolve_device`, `set_device`,
+  `get_device`, `device_info` for CPU / CUDA / Apple-MPS. dtype-aware -- MPS is
+  auto-selected only for float32 (it has no float64), so the default float64
+  high-accuracy regime stays on CPU/CUDA. Override with `set_device(...)` or the
+  `FASTLSQ_DEVICE` environment variable; internal tensor creation respects the
+  active device at call time.
+- **Pluggable linear solver** `solve_lstsq(..., method=...)`:
+  - `"svd"` -- rank-revealing truncated SVD (LAPACK `gelsd` fast path on CPU);
+  - `"cholesky"` -- fast normal-equations solve for well-conditioned systems;
+  - `"rsvd"` -- torch-native randomized SVD (`O(MNk)`) for strongly low-rank `A`;
+  - `"auto"` (default) -- Cholesky with a cheap conditioning probe, falling back
+    to SVD when ill-conditioned (recovers the fast path without losing accuracy).
+  MPS factorizations run on CPU (no robust `svd`/`lstsq` there) and move back.
+- **Working anisotropic Sigma = L Lᵀ learner**: `LearnableFastLSQ` (diagonal &
+  cholesky modes) now converges -- `solve_inner` uses a differentiable
+  *rank-revealing* solve, the Cholesky factor is log-parameterized (clamped,
+  positive-definite), and `train_bandwidth` is robust (gradient clipping,
+  best-iterate restore, graceful SVD/gradient-failure handling). Chainable
+  `LearnableFastLSQ.fit(problem, ...)` for one-line learn-then-predict.
+- **Vector-valued solutions (`u: ℝᵈ → ℝᵏ`)**: first-class support for coupled and
+  decoupled multi-output PDEs via the new `fastlsq.block` module. Problems opt in
+  with `self.n_outputs = k`; `solver.beta` is `(N, k)` and `solver.predict(x)`
+  returns `(M, k)` (scalar `k=1` is bit-for-bit unchanged). `block_concat`
+  assembles a nested list of blocks (`None` = zero block); `pack_beta` /
+  `unpack_beta` convert between `(N, k)` and the block-stacked `(N*k, 1)` solve.
+  The block-stacked LSQ is solved by the rank-revealing solver, and the
+  Σ-learner computes its loss on the flat `_beta_flat` so it stays correct for
+  `k>1`.
+- **Learnable operator coefficients**: `Op` accepts `nn.Parameter` (and tensors)
+  as coefficients, e.g. `Op.laplacian(d=2) + k**2 * Op.identity(d=2)` with
+  `k = nn.Parameter(...)`; gradients flow through the prebuilt linear solve.
+
+### Changed
+
+- `solve_lstsq` defaults to the rank-revealing / `auto` solve instead of forming
+  the normal equations -- several orders of magnitude more accurate on the
+  rank-deficient random-feature systems (at a higher, still one-shot, cost).
+- `solve_linear` is one-shot: the bandwidth-learning hyper-parameters live on
+  `LearnableFastLSQ.fit()` / `train_bandwidth`, not the solve signature.
+- Packaging: `requires-python` lowered to `>=3.9` (+3.9 classifier); description
+  updated.
+
+### Fixed
+
+- The previously dead `LearnableFastLSQ(mode="cholesky")` path, which diverged
+  from an unstable inner `torch.linalg.lstsq` and an unconstrained `L`.
+
+## [0.1.5] - 2026-05-25
+
+### Added
+
+- **Vector-valued features**: new `VectorBasis` and `VectorFastLSQSolver`
+  (`fastlsq/vector.py`) for solving coupled systems where the unknown is a
+  vector field `u(x) = (u_1, ..., u_K)`.  Use cases include
+  streamfunction-vorticity NS `(psi, omega)`, incompressible NS primitive
+  variables `(u, v, p)`, multi-species transport, MHD, etc.
+
+  - `VectorBasis.random(input_dim, n_features, sigma, n_components, sigmas=...)`
+    creates K independent random-Fourier `SinusoidalBasis`es; per-component
+    bandwidth can be tuned via `sigmas`.
+  - Stacked evaluators: `evaluate(x) -> (M, K, N)`, `gradient -> (M, K, d, N)`,
+    `laplacian`, `hessian_diag`, `derivative(alpha)`.
+  - Block-diagonal assembly helpers (`block_diag_evaluate`,
+    `block_diag_laplacian`, `block_diag_derivative`) for systems whose
+    rows are independent per component.
+  - Coefficient packing utilities (`stack_betas`, `unstack_beta`,
+    `predict`) accept per-component lists, stacked columns, or
+    `(N, K)` matrices interchangeably.
+  - `VectorFastLSQSolver(input_dim, n_components, normalize=True)` is the
+    multi-component counterpart of `FastLSQSolver`; `add_block(scale=...)`
+    accepts either a scalar or a list of K scalars for per-component
+    bandwidth.
+  - `component(k)` / `component_solver(k)` give direct access to the k-th
+    scalar basis / solver for ad-hoc per-component work.
+
+### Other
+
+- Synchronised `pyproject.toml`, `fastlsq.__version__`, and CHANGELOG (the
+  package source had drifted to `__version__ = "0.1.0"` against
+  `pyproject.toml = "0.1.4"`; both now read `"0.1.5"`).
 
 ## [0.2.0] - 2026-03-01
 

{fastlsq-0.1.5 → fastlsq-0.2.1/FastLSQ.egg-info}/PKG-INFO

@@ -1,7 +1,7 @@
 Metadata-Version: 2.4
 Name: FastLSQ
-Version: 0.1.5
-Summary: Solving PDEs in one shot via Fourier features with exact analytical derivatives
+Version: 0.2.1
+Summary: One-shot PDE solving via Fourier features with exact analytical derivatives; rank-revealing solvers, learnable anisotropic bandwidth, and CPU/CUDA/MPS support
 Author: Antonin Sulc
 License-Expression: MIT
 Project-URL: Homepage, https://github.com/asulc/FastLSQ
@@ -14,12 +14,13 @@ Classifier: Development Status :: 4 - Beta
 Classifier: Intended Audience :: Science/Research
 Classifier: Operating System :: OS Independent
 Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Topic :: Scientific/Engineering :: Mathematics
 Classifier: Topic :: Scientific/Engineering :: Physics
-Requires-Python: >=3.10
+Requires-Python: >=3.9
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: torch>=2.0

fastlsq-0.2.1/FastLSQ.egg-info/SOURCES.txt

@@ -0,0 +1,116 @@
+CHANGELOG.md
+LICENSE
+MANIFEST.in
+README.md
+pyproject.toml
+requirements.txt
+FastLSQ.egg-info/PKG-INFO
+FastLSQ.egg-info/SOURCES.txt
+FastLSQ.egg-info/dependency_links.txt
+FastLSQ.egg-info/requires.txt
+FastLSQ.egg-info/top_level.txt
+examples/add_your_own_pde.py
+examples/benchmark_comparison.py
+examples/custom_features.py
+examples/grad_shafranov.py
+examples/grid_inverse.py
+examples/grid_rl_control.py
+examples/grid_swing.py
+examples/gs_inverse.py
+examples/gs_rl_control.py
+examples/inverse_heat_source.py
+examples/inverse_magnetostatics.py
+examples/inverse_source_position.py
+examples/learnable_helmholtz.py
+examples/orbit_hill.py
+examples/orbit_inverse.py
+examples/orbit_rl.py
+examples/pde_discovery.py
+examples/run_all_extensions.py
+examples/run_linear.py
+examples/run_nonlinear.py
+examples/tutorial_basic.py
+examples/tutorial_nonlinear.py
+examples/vector_basis_stream_vorticity.py
+examples/extras/fred_sde.py
+examples/extras/fred_sde_fastlsq.py
+examples/extras/gaia_potential.py
+examples/extras/gaia_potential_fastlsq.py
+examples/extras/horizons_ephemeris.py
+examples/extras/numerai_alpha.py
+examples/extras/numerai_alpha_fastlsq.py
+examples/extras/run_all_fastlsq.py
+examples/extras/spectral_expansion.py
+examples/extras/scenarios/__init__.py
+examples/extras/scenarios/_alsu_lattice.py
+examples/extras/scenarios/_common.py
+examples/extras/scenarios/run_all.py
+examples/extras/scenarios/s01_beamloss_ode.py
+examples/extras/scenarios/s01_betatron_tune.py
+examples/extras/scenarios/s01_green_fff.py
+examples/extras/scenarios/s01_hill_ivp.py
+examples/extras/scenarios/s01_observe_fit_act_simulator.py
+examples/extras/scenarios/s01_orbit_inverse.py
+examples/extras/scenarios/s01_passive_loco.py
+examples/extras/scenarios/s01_perturbed_hill.py
+examples/extras/scenarios/s01_sofb_observe_fit_act.py
+examples/extras/scenarios/s01_streaming_archive_growth.py
+examples/extras/scenarios/s01_synchrotron_ode.py
+examples/extras/scenarios/s01_tides_3months.py
+examples/extras/scenarios/s01_topoff_impulse.py
+examples/extras/scenarios/s01_visualize.py
+examples/extras/scenarios/s02_plasma_wakefield.py
+examples/extras/scenarios/s03_synchrobetatron.py
+examples/extras/scenarios/s04_sunspots.py
+examples/extras/scenarios/s05_helioseismology.py
+examples/extras/scenarios/s06_tides.py
+examples/extras/scenarios/s07_iers_earth_rotation.py
+examples/extras/scenarios/s08_mauna_loa_co2.py
+examples/extras/scenarios/s09_enso_qbo.py
+examples/extras/scenarios/s10_pulsar_timing.py
+examples/extras/scenarios/s11_modal_analysis.py
+examples/extras/scenarios/s12_mems_resonator.py
+examples/extras/scenarios/s13_variable_stars_kepler.py
+examples/extras/scenarios/s14_eeg.py
+examples/extras/scenarios/s15_circadian.py
+fastlsq/__init__.py
+fastlsq/api.py
+fastlsq/basis.py
+fastlsq/block.py
+fastlsq/device.py
+fastlsq/diagnostics.py
+fastlsq/export.py
+fastlsq/geometry.py
+fastlsq/learnable.py
+fastlsq/lightning.py
+fastlsq/linalg.py
+fastlsq/newton.py
+fastlsq/plotting.py
+fastlsq/solvers.py
+fastlsq/tuning.py
+fastlsq/utils.py
+fastlsq/vector.py
+fastlsq/viz.py
+fastlsq/problems/__init__.py
+fastlsq/problems/linear.py
+fastlsq/problems/nonlinear.py
+fastlsq/problems/regression.py
+misc/fastlsq_teaser.png
+misc/ideal_quadrupole.png
+misc/inverse_heat_source.gif
+misc/inverse_heat_source.png
+misc/inverse_magnetostatics.png
+misc/inverse_magnetostatics_convergence.png
+misc/quadrupole_convergence.png
+misc/quadrupole_optimization.png
+misc/tutorial_nlpoisson_convergence.png
+misc/tutorial_nlpoisson_solution.png
+tests/test_basic.py
+tests/test_block.py
+tests/test_derivatives.py
+tests/test_device.py
+tests/test_grad_shafranov.py
+tests/test_grid_swing.py
+tests/test_learnable.py
+tests/test_orbit_hill.py
+tests/test_vector_basis.py

{fastlsq-0.1.5/FastLSQ.egg-info → fastlsq-0.2.1}/PKG-INFO

@@ -1,7 +1,7 @@
 Metadata-Version: 2.4
 Name: FastLSQ
-Version: 0.1.5
-Summary: Solving PDEs in one shot via Fourier features with exact analytical derivatives
+Version: 0.2.1
+Summary: One-shot PDE solving via Fourier features with exact analytical derivatives; rank-revealing solvers, learnable anisotropic bandwidth, and CPU/CUDA/MPS support
 Author: Antonin Sulc
 License-Expression: MIT
 Project-URL: Homepage, https://github.com/asulc/FastLSQ
@@ -14,12 +14,13 @@ Classifier: Development Status :: 4 - Beta
 Classifier: Intended Audience :: Science/Research
 Classifier: Operating System :: OS Independent
 Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Topic :: Scientific/Engineering :: Mathematics
 Classifier: Topic :: Scientific/Engineering :: Physics
-Requires-Python: >=3.10
+Requires-Python: >=3.9
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: torch>=2.0

fastlsq-0.2.1/examples/extras/fred_sde.py

@@ -0,0 +1,262 @@
+#!/usr/bin/env python
+"""Discover the drift and diffusion of an interest-rate SDE from data.
+
+A short-rate model has the generic form
+    dr = mu(r,t) dt + sigma(r,t) dW.
+The four classical templates:
+    Vasicek:        mu = kappa(theta - r),       sigma = sigma_0
+    CIR:            mu = kappa(theta - r),       sigma = sigma_0 sqrt(r)
+    CKLS:           mu = kappa(theta - r),       sigma = sigma_0 r^gamma
+    Hull-White:     mu = (theta(t) - a r),       sigma = sigma_0
+For a fixed time-step approximation we have, increment by increment,
+    dr_i = mu(r_i) dt + sigma(r_i) sqrt(dt) z_i,   z_i ~ N(0,1).
+So given a series r_1,...,r_N we can estimate the conditional mean and
+conditional variance of  dr_i  given  r_i  by binning, and fit a symbolic
+mu(r) and sigma^2(r) to the binned estimates.
+
+We try to download the FRED DGS10 series (10-year Treasury constant
+maturity); if not available offline we fall back to synthetic Vasicek
+or CIR data.
+
+Usage:  python fred_sde.py
+"""
+from __future__ import annotations
+
+import time
+import io
+import urllib.request
+import numpy as np
+
+np.set_printoptions(precision=4, suppress=True)
+
+
+# ----------------------------------------------------------------------
+# Data fetching (or synthetic fallback).
+# ----------------------------------------------------------------------
+
+def fetch_fred_dgs10(timeout=10):
+    """Try to download FRED DGS10 (10-year CMT) daily yields.
+    Returns (dates, values_in_decimal) or None if no internet."""
+    url = "https://fred.stlouisfed.org/graph/fredgraph.csv?id=DGS10"
+    try:
+        req = urllib.request.Request(url, headers={"User-Agent": "Mozilla/5.0"})
+        with urllib.request.urlopen(req, timeout=timeout) as resp:
+            raw = resp.read().decode("utf-8")
+    except Exception as e:
+        print(f"   FRED fetch failed ({type(e).__name__}: {e})", flush=True)
+        return None
+    # parse: DATE,DGS10 columns; skip header
+    dates, vals = [], []
+    for line in raw.splitlines()[1:]:
+        try:
+            d, v = line.split(",")
+            f = float(v)
+            dates.append(d); vals.append(f / 100.0)   # percent → decimal
+        except (ValueError, IndexError):
+            continue
+    if len(vals) < 500:
+        return None
+    return np.array(dates), np.array(vals)
+
+
+def synth_vasicek(n=3000, kappa=0.10, theta=0.045, sigma0=0.015, r0=0.04,
+                  dt=1/252, seed=0):
+    rng = np.random.default_rng(seed)
+    r = np.empty(n); r[0] = r0
+    for i in range(1, n):
+        eps = rng.standard_normal()
+        r[i] = r[i-1] + kappa*(theta - r[i-1])*dt + sigma0*np.sqrt(dt)*eps
+    return r
+
+
+def synth_cir(n=3000, kappa=0.08, theta=0.05, sigma0=0.07, r0=0.04,
+              dt=1/252, seed=0):
+    rng = np.random.default_rng(seed)
+    r = np.empty(n); r[0] = r0
+    for i in range(1, n):
+        eps = rng.standard_normal()
+        diff = sigma0 * np.sqrt(max(r[i-1], 1e-8)) * np.sqrt(dt) * eps
+        r[i] = max(r[i-1] + kappa*(theta - r[i-1])*dt + diff, 1e-8)
+    return r
+
+
+# ----------------------------------------------------------------------
+# Conditional moment estimation: bin r and compute E[dr|r], Var[dr|r].
+# ----------------------------------------------------------------------
+
+def conditional_moments(r, n_bins=20, dt=1/252):
+    """Return (r_bin_centres, mu_hat(r_bin), sigma2_hat(r_bin)) per bin."""
+    dr = np.diff(r)
+    r_at_step = r[:-1]
+    bins = np.quantile(r_at_step, np.linspace(0, 1, n_bins + 1))
+    bins[0] -= 1e-9; bins[-1] += 1e-9
+    idx = np.digitize(r_at_step, bins) - 1
+    r_c, m_c, v_c, n_c = [], [], [], []
+    for k in range(n_bins):
+        mask = (idx == k)
+        if mask.sum() < 5:
+            continue
+        r_c.append(r_at_step[mask].mean())
+        m_c.append(dr[mask].mean() / dt)            # ≈ mu(r)
+        v_c.append(dr[mask].var()  / dt)            # ≈ sigma^2(r)
+        n_c.append(mask.sum())
+    return (np.asarray(r_c), np.asarray(m_c),
+            np.asarray(v_c), np.asarray(n_c))
+
+
+# ----------------------------------------------------------------------
+# Symbolic candidates for mu(r) and sigma^2(r).
+# ----------------------------------------------------------------------
+
+def fit_drift(r_c, mu_hat, n_c):
+    """Fit candidate drift forms; report best by AIC.
+    Candidates:
+       (Const)         mu = a
+       (Linear-MR)     mu = a + b*r          (mean-reverting: b<0 means kappa>0)
+       (Quadratic)     mu = a + b*r + c*r^2
+    """
+    w = np.sqrt(n_c)
+    cands = []
+    for name, basis_fn in [
+        ("Const",     lambda r: np.stack([np.ones_like(r)], axis=1)),
+        ("Linear-MR", lambda r: np.stack([np.ones_like(r), r], axis=1)),
+        ("Quadratic", lambda r: np.stack([np.ones_like(r), r, r**2], axis=1)),
+    ]:
+        X = basis_fn(r_c)
+        Xw = X * w[:, None]; yw = mu_hat * w
+        coef, *_ = np.linalg.lstsq(Xw, yw, rcond=None)
+        resid = (X @ coef - mu_hat) * w
+        rss = float(resid @ resid)
+        k = X.shape[1]
+        aic = len(r_c) * np.log(max(rss / max(len(r_c), 1), 1e-30)) + 2 * k
+        cands.append((name, coef, rss, aic))
+    cands.sort(key=lambda x: x[3])
+    return cands
+
+
+def fit_diffusion(r_c, sigma2_hat, n_c):
+    """Fit candidate sigma^2(r) forms; report best by AIC.
+       (Const)         sigma^2 = a
+       (Linear-r)      sigma^2 = a * r           ← CIR (after rescaling)
+       (Linear-r^2)    sigma^2 = a * r^2         ← geometric BM
+       (Power)         sigma^2 = a * r^(2*gamma) — fit gamma in log space
+    """
+    w = np.sqrt(n_c)
+    cands = []
+    # Const, linear-r, linear-r^2 by weighted LSQ
+    for name, basis_fn in [
+        ("Const",      lambda r: np.stack([np.ones_like(r)], axis=1)),
+        ("Linear-r",   lambda r: np.stack([r],                axis=1)),
+        ("Linear-r^2", lambda r: np.stack([r**2],             axis=1)),
+    ]:
+        X = basis_fn(r_c)
+        Xw = X * w[:, None]; yw = sigma2_hat * w
+        coef, *_ = np.linalg.lstsq(Xw, yw, rcond=None)
+        resid = (X @ coef - sigma2_hat) * w
+        rss = float(resid @ resid)
+        k = X.shape[1]
+        aic = len(r_c) * np.log(max(rss / max(len(r_c), 1), 1e-30)) + 2 * k
+        cands.append((name, coef, rss, aic))
+    # Power: log(sigma^2) = log(a) + 2*gamma * log(r), needs r > 0
+    pos = r_c > 0
+    if pos.sum() > 3:
+        lx = np.log(r_c[pos]); ly = np.log(np.maximum(sigma2_hat[pos], 1e-20))
+        ww = w[pos]
+        X = np.stack([np.ones_like(lx), lx], axis=1)
+        Xw = X * ww[:, None]; yw = ly * ww
+        coef_p, *_ = np.linalg.lstsq(Xw, yw, rcond=None)
+        log_a, two_gamma = coef_p
+        # AIC in original space
+        pred = np.exp(log_a) * (r_c[pos] ** two_gamma)
+        resid = (pred - sigma2_hat[pos]) * w[pos]
+        rss = float(resid @ resid)
+        k = 2
+        aic = pos.sum() * np.log(max(rss / max(pos.sum(), 1), 1e-30)) + 2 * k
+        cands.append((f"Power(gamma={two_gamma/2:.2f})", coef_p, rss, aic))
+    cands.sort(key=lambda x: x[3])
+    return cands
+
+
+# ----------------------------------------------------------------------
+# Pretty-print one analysis
+# ----------------------------------------------------------------------
+
+def pretty_drift(name, coef):
+    if name == "Const":
+        return f"mu(r) = {coef[0]:+.4f}"
+    if name == "Linear-MR":
+        a, b = coef
+        # mean reversion: mu = a + b*r = b*(a/b + r), kappa = -b, theta = -a/b
+        kappa = -b
+        theta = -a / b if abs(b) > 1e-12 else float("nan")
+        return (f"mu(r) = {a:+.4f} + {b:+.4f} r   "
+                f"(kappa={kappa:+.3f}, theta={theta:+.4f})")
+    if name == "Quadratic":
+        return (f"mu(r) = {coef[0]:+.4f} + {coef[1]:+.4f} r + {coef[2]:+.4f} r^2")
+    return f"{name}: coefs={coef}"
+
+
+def pretty_diffusion(name, coef):
+    if name == "Const":
+        return f"sigma^2(r) = {coef[0]:.4g}        => sigma = {np.sqrt(max(coef[0],0)):.4f}"
+    if name == "Linear-r":
+        return f"sigma^2(r) = {coef[0]:.4g} * r    (CIR-like)"
+    if name == "Linear-r^2":
+        return f"sigma^2(r) = {coef[0]:.4g} * r^2 (geom-BM-like)"
+    if name.startswith("Power"):
+        log_a, two_gamma = coef
+        return (f"sigma^2(r) = {np.exp(log_a):.4g} * r^{two_gamma:.3f}    "
+                f"({name})")
+    return f"{name}: coefs={coef}"
+
+
+# ----------------------------------------------------------------------
+# Driver
+# ----------------------------------------------------------------------
+
+def analyse(label, r, dt=1/252):
+    print(f"\n=== {label}   (N={len(r)}  observations)")
+    r_c, mu_hat, sig2_hat, n_c = conditional_moments(r, n_bins=20, dt=dt)
+    print(f"   r range:        {r_c.min():.4f} .. {r_c.max():.4f}")
+    print(f"   binned drifts:  range {mu_hat.min():+.4f} .. {mu_hat.max():+.4f}")
+    print(f"   binned vars:    range {sig2_hat.min():.4g} .. {sig2_hat.max():.4g}")
+    drift_cands  = fit_drift(r_c, mu_hat, n_c)
+    diff_cands   = fit_diffusion(r_c, sig2_hat, n_c)
+    print("   Drift candidates (sorted by AIC):")
+    for nm, coef, rss, aic in drift_cands:
+        marker = "←" if nm == drift_cands[0][0] else " "
+        print(f"     {marker} {nm:12s}  AIC={aic:9.1f}  {pretty_drift(nm, coef)}")
+    print("   Diffusion candidates (sorted by AIC):")
+    for nm, coef, rss, aic in diff_cands:
+        marker = "←" if nm == diff_cands[0][0] else " "
+        print(f"     {marker} {nm:20s}  AIC={aic:9.1f}  {pretty_diffusion(nm, coef)}")
+    return drift_cands[0], diff_cands[0]
+
+
+def main():
+    print(">> Short-rate SDE discovery (FRED + synthetic checks)\n", flush=True)
+
+    # --- 1. Synthetic Vasicek
+    r_vas = synth_vasicek(n=3000, kappa=0.10, theta=0.045, sigma0=0.015, seed=0)
+    analyse("Synthetic Vasicek  truth: drift Linear-MR (kappa=0.10, theta=0.045); "
+            "sigma^2 = const = 0.000225", r_vas)
+
+    # --- 2. Synthetic CIR
+    r_cir = synth_cir(n=3000, kappa=0.08, theta=0.05, sigma0=0.07, seed=0)
+    analyse("Synthetic CIR      truth: drift Linear-MR (kappa=0.08, theta=0.050); "
+            "sigma^2 = 0.0049 * r", r_cir)
+
+    # --- 3. Real FRED data, if reachable
+    res = fetch_fred_dgs10()
+    if res is None:
+        print("\n=== FRED DGS10 unavailable (offline) --- skipping real-data run.")
+    else:
+        dates, r_real = res
+        # drop NaN-ish entries (FRED uses '.' for missing)
+        r_real = r_real[~np.isnan(r_real)]
+        analyse(f"FRED DGS10 (10-yr Treasury daily yields, {len(r_real)} obs)",
+                r_real)
+
+
+if __name__ == "__main__":
+    main()