scimba_torch.numerical_solvers.temporal_pde.discrete_pinn

Implement a discrete PINN for solving temporal PDEs.

Classes

DiscretePINN(pde, projector[, scheme])

A discrete PINN for solving a differential equation.

DiscretePINNImplicit(pde, pde_imp, projector)

Feature-branch DiscretePINN_Implicit that extends the base DiscretePINN.

class DiscretePINN(pde, projector, scheme='euler_exp', **kwargs)[source]

Bases: ExplicitTimeDiscreteScheme

A discrete PINN for solving a differential equation.

It uses linear and/or nonlinear spaces.

The class supports initialization of the model with a target function, computation of the right-hand side (RHS) from the model, and stepping through time using a projector.

Parameters:
  • pde (TemporalPDE | KineticPDE) – The PDE model, which can be either a TemporalPDE or KineticPDE.

  • projector (TimeDiscreteCollocationProjector) – The projector for training the model, typically a TimeDiscreteCollocationProjector.

  • scheme (str) – The time scheme used for updating the model. Options include “euler_exp”, “rk2”, and “rk4”.

  • **kwargs – Additional hyperparameters for the scheme.

model

The neural network model used for computing the solution.

Type:

nn.Module

projector

The projector used for training the model.

Type:

ProjectorPINNs

scheme

The time scheme used for updating the model.

Type:

str

construct_rhs_ee(pde_n, t, dt)[source]

Returns \(u_n + \Delta t f(u_n)\) as a function.

Parameters:
  • pde_n (TemporalPDE | KineticPDE) – The PDE model at the current time step.

  • t (float) – The current time.

  • dt (float) – The time step size.

Return type:

Union[Callable[[LabelTensor, LabelTensor], Tensor], Callable[[LabelTensor, LabelTensor, LabelTensor], Tensor]]

Returns:

A function representing the right-hand side (RHS) for the explicit Euler scheme.

construct_rhs_ee_weak_bc(pde_n, t, dt)[source]

Returns \(u_n + \Delta t f(u_n)\) as a function.

Parameters:
  • pde_n (TemporalPDE | KineticPDE) – The PDE model at the current time step.

  • t (float) – The current time.

  • dt (float) – The time step size.

Returns:

one for the right-hand side (RHS) and another for the boundary condition (BC) residuals for the explicit Euler scheme with weak boundary conditions.

Return type:

A tuple containing two functions

construct_rhs_rk2(pde_n, pde_nph, t, dt)[source]

Returns \(u_n + \Delta t f(u_n)\) as a function.

Parameters:
  • pde_n (TemporalPDE | KineticPDE) – The PDE model at the current time step.

  • pde_nph (TemporalPDE | KineticPDE) – The PDE model at the next half time step.

  • t (float) – The current time.

  • dt (float) – The time step size.

Return type:

Union[Callable[[LabelTensor, LabelTensor], Tensor], Callable[[LabelTensor, LabelTensor, LabelTensor], Tensor]]

Returns:

A function representing the right-hand side (RHS) for the second-order Runge-Kutta (RK2) scheme.

construct_rhs_rk2_weak_bc(pde_n, pde_nph, t, dt)[source]

Returns \(u_n + \Delta t f(u_n)\) as a function.

Parameters:
  • pde_n (TemporalPDE | KineticPDE) – The PDE model at the current time step.

  • pde_nph (TemporalPDE | KineticPDE) – The PDE model at the next half time step.

  • t (float) – The current time.

  • dt (float) – The time step size.

Returns:

one for the right-hand side (RHS) and another for the boundary condition (BC) residuals for the second-order Runge-Kutta (RK2) scheme with weak boundary conditions.

Return type:

A tuple containing two functions

construct_rhs_rk4(pde_n, pde_1, pde_2, pde_3, t, dt)[source]

Returns \(u_n + \Delta t f(u_n)\) as a function.

Parameters:
Return type:

Union[Callable[[LabelTensor, LabelTensor], Tensor], Callable[[LabelTensor, LabelTensor, LabelTensor], Tensor]]

Returns:

A function representing the right-hand side (RHS) of the RK4 scheme.

Raises:

NotImplementedError – If the PDE models are not of type TemporalPDE, as RK4 is not implemented for KineticPDE.

update(t, dt, **kwargs)[source]

Compute the new parameters \(\theta_{n+1}\) from \(\theta_n\).

Uses the time step dt and the chosen time scheme.

Parameters:
  • t (float) – The current time.

  • dt (float) – The time step size.

  • **kwargs – Additional arguments for the projector’s solve method.

Raises:

NotImplementedError – If the RK4 scheme is chosen with weak boundary conditions, as it is not implemented.

class DiscretePINNImplicit(pde, pde_imp, projector, type_scheme='implicit', scheme='euler_imp', **kwargs)[source]

Bases: DiscretePINN

Feature-branch DiscretePINN_Implicit that extends the base DiscretePINN.

This subclass keeps all base behaviour (explicit schemes) and adds implicit schemes (Euler-imp, Crank–Nicolson, SDIRK2) using an extra PDE instance pde_imp and a type_scheme switch.

Parameters:
  • pde (TemporalPDE | KineticPDE) – The PDE model, which can be either a TemporalPDE or KineticPDE.

  • pde_imp (TemporalPDE | KineticPDE) – The PDE for implicit scheme.

  • projector (TimeDiscreteCollocationProjector) – The projector for training the model, typically a TimeDiscreteCollocationProjector.

  • type_scheme (str) – a switch explicit/implicit.

  • scheme (str) – The time scheme used for updating the model. Options include “euler_exp”, “rk2”, and “rk4”.

  • **kwargs – Additional hyperparameters for the scheme.

Raises:

ValueError – If the time_scheme is not implicit or explicit.

construct_rhs_ei(pde_n, t, dt)[source]

Construct the right-hand side of the implicit Euler scheme.

Returns u_n + Delta t , S(u_n) as a function, corresponding to the right-hand side (RHS) of the implicit Euler scheme.

Parameters:
  • pde_n (TemporalPDE | KineticPDE) – The PDE model at the current time step.

  • t (float) – The current time.

  • dt (float) – The time step size.

Return type:

Union[Callable[[LabelTensor, LabelTensor], Tensor], Callable[[LabelTensor, LabelTensor, LabelTensor], Tensor]]

Returns:

A function representing the right-hand side (RHS) for the implicit Euler scheme, where S(u_n) denotes the source or reaction term of the PDE.

construct_lhs_ei(pde_np1, t, dt)[source]

Construct the left-hand side of the implicit Euler scheme.

Returns u_{n+1} - Delta t , F(u_{n+1}) as a function, corresponding to the left-hand side (LHS) of the implicit Euler scheme.

Parameters:
  • pde_np1 (TemporalPDE | KineticPDE) – The PDE model evaluated at the next time step \(t_{n+1}\).

  • t (float) – The current time.

  • dt (float) – The time step size.

Return type:

Union[Callable[[LabelTensor, LabelTensor], Tensor], Callable[[LabelTensor, LabelTensor, LabelTensor], Tensor]]

Returns:

A function representing the left-hand side (LHS) for the implicit Euler scheme, where F(u_{n+1}) denotes the spatial operator of the PDE.

construct_rhs_cn(pde_n, t, dt)[source]

Construct the right-hand side of the Crank–Nicolson scheme.

Returns u^n - tfrac{1}{2} Delta t , (F(u^n) + S(u^n)) as a function, corresponding to the right-hand side (RHS) of the Crank–Nicolson scheme.

Parameters:
  • pde_n (TemporalPDE | KineticPDE) – The PDE model at the current time step.

  • t (float) – The current time.

  • dt (float) – The time step size.

Return type:

Union[Callable[[LabelTensor, LabelTensor], Tensor], Callable[[LabelTensor, LabelTensor, LabelTensor], Tensor]]

Returns:

A function representing the right-hand side (RHS) for the Crank–Nicolson scheme, combining both the spatial operator F(u^n) and the source term S(u^n) evaluated at the current time step.

construct_lhs_cn(pde_np1, t, dt)[source]

Construct the left-hand side of the Crank–Nicolson scheme.

Returns u^{n+1} + tfrac{1}{2} Delta t , (F(u^{n+1}) - S(u^{n+1}))

as a function, corresponding to the left-hand side (LHS) of the Crank–Nicolson scheme.

Parameters:
  • pde_np1 (TemporalPDE | KineticPDE) – The PDE model evaluated at the next time step t_{n+1}.

  • t (float) – The current time.

  • dt (float) – The time step size.

Return type:

Union[Callable[[LabelTensor, LabelTensor], Tensor], Callable[[LabelTensor, LabelTensor, LabelTensor], Tensor]]

Returns:

A function representing the left-hand side (LHS) for the Crank–Nicolson scheme, where F(u^{n+1}) and S(u^{n+1}) denote respectively the spatial operator and the source term of the PDE, both evaluated at the next time step.

construct_rhs_sdirk2_step1(pde_n, t, dt)[source]

Construct the right-hand side of the first stage of the SDIRK2 scheme.

Returns u^n as a function, corresponding to the right-hand side (RHS) of the first stage of the SDIRK2 scheme.

Parameters:
  • pde_n (TemporalPDE | KineticPDE) – The PDE model at the current time step.

  • t (float) – The current time.

  • dt (float) – The time step size.

Return type:

Union[Callable[[LabelTensor, LabelTensor], Tensor], Callable[[LabelTensor, LabelTensor, LabelTensor], Tensor]]

Returns:

A function representing the RHS for the first stage of the SDIRK2 scheme, where the initial value \(u^n\) is used as input for the implicit solve.

construct_lhs_sdirk2_step1(pde, t, dt)[source]

Construct the left-hand side of the first stage of the SDIRK2 scheme.

Returns u_1 - Delta t , gamma , (F(u_1) - S(u_1)) as a function, corresponding to the left-hand side (LHS) of the first stage of the SDIRK2 scheme.

Parameters:
  • pde (TemporalPDE | KineticPDE) – The PDE model evaluated at the first stage.

  • t (float) – The current time.

  • dt (float) – The time step size.

Return type:

Union[Callable[[LabelTensor, LabelTensor], Tensor], Callable[[LabelTensor, LabelTensor, LabelTensor], Tensor]]

Returns:

A function representing the LHS for the first stage of the SDIRK2 scheme, where gamma is the SDIRK2 coefficient and F and S denote respectively the spatial and source operators of the PDE.

construct_rhs_sdirk2_step2(pde_n, u1_copy, t, dt)[source]

Construct the right-hand side of the second stage of the SDIRK2 scheme.

Returns u^n - Delta t , (1 - gamma) , (F(u_1) - S(u_1)) as a function, corresponding to the right-hand side (RHS) of the second stage of the SDIRK2 scheme.

Parameters:
  • pde_n (TemporalPDE | KineticPDE) – The PDE model at the current time step.

  • u1_copy (TemporalPDE | KineticPDE) – The PDE model containing the intermediate stage solution u_1.

  • t (float) – The current time.

  • dt (float) – The time step size.

Return type:

Union[Callable[[LabelTensor, LabelTensor], Tensor], Callable[[LabelTensor, LabelTensor, LabelTensor], Tensor]]

Returns:

A function representing the RHS for the second stage of the SDIRK2 scheme, using the already computed intermediate solution u_1.

construct_lhs_sdirk2_step2(pde, t, dt)[source]

Construct the left-hand side of the second stage of the SDIRK2 scheme.

Returns u_2 - Delta t , gamma , (F(u_2) - S(u_2)) as a function, corresponding to the left-hand side (LHS) of the second stage of the SDIRK2 scheme.

Parameters:
  • pde (TemporalPDE | KineticPDE) – The PDE model evaluated at the second stage.

  • t (float) – The current time.

  • dt (float) – The time step size.

Return type:

Union[Callable[[LabelTensor, LabelTensor], Tensor], Callable[[LabelTensor, LabelTensor, LabelTensor], Tensor]]

Returns:

A function representing the LHS for the second stage of the SDIRK2 scheme, where gamma is the SDIRK2 coefficient and F and S denote respectively the spatial and source operators of the PDE.

update(t, dt, **kwargs)[source]

Compute theta_{n+1} from theta_n.

Compute the new parameters theta_{n+1} from theta_n using either explicit or implicit time schemes.

Dispatches to : update (explicit) or update_implicit (implicit) depending on type_scheme.

Parameters:
  • t (float) – The current time.

  • dt (float) – The time step size.

  • **kwargs – Additional arguments for the projector’s solve method.

Raises:

ValueError – If type_scheme is neither explicit nor implicit.

update_implicit(t, dt, **kwargs)[source]

Compute theta_{n+1} from theta_n for implicit schemes.

Compute the new parameters theta_{n+1} from theta_n for implicit schemes, using self.pde_imp`.

Uses the time step dt and the chosen implicit scheme.

Parameters:
  • t (float) – The current time.

  • dt (float) – The time step size.

  • **kwargs – Additional arguments for the projector’s solve method.