Select Model

Select a Model for a PINA Problem

[!IMPORTANT] Read RULES.md before using this skill — it applies to all skills.

Use this skill to choose the right neural network architecture for a PINA problem, configure its parameters, and optionally enhance it with embeddings or blocks.

Step 1 — Understand the problem together

Start by gathering enough context to make a reasoned recommendation. Ask the user these questions conversationally — one at a time, adapting to their answers — rather than firing them all at once.

• Input/output: What variables go in and come out? (e.g., input: x, y, t → output: u, v). If they have a PINA problem object, extract directly: len(problem.input_variables) / len(problem.output_variables).
• Data shape: What does a single sample look like? Scalar values per point? A field on a grid? A graph? Sensor readings + coordinates?
• Data structure (determines what operations fit): A flat table of points (linear / MLP layers), a regular grid (conv1d/conv2d/conv3d, or FNO), an irregular mesh or point cloud (message passing / graph nets, or coordinate-based MLPs like PINNs), or sequences / time series (RNN / transformer)?
• Problem type: Are you solving a PDE (physics-informed), learning an operator (map function to function), fitting a curve/surface, or discovering equations from data?
• Constraints: Roughly how much data? Any compute or time limits? Is accuracy-critical or is a rough baseline enough?

If the user cannot describe their problem at all, reply:

I need at least the input/output dimensions to suggest an architecture. What are the input variables and what are the outputs you want to predict?

Step 2 — Reason about architecture

Take the context from Step 1 and reason about which architecture family fits best. Present the reasoning to the user and invite their input rather than simply prescribing.

Available building blocks

Models live in PINA's model (pina.model), with smaller blocks in pina.model.block. Browsing these modules is the best way to discover what's available — the sections below discuss how to choose.

Reasoning guide

Work through these considerations in dialogue with the user:

1. Operator learning? — mapping between function spaces? That points to FNO (regular grid), DeepONet/MIONet (sensor→field), or neural operators (Averaging/LowRank/Graph). Ask about the data geometry.

2. Message-passing / GNN? — input is a graph with explicit connectivity (mesh, point cloud, network)?

   • PINA has many GNN based blocks in pina.model.block.message_passing.

3. Equation discovery? — SINDy or similar. Make sure the user knows this discovers symbolic expressions, not just a black-box surrogate.

4. Spline fitting? — Spline for 1D, SplineSurface for 2D, or VectorizedSpline for independent per-feature splines. Typically explicit fitting, not deep learning.

5. Standard pointwise mapping? — This is where most PDE problems land. Within this family, consider:

   • Complexity: Start simple (FeedForward). Only escalate if it struggles — PirateNet for multi-scale, ResidualFeedForward for depth.
   • Data regime: Lots of collocation points → deeper/larger MLP works. Very little data → KAN can be parameter-efficient.
   • Multi-scale / high-frequency: FourierFeatureEmbedding (Step 3) is a cheap add-on to any MLP. PirateNet bundles it natively.

6. Custom nn.Module — When none of the above fit, use pure PyTorch. This is common for:

   • Multi-input/multi-output architectures PINA doesn't have
   • Pretrained encoders (CNNs, transformers)
   • Novel research architectures
   • Lightweight models where PINA's wrapping adds unnecessary overhead

   The only requirement: it must be a valid torch.nn.Module that accepts a tensor/graph and returns a tensor/graph. With PINA solvers, keep the constructor parameter names input_dimensions and output_dimensions so labels are handled automatically (or use nn.Module directly without PINA).

Step 3 — Enhancements and embeddings

These are add-ons to improve convergence or enforce structure. Discuss whether the user's problem benefits from them.

Fourier Feature Embedding (multi-scale)

Use when the PDE solution has features at multiple length scales, or when a plain MLP converges slowly (spectral bias).

from pina.model.block import FourierFeatureEmbedding

embedding = FourierFeatureEmbedding(
    input_dimension=input_dim,
    output_dimension=128,  # must be even
    sigma=2.0,
)

Wrap any model manually:

class WithEmbedding(nn.Module):
    def __init__(self, input_dim, output_dim, inner_size=64):
        super().__init__()
        self.embedding = FourierFeatureEmbedding(input_dim, 128)
        self.net = nn.Sequential(
            nn.Linear(128, inner_size), nn.Tanh(),
            nn.Linear(inner_size, output_dim),
        )

    def forward(self, x):
        return self.net(self.embedding(x))

Periodic Boundary Embedding

Hard-code periodicity so the model never needs to learn it:

from pina.model.block import PeriodicBoundaryEmbedding

embedding = PeriodicBoundaryEmbedding(
    input_dimension=2,
    periods={"x": 2 * 3.14159},
)

Residual / Orthogonal blocks

• ResidualBlock(input_dim, output_dim, hidden_dim, activation) — stacks layers with a skip connection. Use inside a custom model.
• OrthogonalBlock(dim=-1) — orthonormalize features along a dimension. Helps training stability.

Both are in pina.model.block.

Message-passing blocks (graph models)

For graph-structured inputs, PINA provides blocks that follow the same pattern: update node features by aggregating neighbor information. Available in pina.model.block (top-level) and pina.model.block.message_passing:

Prefer the built-in GraphNeuralOperator (in pina.model) if it matches your problem — it composes lifting, message passing, and projection into one class. Drop down to these blocks when you need a custom arrangement.

Step 4 — Build and test

Construction pattern

Prefer PINA's built-in models when they match (less boilerplate). Use pure nn.Module when they don't.

# PINA model
from pina.model import FeedForward
model = FeedForward(input_dimensions=3, output_dimensions=3, inner_size=64, n_layers=4, func=nn.Tanh)

# Pure torch custom model
class MyModel(nn.Module):
    def __init__(self, input_dimensions, output_dimensions):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(input_dimensions, 128), nn.Tanh(),
            nn.Linear(128, output_dimensions),
        )
    def forward(self, x):
        return self.net(x)

model = MyModel(input_dimensions=3, output_dimensions=3)

# Graph model: lifting → message passing → projection
from pina.model.block import GNOBlock, FeedForward

class GraphModel(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        super().__init__()
        self.lifting = FeedForward(input_dim, hidden_dim)
        self.gno = GNOBlock(...)
        self.projection = FeedForward(hidden_dim, output_dim)

    def forward(self, x, edge_index, edge_attr):
        x = self.lifting(x)
        x = self.gno(x, edge_index, edge_attr)
        return self.projection(x)

Validation

• Forward pass — tensor model: model(torch.randn(batch, input_dim)) → (batch, output_dim).
• Forward pass — graph model: model(x, edge_index, edge_attr) with dummy tensors of the right shapes, or a torch_geometric.data.Data object if the model expects one. The output should match the expected node-feature shape.
• If using PINA solvers: confirm the solver can call model.forward() with the actual condition tensors (same dtype, device, shapes).
• If using custom nn.Module with PINA: keep input_dimensions and output_dimensions as constructor kwargs for automatic label extraction.

Checklist

• Input/output dimensions determined (from problem object or user Q&A)
• Architecture reasoning discussed with the user, not prescribed
• Model class selected (PINA built-in, or custom nn.Module)
• Embeddings/enhancements added if the problem needs them (multi-scale, periodic boundaries, residual connections)
• Model passes a forward and backward pass smoke test with correct shapes