Training

PDETrainer

PDETrainer(
    model, pde, config, device=None, bandit_sampler=None
)

Trainer for Physics-Informed Neural Networks.

Initialize trainer.

Parameters:

Name	Type	Description	Default
model	Module	PINN model.	required
pde	PDEBase	PDE instance.	required
config	Config	Full configuration object. Optimizer, scheduler, early stopping and loss settings all come from config.training.	required
device	device | None	Device to train on.	None
bandit_sampler		Optional ResidualBanditSampler. When given (or when collocation_distribution is "adaptive") collocation points come from the experimental RL bandit sampler.	None

Source code in pinnrl/training/trainer.py

def __init__(
    self,
    model: nn.Module,
    pde: "PDEBase",
    config: "Config",
    device: torch.device | None = None,
    bandit_sampler=None,
):
    """Initialize trainer.

    Args:
        model: PINN model.
        pde: PDE instance.
        config: Full configuration object. Optimizer, scheduler, early
            stopping and loss settings all come from ``config.training``.
        device: Device to train on.
        bandit_sampler: Optional ``ResidualBanditSampler``. When given (or when
            ``collocation_distribution`` is ``"adaptive"``) collocation points
            come from the experimental RL bandit sampler.
    """
    # An explicit device wins, then config.device, then MPS if available, else CPU
    self.device = device or (
        config.device
        if hasattr(config, "device")
        else torch.device("mps" if torch.backends.mps.is_available() else "cpu")
    )

    self.model = model.to(self.device)
    self.pde = pde
    self.config = config
    self.validation_frequency = config.training.validation_frequency

    # Setup logging up-front so optimizer init can warn through self.logger.
    self._setup_logging()

    # Setup optimizer
    self._num_epochs = config.training.num_epochs
    self._initialize_optimizer_and_scheduler()

    # Training history
    self.history = {
        "train_loss": [],
        "val_loss": [],
        "residual_loss": [],
        "boundary_loss": [],
        "initial_loss": [],
        "learning_rate": [],
        "loss_weights": [],  # adaptive weights at the end of each epoch that updated them
    }
    self._latest_loss_weights = None  # set by every adaptive step, recorded per epoch

    es_cfg = config.training.early_stopping
    self.early_stopping_enabled = es_cfg.enabled
    self.patience = es_cfg.patience
    self.min_delta = float(es_cfg.min_delta)
    self.best_val_loss = float("inf")
    self.patience_counter = 0

    self.use_rl = (
        bandit_sampler is not None
        or getattr(config.training, "collocation_distribution", None) == "adaptive"
    )
    if bandit_sampler is not None:
        pde.bandit_sampler = bandit_sampler

    # Initialize adaptive weights handler if enabled
    self.use_adaptive_weights = config.training.adaptive_weights.enabled
    if self.use_adaptive_weights:
        self.adaptive_weights = AdaptiveLossWeights(
            strategy=config.training.adaptive_weights.strategy,
            alpha=config.training.adaptive_weights.alpha,
            eps=config.training.adaptive_weights.eps,
            initial_weights=config.training.adaptive_weights.initial_weights,
        )
    else:
        self.logger.info("Adaptive weights are disabled")
        self.adaptive_weights = None

setup_experiment_logging

setup_experiment_logging(experiment_dir=None)

Log to <experiment_dir>/experiment.log, replacing earlier file handlers.

Parameters:

Name	Type	Description	Default
experiment_dir		Directory for the experiment; None is a no-op.	None

Source code in pinnrl/training/trainer.py

def setup_experiment_logging(self, experiment_dir=None):
    """Log to ``<experiment_dir>/experiment.log``, replacing earlier file handlers.

    Args:
        experiment_dir: Directory for the experiment; ``None`` is a no-op.
    """
    if experiment_dir:
        # Remove all FileHandlers from the logger
        for handler in self.logger.handlers[:]:
            if isinstance(handler, logging.FileHandler):
                self.logger.removeHandler(handler)

        # Add experiment-specific log file handler
        log_file = os.path.join(experiment_dir, "experiment.log")
        file_handler = logging.FileHandler(log_file)
        file_handler.setFormatter(
            logging.Formatter("%(asctime)s - %(levelname)s - %(message)s")
        )
        file_handler.setLevel(logging.INFO)
        self.logger.addHandler(file_handler)
        # The run log records INFO even when the application logs at WARNING.
        if self.logger.getEffectiveLevel() > logging.INFO:
            self.logger.setLevel(logging.INFO)
        self.logger.info(f"Experiment logs will be saved to {log_file}")

train_steps

train_steps(
    num_steps=None,
    time_budget_s=None,
    sample_fn=None,
    resample_every=1,
    checkpoints=(),
    on_checkpoint=None,
)

Step-based training loop for benchmarks and studies.

Runs :meth:_step (the same update as :meth:train) until num_steps steps or until training time reaches time_budget_s. Training time is the summed wall-clock of sampling plus optimizer steps; checkpoint callbacks are excluded. There is no validation, early stopping, LR warm-up or scheduling, progress bar or disk I/O.

Parameters:

Name	Type	Description	Default
num_steps	int | None	Stop after this many steps.	None
time_budget_s	float | None	Stop once training time reaches this many seconds.	None
sample_fn	Callable | None	sample_fn(model, progress) -> (x, t), where progress in [0, 1) is the elapsed fraction of the time budget (or of num_steps without one); defaults to the configured collocation strategy with config.training.batch_size.	None
resample_every	int	Draw new points every N steps and reuse them between.	1
checkpoints	Sequence[float]	Increasing training-time marks (seconds).	()
on_checkpoint	Callable | None	Called as on_checkpoint(step, train_time_s, model) when training time first passes each mark.	None

Returns:

Type	Description
dict[str, Any]	{"steps", "train_time_s", "loss"} with the per-step total loss.

Raises:

Type	Description
ValueError	If neither num_steps nor time_budget_s is given.

Source code in pinnrl/training/trainer.py

def train_steps(
    self,
    num_steps: int | None = None,
    time_budget_s: float | None = None,
    sample_fn: Callable | None = None,
    resample_every: int = 1,
    checkpoints: Sequence[float] = (),
    on_checkpoint: Callable | None = None,
) -> dict[str, Any]:
    """Step-based training loop for benchmarks and studies.

    Runs :meth:`_step` (the same update as :meth:`train`) until
    ``num_steps`` steps or until training time reaches ``time_budget_s``.
    Training time is the summed wall-clock of sampling plus optimizer steps;
    checkpoint callbacks are excluded. There is no validation, early
    stopping, LR warm-up or scheduling, progress bar or disk I/O.

    Args:
        num_steps: Stop after this many steps.
        time_budget_s: Stop once training time reaches this many seconds.
        sample_fn: ``sample_fn(model, progress) -> (x, t)``, where ``progress``
            in [0, 1) is the elapsed fraction of the time budget (or of
            ``num_steps`` without one); defaults to the configured
            collocation strategy with ``config.training.batch_size``.
        resample_every: Draw new points every N steps and reuse them between.
        checkpoints: Increasing training-time marks (seconds).
        on_checkpoint: Called as ``on_checkpoint(step, train_time_s, model)``
            when training time first passes each mark.

    Returns:
        ``{"steps", "train_time_s", "loss"}`` with the per-step total loss.

    Raises:
        ValueError: If neither ``num_steps`` nor ``time_budget_s`` is given.
    """
    if num_steps is None and time_budget_s is None:
        raise ValueError("Pass num_steps, time_budget_s, or both.")
    if sample_fn is None:
        batch_size = self.config.training.batch_size

        def sample_fn(_model, progress):
            return self._sample_batch(batch_size, progress)

    pending = list(checkpoints)
    loss_curve = []
    clock = 0.0
    step = 0
    x_batch = t_batch = None
    self.model.train()
    while (num_steps is None or step < num_steps) and (
        time_budget_s is None or clock < time_budget_s
    ):
        if time_budget_s is not None:
            progress = clock / time_budget_s
        else:
            progress = step / num_steps
        start = time.perf_counter()
        if step % resample_every == 0:
            x_batch, t_batch = sample_fn(self.model, progress)
        losses = self._step(x_batch, t_batch)
        loss_value = float(losses["total"].item())  # syncs MPS/CUDA before the clock stops
        clock += time.perf_counter() - start
        step += 1
        loss_curve.append(loss_value)
        while pending and clock >= pending[0]:
            pending.pop(0)
            if on_checkpoint is not None:
                on_checkpoint(step, clock, self.model)
                self.model.train()
    return {"steps": step, "train_time_s": clock, "loss": loss_curve}

train

train(
    num_epochs=None,
    batch_size=None,
    num_points=None,
    experiment_dir=None,
    progress=True,
)

Train the model.

Each epoch takes max(1, num_points // batch_size) optimizer steps on freshly sampled batches (one full batch when optimizer is lbfgs). Every validation_frequency epochs (starting at epoch 0) it validates, logs, checks early stopping (patience counted in epochs) and, with experiment_dir, saves history, metadata and the live snapshot.

Each call is planned as its own run: the learning rate restarts at training.learning_rate and the schedule and the Adam -> L-BFGS switch are sized to num_epochs. Weights, optimizer moments and history carry over from earlier calls.

Parameters:

Name	Type	Description	Default
num_epochs	int | None	Defaults to config.training.num_epochs.	None
batch_size	int | None	Defaults to config.training.batch_size; capped at num_points.	None
num_points	int | None	Collocation points per epoch; defaults to config.training.num_collocation_points.	None
experiment_dir	str | None	Run folder for experiment.log, metadata.json, metrics.json/history.json, live_snapshot.npz, final_model.pt and visualizations/; None writes nothing.	None
progress	bool	Show one tqdm progress bar over epochs.	True

Returns:

Type	Description
	The training history dict (self.history).

Source code in pinnrl/training/trainer.py

def train(
    self,
    num_epochs: int | None = None,
    batch_size: int | None = None,
    num_points: int | None = None,
    experiment_dir: str | None = None,
    progress: bool = True,
):
    """Train the model.

    Each epoch takes ``max(1, num_points // batch_size)`` optimizer steps on
    freshly sampled batches (one full batch when ``optimizer`` is ``lbfgs``). Every
    ``validation_frequency`` epochs (starting at epoch 0) it validates,
    logs, checks early stopping (patience counted in epochs) and, with
    ``experiment_dir``, saves history, metadata and the live snapshot.

    Each call is planned as its own run: the learning rate restarts at
    ``training.learning_rate`` and the schedule and the Adam -> L-BFGS
    switch are sized to ``num_epochs``. Weights, optimizer moments and
    history carry over from earlier calls.

    Args:
        num_epochs: Defaults to ``config.training.num_epochs``.
        batch_size: Defaults to ``config.training.batch_size``; capped at
            ``num_points``.
        num_points: Collocation points per epoch; defaults to
            ``config.training.num_collocation_points``.
        experiment_dir: Run folder for ``experiment.log``, ``metadata.json``,
            ``metrics.json``/``history.json``, ``live_snapshot.npz``,
            ``final_model.pt`` and ``visualizations/``; ``None`` writes nothing.
        progress: Show one tqdm progress bar over epochs.

    Returns:
        The training history dict (``self.history``).
    """
    training_cfg = self.config.training
    num_epochs = training_cfg.num_epochs if num_epochs is None else num_epochs
    batch_size = training_cfg.batch_size if batch_size is None else batch_size
    num_points = training_cfg.num_collocation_points if num_points is None else num_points
    if num_points < batch_size:
        batch_size = num_points
    # Each call is its own run: the LR restarts and cosine T_max and the
    # L-BFGS switch follow num_epochs. Weights and optimizer moments carry over.
    self._plan_epochs(num_epochs)
    self.model.train()
    if hasattr(self.model, "architecture_name"):
        architecture_name = self.model.architecture_name
    else:
        architecture_name = self.model.__class__.__name__

    # Log PDE and architecture information
    self.logger.info("=" * 50)
    self.logger.info(f"STARTING TRAINING FOR PDE: {self.pde.__class__.__name__}")
    self.logger.info(f"NEURAL NETWORK ARCHITECTURE: {architecture_name}")
    if hasattr(self.model, "model"):
        self.logger.info(f"MODEL STRUCTURE: {self.model.model}")
    self.logger.info(f"DEVICE: {self.device}")

    # Log adaptive weights configuration if enabled
    if self.use_adaptive_weights:
        self.logger.info("ADAPTIVE WEIGHTS ENABLED")
        self.logger.info(f"  Strategy: {self.config.training.adaptive_weights.strategy}")
        self.logger.info(f"  Alpha: {self.config.training.adaptive_weights.alpha}")
        self.logger.info(f"  Epsilon: {self.config.training.adaptive_weights.eps}")
    else:
        self.logger.info("ADAPTIVE WEIGHTS DISABLED")
        # Convert loss_weights to something printable
        weights_dict = {
            k: float(v) if isinstance(v, torch.Tensor) else v
            for k, v in self.config.training.loss_weights.items()
        }
        self.logger.info(f"  Fixed weights: {weights_dict}")
    self.logger.info("=" * 50)

    # Inverse-problem warm-up: announce trainable parameters and seed history
    # buckets so the dashboard can render trajectories from epoch 0.
    trainable_pde_params = {}
    if hasattr(self.pde, "_trainable_params"):
        trainable_pde_params = dict(self.pde._trainable_params)
    if trainable_pde_params:
        self.logger.info(
            "INVERSE MODE — identifying %d PDE parameter(s): %s",
            len(trainable_pde_params),
            list(trainable_pde_params.keys()),
        )
        for name in trainable_pde_params:
            self.history.setdefault(f"param_{name}", [])

    # L-BFGS requires deterministic full-batch loss; override batch_size
    # (the L-BFGS phase of adam_lbfgs does the same when it starts).
    if getattr(self, "_is_lbfgs", False) and batch_size != num_points:
        self.logger.warning(
            "L-BFGS optimizer requires full-batch updates; "
            f"overriding batch_size {batch_size} -> {num_points}."
        )
        batch_size = num_points

    # Adaptive weights are incompatible with the L-BFGS closure (per-component
    # gradient passes are too expensive); _step skips them during L-BFGS.
    if getattr(self, "_is_lbfgs", False) and self.use_adaptive_weights:
        self.logger.warning("Adaptive loss weighting is disabled while running L-BFGS.")

    # Record start time
    start_time = datetime.now()

    # Fields shared by the initial and the final metadata.json.
    run_info = {
        "pde_type": getattr(self.pde, "pde_type", type(self.pde).__name__),
        "pde_name": getattr(self.pde.config, "name", ""),
        "architecture": getattr(self.model, "architecture_name", type(self.model).__name__),
        "training_params": {
            "num_epochs": num_epochs,
            "batch_size": batch_size,
            "num_points": num_points,
            "validation_frequency": self.validation_frequency,
        },
        "collocation_distribution": (
            "adaptive" if self.use_rl else self.config.training.collocation_distribution
        ),
        "rl_enabled": self.use_rl,
        "optimizer": self._optimizer_type,
        "mode": self.config.training.mode,
        "trainable_parameters": list(trainable_pde_params.keys()),
        "true_parameters": dict(getattr(self.pde, "_true_parameters", {})),
    }

    # Create directories for visualizations and experiment data
    if experiment_dir:
        os.makedirs(experiment_dir, exist_ok=True)

        # Setup experiment-specific logging
        self.setup_experiment_logging(experiment_dir)

        # Create visualization directory
        viz_dir = os.path.join(experiment_dir, "visualizations")
        os.makedirs(viz_dir, exist_ok=True)

        # Create .running marker file
        running_file = os.path.join(experiment_dir, ".running")
        open(running_file, "w").close()

        initial_metadata = {
            "status": "running",
            "start_time": start_time.strftime("%Y-%m-%d %H:%M:%S"),
            "total_epochs": num_epochs,
            "current_epoch": 0,
            **run_info,
        }
        metadata_path = os.path.join(experiment_dir, "metadata.json")
        with open(metadata_path, "w") as f:
            json.dump(initial_metadata, f, indent=2)

    self.logger.info("Starting training...")

    # Initialize points history
    self.points_history = []

    # Emit an initial live snapshot so the Monitor tab has something to show
    # before the first epoch finishes.
    if experiment_dir:
        self._save_live_snapshot(experiment_dir, epoch=0)

    # tqdm.auto renders a widget in notebooks; refresh at most twice a second.
    pbar = tqdm(total=num_epochs, desc="Training", mininterval=0.5) if progress else None
    for epoch in range(num_epochs):
        self.model.train()
        self._apply_warmup(epoch)
        epoch_losses = []
        epoch_points = []  # Store all points for this epoch
        self._latest_loss_weights = None

        for _ in range(max(1, num_points // batch_size)):
            x_batch, t_batch = self._sample_batch(batch_size, epoch / max(1, num_epochs))
            epoch_points.append(torch.cat([x_batch, t_batch], dim=1).detach().cpu().numpy())
            losses = self._step(x_batch, t_batch)
            epoch_losses.append(float(losses["total"].item()))
        # L-BFGS epochs do not weight adaptively and record nothing; they
        # come after the Adam epochs, so the recorded epochs stay a prefix.
        if self._latest_loss_weights is not None:
            self.history["loss_weights"].append(self._latest_loss_weights)
        if pbar is not None:
            pbar.set_postfix(
                loss=f"{epoch_losses[-1]:.3e}",
                res=f"{losses['residual'].item():.3e}",
                refresh=False,
            )
            pbar.update()

        # End of epoch
        avg_epoch_loss = sum(epoch_losses) / len(epoch_losses)
        points_array = np.concatenate(epoch_points, axis=0)
        self.points_history.append(points_array)

        # Update learning rate (the scheduler takes over after warm-up)
        if epoch >= self._warmup_epochs():
            self._update_scheduler(avg_epoch_loss)

        # Log metrics
        current_lr = self.optimizer.param_groups[0]["lr"]

        epoch_data = {
            "epoch": epoch,
            "train_loss": avg_epoch_loss,
            "residual_loss": _as_float(losses["residual"]),
            "boundary_loss": _as_float(losses["boundary"]),
            "initial_loss": _as_float(losses["initial"]),
            "learning_rate": current_lr,
        }
        if "data" in losses:
            epoch_data["data_loss"] = _as_float(losses["data"])
            self.history.setdefault("data_loss", [])

        # Capture each trainable PDE parameter's current value so the
        # dashboard can plot its trajectory in real time.
        if trainable_pde_params:
            for name, p in trainable_pde_params.items():
                epoch_data[f"param_{name}"] = float(p.detach().cpu().item())

        # Add weights info if using adaptive weights
        if self.use_adaptive_weights and len(self.history["loss_weights"]) > 0:
            current_weights = self.history["loss_weights"][-1]
            weights_str = ""
            component_names = ["residual", "boundary", "initial"]
            for i, name in enumerate(component_names):
                if i < len(current_weights):
                    weights_str += f", {name}_weight={float(current_weights[i]):.4f}"

        else:
            weights_str = ""

        # Update history
        for key, value in epoch_data.items():
            if key in self.history:
                self.history[key].append(value)

        # Validation (if applicable)
        if epoch % self.validation_frequency == 0:
            val_losses = self._compute_validation_loss()
            self.history["val_loss"].append(val_losses["total_loss"])
            self.logger.info(
                f"Epoch {epoch + 1}/{num_epochs}, "
                f"Train Loss: {avg_epoch_loss:.6f}, "
                f"Val Loss: {val_losses['total_loss']:.6f}, "
                f"LR: {current_lr:.6f}{weights_str}"
            )

            # Early stopping check
            if self.early_stopping_enabled:
                val_loss = val_losses["total_loss"]
                if val_loss < self.best_val_loss - self.min_delta:
                    self.best_val_loss = val_loss
                    self.patience_counter = 0
                else:
                    # Patience is in epochs; validation runs every validation_frequency.
                    self.patience_counter += self.validation_frequency

                if self.patience_counter >= self.patience:
                    self.logger.info(
                        f"Early stopping triggered at epoch {epoch + 1}, "
                        f"best val loss: {self.best_val_loss:.6f}"
                    )
                    break

        # Periodic save of history and metadata
        if experiment_dir and epoch % self.validation_frequency == 0:
            save_training_metrics(self.history, experiment_dir)
            metadata_path = os.path.join(experiment_dir, "metadata.json")
            try:
                with open(metadata_path) as f:
                    partial_metadata = json.load(f)
            except (FileNotFoundError, json.JSONDecodeError):
                partial_metadata = {}
            partial_metadata.update(
                {
                    "current_epoch": epoch + 1,
                    "status": "running",
                    "final_loss": float(avg_epoch_loss),
                    "best_val_loss": (
                        float(self.best_val_loss)
                        if self.best_val_loss != float("inf")
                        else None
                    ),
                    "last_updated": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
                }
            )
            if trainable_pde_params:
                partial_metadata["current_parameters"] = {
                    name: float(p.detach().cpu().item())
                    for name, p in trainable_pde_params.items()
                }
            with open(metadata_path, "w") as f:
                json.dump(partial_metadata, f, indent=2)
            # Refresh the live 3D snapshot for the dashboard.
            self._save_live_snapshot(experiment_dir, epoch=epoch + 1)

        # adam_lbfgs: hot-swap the optimizer once we cross the switch boundary.
        if (
            getattr(self, "_optimizer_type", "adam") == "adam_lbfgs"
            and not self._is_lbfgs
            and self._switch_epoch is not None
            and (epoch + 1) >= self._switch_epoch
        ):
            self.logger.info(f"Switching optimizer to L-BFGS at epoch {epoch + 1}")
            self._switch_to_lbfgs()
            batch_size = num_points  # L-BFGS needs the full batch
            if self.use_adaptive_weights:
                self.logger.warning("Adaptive loss weighting is disabled while running L-BFGS.")

    if pbar is not None:
        pbar.close()

    # End of training
    train_time = (datetime.now() - start_time).total_seconds() / 60.0
    self.logger.info(f"Training completed in {train_time} minutes")

    # Save final visualizations and metrics
    if experiment_dir:
        # Create metadata with training information
        metadata = {
            **run_info,
            "status": "completed",
            "training_time_minutes": train_time,
            "total_epochs": num_epochs,
            "final_loss": (
                float(self.history["train_loss"][-1]) if self.history["train_loss"] else None
            ),
            "best_val_loss": (
                float(self.best_val_loss) if self.best_val_loss != float("inf") else None
            ),
            "early_stopping_triggered": (
                self.patience_counter >= self.patience if self.early_stopping_enabled else False
            ),
            "current_epoch": len(self.history["train_loss"]),
            "start_time": start_time.strftime("%Y-%m-%d %H:%M:%S"),
            "end_time": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
            "domain": self.pde.domain,
            "time_domain": self.pde.time_domain,
            "boundary_conditions": self.pde.config.boundary_conditions,
            "initial_condition": self.pde.config.initial_condition,
            "pde_parameters": self.pde.config.parameters,
            "identified_parameters": (
                self.pde.get_trainable_parameter_values() if trainable_pde_params else {}
            ),
        }

        # Save training history as metrics.json and history.json
        self.logger.info("Saving training metrics...")
        save_training_metrics(self.history, experiment_dir, metadata=metadata)

        # Final live snapshot reflects the trained model.
        self._save_live_snapshot(experiment_dir, epoch=len(self.history.get("train_loss", [])))

        # Remove .running marker
        running_file = os.path.join(experiment_dir, ".running")
        if os.path.exists(running_file):
            try:
                os.remove(running_file)
            except OSError:
                pass

        # Save the model directly in the experiment directory
        self.logger.info("Saving model...")
        model_path = os.path.join(experiment_dir, "final_model.pt")
        torch.save(self.model.state_dict(), model_path)

        # Save all plots using the save_plots method
        self.logger.info("Generating and saving all plots...")
        self.save_plots(viz_dir)

    return self.history

save_plots

save_plots(save_dir)

Write the post-training figures to save_dir.

Source code in pinnrl/training/trainer.py

def save_plots(self, save_dir):  # pragma: no cover
    """Write the post-training figures to ``save_dir``."""
    plots.save_training_plots(self, save_dir)

save_training_metrics

save_training_metrics(
    history, experiment_dir, metadata=None
)

Save training metrics to JSON files.

Writes the same history to metrics.json and history.json. When metadata is given it is merged into an existing metadata.json and stamped with last_updated.

Parameters:

Name	Type	Description	Default
history	dict[str, Any]	Dictionary containing training metrics	required
experiment_dir	str	Directory to save the metrics	required
metadata	dict[str, Any] | None	Optional metadata about the experiment	None

Returns:

Type	Description
	Path of history.json.

Source code in pinnrl/training/trainer.py

def save_training_metrics(
    history: dict[str, Any],
    experiment_dir: str,
    metadata: dict[str, Any] | None = None,
):
    """Save training metrics to JSON files.

    Writes the same history to ``metrics.json`` and ``history.json``. When
    ``metadata`` is given it is merged into an existing ``metadata.json`` and
    stamped with ``last_updated``.

    Args:
        history: Dictionary containing training metrics
        experiment_dir: Directory to save the metrics
        metadata: Optional metadata about the experiment

    Returns:
        Path of ``history.json``.
    """
    # Create experiment directory if it doesn't exist
    os.makedirs(experiment_dir, exist_ok=True)

    # Convert NumPy arrays to Python lists in history
    def convert_to_serializable(obj):
        if isinstance(obj, np.ndarray):
            return obj.tolist()
        elif isinstance(obj, dict):
            return {k: convert_to_serializable(v) for k, v in obj.items()}
        elif isinstance(obj, list):
            return [convert_to_serializable(item) for item in obj]
        return obj

    serializable_history = convert_to_serializable(history)

    # Save metrics.json directly in the experiment directory
    metrics_file = os.path.join(experiment_dir, "metrics.json")
    with open(metrics_file, "w") as f:
        json.dump(serializable_history, f, indent=2)

    # Save history.json directly in the experiment directory
    history_file = os.path.join(experiment_dir, "history.json")
    with open(history_file, "w") as f:
        json.dump(serializable_history, f, indent=2)

    # Save metadata if provided
    if metadata:
        metadata_file = os.path.join(experiment_dir, "metadata.json")

        # Update existing metadata if it exists
        if os.path.exists(metadata_file):
            try:
                with open(metadata_file) as f:
                    existing_metadata = json.load(f)
                # Update with new metadata
                existing_metadata.update(metadata)
                metadata = existing_metadata
            except Exception:
                pass

        # Add timestamp
        metadata["last_updated"] = datetime.now().strftime("%Y-%m-%d %H:%M:%S")

        # Convert metadata to serializable format
        serializable_metadata = convert_to_serializable(metadata)

        with open(metadata_file, "w") as f:
            json.dump(serializable_metadata, f, indent=2)

    return history_file

AdaptiveLossWeights

AdaptiveLossWeights(
    strategy="rbw",
    alpha=0.9,
    eps=1e-05,
    initial_weights=None,
)

Adaptive weights for the [residual, boundary, initial] loss components.

The first :meth:update call returns initial_weights (or ones) and seeds the running average; later calls return weights that sum to 1.

Initialize adaptive loss weights handler.

Parameters:

Name	Type	Description	Default
strategy	str	Weight adaptation strategy ('lrw' or 'rbw').	'rbw'
alpha	float	EMA retention factor: new = alpha * old + (1 - alpha) * current.	0.9
eps	float	Small constant for numerical stability.	1e-05
initial_weights	list | None	Weights for [residual, boundary, initial] returned by the first update; None means ones.	None

Source code in pinnrl/training/adaptive_weights.py

def __init__(self, strategy="rbw", alpha=0.9, eps=1e-5, initial_weights=None):
    """Initialize adaptive loss weights handler.

    Args:
        strategy (str): Weight adaptation strategy ('lrw' or 'rbw').
        alpha (float): EMA retention factor: ``new = alpha * old + (1 - alpha) * current``.
        eps (float): Small constant for numerical stability.
        initial_weights (list | None): Weights for [residual, boundary, initial]
            returned by the first update; ``None`` means ones.
    """
    self.strategy = strategy.lower()
    self.alpha = alpha
    # Ensure eps is a float, not a string
    self.eps = float(eps) if isinstance(eps, (str, int)) else eps
    self.initial_weights = (
        torch.tensor(initial_weights) if initial_weights is not None else None
    )
    self.weights = None
    self.running_losses = None
    self.running_grads = None
    self.logger = logging.getLogger(__name__)
    self.logger.info(
        f"AdaptiveLossWeights initialized with strategy={strategy}, alpha={alpha}, eps={self.eps}, initial_weights={initial_weights}"
    )
    self.prev_weights = None

update_weights_lrw

update_weights_lrw(gradients)

Update weights using Learning Rate based Weighting (LRW).

After the first call (which only seeds g_bar)::

g_bar_i = alpha * g_bar_i + (1 - alpha) * g_i
w_i = (1 / (g_bar_i + eps)) / sum_{j active} (1 / (g_bar_j + eps))

over the active components (g_bar_i > 0), so components with smaller gradient norms get larger weights. A component whose gradient norm is zero (a constant loss, e.g. a PDE with no boundary conditions) gets weight 0: it cannot be minimised, and 1 / eps would otherwise hand it almost all the weight and switch the other terms off. If no component is active the weights are uniform.

Parameters:

Name	Type	Description	Default
gradients	Tensor	Gradient norm of each loss component, shape (n,).	required

Returns:

Type	Description
	torch.Tensor: The weights, shape (n,).

Source code in pinnrl/training/adaptive_weights.py

def update_weights_lrw(self, gradients):
    """Update weights using Learning Rate based Weighting (LRW).

    After the first call (which only seeds ``g_bar``)::

        g_bar_i = alpha * g_bar_i + (1 - alpha) * g_i
        w_i = (1 / (g_bar_i + eps)) / sum_{j active} (1 / (g_bar_j + eps))

    over the *active* components (``g_bar_i > 0``), so components with smaller
    gradient norms get larger weights. A component whose gradient norm is zero
    (a constant loss, e.g. a PDE with no boundary conditions) gets weight 0:
    it cannot be minimised, and ``1 / eps`` would otherwise hand it almost all
    the weight and switch the other terms off. If no component is active the
    weights are uniform.

    Args:
        gradients (torch.Tensor): Gradient norm of each loss component, shape ``(n,)``.

    Returns:
        torch.Tensor: The weights, shape ``(n,)``.
    """
    if self.running_grads is None:
        self.running_grads = gradients
        self.weights = (
            self.initial_weights.to(gradients.device)
            if self.initial_weights is not None
            else torch.ones_like(gradients)
        )
        self.logger.debug(f"LRW - Initialized weights: {self.weights}")
        return self.weights

    # Update running average of gradients
    self.running_grads = self.alpha * self.running_grads + (1 - self.alpha) * gradients

    # Ensure eps is a tensor on the same device as the gradients
    eps_tensor = torch.tensor(self.eps, device=self.running_grads.device)

    # Inverse-gradient weights over the components that have a gradient at all.
    active = self.running_grads > 0
    if not torch.any(active):
        self.weights = torch.full_like(self.running_grads, 1.0 / len(self.running_grads))
    else:
        inv_grads = torch.where(
            active,
            1.0 / (self.running_grads + eps_tensor),
            torch.zeros_like(self.running_grads),
        )
        self.weights = inv_grads / torch.sum(inv_grads)

    self.logger.debug(
        f"LRW - Updated weights: gradients={gradients.detach().cpu().numpy()}, running_grads={self.running_grads.detach().cpu().numpy()}, weights={self.weights.detach().cpu().numpy()}"
    )

    return self.weights

update_weights_rbw

update_weights_rbw(losses)

Update weights using Relative Error based Weighting (RBW).

Gives more weight to terms with bigger losses. After the first call (which only seeds L_bar)::

L_bar_i = alpha * L_bar_i + (1 - alpha) * L_i
w_i = alpha * w_prev_i + (1 - alpha) * L_bar_i / (sum_j L_bar_j + eps)

where w_prev is the previous RBW result; on the second call there is none yet and w_i is the normalised term alone.

Parameters:

Name	Type	Description	Default
losses	Tensor	Detached value of each loss component, shape (n,).	required

Returns:

Type	Description
	torch.Tensor: The weights, shape (n,).

Source code in pinnrl/training/adaptive_weights.py

def update_weights_rbw(self, losses):
    """Update weights using Relative Error based Weighting (RBW).

    Gives more weight to terms with bigger losses. After the first call
    (which only seeds ``L_bar``)::

        L_bar_i = alpha * L_bar_i + (1 - alpha) * L_i
        w_i = alpha * w_prev_i + (1 - alpha) * L_bar_i / (sum_j L_bar_j + eps)

    where ``w_prev`` is the previous RBW result; on the second call there is
    none yet and ``w_i`` is the normalised term alone.

    Args:
        losses (torch.Tensor): Detached value of each loss component, shape ``(n,)``.

    Returns:
        torch.Tensor: The weights, shape ``(n,)``.
    """
    if self.running_losses is None:
        self.running_losses = losses
        self.weights = (
            self.initial_weights.to(losses.device)
            if self.initial_weights is not None
            else torch.ones_like(losses)
        )
        self.logger.info(f"RBW - Initialized weights: {self.weights}")
        return self.weights

    # Update running average of losses
    self.running_losses = self.alpha * self.running_losses + (1 - self.alpha) * losses

    # Ensure eps is a tensor on the same device as the losses
    eps_tensor = torch.tensor(self.eps, device=self.running_losses.device)

    # Normalize the running losses to get weights
    # Give more weight to higher losses to focus optimization on problematic terms
    normalized_losses = self.running_losses / (self.running_losses.sum() + eps_tensor)
    self.weights = normalized_losses  # Higher loss -> Higher weight

    # Update weights with exponential moving average
    if self.prev_weights is not None:
        self.weights = self.alpha * self.prev_weights + (1 - self.alpha) * self.weights

    self.prev_weights = self.weights.clone()

    self.logger.debug(
        f"RBW - Updated weights: losses={losses.detach().cpu().numpy()}, running_losses={self.running_losses.detach().cpu().numpy()}, weights={self.weights.detach().cpu().numpy()}"
    )

    return self.weights

update

update(losses=None, gradients=None)

Update weights based on chosen strategy.

Parameters:

Name	Type	Description	Default
losses	Tensor | None	Loss component values (required for rbw).	None
gradients	Tensor | None	Gradient norm per component (required for lrw).	None

Returns:

Type	Description
	torch.Tensor: Updated loss weights.

Raises:

Type	Description
ValueError	If the input the strategy needs is missing, or the strategy is unknown.

Source code in pinnrl/training/adaptive_weights.py

def update(self, losses=None, gradients=None):
    """Update weights based on chosen strategy.

    Args:
        losses (torch.Tensor | None): Loss component values (required for ``rbw``).
        gradients (torch.Tensor | None): Gradient norm per component (required for ``lrw``).

    Returns:
        torch.Tensor: Updated loss weights.

    Raises:
        ValueError: If the input the strategy needs is missing, or the strategy is unknown.
    """
    if self.strategy == "lrw" and gradients is not None:
        return self.update_weights_lrw(gradients)
    elif self.strategy == "rbw" and losses is not None:
        return self.update_weights_rbw(losses)
    else:
        raise ValueError(f"Invalid combination of strategy ({self.strategy}) and inputs")

build_model

build_model(config_dict, device)

Build the PINNModel described by a build_config_dict result.

Returns:

Type	Description
	(model, config) with the model on device and the
	class:~pinnrl.config.Config it was built from.

Source code in pinnrl/training/train.py

def build_model(config_dict, device):
    """Build the PINNModel described by a ``build_config_dict`` result.

    Returns:
        ``(model, config)`` with the model on ``device`` and the
        :class:`~pinnrl.config.Config` it was built from.
    """
    arch_type = config_dict["model"]["architecture"]
    config_obj = Config()
    config_obj.device = device
    config_obj.model = ModelConfig.from_architecture(
        arch_type,
        config_dict["architectures"].get(arch_type, {}),
        input_dim=config_dict["model"]["input_dim"],
        output_dim=config_dict["model"]["output_dim"],
    )
    config_obj.training = TrainingConfig.from_dict(config_dict["training"])
    return PINNModel(config=config_obj, device=device).to(device), config_obj

load_run

load_run(run_dir, device=None)

Rebuild the PDE and trained model of a finished run.

Reads the run's config.yaml snapshot and final_model.pt and goes through the same builders as :func:run_training, so the result matches what was trained. The PDE is the forward problem at the configured (true) parameters, so its exact solution is the reference even for inverse runs; observation data is not reloaded (a Well run's snapshot records the time window its observations covered).

Parameters:

Name	Type	Description	Default
run_dir		Run folder written by :func:run_training.	required
device		Torch device; defaults to CPU.	None

Returns:

Type	Description
	(pde, model, config_dict) with the model in eval mode.

Source code in pinnrl/training/train.py

def load_run(run_dir, device=None):
    """Rebuild the PDE and trained model of a finished run.

    Reads the run's ``config.yaml`` snapshot and ``final_model.pt`` and goes
    through the same builders as :func:`run_training`, so the result matches
    what was trained. The PDE is the forward problem at the configured (true)
    parameters, so its exact solution is the reference even for inverse runs;
    observation data is not reloaded (a Well run's snapshot records the time
    window its observations covered).

    Args:
        run_dir: Run folder written by :func:`run_training`.
        device: Torch device; defaults to CPU.

    Returns:
        ``(pde, model, config_dict)`` with the model in eval mode.
    """
    device = device or torch.device("cpu")
    run_dir = Path(run_dir)
    with open(run_dir / "config.yaml") as f:
        config_dict = yaml.safe_load(f)
    pde_only = {
        **config_dict,
        "pde": {
            **config_dict["pde"],
            "observation_data": None,
            "trainable_parameters": [],
            "parameter_initial_guesses": {},
        },
    }
    pde_only["training"] = {**config_dict["training"], "mode": "forward"}
    pde = create_pde(pde_only, device)
    model, _ = build_model(config_dict, device)
    model.load_state_dict(torch.load(run_dir / "final_model.pt", map_location=device))
    model.eval()
    return pde, model, config_dict