Configure Training for a PINA Solver
[!IMPORTANT] Read RULES.md before using this skill — it applies to all skills.
Use this skill to set up and run training once a solver has been chosen.
PINA's Trainer wraps Lightning's Trainer with PINA-specific defaults:
data splitting across conditions, batching strategies, device placement for
inverse parameters, and gradient tracking for physics-informed solvers.
Step 1 — Understand what the trainer needs
The trainer needs three things. If any are missing, ask the user:
- Solver (required) — already chosen and instantiated.
- Domain discretisation — all
DomainEquationConditiondomains must be sampled before the trainer can create dataloaders. - Training config — epochs, batch size, accelerator, device.
Domain discretisation
Physics-informed problems use domain=... in their conditions. The trainer
requires every such domain to have been sampled:
problem.discretise_domain(n=1000, mode="random")
If you forget, the trainer raises a clear error listing which domains are missing.
Step 2 — Configure the trainer
The constructor has two groups of parameters: PINA-specific and Lightning
**kwargs passed through to the parent class.
PINA-specific parameters
| Parameter | Default | What it controls |
|---|---|---|
solver |
— (required) | The solver instance to train |
batch_size |
None |
None = full batch; int = mini-batches |
train_size / val_size / test_size |
1.0 / 0.0 / 0.0 |
Fraction split per condition |
batching_mode |
"common_batch_size" |
How batches are built across conditions (see below) |
automatic_batching |
False |
True = Lightning's default collation; False = direct subset retrieval |
num_workers |
0 |
Dataloader workers |
pin_memory |
False |
Pin memory for faster GPU transfer |
shuffle |
True |
Shuffle before splitting |
Batching modes
The mode controls how condition data is assembled into batches:
| Mode | Behaviour | When to use |
|---|---|---|
"common_batch_size" |
Each condition supplies batch_size points per batch |
Default — works for most cases |
"proportional" |
Batch sizes are scaled by condition dataset sizes | Unbalanced datasets (e.g., many interior points but few boundary points) |
"separate_conditions" |
Iterates through each condition separately | Each condition's data is heterogeneous (e.g., one is pointwise, another is a graph) |
Common Lightning kwargs
These are passed as **kwargs and fully documented by PyTorch Lightning.
Key ones for PINA users:
| Kwarg | What it does |
|---|---|
max_epochs |
Number of training epochs |
accelerator |
"cpu", "gpu", "mps" (Apple Silicon) |
devices |
Device index or count (e.g., 1, [0], "auto") |
precision |
"16-mixed", "32", "64", "bf16-mixed" |
enable_progress_bar |
True / False |
gradient_clip_val |
Gradient clipping threshold |
callbacks |
List of Lightning callbacks (e.g., ModelCheckpoint, EarlyStopping) |
Step 3 — Train and test
Usage is uniform regardless of solver type:
from pina import Trainer
trainer = Trainer(
solver=solver,
max_epochs=1000,
batch_size=32,
accelerator="cpu",
train_size=0.8,
val_size=0.1,
test_size=0.1,
)
trainer.train() # Lightning fit()
trainer.test() # Lightning test() — optional
Templates
Full-batch physics-informed
problem.discretise_domain(n=5000, mode="random")
trainer = Trainer(
solver=solver,
max_epochs=10000,
batch_size=None, # full batch — no mini-batching
accelerator="cpu",
)
trainer.train()
Mini-batch supervised
trainer = Trainer(
solver=solver,
max_epochs=500,
batch_size=64,
train_size=0.8,
val_size=0.1,
test_size=0.1,
accelerator="gpu",
devices=1,
num_workers=4,
shuffle=True,
)
trainer.train()
trainer.test()
Unbalanced conditions (proportional batching)
trainer = Trainer(
solver=solver,
batch_size=512,
batching_mode="proportional",
max_epochs=1000,
)
Checklist
- Solver already chosen and instantiated (see select-solver skill)
- All domains discretised:
problem.discretise_domain(n=..., mode="...") - Batch size decided:
Nonefor full-batch, an int for mini-batches - Batching mode selected based on condition balance
- Train/val/test split configured
- Accelerator and device chosen
- Training started:
trainer.train()