Back to PINNs

All the ingredients are ready for a basic implementation of a PINN solver. The goal of this section is to gradually build and test this solver. A structure very similar to the one used for the FNN implementation is used.

The LayerNode class:

src
 1@dataclass
 2class LayerNode:
 3    """A layer has a certain number of nodes on the left,
 4    and another on the right. Every node on the left is
 5    fully connected to every node on the right.
 6
 7    Here, contrary to the implementation of the FNN class,
 8    the intermediate z and a values are not stored.
 9
10    The weights and biases are initialized randomly,
11    potentially using a seed for repeatability.
12
13    The NodeVectorized class is used for
14    vectorized auto-differentiation"""
15
16    nb_nodes_left: int
17    nb_nodes_right: int
18    ind_layer: int
19
20    seed: int | None = None
21
22    def __post_init__(self) -> None:
23        self.init_parameters_random()
24
25    def init_parameters_random(self) -> None:
26        rng = np.random.default_rng(self.seed)
27        weights: Array2D = (
28            rng.random((self.nb_nodes_right, self.nb_nodes_left), dtype=FLOAT) - 0.5
29        )
30        biases: Vector = rng.random(self.nb_nodes_right, dtype=FLOAT) - 0.5
31        self.weights = NodeVectorized(weights, name=f"w_L{self.ind_layer}")
32        self.biases = NodeVectorized(biases, name=f"b_L{self.ind_layer}")

Let us focus on the modfication FNN class into the PINN class. Most of the new code appears in the losses and gradient calculations.

Exercise (intermediate)

Implement and test the suggested methods.

src
  1@dataclass
  2class PINN:
  3    """A PINN has a sequence of layers. It owns all the hidden layers
  4    and the output layer, not the input layer. However, to build
  5    the first layer, the number of nodes of the input layer must be known.
  6
  7    Every layer has an activation.
  8
  9    The two main methods are:
 10    - infer: forward execution of the FNN from input to output
 11      with the current weights and biases
 12    - fit: learn all the weights and biases on a train set
 13    """
 14
 15    nb_nodes_per_layer_including_input: Sequence[int]
 16    activations: Sequence[Callable[[NodeVectorized], NodeVectorized]]
 17
 18    seed: int | None = None
 19
 20    def __post_init__(self) -> None:
 21        """creates a sequence of layers."""
 22        self.layers: list[LayerNode] = []
 23        self.nb_layers: int = len(self.nb_nodes_per_layer_including_input) - 1
 24
 25        if self.nb_layers != len(self.activations):
 26            raise IncompatibleDimensions(
 27                "The number of activation functions is not compatible with the number of nodes."
 28            )
 29
 30        for ind in range(self.nb_layers):
 31            layer = LayerNode(
 32                self.nb_nodes_per_layer_including_input[ind],
 33                self.nb_nodes_per_layer_including_input[ind + 1],
 34                ind,
 35                self.seed + ind if self.seed is not None else None,
 36                # to avoid using the same sequence of random numbers
 37                # for every layer which could introduce unwanted symetries
 38            )
 39            self.layers.append(layer)
 40
 41    def infer(self, input_: NodeVectorized) -> NodeVectorized:
 42        """forward sweep across the FNN,
 43        z_right = w . a_left
 44        a_right = sigma(z_right)
 45
 46        computes the attributes z and a of every layer,
 47        returns the last one.
 48
 49        Note that the vectorized node forward operations are used."""
 50        ...
 51
 52    def init_parameters_gradient_zero(self) -> tuple[list[Array2D], list[Vector]]:
 53        """initialize the arrays for the weight and biases gradients to zero
 54        with the correct dimension to be able to compute the mean gradients
 55        accross the train set."""
 56        biases_all: list[Vector] = []
 57        weights_all: list[Array2D] = []
 58        # init to zero for cumsum over training examples
 59        for ind_layer in range(0, self.nb_layers):
 60            layer = self.layers[ind_layer]
 61            weights_all.append(
 62                np.zeros((layer.nb_nodes_right, layer.nb_nodes_left), dtype=FLOAT)
 63            )
 64            biases_all.append(np.zeros(layer.nb_nodes_right, dtype=FLOAT))
 65        return weights_all, biases_all
 66
 67    def parameter_gradient_residual(
 68        self,
 69        inputs: Sequence[Vector],
 70        get_residual: Callable[
 71            [NodeVectorized, NodeVectorized], Sequence[NodeVectorized]
 72        ],
 73    ) -> tuple[list[Array2D], list[Vector]]:
 74        """This method builds the sequence of gradients of the weights and biases
 75        of every layer for the residual.
 76
 77        Since the residual N[u] is arbitrary because it depends on the differential
 78        problem to solve, it is given as a callable input: get_residual
 79
 80        The get_residual takes: the input node, the output node
 81        returns: a sequence of nodes representing every row of the differential system
 82
 83        Here the loss is built by cumulating for every train case
 84        the residual loss.
 85
 86        The inputs represent the colocation points.
 87
 88        The function parameter_gradient_from_loss takes care of the differentiation
 89        with respect to the weights and biases and the averaging.
 90        """
 91        ...
 92
 93    def parameter_gradient_dirichlet(
 94        self,
 95        inputs: Sequence[Vector],
 96        targets: Sequence[Vector],
 97    ) -> tuple[list[Array2D], list[Vector]]:
 98        """This method builds the sequence of gradients of the weights and biases
 99        of every layer for the Dirichlet conditions.
100
101        The inputs represent the points where a Dirichlet condition is imposed,
102        the targets are the value.
103
104        Here the loss is built by cumulating for every train case
105        the Dirchlet loss (which is the MSE between target and current output value).
106
107        The function parameter_gradient_from_loss takes care of the differentiation
108        with respect to the weights and biases and the averaging.
109        """
110        ...
111
112    def parameter_gradient_from_loss(
113        self,
114        nb_training: int,
115        loss: NodeVectorized,
116    ) -> tuple[list[Array2D], list[Vector]]:
117        """differentiates for every layer the loss and produces
118        the mean gradients over the train set.
119
120        The loss attribute is already the sum of a given loss type
121        over the train set."""
122        # first store the sum for all training example of gradients
123        # then at the end, the mean gradient
124        ...
125
126    @staticmethod
127    def bad_input_or_target_size(expected_size: int, vectors: Sequence[Vector]) -> bool:
128        return any([vector.size != expected_size for vector in vectors])
129
130    def fit(
131        self,
132        inputs_dirichlet: Sequence[Vector],
133        targets_dirichlet: Sequence[Vector],
134        inputs_residual: Sequence[Vector],
135        get_residual: Callable[
136            [NodeVectorized, NodeVectorized], Sequence[NodeVectorized]
137        ],
138        learning_rate: float = 0.1,
139        nb_iter_gradient: int = 1_000,
140        display_interval: int | None = 100,
141    ) -> None:
142        """Gradient descent algorithm on both the weights and biases, no stopping criteria,
143        no stochastic pooling.
144
145        The learning_rate is for the gradient from the train set.
146
147        No regularization is implemented.
148
149        A progressbar indicates the progress and the different gradient norms
150        to show a convergence (or not).
151        """
152        if self.bad_input_or_target_size(
153            self.nb_nodes_per_layer_including_input[0], inputs_dirichlet
154        ):
155            raise IncompatibleDimensions("The size of the input is incompatible.")
156        if self.bad_input_or_target_size(
157            self.nb_nodes_per_layer_including_input[0], inputs_residual
158        ):
159            raise IncompatibleDimensions("The size of the input is incompatible.")
160        if self.bad_input_or_target_size(
161            self.nb_nodes_per_layer_including_input[-1], targets_dirichlet
162        ):
163            raise IncompatibleDimensions("The size of the target is incompatible.")
164
165        for ind in tqdm(range(nb_iter_gradient), desc="pinn fit"):
166            gradients_weights_residual, gradients_biases_residual = (
167                self.parameter_gradient_residual(inputs_residual, get_residual)
168            )
169            gradients_weights_dirichlet, gradients_biases_dirichlet = (
170                self.parameter_gradient_dirichlet(inputs_dirichlet, targets_dirichlet)
171            )
172
173            message: list = []
174            for ind_layer, layer in enumerate(self.layers):
175                gradient_weights = (
176                    gradients_weights_residual[ind_layer]
177                    + gradients_weights_dirichlet[ind_layer]
178                )
179                gradient_biases = (
180                    gradients_biases_residual[ind_layer]
181                    + gradients_biases_dirichlet[ind_layer]
182                )
183
184                if display_interval is not None and ind % display_interval == 0:
185                    message.append(
186                        [
187                            ind_layer,
188                            np.linalg.norm(gradients_weights_residual[ind_layer]),
189                            np.linalg.norm(gradients_weights_dirichlet[ind_layer]),
190                            np.linalg.norm(gradients_biases_residual[ind_layer]),
191                            np.linalg.norm(gradients_biases_dirichlet[ind_layer]),
192                        ]
193                    )
194
195                layer.weights.val -= learning_rate * gradient_weights
196                layer.biases.val -= learning_rate * gradient_biases
197
198            if display_interval is not None and ind % display_interval == 0:
199                tqdm.write(tabulate(message, headers=["L", "w N", "w D", "b N", "b D"]))
200                tqdm.write("\n")