rippl 0.1.3__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.
- rippl-0.1.3/LICENSE +21 -0
- rippl-0.1.3/PKG-INFO +115 -0
- rippl-0.1.3/README.md +102 -0
- rippl-0.1.3/pyproject.toml +27 -0
- rippl-0.1.3/rippl/__init__.py +7 -0
- rippl-0.1.3/rippl/api.py +146 -0
- rippl-0.1.3/rippl/cli.py +120 -0
- rippl-0.1.3/rippl/config/models.py +35 -0
- rippl-0.1.3/rippl/core/__init__.py +14 -0
- rippl-0.1.3/rippl/core/api.py +86 -0
- rippl-0.1.3/rippl/core/artifact.py +49 -0
- rippl-0.1.3/rippl/core/config.py +87 -0
- rippl-0.1.3/rippl/core/engine.py +111 -0
- rippl-0.1.3/rippl/core/equation.py +63 -0
- rippl-0.1.3/rippl/core/equation_system.py +55 -0
- rippl-0.1.3/rippl/core/exceptions.py +11 -0
- rippl-0.1.3/rippl/core/experiment.py +225 -0
- rippl-0.1.3/rippl/core/inverse.py +175 -0
- rippl-0.1.3/rippl/core/nondim.py +174 -0
- rippl-0.1.3/rippl/core/operator_experiment.py +69 -0
- rippl-0.1.3/rippl/core/simulation.py +176 -0
- rippl-0.1.3/rippl/core/solver_registry.py +73 -0
- rippl-0.1.3/rippl/core/system.py +303 -0
- rippl-0.1.3/rippl/datasets/__init__.py +6 -0
- rippl-0.1.3/rippl/datasets/generators.py +59 -0
- rippl-0.1.3/rippl/diagnostics/__init__.py +8 -0
- rippl-0.1.3/rippl/diagnostics/energy.py +48 -0
- rippl-0.1.3/rippl/diagnostics/metrics.py +31 -0
- rippl-0.1.3/rippl/diagnostics/physics_validator.py +97 -0
- rippl-0.1.3/rippl/diagnostics/spectral.py +26 -0
- rippl-0.1.3/rippl/io/__init__.py +13 -0
- rippl-0.1.3/rippl/io/checkpoint.py +73 -0
- rippl-0.1.3/rippl/io/export.py +68 -0
- rippl-0.1.3/rippl/nn/__init__.py +30 -0
- rippl-0.1.3/rippl/nn/adaptivesampler.py +117 -0
- rippl-0.1.3/rippl/nn/boundary_block.py +138 -0
- rippl-0.1.3/rippl/nn/conservation_block.py +129 -0
- rippl-0.1.3/rippl/nn/embedding.py +140 -0
- rippl-0.1.3/rippl/nn/energy.py +151 -0
- rippl-0.1.3/rippl/nn/fno.py +224 -0
- rippl-0.1.3/rippl/nn/fourier_mlp.py +66 -0
- rippl-0.1.3/rippl/nn/gradient.py +116 -0
- rippl-0.1.3/rippl/nn/hamiltonian.py +155 -0
- rippl-0.1.3/rippl/nn/hybrid_stepper.py +135 -0
- rippl-0.1.3/rippl/nn/laplacian.py +239 -0
- rippl-0.1.3/rippl/nn/mlp.py +73 -0
- rippl-0.1.3/rippl/nn/multi_field_mlp.py +34 -0
- rippl-0.1.3/rippl/nn/multiscale_ff.py +126 -0
- rippl-0.1.3/rippl/nn/nn_operator_wrapper.py +192 -0
- rippl-0.1.3/rippl/nn/oscillator.py +167 -0
- rippl-0.1.3/rippl/nn/residual.py +115 -0
- rippl-0.1.3/rippl/nn/siren.py +91 -0
- rippl-0.1.3/rippl/nn/spectral.py +184 -0
- rippl-0.1.3/rippl/nn/spectral_conv.py +127 -0
- rippl-0.1.3/rippl/nn/spectral_reg.py +99 -0
- rippl-0.1.3/rippl/operators/__init__.py +7 -0
- rippl-0.1.3/rippl/operators/grid_utils.py +47 -0
- rippl-0.1.3/rippl/operators/operator_mode.py +69 -0
- rippl-0.1.3/rippl/physics/__init__.py +16 -0
- rippl-0.1.3/rippl/physics/boundary.py +123 -0
- rippl-0.1.3/rippl/physics/conservation.py +54 -0
- rippl-0.1.3/rippl/physics/conservative.py +105 -0
- rippl-0.1.3/rippl/physics/derivatives.py +104 -0
- rippl-0.1.3/rippl/physics/distance.py +140 -0
- rippl-0.1.3/rippl/physics/elasticity.py +53 -0
- rippl-0.1.3/rippl/physics/equation.py +7 -0
- rippl-0.1.3/rippl/physics/fractional.py +118 -0
- rippl-0.1.3/rippl/physics/hamilton_jacobi.py +107 -0
- rippl-0.1.3/rippl/physics/navier_stokes.py +88 -0
- rippl-0.1.3/rippl/physics/operators.py +672 -0
- rippl-0.1.3/rippl/physics/pde.py +24 -0
- rippl-0.1.3/rippl/physics/phase_field.py +131 -0
- rippl-0.1.3/rippl/physics/reaction_diffusion.py +136 -0
- rippl-0.1.3/rippl/physics/residuals.py +125 -0
- rippl-0.1.3/rippl/physics/schrodinger.py +79 -0
- rippl-0.1.3/rippl/physics/shock.py +59 -0
- rippl-0.1.3/rippl/solvers/__init__.py +7 -0
- rippl-0.1.3/rippl/solvers/fd_solver.py +363 -0
- rippl-0.1.3/rippl/solvers/spectral_solver.py +65 -0
- rippl-0.1.3/rippl/training/__init__.py +8 -0
- rippl-0.1.3/rippl/training/adaptive_sampler.py +74 -0
- rippl-0.1.3/rippl/training/callbacks.py +91 -0
- rippl-0.1.3/rippl/training/causal.py +78 -0
- rippl-0.1.3/rippl/training/fno_flywheel.py +152 -0
- rippl-0.1.3/rippl/training/lbfgs_config.py +20 -0
- rippl-0.1.3/rippl/training/lightning.py +51 -0
- rippl-0.1.3/rippl/training/lightning_engine.py +85 -0
- rippl-0.1.3/rippl/training/ntk_weighting.py +134 -0
- rippl-0.1.3/rippl/training/pinn_recipe.py +211 -0
- rippl-0.1.3/rippl/training/uq.py +213 -0
- rippl-0.1.3/rippl/visualize.py +49 -0
- rippl-0.1.3/rippl.egg-info/PKG-INFO +115 -0
- rippl-0.1.3/rippl.egg-info/SOURCES.txt +96 -0
- rippl-0.1.3/rippl.egg-info/dependency_links.txt +1 -0
- rippl-0.1.3/rippl.egg-info/entry_points.txt +2 -0
- rippl-0.1.3/rippl.egg-info/requires.txt +5 -0
- rippl-0.1.3/rippl.egg-info/top_level.txt +1 -0
- rippl-0.1.3/setup.cfg +4 -0
rippl-0.1.3/LICENSE
ADDED
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2026 NeuralWave Core Contributors
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|
rippl-0.1.3/PKG-INFO
ADDED
|
@@ -0,0 +1,115 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: rippl
|
|
3
|
+
Version: 0.1.3
|
|
4
|
+
Summary: High-performance physics compiler and PINN engine.
|
|
5
|
+
Requires-Python: >=3.10
|
|
6
|
+
Description-Content-Type: text/markdown
|
|
7
|
+
License-File: LICENSE
|
|
8
|
+
Requires-Dist: torch>=2.0
|
|
9
|
+
Requires-Dist: pydantic
|
|
10
|
+
Provides-Extra: distributed
|
|
11
|
+
Requires-Dist: pytorch-lightning>=2.0; extra == "distributed"
|
|
12
|
+
Dynamic: license-file
|
|
13
|
+
|
|
14
|
+
# Rippl
|
|
15
|
+
### Differentiable Physics Engine for PyTorch
|
|
16
|
+
|
|
17
|
+

|
|
18
|
+

|
|
19
|
+

|
|
20
|
+

|
|
21
|
+
|
|
22
|
+
---
|
|
23
|
+
|
|
24
|
+
**Rippl** is a PyTorch-native, C++-less scientific machine learning framework designed for researchers who are tired of bloated dependencies and slow development cycles. It bridges the gap between **Physics-Informed Neural Networks (PINNs)** and **Discrete Numerical Simulation (FDM)** into a unified, Pydantic-validated API.
|
|
25
|
+
|
|
26
|
+
## Why Rippl? (The Hook)
|
|
27
|
+
|
|
28
|
+
Most SciML libraries suffer from "Academic Bloat": opaque C++ bindings, fragile custom CUDA kernels, and gradient pathologies that make training feel like alchemy. Rippl cuts through the noise:
|
|
29
|
+
|
|
30
|
+
- **Strictly Native PyTorch**: No custom C++ or CUDA. We use `F.conv2d` discrete stencils for performance, meaning it runs anywhere PyTorch does—from your MacBook to an H100 cluster.
|
|
31
|
+
- **The $O(1)$ Autograd Cache**: Our `requires_derived` system pre-caches tensors and graph paths. Stop wasting VRAM on redundant graph traversals during complex PDE residual computations.
|
|
32
|
+
- **Algebraic Hard Constraints**: Forget "soft" boundary losses that clash with your physics. Rippl enforces Dirichlet and Neumann conditions algebraically using distance-field masking, ensuring zero-error boundaries from step one.
|
|
33
|
+
- **Differentiable Simulation**: Backprop through time (BPTT) using standard FDM stencils mapped directly to convolution operations.
|
|
34
|
+
|
|
35
|
+
---
|
|
36
|
+
|
|
37
|
+
## Installation
|
|
38
|
+
|
|
39
|
+
```bash
|
|
40
|
+
pip install git+https://github.com/esranow/rippleml.git
|
|
41
|
+
```
|
|
42
|
+
|
|
43
|
+
*For local development:*
|
|
44
|
+
```bash
|
|
45
|
+
git clone https://github.com/esranow/rippl.git
|
|
46
|
+
cd rippl && pip install -e .
|
|
47
|
+
```
|
|
48
|
+
|
|
49
|
+
---
|
|
50
|
+
|
|
51
|
+
## The 10-Line Quickstart
|
|
52
|
+
|
|
53
|
+
Solve a 1D Wave Equation ($u_{tt} - c^2 u_{xx} = 0$) with zero-gradient overhead:
|
|
54
|
+
|
|
55
|
+
```python
|
|
56
|
+
import torch
|
|
57
|
+
from rippl.core import System, Domain, Experiment
|
|
58
|
+
from rippl.physics.operators import Laplacian, TimeDerivative
|
|
59
|
+
from rippl.physics.equation import Equation
|
|
60
|
+
|
|
61
|
+
# 1. Define the Physics (u_tt - 1.0 * u_xx = 0)
|
|
62
|
+
wave_eq = Equation([
|
|
63
|
+
(1.0, TimeDerivative(order=2)),
|
|
64
|
+
(-1.0, Laplacian())
|
|
65
|
+
])
|
|
66
|
+
|
|
67
|
+
# 2. Setup the System and Domain
|
|
68
|
+
domain = Domain(spatial_dims=1, bounds=[(0, 1)], resolution=(100,))
|
|
69
|
+
system = System(equation=wave_eq, domain=domain, fields=["u"])
|
|
70
|
+
|
|
71
|
+
# 3. Train the PINN
|
|
72
|
+
model = torch.nn.Sequential(torch.nn.Linear(2, 50), torch.nn.Tanh(), torch.nn.Linear(50, 1))
|
|
73
|
+
exp = Experiment(system, model, torch.optim.Adam(model.parameters()))
|
|
74
|
+
exp.train(epochs=100)
|
|
75
|
+
```
|
|
76
|
+
|
|
77
|
+
---
|
|
78
|
+
|
|
79
|
+
## Deep Dive: Core Architecture
|
|
80
|
+
|
|
81
|
+
### 1. The Autograd Cache (`requires_derived`)
|
|
82
|
+
In standard PINNs, calculating $u_{xxt}$ requires multiple passes through the autograd graph. Rippl's engine analyzes the `Operator` signatures in your `EquationSystem` and builds a directed acyclic graph (DAG) of required derivatives. These are computed in a single optimized pass and cached in the `derived` dictionary, slashing the memory footprint of high-order PDEs.
|
|
83
|
+
|
|
84
|
+
### 2. Algebraic Boundaries
|
|
85
|
+
Soft boundary losses are the primary cause of PINN instability. Rippl uses a `HardConstraintWrapper` combined with a `DistanceFunction` $D(x)$:
|
|
86
|
+
$$u_{pred}(x) = G(x) \cdot u_{model}(x) + B(x)$$
|
|
87
|
+
Where $G(x)$ vanishes at the boundary and $B(x)$ enforces the exact value. The result? **The optimizer never sees the boundary; it only sees the residual.**
|
|
88
|
+
|
|
89
|
+
### 3. Differentiable FDM Convolutions
|
|
90
|
+
Standard Finite Difference Methods (FDM) are usually isolated from ML training. Rippl maps these stencils to `torch.nn.functional.conv2d`. This allows you to use a traditional solver as a "Layer" in your network, enabling hybrid architectures that are both numerically stable and learnable.
|
|
91
|
+
|
|
92
|
+
---
|
|
93
|
+
|
|
94
|
+
## Advanced Physics Capabilities
|
|
95
|
+
|
|
96
|
+
- **Spectral Bias Mitigation**: Native support for **SIRENs** (Sinusoidal Representation Networks) and **Fourier Neural Operators (FNOs)** to capture high-frequency shock fronts.
|
|
97
|
+
- **Conservation Law Enforcement**: Wrappers for enforcing flux conservation and divergence-free fields (e.g., in Navier-Stokes).
|
|
98
|
+
- **Inverse Problems**: Parameters (like viscosity $\nu$ or wave speed $c$) can be marked as `requires_grad`, allowing the system to discover physics from raw data automatically.
|
|
99
|
+
|
|
100
|
+
---
|
|
101
|
+
|
|
102
|
+
## "Vibecoding" & Contribution Philosophy
|
|
103
|
+
|
|
104
|
+
Rippl is built for **speed and transparency**. Our codebase is intentionally monolithic and terse. We prioritize architectural elegance over feature-creep.
|
|
105
|
+
|
|
106
|
+
**The Contract:**
|
|
107
|
+
1. **Strict Perimeter**: All inputs are validated via Pydantic. If your dimensions don't match your physics, Rippl fails fast with a clear error message.
|
|
108
|
+
2. **Raw Core**: Inside the loop, it's pure, unadulterated PyTorch. No abstraction leaks. No hidden overhead.
|
|
109
|
+
|
|
110
|
+
If you want to contribute, keep it dry. No boilerplate. No unnecessary classes. If it can be done with a tensor operation, do it with a tensor operation.
|
|
111
|
+
|
|
112
|
+
---
|
|
113
|
+
|
|
114
|
+
## License
|
|
115
|
+
MIT. Go build something cool.
|
rippl-0.1.3/README.md
ADDED
|
@@ -0,0 +1,102 @@
|
|
|
1
|
+
# Rippl
|
|
2
|
+
### Differentiable Physics Engine for PyTorch
|
|
3
|
+
|
|
4
|
+

|
|
5
|
+

|
|
6
|
+

|
|
7
|
+

|
|
8
|
+
|
|
9
|
+
---
|
|
10
|
+
|
|
11
|
+
**Rippl** is a PyTorch-native, C++-less scientific machine learning framework designed for researchers who are tired of bloated dependencies and slow development cycles. It bridges the gap between **Physics-Informed Neural Networks (PINNs)** and **Discrete Numerical Simulation (FDM)** into a unified, Pydantic-validated API.
|
|
12
|
+
|
|
13
|
+
## Why Rippl? (The Hook)
|
|
14
|
+
|
|
15
|
+
Most SciML libraries suffer from "Academic Bloat": opaque C++ bindings, fragile custom CUDA kernels, and gradient pathologies that make training feel like alchemy. Rippl cuts through the noise:
|
|
16
|
+
|
|
17
|
+
- **Strictly Native PyTorch**: No custom C++ or CUDA. We use `F.conv2d` discrete stencils for performance, meaning it runs anywhere PyTorch does—from your MacBook to an H100 cluster.
|
|
18
|
+
- **The $O(1)$ Autograd Cache**: Our `requires_derived` system pre-caches tensors and graph paths. Stop wasting VRAM on redundant graph traversals during complex PDE residual computations.
|
|
19
|
+
- **Algebraic Hard Constraints**: Forget "soft" boundary losses that clash with your physics. Rippl enforces Dirichlet and Neumann conditions algebraically using distance-field masking, ensuring zero-error boundaries from step one.
|
|
20
|
+
- **Differentiable Simulation**: Backprop through time (BPTT) using standard FDM stencils mapped directly to convolution operations.
|
|
21
|
+
|
|
22
|
+
---
|
|
23
|
+
|
|
24
|
+
## Installation
|
|
25
|
+
|
|
26
|
+
```bash
|
|
27
|
+
pip install git+https://github.com/esranow/rippleml.git
|
|
28
|
+
```
|
|
29
|
+
|
|
30
|
+
*For local development:*
|
|
31
|
+
```bash
|
|
32
|
+
git clone https://github.com/esranow/rippl.git
|
|
33
|
+
cd rippl && pip install -e .
|
|
34
|
+
```
|
|
35
|
+
|
|
36
|
+
---
|
|
37
|
+
|
|
38
|
+
## The 10-Line Quickstart
|
|
39
|
+
|
|
40
|
+
Solve a 1D Wave Equation ($u_{tt} - c^2 u_{xx} = 0$) with zero-gradient overhead:
|
|
41
|
+
|
|
42
|
+
```python
|
|
43
|
+
import torch
|
|
44
|
+
from rippl.core import System, Domain, Experiment
|
|
45
|
+
from rippl.physics.operators import Laplacian, TimeDerivative
|
|
46
|
+
from rippl.physics.equation import Equation
|
|
47
|
+
|
|
48
|
+
# 1. Define the Physics (u_tt - 1.0 * u_xx = 0)
|
|
49
|
+
wave_eq = Equation([
|
|
50
|
+
(1.0, TimeDerivative(order=2)),
|
|
51
|
+
(-1.0, Laplacian())
|
|
52
|
+
])
|
|
53
|
+
|
|
54
|
+
# 2. Setup the System and Domain
|
|
55
|
+
domain = Domain(spatial_dims=1, bounds=[(0, 1)], resolution=(100,))
|
|
56
|
+
system = System(equation=wave_eq, domain=domain, fields=["u"])
|
|
57
|
+
|
|
58
|
+
# 3. Train the PINN
|
|
59
|
+
model = torch.nn.Sequential(torch.nn.Linear(2, 50), torch.nn.Tanh(), torch.nn.Linear(50, 1))
|
|
60
|
+
exp = Experiment(system, model, torch.optim.Adam(model.parameters()))
|
|
61
|
+
exp.train(epochs=100)
|
|
62
|
+
```
|
|
63
|
+
|
|
64
|
+
---
|
|
65
|
+
|
|
66
|
+
## Deep Dive: Core Architecture
|
|
67
|
+
|
|
68
|
+
### 1. The Autograd Cache (`requires_derived`)
|
|
69
|
+
In standard PINNs, calculating $u_{xxt}$ requires multiple passes through the autograd graph. Rippl's engine analyzes the `Operator` signatures in your `EquationSystem` and builds a directed acyclic graph (DAG) of required derivatives. These are computed in a single optimized pass and cached in the `derived` dictionary, slashing the memory footprint of high-order PDEs.
|
|
70
|
+
|
|
71
|
+
### 2. Algebraic Boundaries
|
|
72
|
+
Soft boundary losses are the primary cause of PINN instability. Rippl uses a `HardConstraintWrapper` combined with a `DistanceFunction` $D(x)$:
|
|
73
|
+
$$u_{pred}(x) = G(x) \cdot u_{model}(x) + B(x)$$
|
|
74
|
+
Where $G(x)$ vanishes at the boundary and $B(x)$ enforces the exact value. The result? **The optimizer never sees the boundary; it only sees the residual.**
|
|
75
|
+
|
|
76
|
+
### 3. Differentiable FDM Convolutions
|
|
77
|
+
Standard Finite Difference Methods (FDM) are usually isolated from ML training. Rippl maps these stencils to `torch.nn.functional.conv2d`. This allows you to use a traditional solver as a "Layer" in your network, enabling hybrid architectures that are both numerically stable and learnable.
|
|
78
|
+
|
|
79
|
+
---
|
|
80
|
+
|
|
81
|
+
## Advanced Physics Capabilities
|
|
82
|
+
|
|
83
|
+
- **Spectral Bias Mitigation**: Native support for **SIRENs** (Sinusoidal Representation Networks) and **Fourier Neural Operators (FNOs)** to capture high-frequency shock fronts.
|
|
84
|
+
- **Conservation Law Enforcement**: Wrappers for enforcing flux conservation and divergence-free fields (e.g., in Navier-Stokes).
|
|
85
|
+
- **Inverse Problems**: Parameters (like viscosity $\nu$ or wave speed $c$) can be marked as `requires_grad`, allowing the system to discover physics from raw data automatically.
|
|
86
|
+
|
|
87
|
+
---
|
|
88
|
+
|
|
89
|
+
## "Vibecoding" & Contribution Philosophy
|
|
90
|
+
|
|
91
|
+
Rippl is built for **speed and transparency**. Our codebase is intentionally monolithic and terse. We prioritize architectural elegance over feature-creep.
|
|
92
|
+
|
|
93
|
+
**The Contract:**
|
|
94
|
+
1. **Strict Perimeter**: All inputs are validated via Pydantic. If your dimensions don't match your physics, Rippl fails fast with a clear error message.
|
|
95
|
+
2. **Raw Core**: Inside the loop, it's pure, unadulterated PyTorch. No abstraction leaks. No hidden overhead.
|
|
96
|
+
|
|
97
|
+
If you want to contribute, keep it dry. No boilerplate. No unnecessary classes. If it can be done with a tensor operation, do it with a tensor operation.
|
|
98
|
+
|
|
99
|
+
---
|
|
100
|
+
|
|
101
|
+
## License
|
|
102
|
+
MIT. Go build something cool.
|
|
@@ -0,0 +1,27 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["setuptools>=68", "wheel"]
|
|
3
|
+
build-backend = "setuptools.build_meta"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "rippl"
|
|
7
|
+
version = "0.1.3"
|
|
8
|
+
description = "High-performance physics compiler and PINN engine."
|
|
9
|
+
readme = "README.md"
|
|
10
|
+
requires-python = ">=3.10"
|
|
11
|
+
dependencies = [
|
|
12
|
+
"torch>=2.0",
|
|
13
|
+
"pydantic"
|
|
14
|
+
]
|
|
15
|
+
|
|
16
|
+
[project.optional-dependencies]
|
|
17
|
+
distributed = ["pytorch-lightning>=2.0"]
|
|
18
|
+
|
|
19
|
+
[project.scripts]
|
|
20
|
+
rippl = "rippl.cli:main"
|
|
21
|
+
|
|
22
|
+
[tool.setuptools]
|
|
23
|
+
packages = ["rippl", "rippl.core", "rippl.physics", "rippl.solvers", "rippl.training", "rippl.nn", "rippl.config", "rippl.operators", "rippl.io", "rippl.datasets", "rippl.diagnostics"]
|
|
24
|
+
|
|
25
|
+
[tool.pytest.ini_options]
|
|
26
|
+
testpaths = ["rippl/tests"]
|
|
27
|
+
python_files = "test_*.py"
|
|
@@ -0,0 +1,7 @@
|
|
|
1
|
+
__version__ = "0.0.1"
|
|
2
|
+
# from rippl.api import train, simulate, identify
|
|
3
|
+
from rippl.core.system import System, Domain, Constraint
|
|
4
|
+
from rippl.core.config import register_operator, register_solver
|
|
5
|
+
from rippl.core.api import compile, run
|
|
6
|
+
|
|
7
|
+
__all__ = ["System", "Domain", "Constraint", "register_operator", "register_solver", "compile", "run"]
|
rippl-0.1.3/rippl/api.py
ADDED
|
@@ -0,0 +1,146 @@
|
|
|
1
|
+
import torch
|
|
2
|
+
import os
|
|
3
|
+
from typing import Union, Dict, Any, Optional, Tuple
|
|
4
|
+
from rippl.core.config import ConfigParser
|
|
5
|
+
from rippl.core.system import System
|
|
6
|
+
from rippl.core.experiment import Experiment
|
|
7
|
+
from rippl.nn.registry import build_model, load_model
|
|
8
|
+
from rippl.export.exporter import export_model
|
|
9
|
+
|
|
10
|
+
def train(config_path_or_dict: Union[str, Dict[str, Any]]) -> Tuple[torch.nn.Module, Dict[str, Any]]:
|
|
11
|
+
"""
|
|
12
|
+
Train a physics-informed neural network model based on the provided configuration.
|
|
13
|
+
|
|
14
|
+
Args:
|
|
15
|
+
config_path_or_dict: Path to a YAML/JSON config file or a configuration dictionary.
|
|
16
|
+
|
|
17
|
+
Returns:
|
|
18
|
+
A tuple of (trained_model, training_results).
|
|
19
|
+
"""
|
|
20
|
+
config = ConfigParser.load(config_path_or_dict)
|
|
21
|
+
system = System.from_config(config)
|
|
22
|
+
|
|
23
|
+
# Model configuration
|
|
24
|
+
model_cfg = config.get("model", {})
|
|
25
|
+
if not model_cfg:
|
|
26
|
+
raise ValueError("Configuration missing 'model' section.")
|
|
27
|
+
|
|
28
|
+
model = build_model(model_cfg["name"], model_cfg["config"])
|
|
29
|
+
|
|
30
|
+
# Training configuration
|
|
31
|
+
train_cfg = config.get("training", {})
|
|
32
|
+
if not train_cfg:
|
|
33
|
+
raise ValueError("Configuration missing 'training' section.")
|
|
34
|
+
|
|
35
|
+
optimizer_name = train_cfg.get("optimizer", "Adam")
|
|
36
|
+
learning_rate = train_cfg.get("lr", 1e-3)
|
|
37
|
+
optimizer_class = getattr(torch.optim, optimizer_name)
|
|
38
|
+
optimizer = optimizer_class(model.parameters(), lr=learning_rate)
|
|
39
|
+
|
|
40
|
+
# Experiment orchestration
|
|
41
|
+
experiment_kwargs = train_cfg.get("experiment_kwargs", {})
|
|
42
|
+
experiment = Experiment(system, model, optimizer, **experiment_kwargs)
|
|
43
|
+
|
|
44
|
+
# Coordinates for training
|
|
45
|
+
# Default to a dense grid from the domain if not specified
|
|
46
|
+
if "data" in config and "coords" in config["data"]:
|
|
47
|
+
coords = torch.tensor(config["data"]["coords"], dtype=torch.float32)
|
|
48
|
+
else:
|
|
49
|
+
coords, _ = system.domain.build_grid()
|
|
50
|
+
coords = coords.reshape(-1, coords.shape[-1])
|
|
51
|
+
|
|
52
|
+
# Optional logger
|
|
53
|
+
logger = None
|
|
54
|
+
if "logging" in config:
|
|
55
|
+
from rippl.logging.logger import RipplLogger
|
|
56
|
+
logger = RipplLogger(**config["logging"])
|
|
57
|
+
|
|
58
|
+
results = experiment.train(
|
|
59
|
+
coords,
|
|
60
|
+
epochs=train_cfg.get("epochs", 100),
|
|
61
|
+
ntk_freq=train_cfg.get("ntk_freq", 500),
|
|
62
|
+
patience=train_cfg.get("patience", 200),
|
|
63
|
+
logger=logger
|
|
64
|
+
)
|
|
65
|
+
|
|
66
|
+
return model, results
|
|
67
|
+
|
|
68
|
+
def simulate(config_path_or_dict: Union[str, Dict[str, Any]]) -> torch.Tensor:
|
|
69
|
+
"""
|
|
70
|
+
Run a trained model or a numerical solver for inference.
|
|
71
|
+
|
|
72
|
+
Args:
|
|
73
|
+
config_path_or_dict: Path to a configuration file or a dictionary.
|
|
74
|
+
|
|
75
|
+
Returns:
|
|
76
|
+
A torch.Tensor containing the simulated field values.
|
|
77
|
+
"""
|
|
78
|
+
config = ConfigParser.load(config_path_or_dict)
|
|
79
|
+
|
|
80
|
+
# If a model path is provided in config, load it.
|
|
81
|
+
# Otherwise, check if an FD solver is requested.
|
|
82
|
+
if "model_path" in config:
|
|
83
|
+
model = load_model(config["model_path"])
|
|
84
|
+
system = System.from_config(config)
|
|
85
|
+
coords, _ = system.domain.build_grid()
|
|
86
|
+
# Flatten and evaluate
|
|
87
|
+
original_shape = coords.shape[:-1]
|
|
88
|
+
coords_flat = coords.reshape(-1, coords.shape[-1])
|
|
89
|
+
with torch.no_grad():
|
|
90
|
+
output = model(coords_flat)
|
|
91
|
+
# Reshape back to grid
|
|
92
|
+
if isinstance(output, dict):
|
|
93
|
+
return {k: v.view(*original_shape, 1) for k, v in output.items()}
|
|
94
|
+
return output.view(*original_shape, 1)
|
|
95
|
+
|
|
96
|
+
elif "solver" in config:
|
|
97
|
+
from rippl.core.config import get_solver_fn
|
|
98
|
+
solver_name = config["solver"]["name"]
|
|
99
|
+
solver_fn = get_solver_fn(solver_name)
|
|
100
|
+
return solver_fn(**config["solver"].get("kwargs", {}))
|
|
101
|
+
|
|
102
|
+
raise ValueError("Simulation requires either 'model_path' or 'solver' in config.")
|
|
103
|
+
|
|
104
|
+
def identify(config_path_or_dict: Union[str, Dict[str, Any]], data_path: str) -> Dict[str, Any]:
|
|
105
|
+
"""
|
|
106
|
+
Run parameter discovery (Digital Twin) using sensor data.
|
|
107
|
+
|
|
108
|
+
Args:
|
|
109
|
+
config_path_or_dict: Path to a configuration file or a dictionary.
|
|
110
|
+
data_path: Path to the sensor data file (e.g., CSV).
|
|
111
|
+
|
|
112
|
+
Returns:
|
|
113
|
+
A dictionary containing identified parameters and their history.
|
|
114
|
+
"""
|
|
115
|
+
config = ConfigParser.load(config_path_or_dict)
|
|
116
|
+
system = System.from_config(config)
|
|
117
|
+
|
|
118
|
+
# Load model
|
|
119
|
+
model_cfg = config["model"]
|
|
120
|
+
model = build_model(model_cfg["name"], model_cfg["config"])
|
|
121
|
+
|
|
122
|
+
# Digital Twin setup
|
|
123
|
+
from rippl.core.inverse import DigitalTwin, InverseParameter
|
|
124
|
+
params_data = config["physics"].get("inverse_parameters", [])
|
|
125
|
+
parameters = []
|
|
126
|
+
for p_data in params_data:
|
|
127
|
+
parameters.append(InverseParameter(**p_data))
|
|
128
|
+
|
|
129
|
+
# Sensor data columns
|
|
130
|
+
data_cfg = config["data"]
|
|
131
|
+
dt = DigitalTwin.from_csv(
|
|
132
|
+
system, model, parameters,
|
|
133
|
+
csv_path=data_path,
|
|
134
|
+
coord_cols=data_cfg["coord_cols"],
|
|
135
|
+
field_cols=data_cfg["field_cols"],
|
|
136
|
+
data_weight=data_cfg.get("data_weight", 1.0),
|
|
137
|
+
physics_weight=data_cfg.get("physics_weight", 1.0)
|
|
138
|
+
)
|
|
139
|
+
|
|
140
|
+
train_cfg = config.get("training", {})
|
|
141
|
+
results = dt.train(
|
|
142
|
+
epochs=train_cfg.get("epochs", 5000),
|
|
143
|
+
lr=train_cfg.get("lr", 1e-3)
|
|
144
|
+
)
|
|
145
|
+
|
|
146
|
+
return results
|
rippl-0.1.3/rippl/cli.py
ADDED
|
@@ -0,0 +1,120 @@
|
|
|
1
|
+
import argparse
|
|
2
|
+
import sys
|
|
3
|
+
import os
|
|
4
|
+
import yaml
|
|
5
|
+
import torch
|
|
6
|
+
from rippl.api import train, simulate
|
|
7
|
+
from rippl.nn.registry import load_model
|
|
8
|
+
from rippl.export.exporter import export_model
|
|
9
|
+
from rippl.core.config import ConfigParser
|
|
10
|
+
|
|
11
|
+
def main():
|
|
12
|
+
"""
|
|
13
|
+
Rippl CLI Entrypoint.
|
|
14
|
+
Provides commands for training, simulation, and model export.
|
|
15
|
+
"""
|
|
16
|
+
parser = argparse.ArgumentParser(
|
|
17
|
+
description="Rippl: Modular Physics-ML Framework for Enterprise Wave PDEs",
|
|
18
|
+
formatter_class=argparse.RawDescriptionHelpFormatter
|
|
19
|
+
)
|
|
20
|
+
|
|
21
|
+
subparsers = parser.add_subparsers(dest="command", help="Available subcommands")
|
|
22
|
+
|
|
23
|
+
# Train command
|
|
24
|
+
train_parser = subparsers.add_parser("train", help="Train a PINN model from a configuration file")
|
|
25
|
+
train_parser.add_argument("config", help="Path to the YAML or JSON configuration file")
|
|
26
|
+
train_parser.add_argument("--output", "-o", default="checkpoint", help="Directory to save the trained model")
|
|
27
|
+
|
|
28
|
+
# Simulate command
|
|
29
|
+
sim_parser = subparsers.add_parser("simulate", help="Run inference or a numerical solver")
|
|
30
|
+
sim_parser.add_argument("config", help="Path to the configuration file")
|
|
31
|
+
|
|
32
|
+
# Export command
|
|
33
|
+
export_parser = subparsers.add_parser("export", help="Export a trained model to TorchScript or ONNX")
|
|
34
|
+
export_parser.add_argument("config", help="Path to the configuration file used for training")
|
|
35
|
+
export_parser.add_argument("model_dir", help="Directory containing the weights.pt and config.json")
|
|
36
|
+
export_parser.add_argument("--format", default="torchscript", choices=["torchscript", "onnx"],
|
|
37
|
+
help="Target export format (default: torchscript)")
|
|
38
|
+
|
|
39
|
+
if len(sys.argv) == 1:
|
|
40
|
+
parser.print_help()
|
|
41
|
+
sys.exit(0)
|
|
42
|
+
|
|
43
|
+
args = parser.parse_args()
|
|
44
|
+
|
|
45
|
+
try:
|
|
46
|
+
if args.command == "train":
|
|
47
|
+
print(f"[Rippl] Initializing training from {args.config}...")
|
|
48
|
+
|
|
49
|
+
with open(args.config, 'r') as f:
|
|
50
|
+
raw_cfg = yaml.safe_load(f)
|
|
51
|
+
|
|
52
|
+
from rippl.config.models import DomainConfig, EquationConfig, OperatorConfig
|
|
53
|
+
|
|
54
|
+
# 2. Enforce strict Pydantic validation before PyTorch logic
|
|
55
|
+
if "domain" in raw_cfg:
|
|
56
|
+
DomainConfig(**raw_cfg["domain"])
|
|
57
|
+
if "equation" in raw_cfg:
|
|
58
|
+
EquationConfig(**raw_cfg["equation"])
|
|
59
|
+
if "operators" in raw_cfg:
|
|
60
|
+
for op in raw_cfg["operators"]:
|
|
61
|
+
OperatorConfig(**op)
|
|
62
|
+
|
|
63
|
+
# 3 & 4. Handle Neural Network IoC
|
|
64
|
+
model_cfg = raw_cfg.get("model", {})
|
|
65
|
+
if "script" in model_cfg and "class_name" in model_cfg:
|
|
66
|
+
import importlib.util
|
|
67
|
+
script_path = model_cfg["script"]
|
|
68
|
+
class_name = model_cfg["class_name"]
|
|
69
|
+
|
|
70
|
+
spec = importlib.util.spec_from_file_location("custom_model", script_path)
|
|
71
|
+
custom_module = importlib.util.module_from_spec(spec)
|
|
72
|
+
sys.modules["custom_model"] = custom_module
|
|
73
|
+
spec.loader.exec_module(custom_module)
|
|
74
|
+
|
|
75
|
+
model_cls = getattr(custom_module, class_name)
|
|
76
|
+
kwargs = {k: v for k, v in model_cfg.items() if k not in ["script", "class_name", "type"]}
|
|
77
|
+
net = model_cls(**kwargs)
|
|
78
|
+
else:
|
|
79
|
+
import rippl.nn as rnn
|
|
80
|
+
model_type = model_cfg.get("type", "MLP")
|
|
81
|
+
model_cls = getattr(rnn, model_type)
|
|
82
|
+
kwargs = {k: v for k, v in model_cfg.items() if k != "type"}
|
|
83
|
+
net = model_cls(**kwargs)
|
|
84
|
+
|
|
85
|
+
from rippl.core.engine import Engine
|
|
86
|
+
engine = Engine(net)
|
|
87
|
+
|
|
88
|
+
# 5. Compile, fit, and save
|
|
89
|
+
engine.compile()
|
|
90
|
+
|
|
91
|
+
epochs = raw_cfg.get("training", {}).get("epochs", 10)
|
|
92
|
+
engine.fit(epochs=epochs, **raw_cfg.get("training", {}))
|
|
93
|
+
|
|
94
|
+
engine.validate()
|
|
95
|
+
|
|
96
|
+
engine.save("output.rpx")
|
|
97
|
+
|
|
98
|
+
print(f"[Rippl] Training successful. Model saved to output.rpx")
|
|
99
|
+
|
|
100
|
+
elif args.command == "simulate":
|
|
101
|
+
print(f"[Rippl] Running simulation from {args.config}...")
|
|
102
|
+
result = simulate(args.config)
|
|
103
|
+
print(f"[Rippl] Simulation complete. Output tensor shape: {result.shape}")
|
|
104
|
+
|
|
105
|
+
elif args.command == "export":
|
|
106
|
+
print(f"[Rippl] Exporting model from {args.model_dir} to {args.format} format...")
|
|
107
|
+
model = load_model(args.model_dir)
|
|
108
|
+
|
|
109
|
+
# Load metadata from training config
|
|
110
|
+
config = ConfigParser.load(args.config)
|
|
111
|
+
export_model(model, args.model_dir, format=args.format, metadata=config)
|
|
112
|
+
|
|
113
|
+
print(f"[Rippl] Export successful. Artifacts located in {args.model_dir}/")
|
|
114
|
+
|
|
115
|
+
except Exception as e:
|
|
116
|
+
print(f"[Rippl ERROR] {str(e)}")
|
|
117
|
+
sys.exit(1)
|
|
118
|
+
|
|
119
|
+
if __name__ == "__main__":
|
|
120
|
+
main()
|
|
@@ -0,0 +1,35 @@
|
|
|
1
|
+
import os
|
|
2
|
+
import json
|
|
3
|
+
from pathlib import Path
|
|
4
|
+
from typing import Dict, Any, List, Optional
|
|
5
|
+
from pydantic import BaseModel
|
|
6
|
+
|
|
7
|
+
class OperatorConfig(BaseModel):
|
|
8
|
+
type: str
|
|
9
|
+
params: Dict[str, Any] = {}
|
|
10
|
+
|
|
11
|
+
class EquationConfig(BaseModel):
|
|
12
|
+
operators: List[OperatorConfig] = []
|
|
13
|
+
params: Dict[str, Any] = {}
|
|
14
|
+
|
|
15
|
+
class DomainConfig(BaseModel):
|
|
16
|
+
spatial_dims: int = 1
|
|
17
|
+
time_dependent: bool = False
|
|
18
|
+
bounds: List[List[float]] = []
|
|
19
|
+
resolution: List[int] = []
|
|
20
|
+
|
|
21
|
+
class NetworkConfig(BaseModel):
|
|
22
|
+
params: Dict[str, Any] = {}
|
|
23
|
+
|
|
24
|
+
def generate_json_schemas(out_dir="schemas/"):
|
|
25
|
+
Path(out_dir).mkdir(parents=True, exist_ok=True)
|
|
26
|
+
models = {
|
|
27
|
+
"OperatorConfig": OperatorConfig,
|
|
28
|
+
"EquationConfig": EquationConfig,
|
|
29
|
+
"DomainConfig": DomainConfig,
|
|
30
|
+
"NetworkConfig": NetworkConfig
|
|
31
|
+
}
|
|
32
|
+
for name, model in models.items():
|
|
33
|
+
schema = model.model_json_schema()
|
|
34
|
+
with open(os.path.join(out_dir, f"{name}.json"), "w") as f:
|
|
35
|
+
json.dump(schema, f, indent=2)
|
|
@@ -0,0 +1,14 @@
|
|
|
1
|
+
"""rippl.core — System, Simulation, Experiment."""
|
|
2
|
+
from rippl.core.system import System, Domain, Constraint, NeumannConstraint
|
|
3
|
+
from rippl.core.equation import Equation
|
|
4
|
+
from rippl.core.equation_system import EquationSystem
|
|
5
|
+
from rippl.core.experiment import Experiment
|
|
6
|
+
from rippl.core.inverse import InverseProblem, InverseParameter, DigitalTwin
|
|
7
|
+
from rippl.core.nondim import ReferenceScales, NondimSystem, NondimModelWrapper
|
|
8
|
+
|
|
9
|
+
__all__ = [
|
|
10
|
+
"System", "Domain", "Constraint", "NeumannConstraint",
|
|
11
|
+
"Equation", "EquationSystem",
|
|
12
|
+
"Experiment", "InverseProblem", "InverseParameter", "DigitalTwin",
|
|
13
|
+
"ReferenceScales", "NondimSystem", "NondimModelWrapper"
|
|
14
|
+
]
|
|
@@ -0,0 +1,86 @@
|
|
|
1
|
+
import torch
|
|
2
|
+
from typing import Any, Dict, Union
|
|
3
|
+
|
|
4
|
+
def compile(model: torch.nn.Module, backend: str = "inductor",
|
|
5
|
+
mode: str = "max-autotune") -> torch.nn.Module:
|
|
6
|
+
try:
|
|
7
|
+
return torch.compile(model, backend=backend, mode=mode)
|
|
8
|
+
except Exception:
|
|
9
|
+
# Graceful fallback for unsupported platforms (Windows, older CUDA)
|
|
10
|
+
import warnings
|
|
11
|
+
warnings.warn(f"torch.compile backend='{backend}' unavailable. Running eager mode.")
|
|
12
|
+
return model
|
|
13
|
+
|
|
14
|
+
def run(domain: Any, equation: Any, model: torch.nn.Module,
|
|
15
|
+
strategy: str = "auto", devices: Union[int, str] = "auto",
|
|
16
|
+
precision: str = "bf16-mixed", epochs: int = 10000,
|
|
17
|
+
**kwargs) -> Dict[str, Any]:
|
|
18
|
+
try:
|
|
19
|
+
import pytorch_lightning as pl
|
|
20
|
+
from rippl.training.lightning_engine import LightningEngine
|
|
21
|
+
return _run_lightning(domain, equation, model, strategy, devices,
|
|
22
|
+
precision, epochs, kwargs)
|
|
23
|
+
except ImportError:
|
|
24
|
+
from rippl.training.pinn_recipe import PINNTrainingRecipe
|
|
25
|
+
return _run_native(domain, equation, model, epochs, kwargs)
|
|
26
|
+
|
|
27
|
+
def _run_lightning(domain, equation, model, strategy, devices,
|
|
28
|
+
precision, epochs, kwargs):
|
|
29
|
+
import pytorch_lightning as pl
|
|
30
|
+
from rippl.training.lightning_engine import LightningEngine
|
|
31
|
+
from rippl.core.nondim import AutoScaler
|
|
32
|
+
|
|
33
|
+
scaler = AutoScaler.from_domain_equation(domain, equation)
|
|
34
|
+
engine = LightningEngine(model=model, equation=equation, scaler=scaler,
|
|
35
|
+
lr=kwargs.get("lr", 1e-3),
|
|
36
|
+
constraint_weight=kwargs.get("constraint_weight", 100.0),
|
|
37
|
+
lbfgs_steps=kwargs.get("lbfgs_steps", 500),
|
|
38
|
+
causal=kwargs.get("causal", False))
|
|
39
|
+
trainer = pl.Trainer(max_epochs=epochs, strategy=strategy, devices=devices,
|
|
40
|
+
precision=precision, enable_checkpointing=False, logger=False,
|
|
41
|
+
callbacks=kwargs.get("callbacks", []))
|
|
42
|
+
loader = domain.generate_loader(batch_size=kwargs.get("batch_size", 2048))
|
|
43
|
+
trainer.fit(engine, train_dataloaders=loader)
|
|
44
|
+
|
|
45
|
+
return {"model_state": model.state_dict(),
|
|
46
|
+
"scaler_state": scaler.get_state(),
|
|
47
|
+
"final_loss": engine.final_loss}
|
|
48
|
+
|
|
49
|
+
def _run_native(domain, equation, model, epochs, kwargs):
|
|
50
|
+
# Fallback: uses existing PINNTrainingRecipe directly
|
|
51
|
+
from rippl.training.pinn_recipe import PINNTrainingRecipe
|
|
52
|
+
from rippl.core.nondim import AutoScaler
|
|
53
|
+
from rippl.core.system import System
|
|
54
|
+
import torch
|
|
55
|
+
|
|
56
|
+
device = next(model.parameters()).device
|
|
57
|
+
scaler = AutoScaler.from_domain_equation(domain, equation)
|
|
58
|
+
loader = domain.generate_loader(batch_size=kwargs.get("batch_size", 2048))
|
|
59
|
+
|
|
60
|
+
def loss_fn():
|
|
61
|
+
batch = next(iter(loader))[0].to(device)
|
|
62
|
+
c = batch.requires_grad_(True)
|
|
63
|
+
scaled = scaler.scale_inputs(c)
|
|
64
|
+
u = model(scaled)
|
|
65
|
+
u_phys = scaler.scale_outputs(u)
|
|
66
|
+
pde_loss = equation.compute_loss({"u": u_phys} if not isinstance(u_phys, dict) else u_phys, c)
|
|
67
|
+
return pde_loss, pde_loss
|
|
68
|
+
|
|
69
|
+
def constraint_loss_fn():
|
|
70
|
+
return torch.tensor(0.0, device=device, requires_grad=True), {}
|
|
71
|
+
|
|
72
|
+
recipe = PINNTrainingRecipe(
|
|
73
|
+
model=model,
|
|
74
|
+
loss_fn=loss_fn,
|
|
75
|
+
constraint_loss_fn=constraint_loss_fn,
|
|
76
|
+
device=device,
|
|
77
|
+
phase_b_epochs=epochs,
|
|
78
|
+
causal=kwargs.get("causal", False)
|
|
79
|
+
)
|
|
80
|
+
result = recipe.run()
|
|
81
|
+
|
|
82
|
+
return {
|
|
83
|
+
"model_state": model.state_dict(),
|
|
84
|
+
"scaler_state": scaler.get_state(),
|
|
85
|
+
"final_loss": result.get("phase_c_final_loss", result.get("phase_b_final_loss", 0.0))
|
|
86
|
+
}
|