FastLSQ 0.4.0__tar.gz → 0.4.1__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 (153) hide show
  1. {fastlsq-0.4.0 → fastlsq-0.4.1}/CHANGELOG.md +29 -0
  2. {fastlsq-0.4.0 → fastlsq-0.4.1}/FastLSQ.egg-info/PKG-INFO +1 -1
  3. {fastlsq-0.4.0 → fastlsq-0.4.1}/FastLSQ.egg-info/SOURCES.txt +2 -0
  4. {fastlsq-0.4.0 → fastlsq-0.4.1}/PKG-INFO +1 -1
  5. {fastlsq-0.4.0 → fastlsq-0.4.1}/fastlsq/__init__.py +2 -1
  6. {fastlsq-0.4.0 → fastlsq-0.4.1}/fastlsq/api.py +31 -3
  7. {fastlsq-0.4.0 → fastlsq-0.4.1}/fastlsq/basis.py +37 -0
  8. fastlsq-0.4.1/fastlsq/benchmark.py +100 -0
  9. {fastlsq-0.4.0 → fastlsq-0.4.1}/fastlsq/export.py +34 -3
  10. {fastlsq-0.4.0 → fastlsq-0.4.1}/fastlsq/learnable.py +14 -0
  11. {fastlsq-0.4.0 → fastlsq-0.4.1}/fastlsq/linalg.py +84 -19
  12. {fastlsq-0.4.0 → fastlsq-0.4.1}/pyproject.toml +1 -1
  13. fastlsq-0.4.1/tests/test_diagnostics_timing.py +197 -0
  14. {fastlsq-0.4.0 → fastlsq-0.4.1}/tests/test_vector_basis.py +1 -1
  15. {fastlsq-0.4.0 → fastlsq-0.4.1}/FastLSQ.egg-info/dependency_links.txt +0 -0
  16. {fastlsq-0.4.0 → fastlsq-0.4.1}/FastLSQ.egg-info/requires.txt +0 -0
  17. {fastlsq-0.4.0 → fastlsq-0.4.1}/FastLSQ.egg-info/top_level.txt +0 -0
  18. {fastlsq-0.4.0 → fastlsq-0.4.1}/LICENSE +0 -0
  19. {fastlsq-0.4.0 → fastlsq-0.4.1}/MANIFEST.in +0 -0
  20. {fastlsq-0.4.0 → fastlsq-0.4.1}/README.md +0 -0
  21. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/add_your_own_pde.py +0 -0
  22. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/benchmark_comparison.py +0 -0
  23. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/custom_features.py +0 -0
  24. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/digital_twins/darcy_heat.py +0 -0
  25. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/digital_twins/pendulum.py +0 -0
  26. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/digital_twins/pendulum_benchmark.py +0 -0
  27. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/digital_twins/plasma_wakefield.py +0 -0
  28. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/digital_twins/plasma_wakefield_2D_1.py +0 -0
  29. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/digital_twins/plasma_wakefield_2D_2.py +0 -0
  30. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/digital_twins/plasma_wakefield_2d_3.py +0 -0
  31. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/digital_twins/plasma_wakefield_parameteric.py +0 -0
  32. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/digital_twins/plot_utils.py +0 -0
  33. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/digital_twins/structural_health_simple.py +0 -0
  34. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/digital_twins/turbulence_gravity_cooling.py +0 -0
  35. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/extras/fred_sde.py +0 -0
  36. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/extras/fred_sde_fastlsq.py +0 -0
  37. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/extras/gaia_potential.py +0 -0
  38. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/extras/gaia_potential_fastlsq.py +0 -0
  39. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/extras/horizons_ephemeris.py +0 -0
  40. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/extras/numerai_alpha.py +0 -0
  41. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/extras/numerai_alpha_fastlsq.py +0 -0
  42. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/extras/run_all_fastlsq.py +0 -0
  43. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/extras/scenarios/__init__.py +0 -0
  44. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/extras/scenarios/_alsu_lattice.py +0 -0
  45. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/extras/scenarios/_common.py +0 -0
  46. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/extras/scenarios/run_all.py +0 -0
  47. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/extras/scenarios/s01_beamloss_ode.py +0 -0
  48. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/extras/scenarios/s01_betatron_tune.py +0 -0
  49. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/extras/scenarios/s01_green_fff.py +0 -0
  50. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/extras/scenarios/s01_hill_ivp.py +0 -0
  51. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/extras/scenarios/s01_observe_fit_act_simulator.py +0 -0
  52. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/extras/scenarios/s01_orbit_inverse.py +0 -0
  53. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/extras/scenarios/s01_passive_loco.py +0 -0
  54. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/extras/scenarios/s01_perturbed_hill.py +0 -0
  55. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/extras/scenarios/s01_sofb_observe_fit_act.py +0 -0
  56. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/extras/scenarios/s01_streaming_archive_growth.py +0 -0
  57. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/extras/scenarios/s01_synchrotron_ode.py +0 -0
  58. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/extras/scenarios/s01_tides_3months.py +0 -0
  59. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/extras/scenarios/s01_topoff_impulse.py +0 -0
  60. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/extras/scenarios/s01_visualize.py +0 -0
  61. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/extras/scenarios/s02_plasma_wakefield.py +0 -0
  62. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/extras/scenarios/s03_synchrobetatron.py +0 -0
  63. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/extras/scenarios/s04_sunspots.py +0 -0
  64. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/extras/scenarios/s05_helioseismology.py +0 -0
  65. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/extras/scenarios/s06_tides.py +0 -0
  66. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/extras/scenarios/s07_iers_earth_rotation.py +0 -0
  67. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/extras/scenarios/s08_mauna_loa_co2.py +0 -0
  68. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/extras/scenarios/s09_enso_qbo.py +0 -0
  69. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/extras/scenarios/s10_pulsar_timing.py +0 -0
  70. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/extras/scenarios/s11_modal_analysis.py +0 -0
  71. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/extras/scenarios/s12_mems_resonator.py +0 -0
  72. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/extras/scenarios/s13_variable_stars_kepler.py +0 -0
  73. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/extras/scenarios/s14_eeg.py +0 -0
  74. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/extras/scenarios/s15_circadian.py +0 -0
  75. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/extras/spectral_expansion.py +0 -0
  76. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/grad_shafranov.py +0 -0
  77. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/grid_inverse.py +0 -0
  78. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/grid_rl_control.py +0 -0
  79. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/grid_swing.py +0 -0
  80. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/gs_inverse.py +0 -0
  81. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/gs_rl_control.py +0 -0
  82. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/integro_differential_demo.py +0 -0
  83. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/inverse/aero_.py +0 -0
  84. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/inverse/denoising_parameter_estimation.py +0 -0
  85. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/inverse/elastic_wave_animation.py +0 -0
  86. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/inverse/heat_from_video.py +0 -0
  87. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/inverse/inverse_memory_kernel.py +0 -0
  88. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/inverse/inverse_source_integral_sensors.py +0 -0
  89. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/inverse/inverse_turbulence.py +0 -0
  90. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/inverse/rlc_system_id.py +0 -0
  91. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/inverse/shape_ns.py +0 -0
  92. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/inverse/subsurface_imaging.py +0 -0
  93. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/inverse/tomography_projection.py +0 -0
  94. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/inverse/wing_optimize_simple.py +0 -0
  95. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/inverse_heat_source.py +0 -0
  96. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/inverse_magnetostatics.py +0 -0
  97. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/inverse_source_position.py +0 -0
  98. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/learnable_helmholtz.py +0 -0
  99. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/memory_diffusion.py +0 -0
  100. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/orbit_hill.py +0 -0
  101. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/orbit_inverse.py +0 -0
  102. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/orbit_rl.py +0 -0
  103. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/pde_discovery.py +0 -0
  104. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/rlc_integro_differential.py +0 -0
  105. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/run_all_extensions.py +0 -0
  106. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/run_linear.py +0 -0
  107. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/run_nonlinear.py +0 -0
  108. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/sindy/compare_sindy_methods.py +0 -0
  109. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/sindy/sindy_benchmarks.py +0 -0
  110. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/sindy/sindy_differentiable.py +0 -0
  111. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/sindy/sindy_minimal_diff.py +0 -0
  112. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/tutorial_basic.py +0 -0
  113. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/tutorial_nonlinear.py +0 -0
  114. {fastlsq-0.4.0 → fastlsq-0.4.1}/examples/vector_basis_stream_vorticity.py +0 -0
  115. {fastlsq-0.4.0 → fastlsq-0.4.1}/fastlsq/block.py +0 -0
  116. {fastlsq-0.4.0 → fastlsq-0.4.1}/fastlsq/device.py +0 -0
  117. {fastlsq-0.4.0 → fastlsq-0.4.1}/fastlsq/diagnostics.py +0 -0
  118. {fastlsq-0.4.0 → fastlsq-0.4.1}/fastlsq/geometry.py +0 -0
  119. {fastlsq-0.4.0 → fastlsq-0.4.1}/fastlsq/lightning.py +0 -0
  120. {fastlsq-0.4.0 → fastlsq-0.4.1}/fastlsq/newton.py +0 -0
  121. {fastlsq-0.4.0 → fastlsq-0.4.1}/fastlsq/plotting.py +0 -0
  122. {fastlsq-0.4.0 → fastlsq-0.4.1}/fastlsq/problems/__init__.py +0 -0
  123. {fastlsq-0.4.0 → fastlsq-0.4.1}/fastlsq/problems/linear.py +0 -0
  124. {fastlsq-0.4.0 → fastlsq-0.4.1}/fastlsq/problems/nonlinear.py +0 -0
  125. {fastlsq-0.4.0 → fastlsq-0.4.1}/fastlsq/problems/regression.py +0 -0
  126. {fastlsq-0.4.0 → fastlsq-0.4.1}/fastlsq/solvers.py +0 -0
  127. {fastlsq-0.4.0 → fastlsq-0.4.1}/fastlsq/tuning.py +0 -0
  128. {fastlsq-0.4.0 → fastlsq-0.4.1}/fastlsq/utils.py +0 -0
  129. {fastlsq-0.4.0 → fastlsq-0.4.1}/fastlsq/vector.py +0 -0
  130. {fastlsq-0.4.0 → fastlsq-0.4.1}/fastlsq/viz.py +0 -0
  131. {fastlsq-0.4.0 → fastlsq-0.4.1}/misc/fastlsq_teaser.png +0 -0
  132. {fastlsq-0.4.0 → fastlsq-0.4.1}/misc/ideal_quadrupole.png +0 -0
  133. {fastlsq-0.4.0 → fastlsq-0.4.1}/misc/inverse_heat_source.gif +0 -0
  134. {fastlsq-0.4.0 → fastlsq-0.4.1}/misc/inverse_heat_source.png +0 -0
  135. {fastlsq-0.4.0 → fastlsq-0.4.1}/misc/inverse_magnetostatics.png +0 -0
  136. {fastlsq-0.4.0 → fastlsq-0.4.1}/misc/inverse_magnetostatics_convergence.png +0 -0
  137. {fastlsq-0.4.0 → fastlsq-0.4.1}/misc/quadrupole_convergence.png +0 -0
  138. {fastlsq-0.4.0 → fastlsq-0.4.1}/misc/quadrupole_optimization.png +0 -0
  139. {fastlsq-0.4.0 → fastlsq-0.4.1}/misc/tutorial_nlpoisson_convergence.png +0 -0
  140. {fastlsq-0.4.0 → fastlsq-0.4.1}/misc/tutorial_nlpoisson_solution.png +0 -0
  141. {fastlsq-0.4.0 → fastlsq-0.4.1}/requirements.txt +0 -0
  142. {fastlsq-0.4.0 → fastlsq-0.4.1}/setup.cfg +0 -0
  143. {fastlsq-0.4.0 → fastlsq-0.4.1}/tests/test_basic.py +0 -0
  144. {fastlsq-0.4.0 → fastlsq-0.4.1}/tests/test_benchmarks_inverse.py +0 -0
  145. {fastlsq-0.4.0 → fastlsq-0.4.1}/tests/test_block.py +0 -0
  146. {fastlsq-0.4.0 → fastlsq-0.4.1}/tests/test_derivatives.py +0 -0
  147. {fastlsq-0.4.0 → fastlsq-0.4.1}/tests/test_device.py +0 -0
  148. {fastlsq-0.4.0 → fastlsq-0.4.1}/tests/test_grad_shafranov.py +0 -0
  149. {fastlsq-0.4.0 → fastlsq-0.4.1}/tests/test_grid_swing.py +0 -0
  150. {fastlsq-0.4.0 → fastlsq-0.4.1}/tests/test_integral.py +0 -0
  151. {fastlsq-0.4.0 → fastlsq-0.4.1}/tests/test_learnable.py +0 -0
  152. {fastlsq-0.4.0 → fastlsq-0.4.1}/tests/test_orbit_hill.py +0 -0
  153. {fastlsq-0.4.0 → fastlsq-0.4.1}/tests/test_projection.py +0 -0
@@ -2,6 +2,35 @@
2
2
 
3
3
  All notable changes to FastLSQ will be documented in this file.
4
4
 
5
+ ## [0.4.1] - 2026-06-23
6
+
7
+ ### Added
8
+
9
+ - **Solve-time diagnostics — `solve_lstsq(..., return_info=True)`.** Returns
10
+ `(x, info)` with `info = {t_solve, rank_used, residual, cond_estimate}`. `t_solve`
11
+ is the **device-synced wall-time of the solve step alone** (singular values for
12
+ the rank/cond diagnostics are computed *outside* the timed region), `rank_used` is
13
+ the rank-revealing effective numerical rank, and `cond_estimate` is `s_max/s_min`
14
+ over the retained subspace. The default `return_info=False` path is unchanged.
15
+ - **Phased breakdown in `solve_linear` metrics.** `metrics` now reports
16
+ `scale_search_s`, `assemble_s`, `solve_s` (plus `rank_used`, `residual`,
17
+ `cond_estimate`) so the headline time is no longer dominated by the auto-scale
18
+ search — the reported solve time is a number the library hands you, not one a
19
+ benchmark reconstructs by hand.
20
+ - **`fastlsq.benchmark.time_solve(fn, reps, warmup, device)`** — a device-correct
21
+ timing primitive (`synchronize` bracketing + warm-up + min-of-reps) returning the
22
+ reproducible solve-time *floor* in seconds (`return_all=True` for full stats). Also
23
+ exposes `fastlsq.benchmark.synchronize`.
24
+ - **`SinusoidalBasis.random_covariance(d, N, Sigma=… | L=…)`** — fixed full-`Sigma`
25
+ constructor (`W = L @ W_hat`, `Sigma = L Lᵀ`), the symmetric counterpart to
26
+ `random` / `random_anisotropic`.
27
+ - **`LearnableFastLSQ.freeze() -> SinusoidalBasis`** — freezes the learned bandwidth
28
+ into a plain, detached basis for one clean, timed one-shot deployment solve (the
29
+ `.basis` property still reconstructs-with-grad each access).
30
+ - **Checkpoint provenance.** `save_checkpoint` auto-records a `provenance` block
31
+ (library version, device, dtype, timestamp, and the realized scale / `Sigma`) into
32
+ the metadata.
33
+
5
34
  ## [0.4.0] - 2026-06-22
6
35
 
7
36
  ### Added
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: FastLSQ
3
- Version: 0.4.0
3
+ Version: 0.4.1
4
4
  Summary: One-shot PDE solving via Fourier features with exact analytical derivatives; rank-revealing solvers, learnable anisotropic bandwidth, and CPU/CUDA/MPS support
5
5
  Author: Antonin Sulc
6
6
  License-Expression: MIT
@@ -106,6 +106,7 @@ examples/sindy/sindy_minimal_diff.py
106
106
  fastlsq/__init__.py
107
107
  fastlsq/api.py
108
108
  fastlsq/basis.py
109
+ fastlsq/benchmark.py
109
110
  fastlsq/block.py
110
111
  fastlsq/device.py
111
112
  fastlsq/diagnostics.py
@@ -140,6 +141,7 @@ tests/test_benchmarks_inverse.py
140
141
  tests/test_block.py
141
142
  tests/test_derivatives.py
142
143
  tests/test_device.py
144
+ tests/test_diagnostics_timing.py
143
145
  tests/test_grad_shafranov.py
144
146
  tests/test_grid_swing.py
145
147
  tests/test_integral.py
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: FastLSQ
3
- Version: 0.4.0
3
+ Version: 0.4.1
4
4
  Summary: One-shot PDE solving via Fourier features with exact analytical derivatives; rank-revealing solvers, learnable anisotropic bandwidth, and CPU/CUDA/MPS support
5
5
  Author: Antonin Sulc
6
6
  License-Expression: MIT
@@ -53,8 +53,9 @@ from fastlsq.export import (
53
53
  load_checkpoint,
54
54
  )
55
55
  from fastlsq import viz
56
+ from fastlsq import benchmark
56
57
 
57
- __version__ = "0.4.0"
58
+ __version__ = "0.4.1"
58
59
  __all__ = [
59
60
  # Device selection (CPU / CUDA / Apple-MPS, dtype-aware)
60
61
  "resolve_device",
@@ -105,16 +105,20 @@ def solve_linear(
105
105
  if n_bc is None:
106
106
  n_bc = max(800, n_pde // 5)
107
107
 
108
- # Auto-select scale if needed
108
+ # Auto-select scale if needed (timed as its own phase: it is a search that can
109
+ # be amortised / skipped at deployment, so it must not inflate the solve time).
110
+ scale_search_s = 0.0
109
111
  if scale is None and auto_scale:
110
112
  if verbose:
111
113
  print("Auto-selecting optimal scale...")
114
+ t_phase = time.perf_counter()
112
115
  scale = auto_select_scale(
113
116
  problem, solver_class=FastLSQSolver,
114
117
  n_blocks=n_blocks, hidden_size=hidden_size,
115
118
  n_pde=n_pde, n_bc=n_bc, n_trials=auto_scale_trials,
116
119
  verbose=verbose,
117
120
  )
121
+ scale_search_s = time.perf_counter() - t_phase
118
122
  if verbose:
119
123
  print(f"Selected scale: {scale:.3f}")
120
124
 
@@ -137,10 +141,25 @@ def solve_linear(
137
141
  x_pde, bcs = data
138
142
  build_args = (bcs,)
139
143
 
140
- # Assemble and solve
144
+ # Assemble (timed phase)
145
+ t_phase = time.perf_counter()
141
146
  A, b = problem.build(solver, x_pde, *build_args)
142
- beta_raw = solve_lstsq(A, b, mu=mu, method=method)
147
+ assemble_s = time.perf_counter() - t_phase
148
+
149
+ # Solve (timed phase). With metrics on, ``solve_lstsq`` hands back the
150
+ # device-synced solve time plus rank/residual/cond diagnostics directly; with
151
+ # metrics off, keep the plain fast solve and just time it externally.
143
152
  n_outputs = getattr(problem, "n_outputs", 1)
153
+ if return_metrics:
154
+ beta_raw, solve_info = solve_lstsq(
155
+ A, b, mu=mu, method=method, return_info=True
156
+ )
157
+ solve_s = solve_info["t_solve"]
158
+ else:
159
+ t_phase = time.perf_counter()
160
+ beta_raw = solve_lstsq(A, b, mu=mu, method=method)
161
+ solve_s = time.perf_counter() - t_phase
162
+ solve_info = None
144
163
  solver.beta = unpack_beta(beta_raw, solver.n_features, n_outputs)
145
164
 
146
165
  runtime = time.time() - t0
@@ -169,6 +188,15 @@ def solve_linear(
169
188
  "val_err": val_err,
170
189
  "grad_err": grad_err,
171
190
  "runtime": runtime,
191
+ # Phased breakdown: the conflated ``runtime`` split into its parts so
192
+ # the headline solve time is no longer dominated by the scale search.
193
+ "scale_search_s": scale_search_s,
194
+ "assemble_s": assemble_s,
195
+ "solve_s": solve_s,
196
+ # Rank-revealing solve diagnostics (the paper's central mechanism).
197
+ "rank_used": solve_info["rank_used"],
198
+ "residual": solve_info["residual"],
199
+ "cond_estimate": solve_info["cond_estimate"],
172
200
  }
173
201
  if verbose:
174
202
  print(f"Value error: {val_err:.2e}, Gradient error: {grad_err:.2e}")
@@ -197,6 +197,43 @@ class SinusoidalBasis:
197
197
  b = torch.rand(1, n_features, device=get_device()) * 2 * np.pi
198
198
  return cls(W, b, normalize=normalize, dc_eps=dc_eps)
199
199
 
200
+ @classmethod
201
+ def random_covariance(
202
+ cls,
203
+ input_dim: int,
204
+ n_features: int,
205
+ Sigma: Optional[Union[list, np.ndarray, torch.Tensor]] = None,
206
+ L: Optional[Union[list, np.ndarray, torch.Tensor]] = None,
207
+ normalize: bool = True,
208
+ dc_eps: float = 1e-8,
209
+ ) -> SinusoidalBasis:
210
+ """Create a basis with full-covariance Gaussian frequencies.
211
+
212
+ Frequencies are drawn ``W[:, j] ~ N(0, Sigma)`` via ``W = L @ W_hat`` with
213
+ ``W_hat ~ N(0, I_d)`` and ``Sigma = L L^T``. Provide **exactly one** of
214
+ ``Sigma`` (the ``d x d`` covariance) or its lower-triangular Cholesky factor
215
+ ``L``. This is the fixed full-``Sigma`` analogue of :meth:`random`
216
+ (isotropic) and :meth:`random_anisotropic` (axis-aligned), and uses the same
217
+ ``W = L @ W_hat`` reparameterisation as
218
+ :class:`~fastlsq.learnable.LearnableFastLSQ` -- so a swept covariance no
219
+ longer needs a hand-rolled ``chol(Sigma) @ randn``.
220
+ """
221
+ if (Sigma is None) == (L is None):
222
+ raise ValueError(
223
+ "random_covariance: provide exactly one of Sigma or L."
224
+ )
225
+ dev = get_device()
226
+ dt = torch.get_default_dtype()
227
+ if L is None:
228
+ Sigma_t = torch.as_tensor(Sigma, device=dev, dtype=dt)
229
+ L_t = torch.linalg.cholesky(Sigma_t)
230
+ else:
231
+ L_t = torch.as_tensor(L, device=dev, dtype=dt)
232
+ W_hat = torch.randn(input_dim, n_features, device=dev, dtype=L_t.dtype)
233
+ W = L_t @ W_hat
234
+ b = torch.rand(1, n_features, device=dev, dtype=L_t.dtype) * 2 * np.pi
235
+ return cls(W, b, normalize=normalize, dc_eps=dc_eps)
236
+
200
237
  # ------------------------------------------------------------------
201
238
  # Caching
202
239
  # ------------------------------------------------------------------
@@ -0,0 +1,100 @@
1
+ # Copyright (c) 2026 Antonin Sulc
2
+ # Licensed under the MIT License. See LICENSE file for details.
3
+
4
+ """Device-correct timing primitives for FastLSQ.
5
+
6
+ Naively wall-clock-timing a CUDA solve is wrong: kernels are launched
7
+ asynchronously, so ``t1 - t0`` measures launch overhead, not compute. Without a
8
+ ``synchronize`` bracket every GPU number is meaningless (and irreproducible).
9
+ :func:`time_solve` adds the missing primitive -- ``synchronize`` bracketing plus
10
+ warm-up and a min-of-reps reduction -- so the reported solve time is the
11
+ reproducible compute *floor* on any device (CPU / CUDA / Apple-MPS).
12
+
13
+ import fastlsq as fl
14
+ t = fl.benchmark.time_solve(lambda: fl.solve_lstsq(A, b)) # seconds, floor
15
+ """
16
+
17
+ import time
18
+
19
+ import torch
20
+
21
+ from fastlsq.device import get_device
22
+
23
+
24
+ def synchronize(device=None):
25
+ """Block until all queued work on ``device`` has finished.
26
+
27
+ No-op on CPU; calls ``torch.cuda.synchronize`` on CUDA and
28
+ ``torch.mps.synchronize`` on Apple-MPS (when available). ``device=None``
29
+ uses the active FastLSQ device.
30
+ """
31
+ device = get_device() if device is None else torch.device(device)
32
+ if device.type == "cuda":
33
+ torch.cuda.synchronize(device)
34
+ elif device.type == "mps":
35
+ sync = getattr(getattr(torch, "mps", None), "synchronize", None)
36
+ if sync is not None:
37
+ sync()
38
+
39
+
40
+ def time_solve(fn, reps=10, warmup=2, device=None, return_all=False):
41
+ """Device-correct timing floor for a callable.
42
+
43
+ Runs ``warmup`` unmeasured calls (allocator / autotune warm-up), then ``reps``
44
+ measured calls, each bracketed by :func:`synchronize` so CUDA/MPS asynchronous
45
+ execution is fully accounted for. Returns the **minimum** elapsed time in
46
+ seconds -- the reproducible solve-time floor -- which is the right summary for
47
+ a deterministic op whose variance is pure system noise.
48
+
49
+ Parameters
50
+ ----------
51
+ fn : callable
52
+ Zero-argument callable performing the work to time, e.g.
53
+ ``lambda: solve_lstsq(A, b)``.
54
+ reps : int
55
+ Number of measured repetitions (>= 1).
56
+ warmup : int
57
+ Number of unmeasured warm-up repetitions.
58
+ device : str | torch.device | None
59
+ Device to synchronize around each call. ``None`` uses the active FastLSQ
60
+ device.
61
+ return_all : bool
62
+ If True, return a stats dict
63
+ ``{"min", "median", "mean", "std", "times", "reps", "warmup", "device"}``
64
+ instead of the bare floor.
65
+
66
+ Returns
67
+ -------
68
+ float
69
+ The minimum solve time in seconds (``return_all=False``).
70
+ dict
71
+ Full timing statistics (``return_all=True``).
72
+ """
73
+ device = get_device() if device is None else torch.device(device)
74
+
75
+ for _ in range(max(0, warmup)):
76
+ fn()
77
+ synchronize(device)
78
+
79
+ times = []
80
+ for _ in range(max(1, reps)):
81
+ synchronize(device)
82
+ t0 = time.perf_counter()
83
+ fn()
84
+ synchronize(device)
85
+ times.append(time.perf_counter() - t0)
86
+
87
+ t = torch.tensor(times, dtype=torch.float64)
88
+ floor = float(t.min())
89
+ if not return_all:
90
+ return floor
91
+ return {
92
+ "min": floor,
93
+ "median": float(t.median()),
94
+ "mean": float(t.mean()),
95
+ "std": float(t.std(unbiased=False)),
96
+ "times": times,
97
+ "reps": len(times),
98
+ "warmup": max(0, warmup),
99
+ "device": str(device),
100
+ }
@@ -3,6 +3,8 @@
3
3
 
4
4
  """Export utilities for FastLSQ solutions (NumPy, VTK, etc.)."""
5
5
 
6
+ import time
7
+
6
8
  import torch
7
9
  import numpy as np
8
10
  from typing import Optional, Union, Dict, Any
@@ -10,6 +12,32 @@ from typing import Optional, Union, Dict, Any
10
12
  from fastlsq.solvers import FastLSQSolver
11
13
 
12
14
 
15
+ def _provenance(solver: FastLSQSolver) -> Dict[str, Any]:
16
+ """Auto-recorded provenance for a saved model: library version, device, dtype,
17
+ and the realized frequency bandwidth (scale / Sigma) the model actually uses."""
18
+ import fastlsq # lazy: avoids a circular import at module load
19
+
20
+ prov: Dict[str, Any] = {
21
+ "fastlsq_version": getattr(fastlsq, "__version__", None),
22
+ "created": time.time(),
23
+ "input_dim": solver.input_dim,
24
+ "n_features": solver.n_features,
25
+ }
26
+ W = torch.cat(solver.W_list, dim=1) if solver.W_list else None
27
+ ref = W if W is not None else solver.beta
28
+ if ref is not None:
29
+ prov["device"] = str(ref.device)
30
+ prov["dtype"] = str(ref.dtype).replace("torch.", "")
31
+ if W is not None:
32
+ # Realized frequency second moment Sigma = (W Wᵀ)/N and per-axis scale.
33
+ freq_cov = (W @ W.transpose(-2, -1)) / W.shape[1]
34
+ prov["freq_cov"] = freq_cov.cpu().numpy()
35
+ prov["freq_std"] = torch.sqrt(
36
+ torch.diagonal(freq_cov).clamp_min(0.0)
37
+ ).cpu().numpy()
38
+ return prov
39
+
40
+
13
41
  def to_numpy(
14
42
  solver: FastLSQSolver,
15
43
  x: Union[torch.Tensor, np.ndarray],
@@ -138,11 +166,14 @@ def save_checkpoint(
138
166
  path : str
139
167
  File path (.pt or .pth extension recommended).
140
168
  metadata : dict, optional
141
- Additional metadata to save.
169
+ Additional metadata to save. A ``provenance`` block (library version,
170
+ device, dtype, timestamp, and realized scale / Sigma) is auto-recorded
171
+ unless the caller supplies its own ``provenance`` key.
142
172
  """
143
173
  state = to_dict(solver, include_weights=True, include_metadata=True)
144
- if metadata:
145
- state["metadata"] = metadata
174
+ meta = dict(metadata) if metadata else {}
175
+ meta.setdefault("provenance", _provenance(solver))
176
+ state["metadata"] = meta
146
177
  torch.save(state, path)
147
178
 
148
179
 
@@ -153,6 +153,20 @@ class LearnableFastLSQ(nn.Module):
153
153
  """
154
154
  return SinusoidalBasis(self._W(), self.b, normalize=self._normalize)
155
155
 
156
+ def freeze(self) -> SinusoidalBasis:
157
+ """Freeze the learned bandwidth into a plain, detached ``SinusoidalBasis``.
158
+
159
+ Unlike the :attr:`basis` property -- which reconstructs *with grad* on every
160
+ access -- this returns a fixed basis built from the current
161
+ ``W = L @ W_hat`` with the weights detached and cloned. Run the (untimed)
162
+ bandwidth search, then ``freeze()`` the learned ``Sigma`` into this plain
163
+ basis for one clean, timed one-shot solve: the deployment artifact.
164
+ """
165
+ with torch.no_grad():
166
+ W = self._W().detach().clone()
167
+ b = self.b.detach().clone()
168
+ return SinusoidalBasis(W, b, normalize=self._normalize)
169
+
156
170
  # ------------------------------------------------------------------
157
171
  # Prediction
158
172
  # ------------------------------------------------------------------
@@ -36,6 +36,7 @@ All back-ends are device/dtype-aware. Apple-MPS lacks a robust ``svd``/``lstsq`
36
36
  so the factorization is run on CPU and the result moved back (one-time warning).
37
37
  """
38
38
 
39
+ import time
39
40
  import warnings
40
41
 
41
42
  import torch
@@ -64,15 +65,18 @@ def _maybe_cpu(A, b):
64
65
 
65
66
 
66
67
  def _svd_solve(A, b, mu, rcond):
68
+ """Returns ``(x, S)`` where ``S`` is the singular values (descending) when an
69
+ explicit SVD is formed, else ``None`` (the fast LAPACK gelsd path)."""
67
70
  # Fast LAPACK rank-revealing driver for the common CPU / mu==0 case.
68
71
  if not mu and A.device.type == "cpu":
69
72
  try:
70
- return torch.linalg.lstsq(A, b, rcond=rcond, driver="gelsd").solution
73
+ return torch.linalg.lstsq(A, b, rcond=rcond, driver="gelsd").solution, None
71
74
  except (RuntimeError, ValueError):
72
75
  pass # fall through to the explicit SVD
73
76
  U, S, Vh = torch.linalg.svd(A, full_matrices=False)
74
77
  filt = torch.where(S > rcond * S[0], S / (S * S + mu), torch.zeros_like(S))
75
- return Vh.transpose(-2, -1) @ (filt.unsqueeze(-1) * (U.transpose(-2, -1) @ b))
78
+ x = Vh.transpose(-2, -1) @ (filt.unsqueeze(-1) * (U.transpose(-2, -1) @ b))
79
+ return x, S
76
80
 
77
81
 
78
82
  def _cholesky_solve(A, b, mu):
@@ -99,7 +103,8 @@ def _rsvd_solve(A, b, mu, rcond, rank, oversample, n_iter):
99
103
  Ub, S, Vh = torch.linalg.svd(B, full_matrices=False)
100
104
  U = Q @ Ub # (m, k)
101
105
  filt = torch.where(S > rcond * S[0], S / (S * S + mu), torch.zeros_like(S))
102
- return Vh.transpose(-2, -1) @ (filt.unsqueeze(-1) * (U.transpose(-2, -1) @ b))
106
+ x = Vh.transpose(-2, -1) @ (filt.unsqueeze(-1) * (U.transpose(-2, -1) @ b))
107
+ return x, S # S is the truncated (rank-k) spectrum
103
108
 
104
109
 
105
110
  def _qr_solve(A, b, mu):
@@ -116,13 +121,15 @@ def _qr_solve(A, b, mu):
116
121
 
117
122
 
118
123
  def _auto_solve(A, b, mu, rcond):
124
+ """Returns ``(x, S)``; ``S`` is the singular values only when the SVD safety
125
+ net is taken (the Cholesky / QR fast paths return ``None``)."""
119
126
  # Cheap conditioning probe: cond(A) ~ max/min Cholesky pivot. If well within
120
127
  # float64's reach use the fast Cholesky.
121
128
  try:
122
129
  x, L = _cholesky_solve(A, b, mu)
123
130
  d = torch.diagonal(L).abs()
124
131
  if torch.isfinite(d).all() and d.min() > (rcond ** 0.25) * d.max():
125
- return x
132
+ return x, None
126
133
  except torch.linalg.LinAlgError:
127
134
  pass
128
135
  # Ill-conditioned. On CPU with no ridge the LAPACK gelsd driver is both
@@ -137,12 +144,43 @@ def _auto_solve(A, b, mu, rcond):
137
144
  x = _qr_solve(A, b, mu)
138
145
  nx = torch.linalg.vector_norm(x)
139
146
  if torch.isfinite(nx) and nx <= _QR_AUTO_NORM_GUARD * (1.0 + torch.linalg.vector_norm(b)):
140
- return x
147
+ return x, None
141
148
  return _svd_solve(A, b, mu, rcond)
142
149
 
143
150
 
151
+ def _dispatch(A, b, mu, rcond, method, rank, oversample, n_iter):
152
+ """Run the requested back-end; returns ``(x, S)`` where ``S`` is the singular
153
+ values when the back-end already computed an SVD, else ``None``."""
154
+ if method == "auto":
155
+ return _auto_solve(A, b, mu, rcond)
156
+ elif method == "svd":
157
+ return _svd_solve(A, b, mu, rcond)
158
+ elif method == "qr":
159
+ return _qr_solve(A, b, mu), None
160
+ elif method == "cholesky":
161
+ return _cholesky_solve(A, b, mu)[0], None
162
+ elif method == "rsvd":
163
+ return _rsvd_solve(A, b, mu, rcond, rank, oversample, n_iter)
164
+ else:
165
+ raise ValueError(f"Unknown method {method!r}; "
166
+ "choose 'auto', 'qr', 'svd', 'cholesky', or 'rsvd'.")
167
+
168
+
169
+ def _sync_device(device):
170
+ """Block until queued work on ``device`` finishes (no-op off CUDA/MPS).
171
+
172
+ Without this an async CUDA solve's wall-clock measures only kernel-launch
173
+ overhead, not compute -- the same primitive ``fastlsq.benchmark`` uses."""
174
+ if device.type == "cuda":
175
+ torch.cuda.synchronize(device)
176
+ elif device.type == "mps":
177
+ sync = getattr(getattr(torch, "mps", None), "synchronize", None)
178
+ if sync is not None:
179
+ sync()
180
+
181
+
144
182
  def solve_lstsq(A, b, mu=0.0, rcond=1e-12, method="auto",
145
- rank=None, oversample=10, n_iter=4):
183
+ rank=None, oversample=10, n_iter=4, *, return_info=False):
146
184
  """Solve min ||A x - b||^2 + mu ||x||^2.
147
185
 
148
186
  Parameters
@@ -159,23 +197,50 @@ def solve_lstsq(A, b, mu=0.0, rcond=1e-12, method="auto",
159
197
  rank, oversample, n_iter : int
160
198
  Randomized-SVD parameters (``method="rsvd"`` only). Set ``rank`` << N for
161
199
  the speed-up; ``None`` uses the full rank (correct but no acceleration).
200
+ return_info : bool, optional
201
+ If True, return ``(x, info)`` with a per-solve diagnostics dict
202
+ ``{"t_solve", "rank_used", "residual", "cond_estimate"}`` instead of just
203
+ ``x``. ``t_solve`` is the device-synced wall-time of the **solve step
204
+ only** (not assembly or scale search); ``rank_used`` is the rank-revealing
205
+ effective numerical rank (singular values above ``rcond``); ``residual`` is
206
+ the data residual ``||A x - b||``; ``cond_estimate`` is ``s_max / s_min``
207
+ over the retained subspace. The diagnostic singular values are computed
208
+ *outside* the timed region (one extra ``svdvals`` when the chosen back-end
209
+ did not already form an SVD), so ``t_solve`` stays honest.
162
210
 
163
211
  Returns
164
212
  -------
165
213
  x : Tensor, shape (N, K)
214
+ info : dict, only if ``return_info=True``
166
215
  """
167
216
  A2, b2, mps_dev = _maybe_cpu(A, b)
168
- if method == "auto":
169
- x = _auto_solve(A2, b2, mu, rcond)
170
- elif method == "svd":
171
- x = _svd_solve(A2, b2, mu, rcond)
172
- elif method == "qr":
173
- x = _qr_solve(A2, b2, mu)
174
- elif method == "cholesky":
175
- x = _cholesky_solve(A2, b2, mu)[0]
176
- elif method == "rsvd":
177
- x = _rsvd_solve(A2, b2, mu, rcond, rank, oversample, n_iter)
217
+
218
+ if not return_info:
219
+ x, _ = _dispatch(A2, b2, mu, rcond, method, rank, oversample, n_iter)
220
+ return x.to(mps_dev) if mps_dev is not None else x
221
+
222
+ _sync_device(A2.device)
223
+ t0 = time.perf_counter()
224
+ x, S = _dispatch(A2, b2, mu, rcond, method, rank, oversample, n_iter)
225
+ _sync_device(A2.device)
226
+ t_solve = time.perf_counter() - t0
227
+
228
+ # Diagnostics (untimed): reuse the back-end's spectrum, else one extra SVD.
229
+ if S is None:
230
+ S = torch.linalg.svdvals(A2)
231
+ smax = S[0]
232
+ keep = S > rcond * smax
233
+ rank_used = int(keep.sum().item())
234
+ if rank_used > 0:
235
+ cond_estimate = float((smax / S[keep][-1]).item())
178
236
  else:
179
- raise ValueError(f"Unknown method {method!r}; "
180
- "choose 'auto', 'qr', 'svd', 'cholesky', or 'rsvd'.")
181
- return x.to(mps_dev) if mps_dev is not None else x
237
+ cond_estimate = float("inf")
238
+ residual = float((A2 @ x - b2).norm().item())
239
+ info = {
240
+ "t_solve": t_solve,
241
+ "rank_used": rank_used,
242
+ "residual": residual,
243
+ "cond_estimate": cond_estimate,
244
+ }
245
+ x_out = x.to(mps_dev) if mps_dev is not None else x
246
+ return x_out, info
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
4
4
 
5
5
  [project]
6
6
  name = "FastLSQ"
7
- version = "0.4.0"
7
+ version = "0.4.1"
8
8
  description = "One-shot PDE solving via Fourier features with exact analytical derivatives; rank-revealing solvers, learnable anisotropic bandwidth, and CPU/CUDA/MPS support"
9
9
  readme = "README.md"
10
10
  license = "MIT"