pinnrl 0.1.0__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.
- pinnrl-0.1.0/.gitignore +41 -0
- pinnrl-0.1.0/LICENSE +21 -0
- pinnrl-0.1.0/PKG-INFO +276 -0
- pinnrl-0.1.0/README.md +199 -0
- pinnrl-0.1.0/pinnrl/__init__.py +1 -0
- pinnrl-0.1.0/pinnrl/components/adaptive_weights.py +134 -0
- pinnrl-0.1.0/pinnrl/config/__init__.py +725 -0
- pinnrl-0.1.0/pinnrl/config/config.yaml +382 -0
- pinnrl-0.1.0/pinnrl/dashboard.py +1919 -0
- pinnrl-0.1.0/pinnrl/main.py +86 -0
- pinnrl-0.1.0/pinnrl/neural_networks/__init__.py +154 -0
- pinnrl-0.1.0/pinnrl/neural_networks/attention.py +183 -0
- pinnrl-0.1.0/pinnrl/neural_networks/autoencoder.py +100 -0
- pinnrl-0.1.0/pinnrl/neural_networks/base_network.py +104 -0
- pinnrl-0.1.0/pinnrl/neural_networks/feedforward.py +73 -0
- pinnrl-0.1.0/pinnrl/neural_networks/fno.py +146 -0
- pinnrl-0.1.0/pinnrl/neural_networks/fourier.py +124 -0
- pinnrl-0.1.0/pinnrl/neural_networks/resnet.py +142 -0
- pinnrl-0.1.0/pinnrl/neural_networks/siren.py +90 -0
- pinnrl-0.1.0/pinnrl/pdes/__init__.py +49 -0
- pinnrl-0.1.0/pinnrl/pdes/allen_cahn.py +169 -0
- pinnrl-0.1.0/pinnrl/pdes/black_scholes.py +168 -0
- pinnrl-0.1.0/pinnrl/pdes/burgers_equation.py +177 -0
- pinnrl-0.1.0/pinnrl/pdes/cahn_hilliard.py +221 -0
- pinnrl-0.1.0/pinnrl/pdes/convection_equation.py +137 -0
- pinnrl-0.1.0/pinnrl/pdes/heat_equation.py +594 -0
- pinnrl-0.1.0/pinnrl/pdes/kdv_equation.py +159 -0
- pinnrl-0.1.0/pinnrl/pdes/pde_base.py +1302 -0
- pinnrl-0.1.0/pinnrl/pdes/pendulum_equation.py +289 -0
- pinnrl-0.1.0/pinnrl/pdes/wave_equation.py +187 -0
- pinnrl-0.1.0/pinnrl/rl/__init__.py +0 -0
- pinnrl-0.1.0/pinnrl/rl/rl_agent.py +748 -0
- pinnrl-0.1.0/pinnrl/training/__init__.py +5 -0
- pinnrl-0.1.0/pinnrl/training/train.py +390 -0
- pinnrl-0.1.0/pinnrl/training/trainer.py +1398 -0
- pinnrl-0.1.0/pinnrl/utils/__init__.py +23 -0
- pinnrl-0.1.0/pinnrl/utils/utils.py +795 -0
- pinnrl-0.1.0/pyproject.toml +164 -0
- pinnrl-0.1.0/tests/__init__.py +1 -0
- pinnrl-0.1.0/tests/unit_tests/__init__.py +0 -0
- pinnrl-0.1.0/tests/unit_tests/test_config.py +1049 -0
- pinnrl-0.1.0/tests/unit_tests/test_fno_all_pdes.py +147 -0
- pinnrl-0.1.0/tests/unit_tests/test_neural_networks.py +492 -0
- pinnrl-0.1.0/tests/unit_tests/test_pde_arch_matrix.py +314 -0
- pinnrl-0.1.0/tests/unit_tests/test_pde_base_coverage.py +707 -0
- pinnrl-0.1.0/tests/unit_tests/test_pde_sampling.py +389 -0
- pinnrl-0.1.0/tests/unit_tests/test_pdes.py +1251 -0
- pinnrl-0.1.0/tests/unit_tests/test_rl_agent.py +186 -0
- pinnrl-0.1.0/tests/unit_tests/test_rl_coverage.py +249 -0
- pinnrl-0.1.0/tests/unit_tests/test_trainer.py +316 -0
- pinnrl-0.1.0/tests/unit_tests/test_utils.py +116 -0
- pinnrl-0.1.0/tests/unit_tests/test_utils_coverage.py +457 -0
pinnrl-0.1.0/.gitignore
ADDED
|
@@ -0,0 +1,41 @@
|
|
|
1
|
+
# Python
|
|
2
|
+
__pycache__/
|
|
3
|
+
*.py[cod]
|
|
4
|
+
*.egg-info/
|
|
5
|
+
dist/
|
|
6
|
+
build/
|
|
7
|
+
|
|
8
|
+
# Virtual environments
|
|
9
|
+
pinn/
|
|
10
|
+
.venv/
|
|
11
|
+
|
|
12
|
+
# IDE
|
|
13
|
+
.vscode/
|
|
14
|
+
.idea/
|
|
15
|
+
.DS_Store
|
|
16
|
+
.ruff_cache
|
|
17
|
+
|
|
18
|
+
# Testing
|
|
19
|
+
.pytest_cache/
|
|
20
|
+
.coverage
|
|
21
|
+
htmlcov/
|
|
22
|
+
|
|
23
|
+
# Training artifacts
|
|
24
|
+
experiments/
|
|
25
|
+
checkpoints/
|
|
26
|
+
*.log
|
|
27
|
+
|
|
28
|
+
# Model weights
|
|
29
|
+
*.pth
|
|
30
|
+
*.pt
|
|
31
|
+
|
|
32
|
+
# Notebooks
|
|
33
|
+
notebooks/images/
|
|
34
|
+
.ipynb_checkpoints/
|
|
35
|
+
|
|
36
|
+
# Docs
|
|
37
|
+
site/
|
|
38
|
+
docs/LINKEDIN.md
|
|
39
|
+
|
|
40
|
+
# Sensitive
|
|
41
|
+
.env
|
pinnrl-0.1.0/LICENSE
ADDED
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2024 Jose Garcia
|
|
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.
|
pinnrl-0.1.0/PKG-INFO
ADDED
|
@@ -0,0 +1,276 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: pinnrl
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Physics-Informed Neural Networks with RL-based adaptive sampling for solving PDEs
|
|
5
|
+
Project-URL: Homepage, https://github.com/josegarciav/PINNs-RL-PDE
|
|
6
|
+
Project-URL: Documentation, https://github.com/josegarciav/PINNs-RL-PDE/tree/main/docs
|
|
7
|
+
Project-URL: Repository, https://github.com/josegarciav/PINNs-RL-PDE
|
|
8
|
+
Project-URL: Issues, https://github.com/josegarciav/PINNs-RL-PDE/issues
|
|
9
|
+
Project-URL: Changelog, https://github.com/josegarciav/PINNs-RL-PDE/blob/main/CHANGELOG.md
|
|
10
|
+
Author-email: Jose Garcia <jgarciaventura23@gmail.com>
|
|
11
|
+
License: MIT
|
|
12
|
+
License-File: LICENSE
|
|
13
|
+
Keywords: PDE,PINNs,deep learning,partial differential equations,physics-informed neural networks,reinforcement learning,scientific machine learning,simulation
|
|
14
|
+
Classifier: Development Status :: 3 - Alpha
|
|
15
|
+
Classifier: Intended Audience :: Science/Research
|
|
16
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
17
|
+
Classifier: Programming Language :: Python :: 3
|
|
18
|
+
Classifier: Programming Language :: Python :: 3.10
|
|
19
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
20
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
21
|
+
Classifier: Topic :: Scientific/Engineering
|
|
22
|
+
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
|
|
23
|
+
Classifier: Topic :: Scientific/Engineering :: Mathematics
|
|
24
|
+
Classifier: Topic :: Scientific/Engineering :: Physics
|
|
25
|
+
Requires-Python: >=3.10
|
|
26
|
+
Requires-Dist: gymnasium>=0.29.0
|
|
27
|
+
Requires-Dist: matplotlib>=3.7.0
|
|
28
|
+
Requires-Dist: numpy>=1.24.0
|
|
29
|
+
Requires-Dist: pandas>=2.0.0
|
|
30
|
+
Requires-Dist: plotly>=5.15.0
|
|
31
|
+
Requires-Dist: pyyaml>=6.0
|
|
32
|
+
Requires-Dist: scipy>=1.10.0
|
|
33
|
+
Requires-Dist: torch>=2.0.0
|
|
34
|
+
Requires-Dist: tqdm>=4.65.0
|
|
35
|
+
Provides-Extra: all
|
|
36
|
+
Requires-Dist: black>=23.0.0; extra == 'all'
|
|
37
|
+
Requires-Dist: dash-bootstrap-components>=1.5.0; extra == 'all'
|
|
38
|
+
Requires-Dist: dash>=3.0.0; extra == 'all'
|
|
39
|
+
Requires-Dist: ipywidgets>=8.0.0; extra == 'all'
|
|
40
|
+
Requires-Dist: jupyter>=1.0.0; extra == 'all'
|
|
41
|
+
Requires-Dist: mkdocs-material>=9.4.0; extra == 'all'
|
|
42
|
+
Requires-Dist: mkdocs>=1.5.0; extra == 'all'
|
|
43
|
+
Requires-Dist: mkdocstrings[python]>=0.23.0; extra == 'all'
|
|
44
|
+
Requires-Dist: mypy>=1.5.0; extra == 'all'
|
|
45
|
+
Requires-Dist: nbformat>=5.9.0; extra == 'all'
|
|
46
|
+
Requires-Dist: pre-commit>=3.4.0; extra == 'all'
|
|
47
|
+
Requires-Dist: pydocstyle>=6.3.0; extra == 'all'
|
|
48
|
+
Requires-Dist: pytest-cov>=4.1.0; extra == 'all'
|
|
49
|
+
Requires-Dist: pytest>=7.4.0; extra == 'all'
|
|
50
|
+
Requires-Dist: ruff>=0.1.0; extra == 'all'
|
|
51
|
+
Requires-Dist: scikit-image>=0.21.0; extra == 'all'
|
|
52
|
+
Requires-Dist: seaborn>=0.12.0; extra == 'all'
|
|
53
|
+
Requires-Dist: twine>=5.0.0; extra == 'all'
|
|
54
|
+
Provides-Extra: dev
|
|
55
|
+
Requires-Dist: black>=23.0.0; extra == 'dev'
|
|
56
|
+
Requires-Dist: mypy>=1.5.0; extra == 'dev'
|
|
57
|
+
Requires-Dist: pre-commit>=3.4.0; extra == 'dev'
|
|
58
|
+
Requires-Dist: pydocstyle>=6.3.0; extra == 'dev'
|
|
59
|
+
Requires-Dist: pytest-cov>=4.1.0; extra == 'dev'
|
|
60
|
+
Requires-Dist: pytest>=7.4.0; extra == 'dev'
|
|
61
|
+
Requires-Dist: ruff>=0.1.0; extra == 'dev'
|
|
62
|
+
Requires-Dist: twine>=5.0.0; extra == 'dev'
|
|
63
|
+
Provides-Extra: docs
|
|
64
|
+
Requires-Dist: mkdocs-material>=9.4.0; extra == 'docs'
|
|
65
|
+
Requires-Dist: mkdocs>=1.5.0; extra == 'docs'
|
|
66
|
+
Requires-Dist: mkdocstrings[python]>=0.23.0; extra == 'docs'
|
|
67
|
+
Provides-Extra: notebooks
|
|
68
|
+
Requires-Dist: ipywidgets>=8.0.0; extra == 'notebooks'
|
|
69
|
+
Requires-Dist: jupyter>=1.0.0; extra == 'notebooks'
|
|
70
|
+
Requires-Dist: nbformat>=5.9.0; extra == 'notebooks'
|
|
71
|
+
Provides-Extra: viz
|
|
72
|
+
Requires-Dist: dash-bootstrap-components>=1.5.0; extra == 'viz'
|
|
73
|
+
Requires-Dist: dash>=3.0.0; extra == 'viz'
|
|
74
|
+
Requires-Dist: scikit-image>=0.21.0; extra == 'viz'
|
|
75
|
+
Requires-Dist: seaborn>=0.12.0; extra == 'viz'
|
|
76
|
+
Description-Content-Type: text/markdown
|
|
77
|
+
|
|
78
|
+
<p align="center">
|
|
79
|
+
<img src="docs/assets/pinnrl-logo.png" alt="pinnrl" width="420"/>
|
|
80
|
+
</p>
|
|
81
|
+
|
|
82
|
+
# pinnrl — PINNs + RL for PDEs
|
|
83
|
+
|
|
84
|
+
[](https://opensource.org/licenses/MIT)
|
|
85
|
+
[](https://www.python.org/downloads/)
|
|
86
|
+
[](https://github.com/josegarciav/PINNs-RL-PDE)
|
|
87
|
+
[](https://github.com/josegarciav/PINNs-RL-PDE)
|
|
88
|
+
|
|
89
|
+
**Solve PDEs with neural networks that learn where to look.**
|
|
90
|
+
|
|
91
|
+
`pinnrl` combines Physics-Informed Neural Networks (PINNs) with Deep Q-Network reinforcement learning to place collocation points adaptively — concentrating compute where the residual is highest, not where a grid happens to land. Integrate your chosen architecture, your custom equations, a Dashboard trainer, and add exact/quasi-exact analytical solutions for validation.
|
|
92
|
+
|
|
93
|
+
---
|
|
94
|
+
|
|
95
|
+
## Why pinnrl?
|
|
96
|
+
|
|
97
|
+
Traditional PDE solvers (FEM, FDM) require meshes that become expensive in high dimensions or on irregular geometries. PINNs offer a mesh-free alternative, but their accuracy depends critically on *where* training points are sampled. Random or uniform sampling wastes capacity on easy regions while undersampling shocks, steep gradients, and phase boundaries.
|
|
98
|
+
|
|
99
|
+
`pinnrl` treats collocation point placement as a sequential decision problem. A DQN agent observes the current residual field and selects sampling regions to maximize physics compliance. The result is faster convergence and better accuracy on problems that defeat uniform sampling — Burgers shocks, KdV solitons, Cahn-Hilliard interfaces.
|
|
100
|
+
|
|
101
|
+
**Real use cases this targets:**
|
|
102
|
+
|
|
103
|
+
- **Fluid dynamics** — advection-dominated flows, shock formation (Burgers, Convection)
|
|
104
|
+
- **Materials science** — phase-field modeling, spinodal decomposition (Allen-Cahn, Cahn-Hilliard)
|
|
105
|
+
- **Quantitative finance** — option pricing on non-standard payoffs (Black-Scholes)
|
|
106
|
+
- **Nonlinear wave propagation** — soliton dynamics (KdV), resonance (Wave, Pendulum)
|
|
107
|
+
|
|
108
|
+
---
|
|
109
|
+
|
|
110
|
+
## Supported PDEs
|
|
111
|
+
|
|
112
|
+
| PDE | Domain | Key Parameters | Recommended Architecture |
|
|
113
|
+
|-----|--------|----------------|--------------------------|
|
|
114
|
+
| Heat | `[0,1] x [0,T]` | thermal diffusivity `alpha` | Fourier Features |
|
|
115
|
+
| Wave | `[0,1] x [0,T]` | wave speed `c` | SIREN |
|
|
116
|
+
| Burgers | `[-1,1] x [0,T]` | viscosity `nu` | Fourier Features |
|
|
117
|
+
| KdV | `[-pi,pi] x [0,T]` | dispersion `delta` | ResNet |
|
|
118
|
+
| Convection | `[0,1] x [0,T]` | advection velocity `beta` | Fourier Features |
|
|
119
|
+
| Allen-Cahn | `[-1,1] x [0,T]` | interface width `epsilon` | ResNet |
|
|
120
|
+
| Cahn-Hilliard | `[0,1]^2 x [0,T]` | mobility `M`, interface `epsilon` | Self-Attention |
|
|
121
|
+
| Black-Scholes | `[S_min,S_max] x [0,T]` | volatility `sigma`, rate `r` | FeedForward |
|
|
122
|
+
| Pendulum | `[0,2pi] x [0,T]` | damping `gamma`, forcing `F` | SIREN |
|
|
123
|
+
|
|
124
|
+
---
|
|
125
|
+
|
|
126
|
+
## Supported Architectures
|
|
127
|
+
|
|
128
|
+
| Architecture | Best For | Key Hyperparameters |
|
|
129
|
+
|---|---|---|
|
|
130
|
+
| `feedforward` | Smooth, low-frequency solutions | `hidden_dim`, `num_layers`, `activation` |
|
|
131
|
+
| `resnet` | Deep networks, stiff equations | `num_blocks`, `hidden_dim` |
|
|
132
|
+
| `siren` | Oscillatory, wave-type problems | `omega_0` (frequency scale) |
|
|
133
|
+
| `fourier` | High-frequency, spectral accuracy | `mapping_size`, `scale` |
|
|
134
|
+
| `fno` | Operator learning, multi-scale PDEs | `modes` (frequency cutoff), `num_blocks` |
|
|
135
|
+
| `attention` | Multi-scale, complex geometry | `num_heads`, `num_blocks` |
|
|
136
|
+
| `autoencoder` | Latent-space compression tasks | `latent_dim`, `hidden_dim` |
|
|
137
|
+
|
|
138
|
+
---
|
|
139
|
+
|
|
140
|
+
## Quick Start
|
|
141
|
+
|
|
142
|
+
### Installation
|
|
143
|
+
|
|
144
|
+
```bash
|
|
145
|
+
# Recommended (uv)
|
|
146
|
+
uv add pinnrl
|
|
147
|
+
|
|
148
|
+
# pip
|
|
149
|
+
pip install pinnrl
|
|
150
|
+
```
|
|
151
|
+
|
|
152
|
+
Available PDE keys: `heat`, `wave`, `burgers`, `kdv`, `convection`, `allen_cahn`, `cahn_hilliard`, `black_scholes`, `pendulum`.
|
|
153
|
+
|
|
154
|
+
### Python API — solving the heat equation
|
|
155
|
+
|
|
156
|
+
```python
|
|
157
|
+
import torch
|
|
158
|
+
from pinnrl.config import Config, ModelConfig, TrainingConfig, EarlyStoppingConfig
|
|
159
|
+
from pinnrl.neural_networks import PINNModel
|
|
160
|
+
from pinnrl.pdes.heat_equation import HeatEquation
|
|
161
|
+
from pinnrl.pdes.pde_base import PDEConfig
|
|
162
|
+
from pinnrl.training.trainer import PDETrainer
|
|
163
|
+
|
|
164
|
+
device = torch.device("cpu")
|
|
165
|
+
|
|
166
|
+
training = TrainingConfig(
|
|
167
|
+
num_epochs=3000, batch_size=2048,
|
|
168
|
+
num_collocation_points=5000, num_boundary_points=500, num_initial_points=500,
|
|
169
|
+
learning_rate=1e-3, weight_decay=1e-4, gradient_clipping=1.0,
|
|
170
|
+
early_stopping=EarlyStoppingConfig(enabled=True, patience=100, min_delta=1e-7),
|
|
171
|
+
learning_rate_scheduler=None,
|
|
172
|
+
)
|
|
173
|
+
pde_config = PDEConfig(
|
|
174
|
+
name="heat", domain=[[0.0, 1.0]], time_domain=[0.0, 1.0],
|
|
175
|
+
parameters={"alpha": 0.01},
|
|
176
|
+
boundary_conditions={"type": "dirichlet"},
|
|
177
|
+
initial_condition={"type": "sine", "amplitude": 1.0, "frequency": 2.0},
|
|
178
|
+
exact_solution={"type": "sin_exp_decay", "amplitude": 1.0, "frequency": 2.0},
|
|
179
|
+
dimension=1, device=device, training=training,
|
|
180
|
+
)
|
|
181
|
+
config = Config.__new__(Config)
|
|
182
|
+
config.device = device
|
|
183
|
+
config.model = ModelConfig(input_dim=2, hidden_dim=128, output_dim=1,
|
|
184
|
+
num_layers=4, activation="tanh", architecture="fourier")
|
|
185
|
+
config.training = training
|
|
186
|
+
config.pde_config = pde_config
|
|
187
|
+
|
|
188
|
+
pde = HeatEquation(config=pde_config)
|
|
189
|
+
model = PINNModel(config=config, device=device)
|
|
190
|
+
trainer = PDETrainer(model=model, pde=pde,
|
|
191
|
+
optimizer_config={"learning_rate": 1e-3, "weight_decay": 1e-4},
|
|
192
|
+
config=config, device=device)
|
|
193
|
+
trainer.train(num_epochs=3000)
|
|
194
|
+
|
|
195
|
+
# Validate against exact solution u(x,t) = A * exp(-alpha*(2pi*f/L)^2 * t) * sin(2pi*f*x/L)
|
|
196
|
+
metrics = pde.validate(model, num_points=5000)
|
|
197
|
+
print(f"L2 error: {metrics['l2_error']:.2e} | Max error: {metrics['max_error']:.2e}")
|
|
198
|
+
```
|
|
199
|
+
|
|
200
|
+
|
|
201
|
+
---
|
|
202
|
+
|
|
203
|
+
## Features
|
|
204
|
+
|
|
205
|
+
- **RL-adaptive collocation** — a DQN agent (`src/rl/rl_agent.py`) observes the current residual field and steers point sampling toward high-error regions, reducing wasted forward passes on low-residual areas.
|
|
206
|
+
- **Seven neural architectures** — FeedForward, ResNet, SIREN (periodic activations), Fourier Features (random Fourier mapping), Fourier Neural Operator (spectral convolutions), Self-Attention Transformer, and Autoencoder, all exposed through a single `PINNModel` factory class.
|
|
207
|
+
- **Nine PDEs** — linear to nonlinear, parabolic to hyperbolic, 1D to 2D, spanning thermal diffusion, soliton dynamics, phase-field models, and financial derivatives.
|
|
208
|
+
- **Exact analytical solutions** — every PDE ships with an `exact_solution` method for L2 and max-error validation during and after training; no external reference solver is required.
|
|
209
|
+
- **Adaptive loss weighting** — residual-balancing weights (RBW) and learning-rate-based weights (LRW) automatically re-balance physics, boundary, and initial condition loss terms.
|
|
210
|
+
- **Dash dashboard** — configure and launch experiments, monitor convergence metrics, and compare runs without writing plotting code.
|
|
211
|
+
- **Hardware-aware** — automatically selects CUDA, Apple MPS, or CPU; configurable via the `--device` flag or `config.yaml`.
|
|
212
|
+
|
|
213
|
+
---
|
|
214
|
+
|
|
215
|
+
## Comparison with Related Libraries
|
|
216
|
+
|
|
217
|
+
| Feature | pinnrl | DeepXDE | NeuralPDE.jl |
|
|
218
|
+
|---|---|---|---|
|
|
219
|
+
| RL-based adaptive sampling | Yes | No | No |
|
|
220
|
+
| Architectures bundled | 7 | 3 | 2 |
|
|
221
|
+
| Dashboard trainer | Yes | No | No |
|
|
222
|
+
| Exact solutions built in | Yes | Partial | Partial |
|
|
223
|
+
| Language | Python | Python | Julia |
|
|
224
|
+
| Install | `pip install pinnrl` | `pip install deepxde` | Julia pkg manager |
|
|
225
|
+
|
|
226
|
+
---
|
|
227
|
+
|
|
228
|
+
## Roadmap
|
|
229
|
+
|
|
230
|
+
| Version | Status | Milestones |
|
|
231
|
+
|---|---|---|
|
|
232
|
+
| v0.1 | Now | 9 PDEs, 7 architectures, DQN adaptive sampling, Dash dashboard, exact-solution validation |
|
|
233
|
+
| v0.2 | Planned | PyPI release, 2D/3D geometry, PPO agent option, YAML config schema validation |
|
|
234
|
+
| v0.3 | Planned | Operator learning (DeepONet-style), uncertainty quantification, JAX backend |
|
|
235
|
+
| v1.0 | Future | Stable public API, community PDE registry, full documentation site |
|
|
236
|
+
|
|
237
|
+
---
|
|
238
|
+
|
|
239
|
+
## Who Is It For?
|
|
240
|
+
|
|
241
|
+
- **PhD students and researchers** prototyping mesh-free PDE solvers without implementing adaptive sampling from scratch.
|
|
242
|
+
- **Computational scientists** in fluid dynamics, materials science, or plasma physics benchmarking PINNs against FDM/FEM on standard problems.
|
|
243
|
+
- **Quantitative analysts** exploring neural PDE solvers for high-dimensional option pricing beyond Black-Scholes closed forms.
|
|
244
|
+
- **ML engineers** studying how RL exploration strategies transfer to scientific computing.
|
|
245
|
+
|
|
246
|
+
---
|
|
247
|
+
|
|
248
|
+
## Contributing
|
|
249
|
+
|
|
250
|
+
Contributions are welcome. Please open an issue before submitting a large pull request.
|
|
251
|
+
|
|
252
|
+
**Core development:** Jose Garcia Ventura ([@josegarciav](https://github.com/josegarciav))
|
|
253
|
+
|
|
254
|
+
**Community contributions** in scope: new PDEs, new architectures, bug fixes, documentation improvements, and example notebooks. See [CONTRIBUTING.md](CONTRIBUTING.md) for the full guide.
|
|
255
|
+
|
|
256
|
+
---
|
|
257
|
+
|
|
258
|
+
## Citation
|
|
259
|
+
|
|
260
|
+
If you use `pinnrl` in published work, please cite:
|
|
261
|
+
|
|
262
|
+
```bibtex
|
|
263
|
+
@software{garcia2024pinnrl,
|
|
264
|
+
author = {Garcia Ventura, Jose},
|
|
265
|
+
title = {pinnrl: Physics-Informed Neural Networks with Reinforcement Learning Adaptive Sampling},
|
|
266
|
+
year = {2024},
|
|
267
|
+
url = {https://github.com/josegarciav/PINNs-RL-PDE},
|
|
268
|
+
note = {Open-source Python library for mesh-free PDE solving}
|
|
269
|
+
}
|
|
270
|
+
```
|
|
271
|
+
|
|
272
|
+
---
|
|
273
|
+
|
|
274
|
+
## License
|
|
275
|
+
|
|
276
|
+
[MIT License](LICENSE) © 2024 Jose Garcia Ventura
|
pinnrl-0.1.0/README.md
ADDED
|
@@ -0,0 +1,199 @@
|
|
|
1
|
+
<p align="center">
|
|
2
|
+
<img src="docs/assets/pinnrl-logo.png" alt="pinnrl" width="420"/>
|
|
3
|
+
</p>
|
|
4
|
+
|
|
5
|
+
# pinnrl — PINNs + RL for PDEs
|
|
6
|
+
|
|
7
|
+
[](https://opensource.org/licenses/MIT)
|
|
8
|
+
[](https://www.python.org/downloads/)
|
|
9
|
+
[](https://github.com/josegarciav/PINNs-RL-PDE)
|
|
10
|
+
[](https://github.com/josegarciav/PINNs-RL-PDE)
|
|
11
|
+
|
|
12
|
+
**Solve PDEs with neural networks that learn where to look.**
|
|
13
|
+
|
|
14
|
+
`pinnrl` combines Physics-Informed Neural Networks (PINNs) with Deep Q-Network reinforcement learning to place collocation points adaptively — concentrating compute where the residual is highest, not where a grid happens to land. Integrate your chosen architecture, your custom equations, a Dashboard trainer, and add exact/quasi-exact analytical solutions for validation.
|
|
15
|
+
|
|
16
|
+
---
|
|
17
|
+
|
|
18
|
+
## Why pinnrl?
|
|
19
|
+
|
|
20
|
+
Traditional PDE solvers (FEM, FDM) require meshes that become expensive in high dimensions or on irregular geometries. PINNs offer a mesh-free alternative, but their accuracy depends critically on *where* training points are sampled. Random or uniform sampling wastes capacity on easy regions while undersampling shocks, steep gradients, and phase boundaries.
|
|
21
|
+
|
|
22
|
+
`pinnrl` treats collocation point placement as a sequential decision problem. A DQN agent observes the current residual field and selects sampling regions to maximize physics compliance. The result is faster convergence and better accuracy on problems that defeat uniform sampling — Burgers shocks, KdV solitons, Cahn-Hilliard interfaces.
|
|
23
|
+
|
|
24
|
+
**Real use cases this targets:**
|
|
25
|
+
|
|
26
|
+
- **Fluid dynamics** — advection-dominated flows, shock formation (Burgers, Convection)
|
|
27
|
+
- **Materials science** — phase-field modeling, spinodal decomposition (Allen-Cahn, Cahn-Hilliard)
|
|
28
|
+
- **Quantitative finance** — option pricing on non-standard payoffs (Black-Scholes)
|
|
29
|
+
- **Nonlinear wave propagation** — soliton dynamics (KdV), resonance (Wave, Pendulum)
|
|
30
|
+
|
|
31
|
+
---
|
|
32
|
+
|
|
33
|
+
## Supported PDEs
|
|
34
|
+
|
|
35
|
+
| PDE | Domain | Key Parameters | Recommended Architecture |
|
|
36
|
+
|-----|--------|----------------|--------------------------|
|
|
37
|
+
| Heat | `[0,1] x [0,T]` | thermal diffusivity `alpha` | Fourier Features |
|
|
38
|
+
| Wave | `[0,1] x [0,T]` | wave speed `c` | SIREN |
|
|
39
|
+
| Burgers | `[-1,1] x [0,T]` | viscosity `nu` | Fourier Features |
|
|
40
|
+
| KdV | `[-pi,pi] x [0,T]` | dispersion `delta` | ResNet |
|
|
41
|
+
| Convection | `[0,1] x [0,T]` | advection velocity `beta` | Fourier Features |
|
|
42
|
+
| Allen-Cahn | `[-1,1] x [0,T]` | interface width `epsilon` | ResNet |
|
|
43
|
+
| Cahn-Hilliard | `[0,1]^2 x [0,T]` | mobility `M`, interface `epsilon` | Self-Attention |
|
|
44
|
+
| Black-Scholes | `[S_min,S_max] x [0,T]` | volatility `sigma`, rate `r` | FeedForward |
|
|
45
|
+
| Pendulum | `[0,2pi] x [0,T]` | damping `gamma`, forcing `F` | SIREN |
|
|
46
|
+
|
|
47
|
+
---
|
|
48
|
+
|
|
49
|
+
## Supported Architectures
|
|
50
|
+
|
|
51
|
+
| Architecture | Best For | Key Hyperparameters |
|
|
52
|
+
|---|---|---|
|
|
53
|
+
| `feedforward` | Smooth, low-frequency solutions | `hidden_dim`, `num_layers`, `activation` |
|
|
54
|
+
| `resnet` | Deep networks, stiff equations | `num_blocks`, `hidden_dim` |
|
|
55
|
+
| `siren` | Oscillatory, wave-type problems | `omega_0` (frequency scale) |
|
|
56
|
+
| `fourier` | High-frequency, spectral accuracy | `mapping_size`, `scale` |
|
|
57
|
+
| `fno` | Operator learning, multi-scale PDEs | `modes` (frequency cutoff), `num_blocks` |
|
|
58
|
+
| `attention` | Multi-scale, complex geometry | `num_heads`, `num_blocks` |
|
|
59
|
+
| `autoencoder` | Latent-space compression tasks | `latent_dim`, `hidden_dim` |
|
|
60
|
+
|
|
61
|
+
---
|
|
62
|
+
|
|
63
|
+
## Quick Start
|
|
64
|
+
|
|
65
|
+
### Installation
|
|
66
|
+
|
|
67
|
+
```bash
|
|
68
|
+
# Recommended (uv)
|
|
69
|
+
uv add pinnrl
|
|
70
|
+
|
|
71
|
+
# pip
|
|
72
|
+
pip install pinnrl
|
|
73
|
+
```
|
|
74
|
+
|
|
75
|
+
Available PDE keys: `heat`, `wave`, `burgers`, `kdv`, `convection`, `allen_cahn`, `cahn_hilliard`, `black_scholes`, `pendulum`.
|
|
76
|
+
|
|
77
|
+
### Python API — solving the heat equation
|
|
78
|
+
|
|
79
|
+
```python
|
|
80
|
+
import torch
|
|
81
|
+
from pinnrl.config import Config, ModelConfig, TrainingConfig, EarlyStoppingConfig
|
|
82
|
+
from pinnrl.neural_networks import PINNModel
|
|
83
|
+
from pinnrl.pdes.heat_equation import HeatEquation
|
|
84
|
+
from pinnrl.pdes.pde_base import PDEConfig
|
|
85
|
+
from pinnrl.training.trainer import PDETrainer
|
|
86
|
+
|
|
87
|
+
device = torch.device("cpu")
|
|
88
|
+
|
|
89
|
+
training = TrainingConfig(
|
|
90
|
+
num_epochs=3000, batch_size=2048,
|
|
91
|
+
num_collocation_points=5000, num_boundary_points=500, num_initial_points=500,
|
|
92
|
+
learning_rate=1e-3, weight_decay=1e-4, gradient_clipping=1.0,
|
|
93
|
+
early_stopping=EarlyStoppingConfig(enabled=True, patience=100, min_delta=1e-7),
|
|
94
|
+
learning_rate_scheduler=None,
|
|
95
|
+
)
|
|
96
|
+
pde_config = PDEConfig(
|
|
97
|
+
name="heat", domain=[[0.0, 1.0]], time_domain=[0.0, 1.0],
|
|
98
|
+
parameters={"alpha": 0.01},
|
|
99
|
+
boundary_conditions={"type": "dirichlet"},
|
|
100
|
+
initial_condition={"type": "sine", "amplitude": 1.0, "frequency": 2.0},
|
|
101
|
+
exact_solution={"type": "sin_exp_decay", "amplitude": 1.0, "frequency": 2.0},
|
|
102
|
+
dimension=1, device=device, training=training,
|
|
103
|
+
)
|
|
104
|
+
config = Config.__new__(Config)
|
|
105
|
+
config.device = device
|
|
106
|
+
config.model = ModelConfig(input_dim=2, hidden_dim=128, output_dim=1,
|
|
107
|
+
num_layers=4, activation="tanh", architecture="fourier")
|
|
108
|
+
config.training = training
|
|
109
|
+
config.pde_config = pde_config
|
|
110
|
+
|
|
111
|
+
pde = HeatEquation(config=pde_config)
|
|
112
|
+
model = PINNModel(config=config, device=device)
|
|
113
|
+
trainer = PDETrainer(model=model, pde=pde,
|
|
114
|
+
optimizer_config={"learning_rate": 1e-3, "weight_decay": 1e-4},
|
|
115
|
+
config=config, device=device)
|
|
116
|
+
trainer.train(num_epochs=3000)
|
|
117
|
+
|
|
118
|
+
# Validate against exact solution u(x,t) = A * exp(-alpha*(2pi*f/L)^2 * t) * sin(2pi*f*x/L)
|
|
119
|
+
metrics = pde.validate(model, num_points=5000)
|
|
120
|
+
print(f"L2 error: {metrics['l2_error']:.2e} | Max error: {metrics['max_error']:.2e}")
|
|
121
|
+
```
|
|
122
|
+
|
|
123
|
+
|
|
124
|
+
---
|
|
125
|
+
|
|
126
|
+
## Features
|
|
127
|
+
|
|
128
|
+
- **RL-adaptive collocation** — a DQN agent (`src/rl/rl_agent.py`) observes the current residual field and steers point sampling toward high-error regions, reducing wasted forward passes on low-residual areas.
|
|
129
|
+
- **Seven neural architectures** — FeedForward, ResNet, SIREN (periodic activations), Fourier Features (random Fourier mapping), Fourier Neural Operator (spectral convolutions), Self-Attention Transformer, and Autoencoder, all exposed through a single `PINNModel` factory class.
|
|
130
|
+
- **Nine PDEs** — linear to nonlinear, parabolic to hyperbolic, 1D to 2D, spanning thermal diffusion, soliton dynamics, phase-field models, and financial derivatives.
|
|
131
|
+
- **Exact analytical solutions** — every PDE ships with an `exact_solution` method for L2 and max-error validation during and after training; no external reference solver is required.
|
|
132
|
+
- **Adaptive loss weighting** — residual-balancing weights (RBW) and learning-rate-based weights (LRW) automatically re-balance physics, boundary, and initial condition loss terms.
|
|
133
|
+
- **Dash dashboard** — configure and launch experiments, monitor convergence metrics, and compare runs without writing plotting code.
|
|
134
|
+
- **Hardware-aware** — automatically selects CUDA, Apple MPS, or CPU; configurable via the `--device` flag or `config.yaml`.
|
|
135
|
+
|
|
136
|
+
---
|
|
137
|
+
|
|
138
|
+
## Comparison with Related Libraries
|
|
139
|
+
|
|
140
|
+
| Feature | pinnrl | DeepXDE | NeuralPDE.jl |
|
|
141
|
+
|---|---|---|---|
|
|
142
|
+
| RL-based adaptive sampling | Yes | No | No |
|
|
143
|
+
| Architectures bundled | 7 | 3 | 2 |
|
|
144
|
+
| Dashboard trainer | Yes | No | No |
|
|
145
|
+
| Exact solutions built in | Yes | Partial | Partial |
|
|
146
|
+
| Language | Python | Python | Julia |
|
|
147
|
+
| Install | `pip install pinnrl` | `pip install deepxde` | Julia pkg manager |
|
|
148
|
+
|
|
149
|
+
---
|
|
150
|
+
|
|
151
|
+
## Roadmap
|
|
152
|
+
|
|
153
|
+
| Version | Status | Milestones |
|
|
154
|
+
|---|---|---|
|
|
155
|
+
| v0.1 | Now | 9 PDEs, 7 architectures, DQN adaptive sampling, Dash dashboard, exact-solution validation |
|
|
156
|
+
| v0.2 | Planned | PyPI release, 2D/3D geometry, PPO agent option, YAML config schema validation |
|
|
157
|
+
| v0.3 | Planned | Operator learning (DeepONet-style), uncertainty quantification, JAX backend |
|
|
158
|
+
| v1.0 | Future | Stable public API, community PDE registry, full documentation site |
|
|
159
|
+
|
|
160
|
+
---
|
|
161
|
+
|
|
162
|
+
## Who Is It For?
|
|
163
|
+
|
|
164
|
+
- **PhD students and researchers** prototyping mesh-free PDE solvers without implementing adaptive sampling from scratch.
|
|
165
|
+
- **Computational scientists** in fluid dynamics, materials science, or plasma physics benchmarking PINNs against FDM/FEM on standard problems.
|
|
166
|
+
- **Quantitative analysts** exploring neural PDE solvers for high-dimensional option pricing beyond Black-Scholes closed forms.
|
|
167
|
+
- **ML engineers** studying how RL exploration strategies transfer to scientific computing.
|
|
168
|
+
|
|
169
|
+
---
|
|
170
|
+
|
|
171
|
+
## Contributing
|
|
172
|
+
|
|
173
|
+
Contributions are welcome. Please open an issue before submitting a large pull request.
|
|
174
|
+
|
|
175
|
+
**Core development:** Jose Garcia Ventura ([@josegarciav](https://github.com/josegarciav))
|
|
176
|
+
|
|
177
|
+
**Community contributions** in scope: new PDEs, new architectures, bug fixes, documentation improvements, and example notebooks. See [CONTRIBUTING.md](CONTRIBUTING.md) for the full guide.
|
|
178
|
+
|
|
179
|
+
---
|
|
180
|
+
|
|
181
|
+
## Citation
|
|
182
|
+
|
|
183
|
+
If you use `pinnrl` in published work, please cite:
|
|
184
|
+
|
|
185
|
+
```bibtex
|
|
186
|
+
@software{garcia2024pinnrl,
|
|
187
|
+
author = {Garcia Ventura, Jose},
|
|
188
|
+
title = {pinnrl: Physics-Informed Neural Networks with Reinforcement Learning Adaptive Sampling},
|
|
189
|
+
year = {2024},
|
|
190
|
+
url = {https://github.com/josegarciav/PINNs-RL-PDE},
|
|
191
|
+
note = {Open-source Python library for mesh-free PDE solving}
|
|
192
|
+
}
|
|
193
|
+
```
|
|
194
|
+
|
|
195
|
+
---
|
|
196
|
+
|
|
197
|
+
## License
|
|
198
|
+
|
|
199
|
+
[MIT License](LICENSE) © 2024 Jose Garcia Ventura
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
"""Top-level package for PINNs-RL-PDE framework."""
|
|
@@ -0,0 +1,134 @@
|
|
|
1
|
+
import logging
|
|
2
|
+
|
|
3
|
+
import torch
|
|
4
|
+
|
|
5
|
+
|
|
6
|
+
class AdaptiveLossWeights:
|
|
7
|
+
"""Handles adaptive weighting of different loss components in PINNs."""
|
|
8
|
+
|
|
9
|
+
def __init__(self, strategy="rbw", alpha=0.9, eps=1e-5, initial_weights=None):
|
|
10
|
+
"""
|
|
11
|
+
Initialize adaptive loss weights handler.
|
|
12
|
+
|
|
13
|
+
Args:
|
|
14
|
+
strategy (str): Weight adaptation strategy ('lrw' or 'rbw')
|
|
15
|
+
alpha (float): Moving average factor for weight updates
|
|
16
|
+
eps (float): Small constant for numerical stability
|
|
17
|
+
initial_weights (list): Initial weights for [pde, boundary, initial] components
|
|
18
|
+
"""
|
|
19
|
+
self.strategy = strategy.lower()
|
|
20
|
+
self.alpha = alpha
|
|
21
|
+
# Ensure eps is a float, not a string
|
|
22
|
+
self.eps = float(eps) if isinstance(eps, (str, int)) else eps
|
|
23
|
+
self.initial_weights = (
|
|
24
|
+
torch.tensor(initial_weights) if initial_weights is not None else None
|
|
25
|
+
)
|
|
26
|
+
self.weights = None
|
|
27
|
+
self.running_losses = None
|
|
28
|
+
self.running_grads = None
|
|
29
|
+
self.logger = logging.getLogger(__name__)
|
|
30
|
+
self.logger.info(
|
|
31
|
+
f"AdaptiveLossWeights initialized with strategy={strategy}, alpha={alpha}, eps={self.eps}, initial_weights={initial_weights}"
|
|
32
|
+
)
|
|
33
|
+
self.prev_weights = None
|
|
34
|
+
|
|
35
|
+
def update_weights_lrw(self, gradients):
|
|
36
|
+
"""
|
|
37
|
+
Update weights using Learning Rate based Weighting (LRW).
|
|
38
|
+
|
|
39
|
+
Args:
|
|
40
|
+
gradients (list): List of gradient norms for each loss component
|
|
41
|
+
"""
|
|
42
|
+
if self.running_grads is None:
|
|
43
|
+
self.running_grads = gradients
|
|
44
|
+
self.weights = (
|
|
45
|
+
self.initial_weights.to(gradients.device)
|
|
46
|
+
if self.initial_weights is not None
|
|
47
|
+
else torch.ones_like(gradients)
|
|
48
|
+
)
|
|
49
|
+
self.logger.info(f"LRW - Initialized weights: {self.weights}")
|
|
50
|
+
return self.weights
|
|
51
|
+
|
|
52
|
+
# Update running average of gradients
|
|
53
|
+
self.running_grads = self.alpha * self.running_grads + (1 - self.alpha) * gradients
|
|
54
|
+
|
|
55
|
+
# Ensure eps is a tensor on the same device as the gradients
|
|
56
|
+
eps_tensor = torch.tensor(self.eps, device=self.running_grads.device)
|
|
57
|
+
|
|
58
|
+
# Compute weights inversely proportional to gradient magnitudes
|
|
59
|
+
inv_grads = 1.0 / (self.running_grads + eps_tensor)
|
|
60
|
+
self.weights = inv_grads / torch.sum(inv_grads)
|
|
61
|
+
|
|
62
|
+
self.logger.info(
|
|
63
|
+
f"LRW - Updated weights: gradients={gradients.detach().cpu().numpy()}, running_grads={self.running_grads.detach().cpu().numpy()}, weights={self.weights.detach().cpu().numpy()}"
|
|
64
|
+
)
|
|
65
|
+
|
|
66
|
+
return self.weights
|
|
67
|
+
|
|
68
|
+
def update_weights_rbw(self, losses):
|
|
69
|
+
"""
|
|
70
|
+
Update weights using Relative Error based Weighting (RBW).
|
|
71
|
+
Gives more weight to terms with bigger losses to balance them.
|
|
72
|
+
|
|
73
|
+
Args:
|
|
74
|
+
losses (list): List of individual loss components
|
|
75
|
+
"""
|
|
76
|
+
if self.running_losses is None:
|
|
77
|
+
self.running_losses = losses
|
|
78
|
+
self.weights = (
|
|
79
|
+
self.initial_weights.to(losses.device)
|
|
80
|
+
if self.initial_weights is not None
|
|
81
|
+
else torch.ones_like(losses)
|
|
82
|
+
)
|
|
83
|
+
self.logger.info(f"RBW - Initialized weights: {self.weights}")
|
|
84
|
+
return self.weights
|
|
85
|
+
|
|
86
|
+
# Update running average of losses
|
|
87
|
+
self.running_losses = self.alpha * self.running_losses + (1 - self.alpha) * losses
|
|
88
|
+
|
|
89
|
+
# Ensure eps is a tensor on the same device as the losses
|
|
90
|
+
eps_tensor = torch.tensor(self.eps, device=self.running_losses.device)
|
|
91
|
+
|
|
92
|
+
# Normalize the running losses to get weights
|
|
93
|
+
# Give more weight to higher losses to focus optimization on problematic terms
|
|
94
|
+
normalized_losses = self.running_losses / (self.running_losses.sum() + eps_tensor)
|
|
95
|
+
self.weights = normalized_losses # Higher loss -> Higher weight
|
|
96
|
+
|
|
97
|
+
# Update weights with exponential moving average
|
|
98
|
+
if self.prev_weights is not None:
|
|
99
|
+
self.weights = self.alpha * self.prev_weights + (1 - self.alpha) * self.weights
|
|
100
|
+
|
|
101
|
+
self.prev_weights = self.weights.clone()
|
|
102
|
+
|
|
103
|
+
self.logger.info(
|
|
104
|
+
f"RBW - Updated weights: losses={losses.detach().cpu().numpy()}, running_losses={self.running_losses.detach().cpu().numpy()}, weights={self.weights.detach().cpu().numpy()}"
|
|
105
|
+
)
|
|
106
|
+
|
|
107
|
+
return self.weights
|
|
108
|
+
|
|
109
|
+
def update(self, losses=None, gradients=None):
|
|
110
|
+
"""
|
|
111
|
+
Update weights based on chosen strategy.
|
|
112
|
+
|
|
113
|
+
Args:
|
|
114
|
+
losses (list): List of individual loss components
|
|
115
|
+
gradients (list): List of gradient norms for each loss component
|
|
116
|
+
|
|
117
|
+
Returns:
|
|
118
|
+
torch.Tensor: Updated loss weights
|
|
119
|
+
"""
|
|
120
|
+
if self.strategy == "lrw" and gradients is not None:
|
|
121
|
+
return self.update_weights_lrw(gradients)
|
|
122
|
+
elif self.strategy == "rbw" and losses is not None:
|
|
123
|
+
return self.update_weights_rbw(losses)
|
|
124
|
+
else:
|
|
125
|
+
raise ValueError(f"Invalid combination of strategy ({self.strategy}) and inputs")
|
|
126
|
+
|
|
127
|
+
def get_weights(self):
|
|
128
|
+
"""Return current weights."""
|
|
129
|
+
if self.weights is not None:
|
|
130
|
+
return self.weights
|
|
131
|
+
elif self.initial_weights is not None:
|
|
132
|
+
return self.initial_weights
|
|
133
|
+
else:
|
|
134
|
+
return torch.ones(3) / 3.0 # Default to equal weights for 3 components
|