FastLSQ 0.1.4__tar.gz → 0.1.5__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (79) hide show
  1. {fastlsq-0.1.4 → fastlsq-0.1.5}/CHANGELOG.md +18 -0
  2. {fastlsq-0.1.4 → fastlsq-0.1.5/FastLSQ.egg-info}/PKG-INFO +83 -2
  3. {fastlsq-0.1.4 → fastlsq-0.1.5}/FastLSQ.egg-info/SOURCES.txt +12 -0
  4. {fastlsq-0.1.4 → fastlsq-0.1.5}/FastLSQ.egg-info/requires.txt +5 -0
  5. {fastlsq-0.1.4/FastLSQ.egg-info → fastlsq-0.1.5}/PKG-INFO +83 -2
  6. {fastlsq-0.1.4 → fastlsq-0.1.5}/README.md +78 -1
  7. fastlsq-0.1.5/examples/digital_twins/darcy_heat.py +583 -0
  8. fastlsq-0.1.5/examples/digital_twins/pendulum.py +242 -0
  9. fastlsq-0.1.5/examples/digital_twins/pendulum_benchmark.py +281 -0
  10. fastlsq-0.1.5/examples/digital_twins/plasma_wakefield.py +339 -0
  11. fastlsq-0.1.5/examples/digital_twins/plasma_wakefield_2D_1.py +863 -0
  12. fastlsq-0.1.5/examples/digital_twins/plasma_wakefield_2D_2.py +769 -0
  13. fastlsq-0.1.5/examples/digital_twins/plasma_wakefield_2d_3.py +818 -0
  14. fastlsq-0.1.5/examples/digital_twins/plasma_wakefield_parameteric.py +475 -0
  15. fastlsq-0.1.5/examples/digital_twins/plot_utils.py +42 -0
  16. fastlsq-0.1.5/examples/digital_twins/structural_health_simple.py +343 -0
  17. fastlsq-0.1.5/examples/digital_twins/turbulence_gravity_cooling.py +387 -0
  18. {fastlsq-0.1.4 → fastlsq-0.1.5}/fastlsq/__init__.py +5 -0
  19. {fastlsq-0.1.4 → fastlsq-0.1.5}/fastlsq/api.py +8 -2
  20. fastlsq-0.1.5/fastlsq/block.py +105 -0
  21. {fastlsq-0.1.4 → fastlsq-0.1.5}/fastlsq/learnable.py +33 -7
  22. {fastlsq-0.1.4 → fastlsq-0.1.5}/fastlsq/newton.py +22 -6
  23. {fastlsq-0.1.4 → fastlsq-0.1.5}/fastlsq/problems/linear.py +55 -29
  24. {fastlsq-0.1.4 → fastlsq-0.1.5}/fastlsq/solvers.py +32 -6
  25. {fastlsq-0.1.4 → fastlsq-0.1.5}/fastlsq/tuning.py +9 -2
  26. {fastlsq-0.1.4 → fastlsq-0.1.5}/pyproject.toml +6 -1
  27. {fastlsq-0.1.4 → fastlsq-0.1.5}/requirements.txt +1 -0
  28. {fastlsq-0.1.4 → fastlsq-0.1.5}/FastLSQ.egg-info/dependency_links.txt +0 -0
  29. {fastlsq-0.1.4 → fastlsq-0.1.5}/FastLSQ.egg-info/top_level.txt +0 -0
  30. {fastlsq-0.1.4 → fastlsq-0.1.5}/LICENSE +0 -0
  31. {fastlsq-0.1.4 → fastlsq-0.1.5}/MANIFEST.in +0 -0
  32. {fastlsq-0.1.4 → fastlsq-0.1.5}/examples/add_your_own_pde.py +0 -0
  33. {fastlsq-0.1.4 → fastlsq-0.1.5}/examples/benchmark_comparison.py +0 -0
  34. {fastlsq-0.1.4 → fastlsq-0.1.5}/examples/custom_features.py +0 -0
  35. {fastlsq-0.1.4 → fastlsq-0.1.5}/examples/inverse/aero_.py +0 -0
  36. {fastlsq-0.1.4 → fastlsq-0.1.5}/examples/inverse/denoising_parameter_estimation.py +0 -0
  37. {fastlsq-0.1.4 → fastlsq-0.1.5}/examples/inverse/elastic_wave_animation.py +0 -0
  38. {fastlsq-0.1.4 → fastlsq-0.1.5}/examples/inverse/heat_from_video.py +0 -0
  39. {fastlsq-0.1.4 → fastlsq-0.1.5}/examples/inverse/inverse_turbulence.py +0 -0
  40. {fastlsq-0.1.4 → fastlsq-0.1.5}/examples/inverse/shape_ns.py +0 -0
  41. {fastlsq-0.1.4 → fastlsq-0.1.5}/examples/inverse/subsurface_imaging.py +0 -0
  42. {fastlsq-0.1.4 → fastlsq-0.1.5}/examples/inverse/wing_optimize_simple.py +0 -0
  43. {fastlsq-0.1.4 → fastlsq-0.1.5}/examples/inverse_heat_source.py +0 -0
  44. {fastlsq-0.1.4 → fastlsq-0.1.5}/examples/inverse_magnetostatics.py +0 -0
  45. {fastlsq-0.1.4 → fastlsq-0.1.5}/examples/inverse_source_position.py +0 -0
  46. {fastlsq-0.1.4 → fastlsq-0.1.5}/examples/learnable_helmholtz.py +0 -0
  47. {fastlsq-0.1.4 → fastlsq-0.1.5}/examples/pde_discovery.py +0 -0
  48. {fastlsq-0.1.4 → fastlsq-0.1.5}/examples/run_all_extensions.py +0 -0
  49. {fastlsq-0.1.4 → fastlsq-0.1.5}/examples/run_linear.py +0 -0
  50. {fastlsq-0.1.4 → fastlsq-0.1.5}/examples/run_nonlinear.py +0 -0
  51. {fastlsq-0.1.4 → fastlsq-0.1.5}/examples/sindy/compare_sindy_methods.py +0 -0
  52. {fastlsq-0.1.4 → fastlsq-0.1.5}/examples/sindy/sindy_benchmarks.py +0 -0
  53. {fastlsq-0.1.4 → fastlsq-0.1.5}/examples/sindy/sindy_differentiable.py +0 -0
  54. {fastlsq-0.1.4 → fastlsq-0.1.5}/examples/sindy/sindy_minimal_diff.py +0 -0
  55. {fastlsq-0.1.4 → fastlsq-0.1.5}/examples/tutorial_basic.py +0 -0
  56. {fastlsq-0.1.4 → fastlsq-0.1.5}/examples/tutorial_nonlinear.py +0 -0
  57. {fastlsq-0.1.4 → fastlsq-0.1.5}/fastlsq/basis.py +0 -0
  58. {fastlsq-0.1.4 → fastlsq-0.1.5}/fastlsq/diagnostics.py +0 -0
  59. {fastlsq-0.1.4 → fastlsq-0.1.5}/fastlsq/export.py +0 -0
  60. {fastlsq-0.1.4 → fastlsq-0.1.5}/fastlsq/geometry.py +0 -0
  61. {fastlsq-0.1.4 → fastlsq-0.1.5}/fastlsq/lightning.py +0 -0
  62. {fastlsq-0.1.4 → fastlsq-0.1.5}/fastlsq/linalg.py +0 -0
  63. {fastlsq-0.1.4 → fastlsq-0.1.5}/fastlsq/plotting.py +0 -0
  64. {fastlsq-0.1.4 → fastlsq-0.1.5}/fastlsq/problems/__init__.py +0 -0
  65. {fastlsq-0.1.4 → fastlsq-0.1.5}/fastlsq/problems/nonlinear.py +0 -0
  66. {fastlsq-0.1.4 → fastlsq-0.1.5}/fastlsq/problems/regression.py +0 -0
  67. {fastlsq-0.1.4 → fastlsq-0.1.5}/fastlsq/utils.py +0 -0
  68. {fastlsq-0.1.4 → fastlsq-0.1.5}/misc/fastlsq_teaser.png +0 -0
  69. {fastlsq-0.1.4 → fastlsq-0.1.5}/misc/ideal_quadrupole.png +0 -0
  70. {fastlsq-0.1.4 → fastlsq-0.1.5}/misc/inverse_heat_source.gif +0 -0
  71. {fastlsq-0.1.4 → fastlsq-0.1.5}/misc/inverse_heat_source.png +0 -0
  72. {fastlsq-0.1.4 → fastlsq-0.1.5}/misc/inverse_magnetostatics.png +0 -0
  73. {fastlsq-0.1.4 → fastlsq-0.1.5}/misc/inverse_magnetostatics_convergence.png +0 -0
  74. {fastlsq-0.1.4 → fastlsq-0.1.5}/misc/quadrupole_convergence.png +0 -0
  75. {fastlsq-0.1.4 → fastlsq-0.1.5}/misc/quadrupole_optimization.png +0 -0
  76. {fastlsq-0.1.4 → fastlsq-0.1.5}/misc/tutorial_nlpoisson_convergence.png +0 -0
  77. {fastlsq-0.1.4 → fastlsq-0.1.5}/misc/tutorial_nlpoisson_solution.png +0 -0
  78. {fastlsq-0.1.4 → fastlsq-0.1.5}/setup.cfg +0 -0
  79. {fastlsq-0.1.4 → fastlsq-0.1.5}/tests/test_basic.py +0 -0
@@ -6,6 +6,24 @@ All notable changes to FastLSQ will be documented in this file.
6
6
 
7
7
  ### Added
8
8
 
9
+ - **Vector-valued solutions (`u: ℝᵈ → ℝᵏ`)**: first-class support for coupled
10
+ systems and decoupled multi-output PDEs. Problems opt in with
11
+ `self.n_outputs = k`; `solver.beta` has shape `(N, k)` and
12
+ `solver.predict(x)` returns `(M, k)`. Scalar problems are the `k=1` case
13
+ and remain bit-for-bit identical (Helmholtz 2D / Poisson 5D / Bratu 2D
14
+ regressions verified).
15
+ - New module `fastlsq.block`: `block_concat` assembles a 2-D nested list
16
+ of `(M_i, N_j)` tensors into a block matrix (with `None` for zero
17
+ blocks); `pack_beta` / `unpack_beta` convert between `(N, k)` and the
18
+ block-stacked `(N*k, 1)` solve representation.
19
+ - `predict_with_grad` / `predict_with_laplacian` einsum fixed to
20
+ `"idh,hk->idk"` so the output dim no longer collapses for `k>1`. Squeezes
21
+ back to `(M, d)` when `k=1` for backward compatibility.
22
+ - `ElasticWave2D` refactored: gains `n_outputs = 2`, `build()` uses
23
+ `block_concat`, and now exposes `exact_grad` with shape `(M, d, k)`.
24
+ - `LearnableFastLSQ` accepts `n_outputs`; legacy subclasses that store a
25
+ flat `(Nk, 1)` `beta` (e.g. `ElasticLearnable` in the elastic wave
26
+ example) keep working under the default `n_outputs=1`.
9
27
  - **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`.
10
28
 
11
29
  ## [0.2.0] - 2026-03-01
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: FastLSQ
3
- Version: 0.1.4
3
+ Version: 0.1.5
4
4
  Summary: Solving PDEs in one shot via Fourier features with exact analytical derivatives
5
5
  Author: Antonin Sulc
6
6
  License-Expression: MIT
@@ -25,6 +25,10 @@ License-File: LICENSE
25
25
  Requires-Dist: torch>=2.0
26
26
  Requires-Dist: numpy>=1.24
27
27
  Requires-Dist: matplotlib>=3.7
28
+ Provides-Extra: battery
29
+ Requires-Dist: progpy>=1.3; extra == "battery"
30
+ Requires-Dist: pandas>=2.0; extra == "battery"
31
+ Requires-Dist: scipy>=1.10; extra == "battery"
28
32
  Provides-Extra: dev
29
33
  Requires-Dist: pytest>=7.0; extra == "dev"
30
34
  Requires-Dist: pandas>=2.0; extra == "dev"
@@ -36,8 +40,11 @@ Dynamic: license-file
36
40
 
37
41
  # FastLSQ
38
42
 
43
+ [BerkeleyLab ATAP Talk](https://github.com/sulcantonin/FastLSQ/raw/main/presentations/ATAP_Sulc_20260324.pptx)
44
+
45
+
39
46
  <p align="center">
40
- <img src="misc/fastlsq_teaser.png" alt="FastLSQ method overview" width="800"/>
47
+ <img src="misc/fastlsq_teaser.png" alt="FastLSQ method overview" width="400"/>
41
48
  </p>
42
49
 
43
50
  **Solving PDEs in one shot via Fourier features with exact analytical derivatives.**
@@ -128,6 +135,78 @@ A_pde = helmholtz.apply(basis, x) # (5000, 1500)
128
135
  wave = Op.partial(dim=2, order=2, d=3) - c**2 * Op.laplacian(d=3, dims=[0, 1])
129
136
  ```
130
137
 
138
+ ### Vector-valued solutions
139
+
140
+ `solve_linear` / `solve_nonlinear` support vector-valued **u**: ℝᵈ → ℝᵏ for
141
+ coupled systems (elasticity, Stokes, Maxwell vector potential, …) and for
142
+ decoupled multi-output problems sharing one basis. The math is unchanged; the
143
+ solver just allocates `beta` with shape `(N, k)` so that `solver.predict(x)`
144
+ returns shape `(M, k)` directly.
145
+
146
+ A problem opts in by setting `self.n_outputs = k` and assembling its operator
147
+ in block-stacked form `A ∈ ℝ^{Mk × Nk}`, `b ∈ ℝ^{Mk × 1}`. The helper
148
+ `block_concat` removes the manual `torch.cat` bookkeeping:
149
+
150
+ ```python
151
+ import torch
152
+ from fastlsq import solve_linear, block_concat
153
+
154
+ class Stokes2D:
155
+ n_outputs = 3 # (u, v, p)
156
+ dim = 2
157
+ name = "Stokes 2D"
158
+ # ... exact, exact_grad, get_train_data, get_test_points ...
159
+
160
+ def build(self, slv, x, bcs, f):
161
+ basis = slv.basis
162
+ cache = basis.cache(x)
163
+ dx = basis.derivative(x, (1, 0), cache=cache)
164
+ dy = basis.derivative(x, (0, 1), cache=cache)
165
+ lap = basis.laplacian(x, cache=cache)
166
+
167
+ # Rows = equations (mom_x, mom_y, continuity);
168
+ # columns = coefficient blocks (u, v, p)
169
+ A = block_concat([
170
+ [-lap, None, dx ], # -Δu + ∂p/∂x = f_x
171
+ [ None, -lap, dy ], # -Δv + ∂p/∂y = f_y
172
+ [ dx, dy, None], # ∂u/∂x + ∂v/∂y = 0
173
+ ])
174
+ b = block_concat([[f[:, 0:1]], [f[:, 1:2]], [torch.zeros_like(f[:, 0:1])]])
175
+ # ... add BC blocks the same way ...
176
+ return A, b
177
+
178
+ result = solve_linear(Stokes2D(), scale=5.0)
179
+ u = result["u_fn"](x_test) # shape (M, 3): columns are (u, v, p)
180
+ ```
181
+
182
+ #### Partial derivatives for a vector u
183
+
184
+ The basis-level operators (`basis.derivative`, `basis.gradient`,
185
+ `basis.laplacian`, `DiffOperator.apply`) all return shape `(M, N)` regardless
186
+ of how many components `u` has — vector-ness only enters when you contract
187
+ with `beta`:
188
+
189
+ ```python
190
+ # Full Jacobian, then slice (M, d, k) -> per (component, dim)
191
+ u, J = solver.predict_with_grad(x) # J shape (M, d, k); J[:, j, c] = ∂u_c/∂x_j
192
+
193
+ # Single operator on a single component
194
+ D_y = solver.basis.derivative(x, alpha=(0, 1)) # (M, N): ∂φ/∂y
195
+ du0_dy = D_y @ solver.beta[:, 0:1] # ∂u_0/∂y
196
+
197
+ # Symbolic operator, all components at once
198
+ from fastlsq import Op
199
+ yy = Op.partial(dim=1, order=2, d=2)
200
+ A = yy.apply(solver.basis, x) # (M, N)
201
+ u_yy = A @ solver.beta # (M, k): ∂²u/∂y² per component
202
+ ```
203
+
204
+ Scalar problems are untouched: `n_outputs` defaults to `1`, `solver.beta` keeps
205
+ shape `(N, 1)`, and `predict_with_grad` returns gradient shape `(M, d)` for
206
+ backward compatibility (the trailing component axis is squeezed when k=1).
207
+ `ElasticWave2D` in [fastlsq/problems/linear.py](fastlsq/problems/linear.py) is
208
+ the canonical coupled vector example.
209
+
131
210
  ### Plot solutions
132
211
 
133
212
  ```python
@@ -177,6 +256,7 @@ derivative engine:
177
256
  | `FeatureBasis` | Adapter for non-sinusoidal solvers (e.g. PIELM with tanh) |
178
257
  | `FastLSQSolver` | Manages feature blocks; exposes `.basis` for all derivative computations |
179
258
  | `LearnableFastLSQ` | Differentiable solver with learnable bandwidth via reparameterisation trick |
259
+ | `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 |
180
260
 
181
261
  ### How it works
182
262
 
@@ -253,6 +333,7 @@ See `examples/add_your_own_pde.py` for the complete tutorial.
253
333
 
254
334
  - **Analytical derivative engine**: `SinusoidalBasis` computes arbitrary-order derivatives exactly in O(1) -- the foundation of the entire framework
255
335
  - **Symbolic PDE operators**: Compose differential operators with `Op` (Laplacian, wave, Helmholtz, biharmonic, custom) via intuitive arithmetic; coefficients can be `nn.Parameter` for AdamW optimisation
336
+ - **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
256
337
  - **High-level API**: Solve PDEs in one line with `solve_linear()` and `solve_nonlinear()`
257
338
  - **Learnable bandwidth**: `LearnableFastLSQ` optimises the bandwidth (scalar or anisotropic) via reparameterisation
258
339
  - **Learnable PDE coefficients**: Plug `nn.Parameter` into `Op` (e.g. Helmholtz wavenumber `k`) and optimise via AdamW; gradients flow through the prebuilt linear solve
@@ -22,6 +22,17 @@ examples/run_linear.py
22
22
  examples/run_nonlinear.py
23
23
  examples/tutorial_basic.py
24
24
  examples/tutorial_nonlinear.py
25
+ examples/digital_twins/darcy_heat.py
26
+ examples/digital_twins/pendulum.py
27
+ examples/digital_twins/pendulum_benchmark.py
28
+ examples/digital_twins/plasma_wakefield.py
29
+ examples/digital_twins/plasma_wakefield_2D_1.py
30
+ examples/digital_twins/plasma_wakefield_2D_2.py
31
+ examples/digital_twins/plasma_wakefield_2d_3.py
32
+ examples/digital_twins/plasma_wakefield_parameteric.py
33
+ examples/digital_twins/plot_utils.py
34
+ examples/digital_twins/structural_health_simple.py
35
+ examples/digital_twins/turbulence_gravity_cooling.py
25
36
  examples/inverse/aero_.py
26
37
  examples/inverse/denoising_parameter_estimation.py
27
38
  examples/inverse/elastic_wave_animation.py
@@ -37,6 +48,7 @@ examples/sindy/sindy_minimal_diff.py
37
48
  fastlsq/__init__.py
38
49
  fastlsq/api.py
39
50
  fastlsq/basis.py
51
+ fastlsq/block.py
40
52
  fastlsq/diagnostics.py
41
53
  fastlsq/export.py
42
54
  fastlsq/geometry.py
@@ -2,6 +2,11 @@ torch>=2.0
2
2
  numpy>=1.24
3
3
  matplotlib>=3.7
4
4
 
5
+ [battery]
6
+ progpy>=1.3
7
+ pandas>=2.0
8
+ scipy>=1.10
9
+
5
10
  [dev]
6
11
  pytest>=7.0
7
12
  pandas>=2.0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: FastLSQ
3
- Version: 0.1.4
3
+ Version: 0.1.5
4
4
  Summary: Solving PDEs in one shot via Fourier features with exact analytical derivatives
5
5
  Author: Antonin Sulc
6
6
  License-Expression: MIT
@@ -25,6 +25,10 @@ License-File: LICENSE
25
25
  Requires-Dist: torch>=2.0
26
26
  Requires-Dist: numpy>=1.24
27
27
  Requires-Dist: matplotlib>=3.7
28
+ Provides-Extra: battery
29
+ Requires-Dist: progpy>=1.3; extra == "battery"
30
+ Requires-Dist: pandas>=2.0; extra == "battery"
31
+ Requires-Dist: scipy>=1.10; extra == "battery"
28
32
  Provides-Extra: dev
29
33
  Requires-Dist: pytest>=7.0; extra == "dev"
30
34
  Requires-Dist: pandas>=2.0; extra == "dev"
@@ -36,8 +40,11 @@ Dynamic: license-file
36
40
 
37
41
  # FastLSQ
38
42
 
43
+ [BerkeleyLab ATAP Talk](https://github.com/sulcantonin/FastLSQ/raw/main/presentations/ATAP_Sulc_20260324.pptx)
44
+
45
+
39
46
  <p align="center">
40
- <img src="misc/fastlsq_teaser.png" alt="FastLSQ method overview" width="800"/>
47
+ <img src="misc/fastlsq_teaser.png" alt="FastLSQ method overview" width="400"/>
41
48
  </p>
42
49
 
43
50
  **Solving PDEs in one shot via Fourier features with exact analytical derivatives.**
@@ -128,6 +135,78 @@ A_pde = helmholtz.apply(basis, x) # (5000, 1500)
128
135
  wave = Op.partial(dim=2, order=2, d=3) - c**2 * Op.laplacian(d=3, dims=[0, 1])
129
136
  ```
130
137
 
138
+ ### Vector-valued solutions
139
+
140
+ `solve_linear` / `solve_nonlinear` support vector-valued **u**: ℝᵈ → ℝᵏ for
141
+ coupled systems (elasticity, Stokes, Maxwell vector potential, …) and for
142
+ decoupled multi-output problems sharing one basis. The math is unchanged; the
143
+ solver just allocates `beta` with shape `(N, k)` so that `solver.predict(x)`
144
+ returns shape `(M, k)` directly.
145
+
146
+ A problem opts in by setting `self.n_outputs = k` and assembling its operator
147
+ in block-stacked form `A ∈ ℝ^{Mk × Nk}`, `b ∈ ℝ^{Mk × 1}`. The helper
148
+ `block_concat` removes the manual `torch.cat` bookkeeping:
149
+
150
+ ```python
151
+ import torch
152
+ from fastlsq import solve_linear, block_concat
153
+
154
+ class Stokes2D:
155
+ n_outputs = 3 # (u, v, p)
156
+ dim = 2
157
+ name = "Stokes 2D"
158
+ # ... exact, exact_grad, get_train_data, get_test_points ...
159
+
160
+ def build(self, slv, x, bcs, f):
161
+ basis = slv.basis
162
+ cache = basis.cache(x)
163
+ dx = basis.derivative(x, (1, 0), cache=cache)
164
+ dy = basis.derivative(x, (0, 1), cache=cache)
165
+ lap = basis.laplacian(x, cache=cache)
166
+
167
+ # Rows = equations (mom_x, mom_y, continuity);
168
+ # columns = coefficient blocks (u, v, p)
169
+ A = block_concat([
170
+ [-lap, None, dx ], # -Δu + ∂p/∂x = f_x
171
+ [ None, -lap, dy ], # -Δv + ∂p/∂y = f_y
172
+ [ dx, dy, None], # ∂u/∂x + ∂v/∂y = 0
173
+ ])
174
+ b = block_concat([[f[:, 0:1]], [f[:, 1:2]], [torch.zeros_like(f[:, 0:1])]])
175
+ # ... add BC blocks the same way ...
176
+ return A, b
177
+
178
+ result = solve_linear(Stokes2D(), scale=5.0)
179
+ u = result["u_fn"](x_test) # shape (M, 3): columns are (u, v, p)
180
+ ```
181
+
182
+ #### Partial derivatives for a vector u
183
+
184
+ The basis-level operators (`basis.derivative`, `basis.gradient`,
185
+ `basis.laplacian`, `DiffOperator.apply`) all return shape `(M, N)` regardless
186
+ of how many components `u` has — vector-ness only enters when you contract
187
+ with `beta`:
188
+
189
+ ```python
190
+ # Full Jacobian, then slice (M, d, k) -> per (component, dim)
191
+ u, J = solver.predict_with_grad(x) # J shape (M, d, k); J[:, j, c] = ∂u_c/∂x_j
192
+
193
+ # Single operator on a single component
194
+ D_y = solver.basis.derivative(x, alpha=(0, 1)) # (M, N): ∂φ/∂y
195
+ du0_dy = D_y @ solver.beta[:, 0:1] # ∂u_0/∂y
196
+
197
+ # Symbolic operator, all components at once
198
+ from fastlsq import Op
199
+ yy = Op.partial(dim=1, order=2, d=2)
200
+ A = yy.apply(solver.basis, x) # (M, N)
201
+ u_yy = A @ solver.beta # (M, k): ∂²u/∂y² per component
202
+ ```
203
+
204
+ Scalar problems are untouched: `n_outputs` defaults to `1`, `solver.beta` keeps
205
+ shape `(N, 1)`, and `predict_with_grad` returns gradient shape `(M, d)` for
206
+ backward compatibility (the trailing component axis is squeezed when k=1).
207
+ `ElasticWave2D` in [fastlsq/problems/linear.py](fastlsq/problems/linear.py) is
208
+ the canonical coupled vector example.
209
+
131
210
  ### Plot solutions
132
211
 
133
212
  ```python
@@ -177,6 +256,7 @@ derivative engine:
177
256
  | `FeatureBasis` | Adapter for non-sinusoidal solvers (e.g. PIELM with tanh) |
178
257
  | `FastLSQSolver` | Manages feature blocks; exposes `.basis` for all derivative computations |
179
258
  | `LearnableFastLSQ` | Differentiable solver with learnable bandwidth via reparameterisation trick |
259
+ | `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 |
180
260
 
181
261
  ### How it works
182
262
 
@@ -253,6 +333,7 @@ See `examples/add_your_own_pde.py` for the complete tutorial.
253
333
 
254
334
  - **Analytical derivative engine**: `SinusoidalBasis` computes arbitrary-order derivatives exactly in O(1) -- the foundation of the entire framework
255
335
  - **Symbolic PDE operators**: Compose differential operators with `Op` (Laplacian, wave, Helmholtz, biharmonic, custom) via intuitive arithmetic; coefficients can be `nn.Parameter` for AdamW optimisation
336
+ - **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
256
337
  - **High-level API**: Solve PDEs in one line with `solve_linear()` and `solve_nonlinear()`
257
338
  - **Learnable bandwidth**: `LearnableFastLSQ` optimises the bandwidth (scalar or anisotropic) via reparameterisation
258
339
  - **Learnable PDE coefficients**: Plug `nn.Parameter` into `Op` (e.g. Helmholtz wavenumber `k`) and optimise via AdamW; gradients flow through the prebuilt linear solve
@@ -1,7 +1,10 @@
1
1
  # FastLSQ
2
2
 
3
+ [BerkeleyLab ATAP Talk](https://github.com/sulcantonin/FastLSQ/raw/main/presentations/ATAP_Sulc_20260324.pptx)
4
+
5
+
3
6
  <p align="center">
4
- <img src="misc/fastlsq_teaser.png" alt="FastLSQ method overview" width="800"/>
7
+ <img src="misc/fastlsq_teaser.png" alt="FastLSQ method overview" width="400"/>
5
8
  </p>
6
9
 
7
10
  **Solving PDEs in one shot via Fourier features with exact analytical derivatives.**
@@ -92,6 +95,78 @@ A_pde = helmholtz.apply(basis, x) # (5000, 1500)
92
95
  wave = Op.partial(dim=2, order=2, d=3) - c**2 * Op.laplacian(d=3, dims=[0, 1])
93
96
  ```
94
97
 
98
+ ### Vector-valued solutions
99
+
100
+ `solve_linear` / `solve_nonlinear` support vector-valued **u**: ℝᵈ → ℝᵏ for
101
+ coupled systems (elasticity, Stokes, Maxwell vector potential, …) and for
102
+ decoupled multi-output problems sharing one basis. The math is unchanged; the
103
+ solver just allocates `beta` with shape `(N, k)` so that `solver.predict(x)`
104
+ returns shape `(M, k)` directly.
105
+
106
+ A problem opts in by setting `self.n_outputs = k` and assembling its operator
107
+ in block-stacked form `A ∈ ℝ^{Mk × Nk}`, `b ∈ ℝ^{Mk × 1}`. The helper
108
+ `block_concat` removes the manual `torch.cat` bookkeeping:
109
+
110
+ ```python
111
+ import torch
112
+ from fastlsq import solve_linear, block_concat
113
+
114
+ class Stokes2D:
115
+ n_outputs = 3 # (u, v, p)
116
+ dim = 2
117
+ name = "Stokes 2D"
118
+ # ... exact, exact_grad, get_train_data, get_test_points ...
119
+
120
+ def build(self, slv, x, bcs, f):
121
+ basis = slv.basis
122
+ cache = basis.cache(x)
123
+ dx = basis.derivative(x, (1, 0), cache=cache)
124
+ dy = basis.derivative(x, (0, 1), cache=cache)
125
+ lap = basis.laplacian(x, cache=cache)
126
+
127
+ # Rows = equations (mom_x, mom_y, continuity);
128
+ # columns = coefficient blocks (u, v, p)
129
+ A = block_concat([
130
+ [-lap, None, dx ], # -Δu + ∂p/∂x = f_x
131
+ [ None, -lap, dy ], # -Δv + ∂p/∂y = f_y
132
+ [ dx, dy, None], # ∂u/∂x + ∂v/∂y = 0
133
+ ])
134
+ b = block_concat([[f[:, 0:1]], [f[:, 1:2]], [torch.zeros_like(f[:, 0:1])]])
135
+ # ... add BC blocks the same way ...
136
+ return A, b
137
+
138
+ result = solve_linear(Stokes2D(), scale=5.0)
139
+ u = result["u_fn"](x_test) # shape (M, 3): columns are (u, v, p)
140
+ ```
141
+
142
+ #### Partial derivatives for a vector u
143
+
144
+ The basis-level operators (`basis.derivative`, `basis.gradient`,
145
+ `basis.laplacian`, `DiffOperator.apply`) all return shape `(M, N)` regardless
146
+ of how many components `u` has — vector-ness only enters when you contract
147
+ with `beta`:
148
+
149
+ ```python
150
+ # Full Jacobian, then slice (M, d, k) -> per (component, dim)
151
+ u, J = solver.predict_with_grad(x) # J shape (M, d, k); J[:, j, c] = ∂u_c/∂x_j
152
+
153
+ # Single operator on a single component
154
+ D_y = solver.basis.derivative(x, alpha=(0, 1)) # (M, N): ∂φ/∂y
155
+ du0_dy = D_y @ solver.beta[:, 0:1] # ∂u_0/∂y
156
+
157
+ # Symbolic operator, all components at once
158
+ from fastlsq import Op
159
+ yy = Op.partial(dim=1, order=2, d=2)
160
+ A = yy.apply(solver.basis, x) # (M, N)
161
+ u_yy = A @ solver.beta # (M, k): ∂²u/∂y² per component
162
+ ```
163
+
164
+ Scalar problems are untouched: `n_outputs` defaults to `1`, `solver.beta` keeps
165
+ shape `(N, 1)`, and `predict_with_grad` returns gradient shape `(M, d)` for
166
+ backward compatibility (the trailing component axis is squeezed when k=1).
167
+ `ElasticWave2D` in [fastlsq/problems/linear.py](fastlsq/problems/linear.py) is
168
+ the canonical coupled vector example.
169
+
95
170
  ### Plot solutions
96
171
 
97
172
  ```python
@@ -141,6 +216,7 @@ derivative engine:
141
216
  | `FeatureBasis` | Adapter for non-sinusoidal solvers (e.g. PIELM with tanh) |
142
217
  | `FastLSQSolver` | Manages feature blocks; exposes `.basis` for all derivative computations |
143
218
  | `LearnableFastLSQ` | Differentiable solver with learnable bandwidth via reparameterisation trick |
219
+ | `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 |
144
220
 
145
221
  ### How it works
146
222
 
@@ -217,6 +293,7 @@ See `examples/add_your_own_pde.py` for the complete tutorial.
217
293
 
218
294
  - **Analytical derivative engine**: `SinusoidalBasis` computes arbitrary-order derivatives exactly in O(1) -- the foundation of the entire framework
219
295
  - **Symbolic PDE operators**: Compose differential operators with `Op` (Laplacian, wave, Helmholtz, biharmonic, custom) via intuitive arithmetic; coefficients can be `nn.Parameter` for AdamW optimisation
296
+ - **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
220
297
  - **High-level API**: Solve PDEs in one line with `solve_linear()` and `solve_nonlinear()`
221
298
  - **Learnable bandwidth**: `LearnableFastLSQ` optimises the bandwidth (scalar or anisotropic) via reparameterisation
222
299
  - **Learnable PDE coefficients**: Plug `nn.Parameter` into `Op` (e.g. Helmholtz wavenumber `k`) and optimise via AdamW; gradients flow through the prebuilt linear solve