Skip to content

Neural Networks

Model Factory

PINNModel 

PINNModel(config, device=None, **kwargs)

Bases: BaseNetwork

Physics-Informed Neural Network model factory.

Builds the network named by config.model.architecture (fourier, resnet, siren, attention, autoencoder, fno; anything else gives feedforward) and stores it as self.model.

Initialize the PINN model.

Parameters:

Name Type Description Default
config Config

Configuration object with model and device; config.model.device is set to the chosen device.

 required
device

Optional device to override config device

 None
**kwargs

Ignored.

 {}
Source code in pinnrl/neural_networks/__init__.py 
 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
def __init__(self, config: "Config", device=None, **kwargs):
    """Initialize the PINN model.

    Args:
        config: Configuration object with ``model`` and ``device``;
            ``config.model.device`` is set to the chosen device.
        device: Optional device to override config device
        **kwargs: Ignored.
    """
    # Store configuration
    self.config = config

    # Use provided device or get from config
    self.device = device if device is not None else config.device

    # Override device in config with the determined device
    config_with_device = config.model
    config_with_device.device = self.device

    # BaseNetwork reads the ModelConfig through its dict-like .get
    super().__init__(config_with_device)

    # Store architecture type
    self.architecture = config.model.architecture
    self.architecture_name = config.model.architecture  # More accessible name for metadata

    # Create model based on architecture
    if self.architecture == "fourier":
        self.model = FourierNetwork(config_with_device)
    elif self.architecture == "resnet":
        # Create a specific configuration dictionary for ResNet
        resnet_config = {
            "input_dim": config.model.input_dim,
            "hidden_dim": config.model.hidden_dim,
            "output_dim": config.model.output_dim,
            "activation": config.model.activation,
            "dropout": config.model.dropout,
            "device": self.device,
        }

        # Make sure num_blocks is defined, using num_layers as fallback
        if hasattr(config.model, "num_blocks") and config.model.num_blocks is not None:
            resnet_config["num_blocks"] = config.model.num_blocks
        else:
            resnet_config["num_blocks"] = config.model.num_layers

        # Add hidden_dims if defined
        if hasattr(config.model, "hidden_dims") and config.model.hidden_dims is not None:
            resnet_config["hidden_dims"] = config.model.hidden_dims

        self.model = ResNet(resnet_config)
    elif self.architecture == "siren":
        self.model = SIREN(config_with_device)
    elif self.architecture == "attention":
        self.model = AttentionNetwork(config_with_device)
    elif self.architecture == "autoencoder":
        self.model = AutoEncoder(config_with_device)
    elif self.architecture == "fno":
        fno_config = {
            "input_dim": config.model.input_dim,
            "hidden_dim": config.model.hidden_dim,
            "output_dim": config.model.output_dim,
            "num_blocks": getattr(config.model, "num_blocks", None) or config.model.num_layers,
            "modes": config.model.get("modes", 16),
            "activation": config.model.activation,
            "device": self.device,
        }
        self.model = FNONetwork(fno_config)
    else:  # Default to feedforward
        self.model = FeedForwardNetwork(config_with_device)

    # Explicitly move the model to the specified device
    self.model = self.model.to(self.device)
    self.to(self.device)

forward 

forward(x)

Forward pass of the model.

Parameters:

Name Type Description Default
x

Input tensor of shape (batch_size, input_dim)

 required

Returns:

Type Description

Output tensor of shape (batch_size, output_dim)
Source code in pinnrl/neural_networks/__init__.py 
122
123
124
125
126
127
128
129
130
131
def forward(self, x):
    """Forward pass of the model.

    Args:
        x: Input tensor of shape (batch_size, input_dim)

    Returns:
        Output tensor of shape (batch_size, output_dim)
    """
    return self.model(x)

Base

BaseNetwork 

BaseNetwork(config)

Bases: Module

Base class for all neural network architectures.

Initialize the base network.

Parameters:

Name Type Description Default
config dict

Configuration mapping with .get (a dict or a ModelConfig)

 required
Source code in pinnrl/neural_networks/base_network.py 
18
19
20
21
22
23
24
25
26
27
def __init__(self, config: dict) -> None:
    """Initialize the base network.

    Args:
        config: Configuration mapping with ``.get`` (a dict or a ``ModelConfig``)
    """
    super().__init__()
    self.config = config
    self.device = config.get("device", torch.device("cpu"))
    self.to(self.device)

forward 

forward(x)

Forward pass of the network.

Parameters:

Name Type Description Default
x InputType

Input tensor or array

 required

Returns:

Type Description
OutputType

Output tensor or array
Source code in pinnrl/neural_networks/base_network.py 
29
30
31
32
33
34
35
36
37
38
def forward(self, x: InputType) -> OutputType:
    """Forward pass of the network.

    Args:
        x: Input tensor or array

    Returns:
        Output tensor or array
    """
    raise NotImplementedError

save_state 

save_state(path)

Save model state with architecture information.

Parameters:

Name Type Description Default
path str

Path to save the model

 required
Source code in pinnrl/neural_networks/base_network.py 
59
60
61
62
63
64
65
66
def save_state(self, path: str) -> None:
    """Save model state with architecture information.

    Args:
        path: Path to save the model
    """
    state = {"model_state_dict": self.state_dict(), "config": self.config}
    torch.save(state, path)

load_state 

load_state(path)

Load model state with architecture information.

Parameters:

Name Type Description Default
path str

Path to load the model from

 required
Source code in pinnrl/neural_networks/base_network.py 
68
69
70
71
72
73
74
75
76
def load_state(self, path: str) -> None:
    """Load model state with architecture information.

    Args:
        path: Path to load the model from
    """
    state = torch.load(path, map_location=self.device, weights_only=False)
    self.load_state_dict(state["model_state_dict"])
    self.config = state["config"]

count_parameters 

count_parameters()

Count the number of trainable parameters.

Source code in pinnrl/neural_networks/base_network.py 
78
79
80
def count_parameters(self) -> int:
    """Count the number of trainable parameters."""
    return sum(p.numel() for p in self.parameters() if p.requires_grad)

FeedForward

FeedForwardNetwork 

FeedForwardNetwork(config)

Bases: BaseNetwork

Simple feed-forward neural network with configurable hidden layers.

Initialize the feed-forward network.

Parameters:

Name Type Description Default
config NetworkConfig

Network configuration dictionary containing: - input_dim: Input dimension - hidden_dims: List of hidden layer dimensions - output_dim: Output dimension - activation: Activation function name ('relu', 'tanh', etc.; default 'relu') - dropout: Dropout rate (default 0.1) - layer_norm: Whether to use layer normalization (default True) - device: Device to place the model on

 required
Source code in pinnrl/neural_networks/feedforward.py 
12
13
14
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
def __init__(self, config: NetworkConfig) -> None:
    """Initialize the feed-forward network.

    Args:
        config: Network configuration dictionary containing:
            - input_dim: Input dimension
            - hidden_dims: List of hidden layer dimensions
            - output_dim: Output dimension
            - activation: Activation function name ('relu', 'tanh', etc.; default 'relu')
            - dropout: Dropout rate (default 0.1)
            - layer_norm: Whether to use layer normalization (default True)
            - device: Device to place the model on
    """
    super().__init__(config)
    self.input_dim = config["input_dim"]
    self.hidden_dims = config["hidden_dims"]
    self.output_dim = config["output_dim"]
    self.dropout_rate = config.get("dropout", 0.1)
    self.use_layer_norm = config.get("layer_norm", True)
    self.device = config.get("device", torch.device("cpu"))

    # Get activation function name for per-layer instantiation
    activation_name = config.get("activation", "relu")

    # Build layers
    layers = []
    prev_dim = self.input_dim

    for hidden_dim in self.hidden_dims:
        layers.append(nn.Linear(prev_dim, hidden_dim))
        if self.use_layer_norm:
            layers.append(nn.LayerNorm(hidden_dim))
        # A fresh activation module per layer
        layers.append(self._get_activation_module(activation_name))
        if self.dropout_rate > 0.0:
            layers.append(nn.Dropout(self.dropout_rate))
        prev_dim = hidden_dim

    # Output layer
    layers.append(nn.Linear(prev_dim, self.output_dim))

    self.layers = nn.Sequential(*layers)

    # Move model to device
    self.to(self.device)

forward 

forward(x)

Forward pass of the network.

Parameters:

Name Type Description Default
x InputType

Input tensor or array

 required

Returns:

Type Description
OutputType

Output tensor or array
Source code in pinnrl/neural_networks/feedforward.py 
58
59
60
61
62
63
64
65
66
67
68
69
70
71
def forward(self, x: InputType) -> OutputType:
    """Forward pass of the network.

    Args:
        x: Input tensor or array

    Returns:
        Output tensor or array
    """
    x = self._prepare_input(x)
    # Ensure input is on the same device as the model
    if x.device != self.device:
        x = x.to(self.device)
    return self.layers(x)

ResNet

ResNet 

ResNet(config)

Bases: BaseNetwork

Residual Neural Network architecture built from multiple ResNetBlock components.

Initialize the ResNet architecture.

Parameters:

Name Type Description Default
config NetworkConfig

Network configuration dictionary containing: - input_dim: Input dimension - hidden_dim: Hidden layer dimension (else hidden_dims[0], else 124) - num_blocks: Number of ResNetBlock components (else num_layers, else 4) - output_dim: Output dimension - activation: Activation function name (default relu) - dropout: Dropout rate (default 0.1)

 required
Source code in pinnrl/neural_networks/resnet.py 
 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
def __init__(self, config: NetworkConfig) -> None:
    """Initialize the ResNet architecture.

    Args:
        config: Network configuration dictionary containing:
            - input_dim: Input dimension
            - hidden_dim: Hidden layer dimension (else ``hidden_dims[0]``, else 124)
            - num_blocks: Number of ResNetBlock components (else ``num_layers``, else 4)
            - output_dim: Output dimension
            - activation: Activation function name (default ``relu``)
            - dropout: Dropout rate (default 0.1)
    """
    super().__init__(config)
    self.input_dim = config["input_dim"]

    # Support both hidden_dim (scalar) and hidden_dims (list)
    if "hidden_dim" in config:
        self.hidden_dim = config["hidden_dim"]
    elif (
        "hidden_dims" in config
        and isinstance(config["hidden_dims"], list)
        and len(config["hidden_dims"]) > 0
    ):
        self.hidden_dim = config["hidden_dims"][0]
    else:
        self.hidden_dim = 124  # Default value if nothing is specified

    # Support both num_blocks and num_layers
    if "num_blocks" in config:
        self.num_blocks = config["num_blocks"]
    elif "num_layers" in config:
        self.num_blocks = config["num_layers"]
    else:
        self.num_blocks = 4  # Default value if nothing is specified

    self.output_dim = config["output_dim"]
    activation_name = config.get("activation", "relu")
    self.activation_fn = self._get_activation_module(activation_name)
    self.dropout = config.get("dropout", 0.1)

    # Build the ResNet architecture
    self.input_layer = nn.Linear(self.input_dim, self.hidden_dim)
    # Stack multiple ResNetBlock components
    self.blocks = nn.ModuleList(
        [
            ResNetBlock(
                self.hidden_dim,
                self.hidden_dim,
                config.get("activation", "relu"),
                self.dropout,
            )
            for _ in range(self.num_blocks)
        ]
    )
    self.output_layer = nn.Linear(self.hidden_dim, self.output_dim)

forward 

forward(x)

Forward pass of the network.

Parameters:

Name Type Description Default
x InputType

Input tensor or array

 required

Returns:

Type Description
OutputType

Output tensor or array
Source code in pinnrl/neural_networks/resnet.py 
128
129
130
131
132
133
134
135
136
137
138
139
140
141
def forward(self, x: InputType) -> OutputType:
    """Forward pass of the network.

    Args:
        x: Input tensor or array

    Returns:
        Output tensor or array
    """
    x = self._prepare_input(x)
    x = self.activation_fn(self.input_layer(x))
    for block in self.blocks:
        x = block(x)
    return self.output_layer(x)

ResNetBlock 

ResNetBlock(
    in_dim, hidden_dim, activation="relu", dropout=0.1
)

Bases: Module

Residual block act(x + F(x)).

F is Linear -> LayerNorm -> act -> Dropout -> Linear -> LayerNorm -> Dropout.

Initialize a single residual block component.

Parameters:

Name Type Description Default
in_dim int

Input dimension

 required
hidden_dim int

Hidden dimension

 required
activation str

Activation function name

 'relu'
dropout float

Dropout rate

 0.1
Source code in pinnrl/neural_networks/resnet.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
def __init__(
    self,
    in_dim: int,
    hidden_dim: int,
    activation: str = "relu",
    dropout: float = 0.1,
):
    """Initialize a single residual block component.

    Args:
        in_dim: Input dimension
        hidden_dim: Hidden dimension
        activation: Activation function name
        dropout: Dropout rate
    """
    super().__init__()

    # Get activation function
    if activation == "relu":
        self.activation_fn = nn.ReLU()
    elif activation == "leaky_relu":
        self.activation_fn = nn.LeakyReLU()
    elif activation == "tanh":
        self.activation_fn = nn.Tanh()
    elif activation == "sigmoid":
        self.activation_fn = nn.Sigmoid()
    elif activation == "gelu":
        self.activation_fn = nn.GELU()
    else:
        raise ValueError(f"Unsupported activation: {activation}")

    # Build the residual block
    self.layers = nn.Sequential(
        nn.Linear(in_dim, hidden_dim),
        nn.LayerNorm(hidden_dim),
        self.activation_fn,
        nn.Dropout(dropout),
        nn.Linear(hidden_dim, in_dim),
        nn.LayerNorm(in_dim),
        nn.Dropout(dropout),
    )

forward 

forward(x)

Forward pass with residual connection.

Parameters:

Name Type Description Default
x Tensor

Input tensor

 required

Returns:

Type Description
Tensor

Output tensor
Source code in pinnrl/neural_networks/resnet.py 
57
58
59
60
61
62
63
64
65
66
def forward(self, x: torch.Tensor) -> torch.Tensor:
    """Forward pass with residual connection.

    Args:
        x: Input tensor

    Returns:
        Output tensor
    """
    return self.activation_fn(x + self.layers(x))

SIREN

SIREN 

SIREN(config)

Bases: BaseNetwork

Sinusoidal Representation Network: SIRENLayer stack and a linear output layer.

Initialize the network.

Parameters:

Name Type Description Default
config NetworkConfig

Network configuration dictionary containing: - input_dim: Input dimension - hidden_dims: List of hidden layer dimensions - output_dim: Output dimension - omega_0: Frequency parameter (default 30.0); activation is ignored

 required
Source code in pinnrl/neural_networks/siren.py 
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
def __init__(self, config: NetworkConfig) -> None:
    """Initialize the network.

    Args:
        config: Network configuration dictionary containing:
            - input_dim: Input dimension
            - hidden_dims: List of hidden layer dimensions
            - output_dim: Output dimension
            - omega_0: Frequency parameter (default 30.0); ``activation`` is ignored
    """
    super().__init__(config)
    self.input_dim = config["input_dim"]
    self.hidden_dims = config["hidden_dims"]
    self.output_dim = config["output_dim"]
    self.omega_0 = config.get("omega_0", 30.0)

    # Build layers
    self.layers = nn.ModuleList()
    prev_dim = self.input_dim
    for hidden_dim in self.hidden_dims:
        self.layers.append(SIRENLayer(prev_dim, hidden_dim, omega_0=self.omega_0))
        prev_dim = hidden_dim
    self.layers.append(nn.Linear(prev_dim, self.output_dim))

forward 

forward(x)

Forward pass of the network.

Parameters:

Name Type Description Default
x InputType

Input tensor or array

 required

Returns:

Type Description
OutputType

Output tensor or array
Source code in pinnrl/neural_networks/siren.py 
79
80
81
82
83
84
85
86
87
88
89
90
91
def forward(self, x: InputType) -> OutputType:
    """Forward pass of the network.

    Args:
        x: Input tensor or array

    Returns:
        Output tensor or array
    """
    x = self._prepare_input(x)
    for layer in self.layers[:-1]:
        x = layer(x)
    return self.layers[-1](x)

SIRENLayer 

SIRENLayer(in_features, out_features, omega_0=30.0)

Bases: Module

SIREN layer sin(omega_0 * (W x + b)).

Initialize the layer.

Parameters:

Name Type Description Default
in_features int

Input dimension

 required
out_features int

Output dimension

 required
omega_0 float

Frequency parameter

 30.0
Source code in pinnrl/neural_networks/siren.py 
14
15
16
17
18
19
20
21
22
23
24
25
def __init__(self, in_features: int, out_features: int, omega_0: float = 30.0) -> None:
    """Initialize the layer.

    Args:
        in_features: Input dimension
        out_features: Output dimension
        omega_0: Frequency parameter
    """
    super().__init__()
    self.omega_0 = omega_0
    self.linear = nn.Linear(in_features, out_features)
    self.initialize_weights()

initialize_weights 

initialize_weights()

Draw W ~ U(-sqrt(6 / in) / omega_0, sqrt(6 / in) / omega_0).

This is the SIREN hidden-layer rule; it is used for every layer, including the first (the paper's first-layer U(-1/in, 1/in) is not applied). Biases keep PyTorch's default nn.Linear initialization.

Source code in pinnrl/neural_networks/siren.py 
27
28
29
30
31
32
33
34
35
36
37
38
def initialize_weights(self) -> None:
    """Draw ``W ~ U(-sqrt(6 / in) / omega_0, sqrt(6 / in) / omega_0)``.

    This is the SIREN hidden-layer rule; it is used for every layer,
    including the first (the paper's first-layer ``U(-1/in, 1/in)`` is not
    applied). Biases keep PyTorch's default ``nn.Linear`` initialization.
    """
    with torch.no_grad():
        self.linear.weight.uniform_(
            -np.sqrt(6 / self.linear.in_features) / self.omega_0,
            np.sqrt(6 / self.linear.in_features) / self.omega_0,
        )

forward 

forward(x)

Forward pass of the layer.

Parameters:

Name Type Description Default
x Tensor

Input tensor

 required

Returns:

Type Description
Tensor

Output tensor
Source code in pinnrl/neural_networks/siren.py 
40
41
42
43
44
45
46
47
48
49
def forward(self, x: Tensor) -> Tensor:
    """Forward pass of the layer.

    Args:
        x: Input tensor

    Returns:
        Output tensor
    """
    return torch.sin(self.omega_0 * self.linear(x))

Fourier

FourierNetwork 

FourierNetwork(config)

Bases: BaseNetwork

MLP on random Fourier features: [sin(xB), cos(xB)] -> hidden layers -> output.

Initialize the network.

Parameters:

Name Type Description Default
config NetworkConfig

Network configuration dictionary containing: - input_dim: Input dimension - mapping_size: Size of the Fourier feature mapping (default 32) - hidden_dims: Width of each hidden layer (one layer per entry, like the feedforward network); falls back to [hidden_dim] * num_layers (num_layers default 4) - output_dim: Output dimension - activation: Activation function name (default relu) - scale: Standard deviation of B (default 10.0)

 required
Source code in pinnrl/neural_networks/fourier.py 
 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
def __init__(self, config: NetworkConfig) -> None:
    """Initialize the network.

    Args:
        config: Network configuration dictionary containing:
            - input_dim: Input dimension
            - mapping_size: Size of the Fourier feature mapping (default 32)
            - hidden_dims: Width of each hidden layer (one layer per entry,
              like the feedforward network); falls back to
              ``[hidden_dim] * num_layers`` (``num_layers`` default 4)
            - output_dim: Output dimension
            - activation: Activation function name (default ``relu``)
            - scale: Standard deviation of ``B`` (default 10.0)
    """
    super().__init__(config)
    self.input_dim = config["input_dim"]
    self.mapping_size = config.get("mapping_size", 32)  # Default mapping size
    self.hidden_dims = list(
        config.get("hidden_dims") or [config["hidden_dim"]] * config.get("num_layers", 4)
    )
    self.output_dim = config["output_dim"]
    activation_name = config.get("activation", "relu")
    self.activation_fn = self._get_activation_module(activation_name)
    self.scale = config.get("scale", 10.0)

    # Build layers
    self.fourier = FourierFeatures(
        self.input_dim, self.mapping_size, self.scale, device=self.device
    )
    self.layers = nn.ModuleList()
    prev_dim = 2 * self.mapping_size  # Sine and cosine features

    # One hidden layer per width, then the output layer
    for width in self.hidden_dims:
        self.layers.append(nn.Linear(prev_dim, width))
        prev_dim = width

    # Add output layer
    self.layers.append(nn.Linear(prev_dim, self.output_dim))

    # Move all layers to device
    self.to(self.device)

forward 

forward(x)

Forward pass of the network.

Parameters:

Name Type Description Default
x InputType

Input tensor or array

 required

Returns:

Type Description
OutputType

Output tensor or array
Source code in pinnrl/neural_networks/fourier.py 
109
110
111
112
113
114
115
116
117
118
119
120
121
122
def forward(self, x: InputType) -> OutputType:
    """Forward pass of the network.

    Args:
        x: Input tensor or array

    Returns:
        Output tensor or array
    """
    x = self._prepare_input(x)
    x = self.fourier(x)
    for layer in self.layers[:-1]:
        x = self.activation_fn(layer(x))
    return self.layers[-1](x)

FourierFeatures 

FourierFeatures(
    input_dim, mapping_size, scale=10.0, device=None
)

Bases: Module

Random Fourier feature map [sin(xB), cos(xB)] with fixed B ~ N(0, scale^2).

B has shape (input_dim, mapping_size) and is a buffer, not trained; there is no 2Ï€ factor.

Initialize Fourier features.

Parameters:

Name Type Description Default
input_dim int

Input dimension

 required
mapping_size int

Number of frequencies (output has 2 * mapping_size features)

 required
scale float

Standard deviation of the entries of B

 10.0
device device | None

Device to place the features on (default: MPS if available, else CPU)

 None
Source code in pinnrl/neural_networks/fourier.py 
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
def __init__(
    self,
    input_dim: int,
    mapping_size: int,
    scale: float = 10.0,
    device: torch.device | None = None,
):
    """Initialize Fourier features.

    Args:
        input_dim: Input dimension
        mapping_size: Number of frequencies (output has ``2 * mapping_size`` features)
        scale: Standard deviation of the entries of ``B``
        device: Device to place the features on (default: MPS if available, else CPU)
    """
    super().__init__()
    self.input_dim = input_dim
    self.mapping_size = mapping_size
    self.scale = scale
    self.device = device or torch.device("mps" if torch.backends.mps.is_available() else "cpu")

    # Fixed random frequencies B ~ N(0, scale^2)
    self.register_buffer("B", torch.randn(input_dim, mapping_size, device=self.device) * scale)

    # Pre-compute output dimension
    self.output_dim = mapping_size * 2

forward 

forward(x)

Forward pass with Fourier feature mapping.

Parameters:

Name Type Description Default
x Tensor

Input tensor

 required

Returns:

Type Description
Tensor

[sin(xB), cos(xB)], shape (..., 2 * mapping_size)
Source code in pinnrl/neural_networks/fourier.py 
49
50
51
52
53
54
55
56
57
58
59
60
def forward(self, x: torch.Tensor) -> torch.Tensor:
    """Forward pass with Fourier feature mapping.

    Args:
        x: Input tensor

    Returns:
        ``[sin(xB), cos(xB)]``, shape ``(..., 2 * mapping_size)``
    """
    if self.B.device != x.device:
        self.B = self.B.to(x.device)
    return _fourier_feature_transform(x, self.B)

Fourier Neural Operator (FNO)

FNONetwork 

FNONetwork(config)

Bases: BaseNetwork

Fourier Neural Operator network for point-wise PINN training.

Architecture

Input(input_dim) -> Lift(Linear-GELU-Linear, hidden_dim) -> [FNOBlock(hidden_dim)] x num_blocks -> Project(Linear-GELU-Linear, output_dim)

Config keys: input_dim, output_dim, hidden_dim (default 128), num_blocks (default 4), modes (default 16) and activation of the blocks (default gelu).

Source code in pinnrl/neural_networks/fno.py 
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
146
def __init__(self, config: NetworkConfig) -> None:
    super().__init__(config)

    self.input_dim = config["input_dim"]
    self.output_dim = config["output_dim"]
    self.hidden_dim = config.get("hidden_dim", 128)
    self.num_blocks = config.get("num_blocks", 4)
    self.modes = config.get("modes", 16)
    activation = config.get("activation", "gelu")

    # Lifting layer: project (x, t) to higher-dimensional channel space
    self.lift = nn.Sequential(
        nn.Linear(self.input_dim, self.hidden_dim),
        nn.GELU(),
        nn.Linear(self.hidden_dim, self.hidden_dim),
    )

    # Spectral convolution blocks
    self.blocks = nn.ModuleList(
        [FNOBlock(self.hidden_dim, self.modes, activation) for _ in range(self.num_blocks)]
    )

    # Projection layers: channel space -> solution
    self.project = nn.Sequential(
        nn.Linear(self.hidden_dim, self.hidden_dim),
        nn.GELU(),
        nn.Linear(self.hidden_dim, self.output_dim),
    )

    self.to(self.device)

forward 

forward(x)

Map (x, t) points to the solution through the spectral blocks.

Source code in pinnrl/neural_networks/fno.py 
148
149
150
151
152
153
154
155
156
157
158
159
160
def forward(self, x: InputType) -> OutputType:
    """Map ``(x, t)`` points to the solution through the spectral blocks."""
    x = self._prepare_input(x)

    # Lift to channel space
    h = self.lift(x)

    # Apply FNO blocks
    for block in self.blocks:
        h = block(h)

    # Project to output
    return self.project(h)

FNOBlock 

FNOBlock(channels, modes, activation='gelu')

Bases: Module

Single FNO layer: act(LayerNorm(spectral(x) + linear(x) + x)).

Unknown activation names fall back to GELU.

Source code in pinnrl/neural_networks/fno.py 
75
76
77
78
79
80
81
82
83
84
85
86
87
88
def __init__(self, channels: int, modes: int, activation: str = "gelu"):
    super().__init__()
    self.spectral = SpectralConv1d(channels, modes)
    self.linear = nn.Linear(channels, channels)
    self.norm = nn.LayerNorm(channels)

    activations = {
        "relu": nn.ReLU(),
        "gelu": nn.GELU(),
        "tanh": nn.Tanh(),
        "leaky_relu": nn.LeakyReLU(),
        "sigmoid": nn.Sigmoid(),
    }
    self.activation = activations.get(activation, nn.GELU())

forward 

forward(x)

Forward: spectral path + linear bypass with residual connection.

Parameters:

Name Type Description Default
x Tensor

(batch, channels) tensor

 required

Returns:

Type Description
Tensor

(batch, channels) tensor
Source code in pinnrl/neural_networks/fno.py 
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
def forward(self, x: torch.Tensor) -> torch.Tensor:
    """Forward: spectral path + linear bypass with residual connection.

    Args:
        x: (batch, channels) tensor

    Returns:
        (batch, channels) tensor
    """
    spectral_out = self.spectral(x)
    linear_out = self.linear(x)
    return self.activation(self.norm(spectral_out + linear_out + x))

SpectralConv1d 

SpectralConv1d(channels, modes)

Bases: Module

1D spectral convolution layer operating in Fourier space.

Treats the feature/channel dimension as a 1D signal::

X = rfft(x)                                    # (batch, channels // 2 + 1)
Y[:, n] = sum_{m < M} X[:, m] * W[m, n]        # n < M; Y = 0 for n >= M
y = irfft(Y, n=channels)

with M = min(modes, channels // 2 + 1) and W a learnable complex (M, M) matrix (initialised N(0, 1) / (channels * M) per real and imaginary part), so the kept modes are mixed, not scaled one by one.

Source code in pinnrl/neural_networks/fno.py 
34
35
36
37
38
39
40
41
def __init__(self, channels: int, modes: int):
    super().__init__()
    self.channels = channels
    self.modes = min(modes, channels // 2 + 1)

    # Learnable complex weights: (modes, modes) in complex = (modes, modes, 2) real
    scale = 1.0 / (channels * self.modes)
    self.weights = nn.Parameter(scale * torch.randn(self.modes, self.modes, 2))

forward 

forward(x)

Apply spectral convolution along the feature dimension.

Parameters:

Name Type Description Default
x Tensor

(batch, channels) tensor

 required

Returns:

Type Description
Tensor

(batch, channels) tensor
Source code in pinnrl/neural_networks/fno.py 
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
def forward(self, x: torch.Tensor) -> torch.Tensor:
    """Apply spectral convolution along the feature dimension.

    Args:
        x: (batch, channels) tensor

    Returns:
        (batch, channels) tensor
    """
    # FFT along feature dimension
    x_ft = torch.fft.rfft(x, dim=-1)  # (batch, channels//2+1) complex

    # Truncate to kept modes
    modes = self.modes
    x_ft_trunc = x_ft[:, :modes]  # (batch, modes)

    # Complex weight multiplication: (batch, modes) x (modes, modes) -> (batch, modes)
    w = torch.view_as_complex(self.weights)  # (modes, modes)
    out_ft = torch.einsum("bm,mn->bn", x_ft_trunc, w)

    # Pad back to full frequency size and IFFT
    out_full = torch.zeros_like(x_ft)
    out_full[:, :modes] = out_ft
    return torch.fft.irfft(out_full, n=self.channels, dim=-1)

Attention

AttentionNetwork 

AttentionNetwork(config)

Bases: BaseNetwork

Neural network with self-attention layers.

Initialize the network.

Parameters:

Name Type Description Default
config NetworkConfig

Network configuration dictionary containing: - input_dim: Input dimension - hidden_dim: Hidden dimension - output_dim: Output dimension - num_layers: Number of SelfAttention + FeedForwardBlock pairs - num_heads: Number of attention heads (must divide hidden_dim) - dropout: Dropout rate - activation: Activation after the input projection only

 required
Source code in pinnrl/neural_networks/attention.py 
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
146
147
148
149
150
151
152
153
154
155
def __init__(self, config: NetworkConfig) -> None:
    """Initialize the network.

    Args:
        config: Network configuration dictionary containing:
            - input_dim: Input dimension
            - hidden_dim: Hidden dimension
            - output_dim: Output dimension
            - num_layers: Number of SelfAttention + FeedForwardBlock pairs
            - num_heads: Number of attention heads (must divide hidden_dim)
            - dropout: Dropout rate
            - activation: Activation after the input projection only
    """
    super().__init__(config)
    self.input_dim = config["input_dim"]
    self.hidden_dim = config["hidden_dim"]
    self.output_dim = config["output_dim"]
    self.num_layers = config.get("num_layers", 4)
    self.num_heads = config.get("num_heads", 4)
    self.dropout = config.get("dropout", 0.1)
    activation_name = config.get("activation", "gelu")
    self.activation_fn = self._get_activation_module(activation_name)

    # Input projection
    self.input_proj = nn.Linear(self.input_dim, self.hidden_dim)

    # Attention layers
    self.layers = nn.ModuleList()
    for _ in range(self.num_layers):
        self.layers.append(
            nn.ModuleList(
                [
                    SelfAttention(self.hidden_dim, self.num_heads, self.dropout),
                    FeedForwardBlock(self.hidden_dim, dropout=self.dropout),
                ]
            )
        )

    # Output projection
    self.output_proj = nn.Linear(self.hidden_dim, self.output_dim)

    # Initialize weights
    self.apply(self._init_weights)

forward 

forward(x)

Forward pass of the network.

Parameters:

Name Type Description Default
x InputType

Input tensor or array

 required

Returns:

Type Description
OutputType

Output tensor or array
Source code in pinnrl/neural_networks/attention.py 
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
def forward(self, x: InputType) -> OutputType:
    """Forward pass of the network.

    Args:
        x: Input tensor or array

    Returns:
        Output tensor or array
    """
    x = self._prepare_input(x)
    x = self.activation_fn(self.input_proj(x))

    # Apply attention layers
    for attn, ff in self.layers:
        x = attn(x)
        x = ff(x)

    return self.output_proj(x)

SelfAttention 

SelfAttention(dim, heads=4, dropout=0.1)

Bases: Module

Multi-head self-attention over a length-1 sequence, plus residual and LayerNorm.

Each point is its own sequence, so the softmax weight is always 1 and, without attention dropout, the layer computes LayerNorm(x + proj(value(x))).

Initialize the attention layer.

Parameters:

Name Type Description Default
dim int

Input dimension

 required
heads int

Number of attention heads

 4
dropout float

Dropout rate

 0.1
Source code in pinnrl/neural_networks/attention.py 
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
def __init__(self, dim: int, heads: int = 4, dropout: float = 0.1) -> None:
    """Initialize the attention layer.

    Args:
        dim: Input dimension
        heads: Number of attention heads
        dropout: Dropout rate
    """
    super().__init__()
    self.dim = dim
    self.heads = heads
    self.head_dim = dim // heads
    assert self.head_dim * heads == dim, "Dimension must be divisible by heads"

    self.query = nn.Linear(dim, dim)
    self.key = nn.Linear(dim, dim)
    self.value = nn.Linear(dim, dim)
    self.proj = nn.Linear(dim, dim)
    self.dropout = nn.Dropout(dropout)
    self.scale = self.head_dim**-0.5
    self.layer_norm = nn.LayerNorm(dim)

forward 

forward(x)

Forward pass of the layer.

Parameters:

Name Type Description Default
x Tensor

Input tensor of shape (batch_size, dim)

 required

Returns:

Type Description
Tensor

Output tensor of shape (batch_size, dim)
Source code in pinnrl/neural_networks/attention.py 
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
def forward(self, x: Tensor) -> Tensor:
    """Forward pass of the layer.

    Args:
        x: Input tensor of shape (batch_size, dim)

    Returns:
        Output tensor of shape (batch_size, dim)
    """
    residual = x
    batch_size = x.shape[0]

    # Use unsqueeze to add sequence length dimension of 1
    # Shape: (batch_size, 1, dim)
    x = x.unsqueeze(1)

    # Linear projections
    q = self.query(x).view(batch_size, 1, self.heads, self.head_dim).transpose(1, 2)
    k = self.key(x).view(batch_size, 1, self.heads, self.head_dim).transpose(1, 2)
    v = self.value(x).view(batch_size, 1, self.heads, self.head_dim).transpose(1, 2)

    # Attention
    scores = torch.matmul(q, k.transpose(-2, -1)) * self.scale
    attn = F.softmax(scores, dim=-1)
    attn = self.dropout(attn)

    # Apply attention to values
    out = torch.matmul(attn, v)
    out = out.transpose(1, 2).contiguous().view(batch_size, 1, self.dim)
    out = self.proj(out).squeeze(1)  # Remove sequence dimension

    # Add residual connection and apply layer norm
    out = self.layer_norm(out + residual)

    return out

AutoEncoder

AutoEncoder 

AutoEncoder(config)

Bases: BaseNetwork

Encoder-decoder MLP: input -> hidden_dims -> latent -> reversed(hidden_dims) -> output.

It maps coordinates to the solution; it does not reconstruct its input.

Initialize the network.

Parameters:

Name Type Description Default
config NetworkConfig

Network configuration dictionary containing: - input_dim: Input dimension - latent_dim: Latent space dimension - hidden_dims: Encoder hidden widths; the decoder uses them reversed - output_dim: Output dimension (default 1) - activation: Activation function name - dropout: Dropout rate (default 0.0) - layer_norm: LayerNorm after each hidden linear layer (default True)

 required
Source code in pinnrl/neural_networks/autoencoder.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
def __init__(self, config: NetworkConfig) -> None:
    """Initialize the network.

    Args:
        config: Network configuration dictionary containing:
            - input_dim: Input dimension
            - latent_dim: Latent space dimension
            - hidden_dims: Encoder hidden widths; the decoder uses them reversed
            - output_dim: Output dimension (default 1)
            - activation: Activation function name
            - dropout: Dropout rate (default 0.0)
            - layer_norm: LayerNorm after each hidden linear layer (default True)
    """
    super().__init__(config)
    self.input_dim = config["input_dim"]
    self.latent_dim = config.get("latent_dim", 16)
    self.hidden_dims = config.get("hidden_dims", [32, 64])
    activation_name = config.get("activation", "relu")
    self.dropout_rate = config.get("dropout", 0.0)
    self.use_layer_norm = config.get("layer_norm", True)

    self.output_dim = config.get("output_dim", 1)

    # Encoder maps input -> latent; the decoder maps latent -> output_dim
    # for PINN function approximation.
    self.encoder = nn.ModuleList()
    prev_dim = self.input_dim
    for hidden_dim in self.hidden_dims:
        self.encoder.append(nn.Linear(prev_dim, hidden_dim))
        if self.use_layer_norm:
            self.encoder.append(nn.LayerNorm(hidden_dim))
        self.encoder.append(self._get_activation_module(activation_name))
        if self.dropout_rate > 0.0:
            self.encoder.append(nn.Dropout(self.dropout_rate))
        prev_dim = hidden_dim
    self.encoder.append(nn.Linear(prev_dim, self.latent_dim))

    # Decoder: latent -> output_dim (solution value)
    self.decoder = nn.ModuleList()
    prev_dim = self.latent_dim
    for hidden_dim in reversed(self.hidden_dims):
        self.decoder.append(nn.Linear(prev_dim, hidden_dim))
        if self.use_layer_norm:
            self.decoder.append(nn.LayerNorm(hidden_dim))
        self.decoder.append(self._get_activation_module(activation_name))
        if self.dropout_rate > 0.0:
            self.decoder.append(nn.Dropout(self.dropout_rate))
        prev_dim = hidden_dim
    self.decoder.append(nn.Linear(prev_dim, self.output_dim))

encode 

encode(x)

Encode input to latent space.

Parameters:

Name Type Description Default
x InputType

Input tensor or array

 required

Returns:

Type Description
Tensor

Latent space representation
Source code in pinnrl/neural_networks/autoencoder.py 
65
66
67
68
69
70
71
72
73
74
75
76
77
def encode(self, x: InputType) -> torch.Tensor:
    """Encode input to latent space.

    Args:
        x: Input tensor or array

    Returns:
        Latent space representation
    """
    x = self._prepare_input(x)
    for layer in self.encoder:
        x = layer(x)
    return x

decode 

decode(z)

Decode latent space to output.

Parameters:

Name Type Description Default
z Tensor

Latent space tensor

 required

Returns:

Type Description
Tensor

Decoded output
Source code in pinnrl/neural_networks/autoencoder.py 
79
80
81
82
83
84
85
86
87
88
89
90
def decode(self, z: torch.Tensor) -> torch.Tensor:
    """Decode latent space to output.

    Args:
        z: Latent space tensor

    Returns:
        Decoded output
    """
    for layer in self.decoder:
        z = layer(z)
    return z

forward 

forward(x)

Forward pass of the network.

Parameters:

Name Type Description Default
x InputType

Input tensor or array

 required

Returns:

Type Description
OutputType

Predicted solution, shape (batch, output_dim)
Source code in pinnrl/neural_networks/autoencoder.py 
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
def forward(self, x: InputType) -> OutputType:
    """Forward pass of the network.

    Args:
        x: Input tensor or array

    Returns:
        Predicted solution, shape ``(batch, output_dim)``
    """
    z = self.encode(x)
    return self.decode(z)