FastLSQ 0.1.4__tar.gz → 0.2.1__tar.gz

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

@@ -2,11 +2,90 @@
 
 All notable changes to FastLSQ will be documented in this file.
 
-## [Unreleased]
+## [0.2.1] - 2026-06-02
 
 ### Added
 
-- **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.4 → fastlsq-0.2.1/FastLSQ.egg-info}/PKG-INFO

@@ -1,7 +1,7 @@
 Metadata-Version: 2.4
 Name: FastLSQ
-Version: 0.1.4
-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,17 +14,22 @@ 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
 Requires-Dist: numpy>=1.24
 Requires-Dist: matplotlib>=3.7
+Provides-Extra: battery
+Requires-Dist: progpy>=1.3; extra == "battery"
+Requires-Dist: pandas>=2.0; extra == "battery"
+Requires-Dist: scipy>=1.10; extra == "battery"
 Provides-Extra: dev
 Requires-Dist: pytest>=7.0; extra == "dev"
 Requires-Dist: pandas>=2.0; extra == "dev"
@@ -36,8 +41,11 @@ Dynamic: license-file
 
 # FastLSQ
 
+[BerkeleyLab ATAP Talk](https://github.com/sulcantonin/FastLSQ/raw/main/presentations/ATAP_Sulc_20260324.pptx)
+
+
 <p align="center">
-  <img src="misc/fastlsq_teaser.png" alt="FastLSQ method overview" width="800"/>
+  <img src="misc/fastlsq_teaser.png" alt="FastLSQ method overview" width="400"/>
 </p>
 
 **Solving PDEs in one shot via Fourier features with exact analytical derivatives.**
@@ -128,6 +136,78 @@ A_pde = helmholtz.apply(basis, x)    # (5000, 1500)
 wave = Op.partial(dim=2, order=2, d=3) - c**2 * Op.laplacian(d=3, dims=[0, 1])
 ```
 
+### Vector-valued solutions
+
+`solve_linear` / `solve_nonlinear` support vector-valued **u**: ℝᵈ → ℝᵏ for
+coupled systems (elasticity, Stokes, Maxwell vector potential, …) and for
+decoupled multi-output problems sharing one basis. The math is unchanged; the
+solver just allocates `beta` with shape `(N, k)` so that `solver.predict(x)`
+returns shape `(M, k)` directly.
+
+A problem opts in by setting `self.n_outputs = k` and assembling its operator
+in block-stacked form `A ∈ ℝ^{Mk × Nk}`, `b ∈ ℝ^{Mk × 1}`. The helper
+`block_concat` removes the manual `torch.cat` bookkeeping:
+
+```python
+import torch
+from fastlsq import solve_linear, block_concat
+
+class Stokes2D:
+    n_outputs = 3        # (u, v, p)
+    dim = 2
+    name = "Stokes 2D"
+    # ... exact, exact_grad, get_train_data, get_test_points ...
+
+    def build(self, slv, x, bcs, f):
+        basis = slv.basis
+        cache = basis.cache(x)
+        dx = basis.derivative(x, (1, 0), cache=cache)
+        dy = basis.derivative(x, (0, 1), cache=cache)
+        lap = basis.laplacian(x, cache=cache)
+
+        # Rows = equations (mom_x, mom_y, continuity);
+        # columns = coefficient blocks (u, v, p)
+        A = block_concat([
+            [-lap,  None,  dx  ],   # -Δu + ∂p/∂x = f_x
+            [ None, -lap,  dy  ],   # -Δv + ∂p/∂y = f_y
+            [ dx,   dy,    None],   #  ∂u/∂x + ∂v/∂y = 0
+        ])
+        b = block_concat([[f[:, 0:1]], [f[:, 1:2]], [torch.zeros_like(f[:, 0:1])]])
+        # ... add BC blocks the same way ...
+        return A, b
+
+result = solve_linear(Stokes2D(), scale=5.0)
+u = result["u_fn"](x_test)        # shape (M, 3): columns are (u, v, p)
+```
+
+#### Partial derivatives for a vector u
+
+The basis-level operators (`basis.derivative`, `basis.gradient`,
+`basis.laplacian`, `DiffOperator.apply`) all return shape `(M, N)` regardless
+of how many components `u` has — vector-ness only enters when you contract
+with `beta`:
+
+```python
+# Full Jacobian, then slice (M, d, k) -> per (component, dim)
+u, J = solver.predict_with_grad(x)   # J shape (M, d, k); J[:, j, c] = ∂u_c/∂x_j
+
+# Single operator on a single component
+D_y = solver.basis.derivative(x, alpha=(0, 1))   # (M, N): ∂φ/∂y
+du0_dy = D_y @ solver.beta[:, 0:1]               # ∂u_0/∂y
+
+# Symbolic operator, all components at once
+from fastlsq import Op
+yy = Op.partial(dim=1, order=2, d=2)
+A  = yy.apply(solver.basis, x)                   # (M, N)
+u_yy = A @ solver.beta                           # (M, k): ∂²u/∂y² per component
+```
+
+Scalar problems are untouched: `n_outputs` defaults to `1`, `solver.beta` keeps
+shape `(N, 1)`, and `predict_with_grad` returns gradient shape `(M, d)` for
+backward compatibility (the trailing component axis is squeezed when k=1).
+`ElasticWave2D` in [fastlsq/problems/linear.py](fastlsq/problems/linear.py) is
+the canonical coupled vector example.
+
 ### Plot solutions
 
 ```python
@@ -177,6 +257,7 @@ derivative engine:
 | `FeatureBasis` | Adapter for non-sinusoidal solvers (e.g. PIELM with tanh) |
 | `FastLSQSolver` | Manages feature blocks; exposes `.basis` for all derivative computations |
 | `LearnableFastLSQ` | Differentiable solver with learnable bandwidth via reparameterisation trick |
+| `block_concat`, `pack_beta`, `unpack_beta` | Block-structured assembly helpers for vector-valued **u** (coupled systems). `solver.beta` has shape `(N, k)`; scalar problems are the k=1 case |
 
 ### How it works
 
@@ -253,6 +334,7 @@ See `examples/add_your_own_pde.py` for the complete tutorial.
 
 - **Analytical derivative engine**: `SinusoidalBasis` computes arbitrary-order derivatives exactly in O(1) -- the foundation of the entire framework
 - **Symbolic PDE operators**: Compose differential operators with `Op` (Laplacian, wave, Helmholtz, biharmonic, custom) via intuitive arithmetic; coefficients can be `nn.Parameter` for AdamW optimisation
+- **Vector-valued solutions**: First-class support for **u**: ℝᵈ → ℝᵏ (elasticity, Stokes, Maxwell). Problems declare `n_outputs = k`; `block_concat` assembles coupled block systems; `solver.predict(x)` returns shape `(M, k)`. Scalar problems are the `k=1` case
 - **High-level API**: Solve PDEs in one line with `solve_linear()` and `solve_nonlinear()`
 - **Learnable bandwidth**: `LearnableFastLSQ` optimises the bandwidth (scalar or anisotropic) via reparameterisation
 - **Learnable PDE coefficients**: Plug `nn.Parameter` into `Op` (e.g. Helmholtz wavenumber `k`) and optimise via AdamW; gradients flow through the prebuilt linear solve

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.4 → fastlsq-0.2.1}/FastLSQ.egg-info/requires.txt

@@ -2,6 +2,11 @@ torch>=2.0
 numpy>=1.24
 matplotlib>=3.7
 
+[battery]
+progpy>=1.3
+pandas>=2.0
+scipy>=1.10
+
 [dev]
 pytest>=7.0
 pandas>=2.0

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

@@ -1,7 +1,7 @@
 Metadata-Version: 2.4
 Name: FastLSQ
-Version: 0.1.4
-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,17 +14,22 @@ 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
 Requires-Dist: numpy>=1.24
 Requires-Dist: matplotlib>=3.7
+Provides-Extra: battery
+Requires-Dist: progpy>=1.3; extra == "battery"
+Requires-Dist: pandas>=2.0; extra == "battery"
+Requires-Dist: scipy>=1.10; extra == "battery"
 Provides-Extra: dev
 Requires-Dist: pytest>=7.0; extra == "dev"
 Requires-Dist: pandas>=2.0; extra == "dev"
@@ -36,8 +41,11 @@ Dynamic: license-file
 
 # FastLSQ
 
+[BerkeleyLab ATAP Talk](https://github.com/sulcantonin/FastLSQ/raw/main/presentations/ATAP_Sulc_20260324.pptx)
+
+
 <p align="center">
-  <img src="misc/fastlsq_teaser.png" alt="FastLSQ method overview" width="800"/>
+  <img src="misc/fastlsq_teaser.png" alt="FastLSQ method overview" width="400"/>
 </p>
 
 **Solving PDEs in one shot via Fourier features with exact analytical derivatives.**
@@ -128,6 +136,78 @@ A_pde = helmholtz.apply(basis, x)    # (5000, 1500)
 wave = Op.partial(dim=2, order=2, d=3) - c**2 * Op.laplacian(d=3, dims=[0, 1])
 ```
 
+### Vector-valued solutions
+
+`solve_linear` / `solve_nonlinear` support vector-valued **u**: ℝᵈ → ℝᵏ for
+coupled systems (elasticity, Stokes, Maxwell vector potential, …) and for
+decoupled multi-output problems sharing one basis. The math is unchanged; the
+solver just allocates `beta` with shape `(N, k)` so that `solver.predict(x)`
+returns shape `(M, k)` directly.
+
+A problem opts in by setting `self.n_outputs = k` and assembling its operator
+in block-stacked form `A ∈ ℝ^{Mk × Nk}`, `b ∈ ℝ^{Mk × 1}`. The helper
+`block_concat` removes the manual `torch.cat` bookkeeping:
+
+```python
+import torch
+from fastlsq import solve_linear, block_concat
+
+class Stokes2D:
+    n_outputs = 3        # (u, v, p)
+    dim = 2
+    name = "Stokes 2D"
+    # ... exact, exact_grad, get_train_data, get_test_points ...
+
+    def build(self, slv, x, bcs, f):
+        basis = slv.basis
+        cache = basis.cache(x)
+        dx = basis.derivative(x, (1, 0), cache=cache)
+        dy = basis.derivative(x, (0, 1), cache=cache)
+        lap = basis.laplacian(x, cache=cache)
+
+        # Rows = equations (mom_x, mom_y, continuity);
+        # columns = coefficient blocks (u, v, p)
+        A = block_concat([
+            [-lap,  None,  dx  ],   # -Δu + ∂p/∂x = f_x
+            [ None, -lap,  dy  ],   # -Δv + ∂p/∂y = f_y
+            [ dx,   dy,    None],   #  ∂u/∂x + ∂v/∂y = 0
+        ])
+        b = block_concat([[f[:, 0:1]], [f[:, 1:2]], [torch.zeros_like(f[:, 0:1])]])
+        # ... add BC blocks the same way ...
+        return A, b
+
+result = solve_linear(Stokes2D(), scale=5.0)
+u = result["u_fn"](x_test)        # shape (M, 3): columns are (u, v, p)
+```
+
+#### Partial derivatives for a vector u
+
+The basis-level operators (`basis.derivative`, `basis.gradient`,
+`basis.laplacian`, `DiffOperator.apply`) all return shape `(M, N)` regardless
+of how many components `u` has — vector-ness only enters when you contract
+with `beta`:
+
+```python
+# Full Jacobian, then slice (M, d, k) -> per (component, dim)
+u, J = solver.predict_with_grad(x)   # J shape (M, d, k); J[:, j, c] = ∂u_c/∂x_j
+
+# Single operator on a single component
+D_y = solver.basis.derivative(x, alpha=(0, 1))   # (M, N): ∂φ/∂y
+du0_dy = D_y @ solver.beta[:, 0:1]               # ∂u_0/∂y
+
+# Symbolic operator, all components at once
+from fastlsq import Op
+yy = Op.partial(dim=1, order=2, d=2)
+A  = yy.apply(solver.basis, x)                   # (M, N)
+u_yy = A @ solver.beta                           # (M, k): ∂²u/∂y² per component
+```
+
+Scalar problems are untouched: `n_outputs` defaults to `1`, `solver.beta` keeps
+shape `(N, 1)`, and `predict_with_grad` returns gradient shape `(M, d)` for
+backward compatibility (the trailing component axis is squeezed when k=1).
+`ElasticWave2D` in [fastlsq/problems/linear.py](fastlsq/problems/linear.py) is
+the canonical coupled vector example.
+
 ### Plot solutions
 
 ```python
@@ -177,6 +257,7 @@ derivative engine:
 | `FeatureBasis` | Adapter for non-sinusoidal solvers (e.g. PIELM with tanh) |
 | `FastLSQSolver` | Manages feature blocks; exposes `.basis` for all derivative computations |
 | `LearnableFastLSQ` | Differentiable solver with learnable bandwidth via reparameterisation trick |
+| `block_concat`, `pack_beta`, `unpack_beta` | Block-structured assembly helpers for vector-valued **u** (coupled systems). `solver.beta` has shape `(N, k)`; scalar problems are the k=1 case |
 
 ### How it works
 
@@ -253,6 +334,7 @@ See `examples/add_your_own_pde.py` for the complete tutorial.
 
 - **Analytical derivative engine**: `SinusoidalBasis` computes arbitrary-order derivatives exactly in O(1) -- the foundation of the entire framework
 - **Symbolic PDE operators**: Compose differential operators with `Op` (Laplacian, wave, Helmholtz, biharmonic, custom) via intuitive arithmetic; coefficients can be `nn.Parameter` for AdamW optimisation
+- **Vector-valued solutions**: First-class support for **u**: ℝᵈ → ℝᵏ (elasticity, Stokes, Maxwell). Problems declare `n_outputs = k`; `block_concat` assembles coupled block systems; `solver.predict(x)` returns shape `(M, k)`. Scalar problems are the `k=1` case
 - **High-level API**: Solve PDEs in one line with `solve_linear()` and `solve_nonlinear()`
 - **Learnable bandwidth**: `LearnableFastLSQ` optimises the bandwidth (scalar or anisotropic) via reparameterisation
 - **Learnable PDE coefficients**: Plug `nn.Parameter` into `Op` (e.g. Helmholtz wavenumber `k`) and optimise via AdamW; gradients flow through the prebuilt linear solve

{fastlsq-0.1.4 → fastlsq-0.2.1}/README.md

@@ -1,7 +1,10 @@
 # FastLSQ
 
+[BerkeleyLab ATAP Talk](https://github.com/sulcantonin/FastLSQ/raw/main/presentations/ATAP_Sulc_20260324.pptx)
+
+
 <p align="center">
-  <img src="misc/fastlsq_teaser.png" alt="FastLSQ method overview" width="800"/>
+  <img src="misc/fastlsq_teaser.png" alt="FastLSQ method overview" width="400"/>
 </p>
 
 **Solving PDEs in one shot via Fourier features with exact analytical derivatives.**
@@ -92,6 +95,78 @@ A_pde = helmholtz.apply(basis, x)    # (5000, 1500)
 wave = Op.partial(dim=2, order=2, d=3) - c**2 * Op.laplacian(d=3, dims=[0, 1])
 ```
 
+### Vector-valued solutions
+
+`solve_linear` / `solve_nonlinear` support vector-valued **u**: ℝᵈ → ℝᵏ for
+coupled systems (elasticity, Stokes, Maxwell vector potential, …) and for
+decoupled multi-output problems sharing one basis. The math is unchanged; the
+solver just allocates `beta` with shape `(N, k)` so that `solver.predict(x)`
+returns shape `(M, k)` directly.
+
+A problem opts in by setting `self.n_outputs = k` and assembling its operator
+in block-stacked form `A ∈ ℝ^{Mk × Nk}`, `b ∈ ℝ^{Mk × 1}`. The helper
+`block_concat` removes the manual `torch.cat` bookkeeping:
+
+```python
+import torch
+from fastlsq import solve_linear, block_concat
+
+class Stokes2D:
+    n_outputs = 3        # (u, v, p)
+    dim = 2
+    name = "Stokes 2D"
+    # ... exact, exact_grad, get_train_data, get_test_points ...
+
+    def build(self, slv, x, bcs, f):
+        basis = slv.basis
+        cache = basis.cache(x)
+        dx = basis.derivative(x, (1, 0), cache=cache)
+        dy = basis.derivative(x, (0, 1), cache=cache)
+        lap = basis.laplacian(x, cache=cache)
+
+        # Rows = equations (mom_x, mom_y, continuity);
+        # columns = coefficient blocks (u, v, p)
+        A = block_concat([
+            [-lap,  None,  dx  ],   # -Δu + ∂p/∂x = f_x
+            [ None, -lap,  dy  ],   # -Δv + ∂p/∂y = f_y
+            [ dx,   dy,    None],   #  ∂u/∂x + ∂v/∂y = 0
+        ])
+        b = block_concat([[f[:, 0:1]], [f[:, 1:2]], [torch.zeros_like(f[:, 0:1])]])
+        # ... add BC blocks the same way ...
+        return A, b
+
+result = solve_linear(Stokes2D(), scale=5.0)
+u = result["u_fn"](x_test)        # shape (M, 3): columns are (u, v, p)
+```
+
+#### Partial derivatives for a vector u
+
+The basis-level operators (`basis.derivative`, `basis.gradient`,
+`basis.laplacian`, `DiffOperator.apply`) all return shape `(M, N)` regardless
+of how many components `u` has — vector-ness only enters when you contract
+with `beta`:
+
+```python
+# Full Jacobian, then slice (M, d, k) -> per (component, dim)
+u, J = solver.predict_with_grad(x)   # J shape (M, d, k); J[:, j, c] = ∂u_c/∂x_j
+
+# Single operator on a single component
+D_y = solver.basis.derivative(x, alpha=(0, 1))   # (M, N): ∂φ/∂y
+du0_dy = D_y @ solver.beta[:, 0:1]               # ∂u_0/∂y
+
+# Symbolic operator, all components at once
+from fastlsq import Op
+yy = Op.partial(dim=1, order=2, d=2)
+A  = yy.apply(solver.basis, x)                   # (M, N)
+u_yy = A @ solver.beta                           # (M, k): ∂²u/∂y² per component
+```
+
+Scalar problems are untouched: `n_outputs` defaults to `1`, `solver.beta` keeps
+shape `(N, 1)`, and `predict_with_grad` returns gradient shape `(M, d)` for
+backward compatibility (the trailing component axis is squeezed when k=1).
+`ElasticWave2D` in [fastlsq/problems/linear.py](fastlsq/problems/linear.py) is
+the canonical coupled vector example.
+
 ### Plot solutions
 
 ```python
@@ -141,6 +216,7 @@ derivative engine:
 | `FeatureBasis` | Adapter for non-sinusoidal solvers (e.g. PIELM with tanh) |
 | `FastLSQSolver` | Manages feature blocks; exposes `.basis` for all derivative computations |
 | `LearnableFastLSQ` | Differentiable solver with learnable bandwidth via reparameterisation trick |
+| `block_concat`, `pack_beta`, `unpack_beta` | Block-structured assembly helpers for vector-valued **u** (coupled systems). `solver.beta` has shape `(N, k)`; scalar problems are the k=1 case |
 
 ### How it works
 
@@ -217,6 +293,7 @@ See `examples/add_your_own_pde.py` for the complete tutorial.
 
 - **Analytical derivative engine**: `SinusoidalBasis` computes arbitrary-order derivatives exactly in O(1) -- the foundation of the entire framework
 - **Symbolic PDE operators**: Compose differential operators with `Op` (Laplacian, wave, Helmholtz, biharmonic, custom) via intuitive arithmetic; coefficients can be `nn.Parameter` for AdamW optimisation
+- **Vector-valued solutions**: First-class support for **u**: ℝᵈ → ℝᵏ (elasticity, Stokes, Maxwell). Problems declare `n_outputs = k`; `block_concat` assembles coupled block systems; `solver.predict(x)` returns shape `(M, k)`. Scalar problems are the `k=1` case
 - **High-level API**: Solve PDEs in one line with `solve_linear()` and `solve_nonlinear()`
 - **Learnable bandwidth**: `LearnableFastLSQ` optimises the bandwidth (scalar or anisotropic) via reparameterisation
 - **Learnable PDE coefficients**: Plug `nn.Parameter` into `Op` (e.g. Helmholtz wavenumber `k`) and optimise via AdamW; gradients flow through the prebuilt linear solve