Skip to content

anypinn.lightning.module

PINNModule

Bases: LightningModule

Generic PINN Lightning module. Expects external Problem + Sampler + optimizer config.

Parameters:

Name Type Description Default
problem Problem

The PINN problem definition (constraints, fields, etc.).

 required
hp PINNHyperparameters

Hyperparameters for training.

 required
Source code in src/anypinn/lightning/module.py 
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
class PINNModule(pl.LightningModule):
    """
    Generic PINN Lightning module.
    Expects external Problem + Sampler + optimizer config.

    Args:
        problem: The PINN problem definition (constraints, fields, etc.).
        hp: Hyperparameters for training.
    """

    def __init__(
        self,
        problem: Problem,
        hp: PINNHyperparameters,
    ):
        super().__init__()
        self.save_hyperparameters(ignore=["problem"])

        self.problem = problem
        self.hp = hp

        def _log(key: str, value: Tensor, progress_bar: bool = False) -> None:
            self.log(
                key,
                value,
                on_step=False,
                on_epoch=True,
                prog_bar=progress_bar,
                batch_size=hp.training_data.batch_size,
            )

        self._log = cast(LogFn, _log)

    @override
    def on_fit_start(self) -> None:
        """
        Called when fit begins. Resolves validation sources using loaded data.
        """
        self.problem.inject_context(self.trainer.datamodule.context)  # type: ignore

    @override
    def on_predict_start(self) -> None:
        """
        Called when predict begins. Resolves validation sources using loaded data.
        """
        self.problem.inject_context(self.trainer.datamodule.context)  # type: ignore

    @override
    def training_step(self, batch: TrainingBatch, batch_idx: int) -> Tensor:
        """
        Performs a single training step.
        Calculates total loss from the problem.
        """
        return self.problem.training_loss(batch, self._log)

    @override
    def predict_step(self, batch: PredictionBatch, batch_idx: int) -> Predictions:
        """
        Performs a prediction step.
        """
        x_data, y_data = batch

        (data_batch, predictions) = self.problem.predict((x_data, y_data))
        true_values = self.problem.true_values(x_data)

        return (data_batch, predictions, true_values)

    @override
    def configure_optimizers(self) -> OptimizerLRScheduler:
        """
        Configures the optimizer and learning rate scheduler.
        """
        opt_cfg = self.hp.optimizer
        if isinstance(opt_cfg, LBFGSConfig):
            opt = torch.optim.LBFGS(
                self.parameters(),
                lr=opt_cfg.lr,
                max_iter=opt_cfg.max_iter,
                max_eval=opt_cfg.max_eval,
                history_size=opt_cfg.history_size,
                line_search_fn=opt_cfg.line_search_fn,
            )
        elif isinstance(opt_cfg, AdamConfig):
            opt = torch.optim.Adam(
                self.parameters(),
                lr=opt_cfg.lr,
                betas=opt_cfg.betas,
                weight_decay=opt_cfg.weight_decay,
            )
        else:
            opt = torch.optim.Adam(self.parameters(), lr=self.hp.lr)

        sch_cfg = self.hp.scheduler
        if not sch_cfg:
            return opt

        if isinstance(sch_cfg, CosineAnnealingConfig):
            sch = torch.optim.lr_scheduler.CosineAnnealingLR(
                opt,
                T_max=sch_cfg.T_max,
                eta_min=sch_cfg.eta_min,
            )
            return {
                "optimizer": opt,
                "lr_scheduler": {
                    "name": "lr",
                    "scheduler": sch,
                    "interval": "epoch",
                    "frequency": 1,
                },
            }

        sch = torch.optim.lr_scheduler.ReduceLROnPlateau(
            opt,
            mode=sch_cfg.mode,
            factor=sch_cfg.factor,
            patience=sch_cfg.patience,
            threshold=sch_cfg.threshold,
            min_lr=sch_cfg.min_lr,
        )

        return {
            "optimizer": opt,
            "lr_scheduler": {
                "name": "lr",
                "scheduler": sch,
                "monitor": LOSS_KEY,
                "interval": "epoch",
                "frequency": 1,
            },
        }

hp = hp instance-attribute

problem = problem instance-attribute

__init__(problem: Problem, hp: PINNHyperparameters)

Source code in src/anypinn/lightning/module.py 
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
def __init__(
    self,
    problem: Problem,
    hp: PINNHyperparameters,
):
    super().__init__()
    self.save_hyperparameters(ignore=["problem"])

    self.problem = problem
    self.hp = hp

    def _log(key: str, value: Tensor, progress_bar: bool = False) -> None:
        self.log(
            key,
            value,
            on_step=False,
            on_epoch=True,
            prog_bar=progress_bar,
            batch_size=hp.training_data.batch_size,
        )

    self._log = cast(LogFn, _log)

configure_optimizers() -> OptimizerLRScheduler

Configures the optimizer and learning rate scheduler.

Source code in src/anypinn/lightning/module.py 
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
@override
def configure_optimizers(self) -> OptimizerLRScheduler:
    """
    Configures the optimizer and learning rate scheduler.
    """
    opt_cfg = self.hp.optimizer
    if isinstance(opt_cfg, LBFGSConfig):
        opt = torch.optim.LBFGS(
            self.parameters(),
            lr=opt_cfg.lr,
            max_iter=opt_cfg.max_iter,
            max_eval=opt_cfg.max_eval,
            history_size=opt_cfg.history_size,
            line_search_fn=opt_cfg.line_search_fn,
        )
    elif isinstance(opt_cfg, AdamConfig):
        opt = torch.optim.Adam(
            self.parameters(),
            lr=opt_cfg.lr,
            betas=opt_cfg.betas,
            weight_decay=opt_cfg.weight_decay,
        )
    else:
        opt = torch.optim.Adam(self.parameters(), lr=self.hp.lr)

    sch_cfg = self.hp.scheduler
    if not sch_cfg:
        return opt

    if isinstance(sch_cfg, CosineAnnealingConfig):
        sch = torch.optim.lr_scheduler.CosineAnnealingLR(
            opt,
            T_max=sch_cfg.T_max,
            eta_min=sch_cfg.eta_min,
        )
        return {
            "optimizer": opt,
            "lr_scheduler": {
                "name": "lr",
                "scheduler": sch,
                "interval": "epoch",
                "frequency": 1,
            },
        }

    sch = torch.optim.lr_scheduler.ReduceLROnPlateau(
        opt,
        mode=sch_cfg.mode,
        factor=sch_cfg.factor,
        patience=sch_cfg.patience,
        threshold=sch_cfg.threshold,
        min_lr=sch_cfg.min_lr,
    )

    return {
        "optimizer": opt,
        "lr_scheduler": {
            "name": "lr",
            "scheduler": sch,
            "monitor": LOSS_KEY,
            "interval": "epoch",
            "frequency": 1,
        },
    }

on_fit_start() -> None

Called when fit begins. Resolves validation sources using loaded data.

Source code in src/anypinn/lightning/module.py 
48
49
50
51
52
53
@override
def on_fit_start(self) -> None:
    """
    Called when fit begins. Resolves validation sources using loaded data.
    """
    self.problem.inject_context(self.trainer.datamodule.context)  # type: ignore

on_predict_start() -> None

Called when predict begins. Resolves validation sources using loaded data.

Source code in src/anypinn/lightning/module.py 
55
56
57
58
59
60
@override
def on_predict_start(self) -> None:
    """
    Called when predict begins. Resolves validation sources using loaded data.
    """
    self.problem.inject_context(self.trainer.datamodule.context)  # type: ignore

predict_step(batch: PredictionBatch, batch_idx: int) -> Predictions

Performs a prediction step.

Source code in src/anypinn/lightning/module.py 
70
71
72
73
74
75
76
77
78
79
80
@override
def predict_step(self, batch: PredictionBatch, batch_idx: int) -> Predictions:
    """
    Performs a prediction step.
    """
    x_data, y_data = batch

    (data_batch, predictions) = self.problem.predict((x_data, y_data))
    true_values = self.problem.true_values(x_data)

    return (data_batch, predictions, true_values)

training_step(batch: TrainingBatch, batch_idx: int) -> Tensor

Performs a single training step. Calculates total loss from the problem.

Source code in src/anypinn/lightning/module.py 
62
63
64
65
66
67
68
@override
def training_step(self, batch: TrainingBatch, batch_idx: int) -> Tensor:
    """
    Performs a single training step.
    Calculates total loss from the problem.
    """
    return self.problem.training_loss(batch, self._log)