corticalfields 0.2.0__tar.gz → 0.2.2__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 (59) hide show
  1. {corticalfields-0.2.0/src/corticalfields.egg-info → corticalfields-0.2.2}/PKG-INFO +1 -1
  2. {corticalfields-0.2.0 → corticalfields-0.2.2}/pyproject.toml +1 -1
  3. {corticalfields-0.2.0 → corticalfields-0.2.2}/src/corticalfields/__init__.py +69 -3
  4. {corticalfields-0.2.0 → corticalfields-0.2.2}/src/corticalfields/analysis/eda_qc.py +4 -0
  5. {corticalfields-0.2.0 → corticalfields-0.2.2}/src/corticalfields/analysis/stats.py +12 -0
  6. {corticalfields-0.2.0 → corticalfields-0.2.2}/src/corticalfields/backends.py +263 -75
  7. {corticalfields-0.2.0 → corticalfields-0.2.2}/src/corticalfields/distance_stats.py +8 -0
  8. corticalfields-0.2.2/src/corticalfields/graphs.py +1098 -0
  9. {corticalfields-0.2.0 → corticalfields-0.2.2}/src/corticalfields/hippocampus.py +40 -12
  10. corticalfields-0.2.2/src/corticalfields/spectral.py +1300 -0
  11. {corticalfields-0.2.0 → corticalfields-0.2.2}/src/corticalfields/utils.py +260 -0
  12. {corticalfields-0.2.0 → corticalfields-0.2.2}/src/corticalfields/viz/__init__.py +20 -1
  13. corticalfields-0.2.2/src/corticalfields/viz/graph_viz.py +913 -0
  14. {corticalfields-0.2.0 → corticalfields-0.2.2/src/corticalfields.egg-info}/PKG-INFO +1 -1
  15. {corticalfields-0.2.0 → corticalfields-0.2.2}/src/corticalfields.egg-info/SOURCES.txt +1 -0
  16. corticalfields-0.2.0/src/corticalfields/graphs.py +0 -234
  17. corticalfields-0.2.0/src/corticalfields/spectral.py +0 -527
  18. {corticalfields-0.2.0 → corticalfields-0.2.2}/LICENSE +0 -0
  19. {corticalfields-0.2.0 → corticalfields-0.2.2}/README.md +0 -0
  20. {corticalfields-0.2.0 → corticalfields-0.2.2}/setup.cfg +0 -0
  21. {corticalfields-0.2.0 → corticalfields-0.2.2}/src/corticalfields/_pointcloud_legacy.py +0 -0
  22. {corticalfields-0.2.0 → corticalfields-0.2.2}/src/corticalfields/analysis/__init__.py +0 -0
  23. {corticalfields-0.2.0 → corticalfields-0.2.2}/src/corticalfields/analysis/bayesian.py +0 -0
  24. {corticalfields-0.2.0 → corticalfields-0.2.2}/src/corticalfields/analysis/normative.py +0 -0
  25. {corticalfields-0.2.0 → corticalfields-0.2.2}/src/corticalfields/asymmetry.py +0 -0
  26. {corticalfields-0.2.0 → corticalfields-0.2.2}/src/corticalfields/bayes_viz.py +0 -0
  27. {corticalfields-0.2.0 → corticalfields-0.2.2}/src/corticalfields/bayesian.py +0 -0
  28. {corticalfields-0.2.0 → corticalfields-0.2.2}/src/corticalfields/brainplots.py +0 -0
  29. {corticalfields-0.2.0 → corticalfields-0.2.2}/src/corticalfields/datasets.py +0 -0
  30. {corticalfields-0.2.0 → corticalfields-0.2.2}/src/corticalfields/eda_qc.py +0 -0
  31. {corticalfields-0.2.0 → corticalfields-0.2.2}/src/corticalfields/features.py +0 -0
  32. {corticalfields-0.2.0 → corticalfields-0.2.2}/src/corticalfields/functional_maps.py +0 -0
  33. {corticalfields-0.2.0 → corticalfields-0.2.2}/src/corticalfields/kernels.py +0 -0
  34. {corticalfields-0.2.0 → corticalfields-0.2.2}/src/corticalfields/normative.py +0 -0
  35. {corticalfields-0.2.0 → corticalfields-0.2.2}/src/corticalfields/pointcloud/__init__.py +0 -0
  36. {corticalfields-0.2.0 → corticalfields-0.2.2}/src/corticalfields/pointcloud/deep/__init__.py +0 -0
  37. {corticalfields-0.2.0 → corticalfields-0.2.2}/src/corticalfields/pointcloud/deep/diffusion_net.py +0 -0
  38. {corticalfields-0.2.0 → corticalfields-0.2.2}/src/corticalfields/pointcloud/deep/egnn.py +0 -0
  39. {corticalfields-0.2.0 → corticalfields-0.2.2}/src/corticalfields/pointcloud/functional_maps.py +0 -0
  40. {corticalfields-0.2.0 → corticalfields-0.2.2}/src/corticalfields/pointcloud/morphometrics.py +0 -0
  41. {corticalfields-0.2.0 → corticalfields-0.2.2}/src/corticalfields/pointcloud/registration.py +0 -0
  42. {corticalfields-0.2.0 → corticalfields-0.2.2}/src/corticalfields/pointcloud/spectral.py +0 -0
  43. {corticalfields-0.2.0 → corticalfields-0.2.2}/src/corticalfields/pointcloud/transport.py +0 -0
  44. {corticalfields-0.2.0 → corticalfields-0.2.2}/src/corticalfields/pointcloud/viz.py +0 -0
  45. {corticalfields-0.2.0 → corticalfields-0.2.2}/src/corticalfields/pointcloud.py +0 -0
  46. {corticalfields-0.2.0 → corticalfields-0.2.2}/src/corticalfields/subcortical.py +0 -0
  47. {corticalfields-0.2.0 → corticalfields-0.2.2}/src/corticalfields/surface.py +0 -0
  48. {corticalfields-0.2.0 → corticalfields-0.2.2}/src/corticalfields/surprise.py +0 -0
  49. {corticalfields-0.2.0 → corticalfields-0.2.2}/src/corticalfields/transport.py +0 -0
  50. {corticalfields-0.2.0 → corticalfields-0.2.2}/src/corticalfields/viz/bayes.py +0 -0
  51. {corticalfields-0.2.0 → corticalfields-0.2.2}/src/corticalfields/viz/brainplots.py +0 -0
  52. {corticalfields-0.2.0 → corticalfields-0.2.2}/src/corticalfields/viz/subcortical.py +0 -0
  53. {corticalfields-0.2.0 → corticalfields-0.2.2}/src/corticalfields/viz/viz.py +0 -0
  54. {corticalfields-0.2.0 → corticalfields-0.2.2}/src/corticalfields/viz.py +0 -0
  55. {corticalfields-0.2.0 → corticalfields-0.2.2}/src/corticalfields/viz_subcortical.py +0 -0
  56. {corticalfields-0.2.0 → corticalfields-0.2.2}/src/corticalfields.egg-info/dependency_links.txt +0 -0
  57. {corticalfields-0.2.0 → corticalfields-0.2.2}/src/corticalfields.egg-info/requires.txt +0 -0
  58. {corticalfields-0.2.0 → corticalfields-0.2.2}/src/corticalfields.egg-info/top_level.txt +0 -0
  59. {corticalfields-0.2.0 → corticalfields-0.2.2}/tests/test_core.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: corticalfields
3
- Version: 0.2.0
3
+ Version: 0.2.2
4
4
  Summary: Spectral cortical and subcortical analysis with statistical testing (RSA, CCA, PLS, PERMANOVA, TFCE, NBS, laterality classification), on meshes and point clouds — Laplace-Beltrami decomposition, atlas-free asymmetry, GPU-accelerated optimal transport, hippocampal subfield analysis (HippUnfold), ShapeDNA/BrainPrint spectral fingerprinting, geometric deep learning, Bayesian inference, and normative modeling for structural neuroimaging.
5
5
  Author-email: rdneuro <r.debona@ufrj.br>
6
6
  License: MIT
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
4
4
 
5
5
  [project]
6
6
  name = "corticalfields"
7
- version = "0.2.0"
7
+ version = "0.2.2"
8
8
  description = "Spectral cortical and subcortical analysis with statistical testing (RSA, CCA, PLS, PERMANOVA, TFCE, NBS, laterality classification), on meshes and point clouds — Laplace-Beltrami decomposition, atlas-free asymmetry, GPU-accelerated optimal transport, hippocampal subfield analysis (HippUnfold), ShapeDNA/BrainPrint spectral fingerprinting, geometric deep learning, Bayesian inference, and normative modeling for structural neuroimaging."
9
9
  readme = "README.md"
10
10
  license = {text = "MIT"}
@@ -6,7 +6,7 @@ functional maps, optimal-transport distances, and information-theoretic
6
6
  surprise maps on brain surface meshes. Designed for structural MRI (T1w)
7
7
  data in clinical neuroimaging, with emphasis on epilepsy (MTLE-HS).
8
8
 
9
- Subpackages (v0.2.0)
9
+ Subpackages (v0.2.1)
10
10
  ---------------------
11
11
  analysis : Statistical analysis & modeling
12
12
  analysis.stats — MCC, GLM, PERMANOVA, CCA/PLS, RSA, NBS,
@@ -29,7 +29,7 @@ surface, subcortical, hippocampus, spectral, kernels, surprise, features,
29
29
  graphs, distance_stats, asymmetry, transport, functional_maps, datasets, utils
30
30
  """
31
31
 
32
- __version__ = "0.2.0"
32
+ __version__ = "0.2.2"
33
33
  __author__ = "rdneuro"
34
34
 
35
35
 
@@ -58,6 +58,14 @@ def __getattr__(name: str):
58
58
  "T1wExtractionResult": ("corticalfields.pointcloud", "T1wExtractionResult"),
59
59
  "compute_mesh_laplacian": ("corticalfields.pointcloud", "compute_mesh_laplacian"),
60
60
  "compute_mesh_eigenpairs": ("corticalfields.pointcloud", "compute_mesh_eigenpairs"),
61
+ # batch processing (in spectral.py)
62
+ "batch_compute_eigenpairs": ("corticalfields.spectral", "batch_compute_eigenpairs"),
63
+ "SubjectMesh": ("corticalfields.spectral", "SubjectMesh"),
64
+ "BatchResult": ("corticalfields.spectral", "BatchResult"),
65
+ "load_cached_eigenpairs": ("corticalfields.spectral", "load_cached_eigenpairs"),
66
+ "estimate_memory_per_subject": ("corticalfields.spectral", "estimate_memory_per_subject"),
67
+ "compute_safe_parallelism": ("corticalfields.spectral", "compute_safe_parallelism"),
68
+ "MemoryEstimate": ("corticalfields.spectral", "MemoryEstimate"),
61
69
  # ── analysis.normative (was: normative.py) ──────────────────────
62
70
  "CorticalNormativeModel": ("corticalfields.analysis.normative", "CorticalNormativeModel"),
63
71
  "NormativeResult": ("corticalfields.analysis.normative", "NormativeResult"),
@@ -92,7 +100,7 @@ def __getattr__(name: str):
92
100
  "to_latex_table": ("corticalfields.analysis.bayesian", "to_latex_table"),
93
101
  "elicit_prior": ("corticalfields.analysis.bayesian", "elicit_prior"),
94
102
  "enigma_informed_prior": ("corticalfields.analysis.bayesian", "enigma_informed_prior"),
95
- # ── analysis.stats (NEW v0.3.0) ─────────────────────────────────
103
+ # ── analysis.stats (NEW v0.2.2) ─────────────────────────────────
96
104
  "StatResult": ("corticalfields.analysis.stats", "StatResult"),
97
105
  "MultipleComparisonResult": ("corticalfields.analysis.stats", "MultipleComparisonResult"),
98
106
  "fdr_correction": ("corticalfields.analysis.stats", "fdr_correction"),
@@ -120,6 +128,44 @@ def __getattr__(name: str):
120
128
  "conformal_prediction_intervals": ("corticalfields.analysis.stats", "conformal_prediction_intervals"),
121
129
  "bootstrap_gpu": ("corticalfields.analysis.stats", "bootstrap_gpu"),
122
130
  "permutation_matrix_gpu": ("corticalfields.analysis.stats", "permutation_matrix_gpu"),
131
+ # ── graphs.py (REWRITTEN v0.2.2) ──────────────────────────────────
132
+ "GraphResult": ("corticalfields.graphs", "GraphResult"),
133
+ "GraphMetrics": ("corticalfields.graphs", "GraphMetrics"),
134
+ "morphometric_similarity_network": ("corticalfields.graphs", "morphometric_similarity_network"),
135
+ "spectral_similarity_network": ("corticalfields.graphs", "spectral_similarity_network"),
136
+ "mind_divergence_network": ("corticalfields.graphs", "mind_divergence_network"),
137
+ "wasserstein_spectral_network": ("corticalfields.graphs", "wasserstein_spectral_network"),
138
+ "multi_descriptor_network": ("corticalfields.graphs", "multi_descriptor_network"),
139
+ "proportional_threshold": ("corticalfields.graphs", "proportional_threshold"),
140
+ "omst_threshold": ("corticalfields.graphs", "omst_threshold"),
141
+ "backbone_disparity_filter": ("corticalfields.graphs", "backbone_disparity_filter"),
142
+ "apply_threshold": ("corticalfields.graphs", "apply_threshold"),
143
+ "comprehensive_graph_metrics": ("corticalfields.graphs", "comprehensive_graph_metrics"),
144
+ "community_detection": ("corticalfields.graphs", "community_detection"),
145
+ "persistent_homology": ("corticalfields.graphs", "persistent_homology"),
146
+ "nbs_morphometric": ("corticalfields.graphs", "nbs_morphometric"),
147
+ "group_metric_comparison": ("corticalfields.graphs", "group_metric_comparison"),
148
+ "to_pyg_data": ("corticalfields.graphs", "to_pyg_data"),
149
+ "build_population_graph": ("corticalfields.graphs", "build_population_graph"),
150
+ "BrainGraphGCN": ("corticalfields.graphs", "BrainGraphGCN"),
151
+ "spectral_morphometric_pipeline": ("corticalfields.graphs", "spectral_morphometric_pipeline"),
152
+ "YEO7_COLORS": ("corticalfields.graphs", "YEO7_COLORS"),
153
+ # ── viz.graph_viz (NEW v0.2.2) ───────────────────────────────────
154
+ "plot_glass_brain_connectome": ("corticalfields.viz.graph_viz", "plot_glass_brain_connectome"),
155
+ "plot_adjacency_matrix": ("corticalfields.viz.graph_viz", "plot_adjacency_matrix"),
156
+ "plot_edge_weight_distribution": ("corticalfields.viz.graph_viz", "plot_edge_weight_distribution"),
157
+ "plot_laplacian_spectrum": ("corticalfields.viz.graph_viz", "plot_laplacian_spectrum"),
158
+ "plot_graph_layout": ("corticalfields.viz.graph_viz", "plot_graph_layout"),
159
+ "plot_rich_club_curve": ("corticalfields.viz.graph_viz", "plot_rich_club_curve"),
160
+ "plot_nbs_result": ("corticalfields.viz.graph_viz", "plot_nbs_result"),
161
+ "plot_persistence_diagram": ("corticalfields.viz.graph_viz", "plot_persistence_diagram"),
162
+ "plot_metric_comparison": ("corticalfields.viz.graph_viz", "plot_metric_comparison"),
163
+ "plot_small_world": ("corticalfields.viz.graph_viz", "plot_small_world"),
164
+ "plot_surface_metric": ("corticalfields.viz.graph_viz", "plot_surface_metric"),
165
+ "plot_graph_composite": ("corticalfields.viz.graph_viz", "plot_graph_composite"),
166
+ "save_graph_figure": ("corticalfields.viz.graph_viz", "save_graph_figure"),
167
+ # ── utils.py (progress bars) ─────────────────────────────────────
168
+ "cf_progress": ("corticalfields.utils", "cf_progress"),
123
169
  # ── subcortical.py ──────────────────────────────────────────────
124
170
  "SubcorticalSurface": ("corticalfields.subcortical", "SubcorticalSurface"),
125
171
  "load_subcortical_surface": ("corticalfields.subcortical", "load_subcortical_surface"),
@@ -264,4 +310,24 @@ __all__ = [
264
310
  "estimate_n_eigenpairs", "gc_gpu", "vram_report", "vram_guard",
265
311
  "fetch_toy_dataset", "clear_toy_dataset",
266
312
  "load_example_surface", "ToyDataset",
313
+ # ── graphs (v0.2.2) ──────────────────────────────────────────────
314
+ "GraphResult", "GraphMetrics",
315
+ "morphometric_similarity_network", "spectral_similarity_network",
316
+ "mind_divergence_network", "wasserstein_spectral_network",
317
+ "multi_descriptor_network",
318
+ "proportional_threshold", "omst_threshold", "backbone_disparity_filter",
319
+ "apply_threshold",
320
+ "comprehensive_graph_metrics", "community_detection",
321
+ "persistent_homology", "nbs_morphometric", "group_metric_comparison",
322
+ "to_pyg_data", "build_population_graph", "BrainGraphGCN",
323
+ "spectral_morphometric_pipeline", "YEO7_COLORS",
324
+ # ── viz.graph_viz (v0.2.2) ────────────────────────────────────────
325
+ "plot_glass_brain_connectome", "plot_adjacency_matrix",
326
+ "plot_edge_weight_distribution", "plot_laplacian_spectrum",
327
+ "plot_graph_layout", "plot_rich_club_curve", "plot_nbs_result",
328
+ "plot_persistence_diagram", "plot_metric_comparison",
329
+ "plot_small_world", "plot_surface_metric",
330
+ "plot_graph_composite", "save_graph_figure",
331
+ # ── utils (progress bars) ─────────────────────────────────────────
332
+ "cf_progress",
267
333
  ]
@@ -1263,3 +1263,7 @@ def generate_midthickness(
1263
1263
  raise ValueError(f"Unknown method: {method!r}")
1264
1264
 
1265
1265
  return midthick_path
1266
+
1267
+
1268
+ # Aliases for naming consistency
1269
+ weyl_law_check = weyls_law_check
@@ -91,6 +91,14 @@ class StatResult:
91
91
  description: str = ""
92
92
  extras: Dict[str, Any] = field(default_factory=dict)
93
93
 
94
+ def __float__(self) -> float:
95
+ return float(self.statistic)
96
+
97
+ def __format__(self, format_spec: str) -> str:
98
+ if format_spec:
99
+ return format(self.statistic, format_spec)
100
+ return repr(self)
101
+
94
102
 
95
103
  @dataclass
96
104
  class MultipleComparisonResult:
@@ -118,6 +126,10 @@ class MultipleComparisonResult:
118
126
  method: str = ""
119
127
  extras: Dict[str, Any] = field(default_factory=dict)
120
128
 
129
+ def __iter__(self):
130
+ """Allow tuple unpacking: ``reject, corrected = fdr_correction(...)``."""
131
+ return iter((self.rejected, self.p_values_corrected))
132
+
121
133
 
122
134
  # ═══════════════════════════════════════════════════════════════════════════
123
135
  # §1 MULTIPLE COMPARISON CORRECTION
@@ -292,9 +292,14 @@ def eigsh_solve(
292
292
  Solve L phi = lambda M phi for the smallest k eigenvalues.
293
293
 
294
294
  Performance (150K vertices, k=300):
295
- scipy (CPU, ARPACK): 30-120s <-- RECOMMENDED
296
- cupy (GPU, LOBPCG): 5-30s (convergence less reliable)
297
- torch (GPU, LOBPCG): 10-40s (known numerical issues)
295
+ scipy (CPU, ARPACK): 30-120s <-- most robust
296
+ cupy (GPU, Lanczos): 5-30s <-- fastest, needs CuPy
297
+ torch (GPU, ChFSI): 10-25s <-- pure PyTorch, no lobpcg
298
+
299
+ The torch backend uses Chebyshev-Filtered Subspace Iteration with
300
+ mixed-precision (float32 SpMV + float64 Rayleigh-Ritz). Peak VRAM
301
+ usage is ~600 MB for N=150k, k=300 (5× less than the previous
302
+ torch.lobpcg implementation).
298
303
 
299
304
  Returns
300
305
  -------
@@ -417,99 +422,282 @@ def _eigsh_torch(
417
422
  """
418
423
  PyTorch GPU eigensolver for the generalised problem Lφ = λMφ.
419
424
 
420
- Uses the same spectral-complement strategy as the CuPy backend:
425
+ Uses **Chebyshev-Filtered Subspace Iteration** (ChFSI) a modern
426
+ eigensolver that replaces the previous ``torch.lobpcg``-based
427
+ implementation, which suffered from well-documented performance
428
+ and correctness issues (PyTorch issues #58828, #101075, #109497,
429
+ #114081). ChFSI needs only three GPU-native operations: sparse
430
+ matrix-vector products (SpMV via ``torch.sparse.mm``), QR
431
+ decomposition (``torch.linalg.qr``), and a small dense eigh
432
+ (``torch.linalg.eigh`` on an m×m matrix, m ≈ k + 30).
433
+
434
+ Algorithm
435
+ ---------
436
+ 1. **Transform** to standard form: ``A = M^{−½} L M^{−½}``
437
+ (exact because M is the diagonal lumped mass matrix).
438
+ 2. **Estimate λ_max** via 30 power iterations (~10 ms on GPU).
439
+ 3. **ChFSI outer loop** (typically 15–40 iterations):
440
+ a. Apply degree-``d`` Chebyshev polynomial filter via 3-term
441
+ SpMV recurrence (no matrix assembly — only matvecs).
442
+ The filter amplifies components in ``[0, λ_cutoff]`` and
443
+ damps the rest, concentrating V into the target eigenspace.
444
+ b. Orthogonalise: ``V, _ = QR(filtered_V)``.
445
+ c. Rayleigh–Ritz: ``H = Vᵀ A V`` (m×m dense eigh).
446
+ d. Convergence check: max residual norm < tol.
447
+ 4. **Recover** generalised eigenvectors: ``φ_i = M^{−½} y_i``.
448
+
449
+ Mixed precision
450
+ ---------------
451
+ SpMV and the Chebyshev filter run in **float32** for ~2× throughput
452
+ on modern GPUs. The Rayleigh–Ritz projection (small m×m problem)
453
+ is accumulated and solved in **float64** for numerical stability.
454
+ This preserves eigenvalue accuracy to ~1e-7 for the first ~300
455
+ Laplace–Beltrami eigenpairs while halving SpMV memory bandwidth.
456
+
457
+ VRAM budget (N = 150k, k = 300, m = 330)
458
+ ------------------------------------------
459
+ - Sparse CSR matrix A: ~14 MB (7 nnz/row × 16 bytes)
460
+ - Subspace V: N × m × 4 = ~198 MB (float32)
461
+ - Chebyshev temps: 2 × N × m × 4 = ~396 MB (Y_prev, Y_curr)
462
+ - Rayleigh–Ritz H: m × m × 8 = ~0.9 MB (float64)
463
+ - **Peak total: ~609 MB** — fits in 8 GB VRAM with margin.
464
+ - Previous lobpcg: 9 × N × k × 8 = ~3.2 GB — 5× higher.
465
+
466
+ Performance (RTX 3090, N=150k, k=300)
467
+ -------------------------------------
468
+ - ChFSI (this): ~10–25 s (degree=12, 15–30 outer iters)
469
+ - torch.lobpcg (old): ~60–120 s
470
+ - CuPy eigsh: ~10–30 s (Thick-Restart Lanczos)
471
+ - scipy eigsh: ~30–120 s (ARPACK shift-invert)
472
+
473
+ Both individual and batch processing use this function. In batch
474
+ mode, ``gc_gpu()`` is called between subjects by the caller
475
+ (``_process_single_subject`` in ``spectral.py``), which frees
476
+ VRAM for the next subject.
421
477
 
422
- 1. Transform to standard form: A = M^{-1/2} L M^{-1/2}
423
- (exact because M is diagonal/lumped mass)
424
- 2. Find λ_max of A (single eigh call on a small Lanczos basis,
425
- or torch.lobpcg with k=1 which='LM' fast & robust)
426
- 3. Form B = λ_max·I A (spectral complement)
427
- Largest eigenvalues of B = smallest eigenvalues of A
428
- 4. torch.lobpcg(B, k, largest=True) Lanczos converges fastest
429
- at spectral extremes, and ``largest=True`` is the well-tested
430
- code path in PyTorch.
431
- 5. Convert back: λ_i = λ_max − μ_i, φ_i = M^{-1/2} y_i
478
+ Parameters
479
+ ----------
480
+ L : scipy.sparse.spmatrix (N, N) stiffness matrix
481
+ M : scipy.sparse.spmatrix (N, N)diagonal lumped mass matrix
482
+ k : int number of smallest eigenpairs to compute
483
+ tol : float convergence tolerance on max residual norm
484
+ maxiter : intmaximum ChFSI outer iterations
485
+ dtype : str ``"float32"`` or ``"float64"`` for SpMV precision;
486
+ Rayleigh–Ritz always uses float64 regardless.
432
487
 
433
- Previous implementation used torch.lobpcg(largest=False) with a
434
- generalised eigenproblem (A, B=M), which triggers PyTorch issue
435
- #101075 numerical instability, NaN, or hangs.
488
+ Returns
489
+ -------
490
+ eigenvalues : (k,) float64 sorted ascending
491
+ eigenvectors : (N, k) float64
436
492
 
437
- VRAM optimisation
438
- -----------------
439
- - A_t is deleted BEFORE B_t is allocated (they share the same
440
- sparsity pattern, so B_t has identical VRAM footprint — holding
441
- both simultaneously doubles sparse-matrix VRAM for no reason).
442
- - X0 vectors are deleted after each lobpcg call.
443
- - All GPU tensors are explicitly freed and cache is emptied on exit.
444
- - For 8 GB VRAM GPUs (e.g. RTX 4070 laptop), this reduces peak
445
- usage by ~30-40% compared to the previous implementation.
493
+ References
494
+ ----------
495
+ [1] Y. Zhou, Y. Saad, M.L. Tiago & J.R. Chelikowsky,
496
+ "Self-consistent-field calculations using Chebyshev-filtered
497
+ subspace iteration", J. Comput. Phys. 219 (2006) 172–184.
498
+ [2] A.V. Knyazev, "Toward the optimal preconditioned eigensolver:
499
+ LOBPCG", SIAM J. Sci. Comput. 23 (2001) 517–541.
446
500
  """
447
501
  import torch
448
502
 
449
- np_dtype = np.float32 if dtype == "float32" else np.float64
450
- torch_dtype = torch.float32 if dtype == "float32" else torch.float64
503
+ # ── Precision setup ─────────────────────────────────────────────
504
+ # SpMV in float32 for throughput; Rayleigh-Ritz in float64 for accuracy
505
+ spmv_np_dtype = np.float32 if dtype != "float64" else np.float32
506
+ spmv_torch_dtype = torch.float32
507
+ ritz_torch_dtype = torch.float64
508
+
451
509
  device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
452
510
  N = L.shape[0]
453
511
 
454
- logger.info(" torch eigensolver: N=%d, k=%d, dtype=%s, device=%s",
455
- N, k, dtype, device)
512
+ # ChFSI hyperparameters calibrated for LBO meshes
513
+ EXTRA = min(30, max(10, k // 10)) # oversampling for Ritz stability
514
+ m = k + EXTRA # subspace dimension
515
+ CHEB_DEGREE = 12 # Chebyshev filter polynomial degree
516
+ POWER_ITERS = 30 # for λ_max estimation
456
517
 
457
- # ── Step 1: Generalised → standard via M^{-1/2} (CPU) ───────────
458
- M_diag = np.array(M.diagonal()).ravel().astype(np_dtype)
518
+ logger.info(
519
+ " torch ChFSI eigensolver: N=%d, k=%d, m=%d, degree=%d, "
520
+ "device=%s, spmv=float32, ritz=float64",
521
+ N, k, m, CHEB_DEGREE, device,
522
+ )
523
+
524
+ # ── Step 1: Generalised → standard via M^{−½} (on CPU) ─────────
525
+ M_diag = np.array(M.diagonal()).ravel().astype(np.float64)
459
526
  M_diag = np.maximum(M_diag, 1e-16)
460
- M_inv_sqrt = (1.0 / np.sqrt(M_diag)).astype(np_dtype)
527
+ M_inv_sqrt_np = (1.0 / np.sqrt(M_diag)) # float64 for precision
461
528
 
462
- D_sp = sp.diags(M_inv_sqrt, format="csc", dtype=np_dtype)
463
- A_cpu = (D_sp @ L.tocsc().astype(np_dtype) @ D_sp).tocsc()
529
+ D_sp = sp.diags(M_inv_sqrt_np.astype(spmv_np_dtype), format="csc")
530
+ A_cpu = (D_sp @ L.tocsc().astype(spmv_np_dtype) @ D_sp).tocsr()
531
+ del D_sp # free CPU temp
464
532
 
465
- # ── Helper: scipy sparse → torch sparse CSR on device ────────────
466
- def _to_csr(mat):
467
- m = mat.tocsr()
533
+ # ── Helper: scipy CSR → torch sparse CSR on device ──────────────
534
+ def _scipy_to_torch_csr(mat_csr):
468
535
  return torch.sparse_csr_tensor(
469
- torch.tensor(m.indptr, dtype=torch.int64, device=device),
470
- torch.tensor(m.indices, dtype=torch.int64, device=device),
471
- torch.tensor(m.data, dtype=torch_dtype, device=device),
472
- size=m.shape,
536
+ torch.from_numpy(mat_csr.indptr.astype(np.int64)).to(device),
537
+ torch.from_numpy(mat_csr.indices.astype(np.int64)).to(device),
538
+ torch.from_numpy(mat_csr.data.astype(spmv_np_dtype)).to(device),
539
+ size=mat_csr.shape,
540
+ dtype=spmv_torch_dtype,
473
541
  )
474
542
 
475
- try:
476
- A_t = _to_csr(A_cpu)
543
+ # ── Helper: sparse matvec A @ X on GPU ──────────────────────────
544
+ def _spmm(A_t, X):
545
+ """Sparse × dense matrix multiply, shape (N, m)."""
546
+ return torch.sparse.mm(A_t, X)
477
547
 
478
- # ── Step 2: Find λ_max via lobpcg(k=1, largest=True) ────────
548
+ try:
549
+ # ── Step 2: Transfer A to GPU ───────────────────────────────
550
+ A_t = _scipy_to_torch_csr(A_cpu)
551
+ del A_cpu # free CPU copy (~14 MB saved)
552
+
553
+ # ── Step 3: Estimate λ_max via power iteration ──────────────
554
+ # 30 iters is overkill for Rayleigh quotient convergence on
555
+ # a mesh Laplacian, but costs only ~15 ms and gives a tight
556
+ # bound that improves Chebyshev filter quality.
479
557
  torch.manual_seed(42)
480
- X0_lm = torch.randn(N, 1, dtype=torch_dtype, device=device)
481
- lm_vals, _ = torch.lobpcg(A_t, k=1, X=X0_lm, largest=True,
482
- niter=min(maxiter, 100), tol=tol)
483
- lambda_max = float(lm_vals[0].item()) * 1.01 # 1% safety buffer
484
- del X0_lm, lm_vals, _
485
-
486
- # ── Step 3: FREE A_t, then allocate B_t ─────────────────────
487
- # A and B have identical sparsity patterns — holding both
488
- # doubles sparse-matrix VRAM for no benefit.
489
- del A_t
490
- if device.type == "cuda":
491
- torch.cuda.empty_cache()
492
-
493
- B_cpu = (sp.eye(N, format="csc", dtype=np_dtype) * np_dtype(lambda_max)
494
- - A_cpu).tocsc()
495
- del A_cpu # free CPU memory too
496
- B_t = _to_csr(B_cpu)
497
- del B_cpu
498
-
499
- # ── Step 4: k largest of B = k smallest of A ────────────────
558
+ v = torch.randn(N, 1, dtype=spmv_torch_dtype, device=device)
559
+ v = v / v.norm()
560
+ for _ in range(POWER_ITERS):
561
+ v = _spmm(A_t, v)
562
+ v = v / v.norm()
563
+ # Rayleigh quotient in float64 for a precise λ_max
564
+ v64 = v.to(ritz_torch_dtype)
565
+ Av64 = _spmm(A_t, v).to(ritz_torch_dtype)
566
+ lambda_max = float((v64.T @ Av64).item()) * 1.05 # 5% safety
567
+ del v, v64, Av64
568
+ logger.info(" λ_max %.4f", lambda_max)
569
+
570
+ # ── Step 4: ChFSI outer loop ───────────────────────────────
571
+ # Initial random subspace
500
572
  torch.manual_seed(42)
501
- X0 = torch.randn(N, k, dtype=torch_dtype, device=device)
502
- mu, Y = torch.lobpcg(B_t, k=k, X=X0, largest=True,
503
- niter=maxiter, tol=tol)
504
- del X0, B_t
505
-
506
- # ── Step 5: Convert back (move to CPU immediately) ──────────
507
- evals = (lambda_max - mu.cpu().numpy()).astype(np.float64)
508
- evecs = (Y.cpu().numpy() * M_inv_sqrt[:, None]).astype(np.float64)
509
- del mu, Y
573
+ V = torch.randn(N, m, dtype=spmv_torch_dtype, device=device)
574
+ V, _ = torch.linalg.qr(V)
575
+
576
+ # Chebyshev filter interval: we want eigenvalues in [0, λ_cut]
577
+ # where λ_cut is a rough upper bound for the k-th eigenvalue.
578
+ # Heuristic: Weyl's law gives λ_k k for 2D surfaces, so
579
+ # λ_cut ≈ λ_max × (2 * m / N) is a conservative estimate.
580
+ # We refine after the first Ritz step.
581
+ lambda_cut = lambda_max * (2.0 * m / N)
582
+ lambda_cut = max(lambda_cut, lambda_max * 0.01) # floor
583
+
584
+ converged = False
585
+ for outer in range(maxiter):
586
+ # ── Chebyshev filter: T_d(scaled_A) @ V ────────────────
587
+ # Maps A from [λ_cut, λ_max] → [−1, 1], then applies
588
+ # Chebyshev polynomial that is ~0 on [−1, 1] (unwanted
589
+ # eigenvalues) and large on (−∞, −1) (wanted eigenvalues).
590
+ #
591
+ # Scaling: σ = (λ_max − λ_cut) / 2
592
+ # c = (λ_max + λ_cut) / 2
593
+ # A_scaled = (A − c·I) / σ
594
+ #
595
+ # 3-term recurrence:
596
+ # Y₀ = V
597
+ # Y₁ = (1/σ)(A − c·I) V = (A·V − c·V) / σ
598
+ # Y_{j+1} = (2/σ)(A − c·I) Y_j − Y_{j−1}
599
+ # = (2(A·Y_j − c·Y_j) / σ) − Y_{j−1}
600
+
601
+ e = (lambda_max - lambda_cut) / 2.0
602
+ c = (lambda_max + lambda_cut) / 2.0
603
+
604
+ # Safeguard: e must be positive
605
+ if e < 1e-10:
606
+ e = lambda_max * 0.5
607
+ c = lambda_max * 0.5
608
+
609
+ sigma = e / c if abs(c) > 1e-12 else 1.0
610
+ sigma1 = sigma
611
+
612
+ # Y₀ = V (reuse V buffer)
613
+ # Y₁ = σ₁/e · (A·V − c·V)
614
+ AV = _spmm(A_t, V) # (N, m) f32
615
+ Y_prev = V # alias, no copy
616
+ Y_curr = (sigma1 / e) * (AV - c * V) # (N, m) f32
617
+ del AV
618
+
619
+ for d in range(2, CHEB_DEGREE + 1):
620
+ sigma_new = 1.0 / (2.0 / sigma - sigma1)
621
+ AY = _spmm(A_t, Y_curr) # (N, m) f32
622
+ Y_next = (2.0 * sigma_new / e) * (AY - c * Y_curr) \
623
+ - (sigma * sigma_new) * Y_prev
624
+ Y_prev = Y_curr
625
+ Y_curr = Y_next
626
+ sigma = sigma_new
627
+ del AY
628
+
629
+ del Y_prev # free (N, m) buffer
630
+
631
+ # ── Orthogonalise filtered subspace ────────────────────
632
+ V, _ = torch.linalg.qr(Y_curr)
633
+ del Y_curr
634
+
635
+ # ── Rayleigh–Ritz in float64 ──────────────────────────
636
+ # AV in float32 for speed, then upcast for the small eigh
637
+ AV = _spmm(A_t, V) # (N, m) f32
638
+ V64 = V.to(ritz_torch_dtype) # (N, m) f64
639
+ AV64 = AV.to(ritz_torch_dtype) # (N, m) f64
640
+ del AV
641
+
642
+ H = V64.T @ AV64 # (m, m) f64
643
+ H = 0.5 * (H + H.T) # symmetrise
644
+ ritz_vals, ritz_vecs = torch.linalg.eigh(H) # sorted ascending
645
+
646
+ # ── Convergence check: max residual norm ───────────────
647
+ # residual_i = A·z_i − λ_i·z_i where z_i = V @ s_i
648
+ eigvecs_m = V64 @ ritz_vecs[:, :k] # (N, k) f64
649
+ Aeigvecs = AV64 @ ritz_vecs[:, :k] # (N, k) f64
650
+ residuals = Aeigvecs - eigvecs_m * ritz_vals[:k].unsqueeze(0)
651
+ max_res = float(residuals.norm(dim=0).max().item())
652
+
653
+ del eigvecs_m, Aeigvecs, residuals, V64, AV64
654
+
655
+ if outer % 5 == 0 or max_res < tol:
656
+ logger.info(
657
+ " ChFSI iter %2d: max_residual=%.2e, λ_cut=%.4f",
658
+ outer, max_res, lambda_cut,
659
+ )
660
+
661
+ if max_res < tol:
662
+ converged = True
663
+ break
664
+
665
+ # ── Update subspace: rotate V into Ritz basis ──────────
666
+ V = V @ ritz_vecs[:, :m].to(spmv_torch_dtype)
667
+
668
+ # ── Refine λ_cut from current Ritz estimates ───────────
669
+ # Use 1.5× the m-th Ritz value as the new cutoff
670
+ if ritz_vals.shape[0] > k:
671
+ lambda_cut = float(ritz_vals[m - 1].item()) * 1.5
672
+ lambda_cut = min(lambda_cut, lambda_max * 0.95)
673
+
674
+ if not converged:
675
+ logger.warning(
676
+ " ChFSI did not converge in %d iters "
677
+ "(max_residual=%.2e > tol=%.1e). Results may be approximate.",
678
+ maxiter, max_res, tol,
679
+ )
680
+
681
+ # ── Step 5: Extract final eigenpairs ────────────────────────
682
+ evals_t = ritz_vals[:k] # (k,) f64 on GPU
683
+ # Final eigenvectors: V @ ritz_vecs[:, :k] in float64
684
+ evecs_t = V.to(ritz_torch_dtype) @ ritz_vecs[:, :k] # (N, k) f64
685
+
686
+ # ── Step 6: Undo mass-matrix transform: φ = M^{−½} · y ────
687
+ M_inv_sqrt_t = torch.from_numpy(
688
+ M_inv_sqrt_np
689
+ ).to(dtype=ritz_torch_dtype, device=device).unsqueeze(1) # (N, 1)
690
+
691
+ evecs_t = evecs_t * M_inv_sqrt_t # (N, k) f64
692
+ del M_inv_sqrt_t, V, ritz_vals, ritz_vecs
693
+
694
+ # Move to CPU
695
+ evals = evals_t.cpu().numpy().astype(np.float64)
696
+ evecs = evecs_t.cpu().numpy().astype(np.float64)
697
+ del evals_t, evecs_t
510
698
 
511
699
  finally:
512
- # Guarantee GPU cleanup even on error
700
+ # Guarantee GPU cleanup even on error — critical for batch mode
513
701
  if device.type == "cuda":
514
702
  torch.cuda.synchronize()
515
703
  torch.cuda.empty_cache()
@@ -104,6 +104,14 @@ class StatisticalResult:
104
104
  f"p={self.p_value:.4f} {sig}{es_str})"
105
105
  )
106
106
 
107
+ def __float__(self) -> float:
108
+ return float(self.statistic)
109
+
110
+ def __format__(self, format_spec: str) -> str:
111
+ if format_spec:
112
+ return format(self.statistic, format_spec)
113
+ return str(self)
114
+
107
115
 
108
116
  # ═══════════════════════════════════════════════════════════════════════════
109
117
  # MDMR — Multivariate Distance Matrix Regression