pinn-pytorch-conventions

name	pinn-pytorch-conventions
description	PyTorch and PINA coding conventions, Pydantic state schemas, and PDE mathematical constraints for the LANG-PINN Code Agent.

PINN PyTorch & PINA Coding Skill

This skill provides the Code Agent with domain-specific knowledge for generating correct, modular, PINA-compliant PyTorch code for Physics-Informed Neural Networks.

---

1 — Mandatory Import Conventions

All generated code MUST use these import patterns:

# Core PyTorch
import torch
import torch.nn as nn
from torch.autograd import grad

# PINA framework — VERIFIED import paths (pina v0.2.6)
from pina import Condition, LabelTensor
from pina.problem import SpatialProblem, TimeDependentProblem
from pina.operator import grad as pina_grad, laplacian  # NOT pina.operators
from pina.model import FeedForward  # or DeepONet, etc.
from pina.solver import PINN as PINNSolver            # NOT pina.solvers
from pina.trainer import Trainer
from pina.domain import CartesianDomain                # NOT pina.geometry
from pina.equation import Equation, FixedValue

# PyTorch Lightning (via PINA)
import pytorch_lightning as pl

Rule: Never use raw torch.optim training loops when PINA's Trainer can be used. Raw loops are only acceptable when PINA lacks a required feature.

---

2 — Code Modularity Rules

The Code Agent MUST produce four separate modules, never a monolithic script:

Module	Filename	Responsibility
Model	model.py	nn.Module subclass defining the neural architecture
Loss	loss.py	PDE residual, BC, IC, and data loss functions
Data	data.py	Collocation point sampling, domain definition, Condition objects
Trainer	train.py	PINA Trainer setup, optimizer config, callback hooks

Interface Contracts

# model.py MUST export:
class PINNModel(nn.Module):
    def forward(self, x: LabelTensor) -> LabelTensor: ...

# loss.py MUST export:
def pde_residual(input_: LabelTensor, output_: LabelTensor) -> LabelTensor: ...
def boundary_loss(input_: LabelTensor, output_: LabelTensor) -> LabelTensor: ...
def initial_loss(input_: LabelTensor, output_: LabelTensor) -> LabelTensor: ...

# data.py MUST export:
class PDEProblem(SpatialProblem):  # or TimeDependentProblem
    output_variables: list[str]
    spatial_domain: CartesianDomain
    conditions: dict[str, Condition]

# train.py MUST export:
def train(problem: PDEProblem, model: PINNModel, epochs: int = 1000) -> dict: ...

---

3 — Pydantic State Schemas

All data passed between agents via session.state MUST conform to these schemas:

from pydantic import BaseModel, Field
from typing import Optional

class PDEDefinition(BaseModel):
    """Written to state['pde_definition'] by the PDE Parser Agent."""
    equation_latex: str = Field(description="LaTeX representation of the PDE")
    sympy_expr: str = Field(description="SymPy-parseable string of the PDE operator")
    variables: list[str] = Field(description="Solution and independent variables")
    domain: dict[str, tuple[float, float]] = Field(description="Variable name → (min, max)")
    boundary_conditions: list[dict] = Field(description="List of BC specs with type, location, value")
    initial_conditions: list[dict] = Field(description="List of IC specs with expression")
    pde_family: str = Field(description="Classification: heat|wave|burgers|navier_stokes|poisson|helmholtz|advection|custom")

class ArchitectureSelection(BaseModel):
    """Written to state['architecture'] by the Architecture Agent."""
    topology: str = Field(description="FourierMLP|MLP|CNN|GNN|Transformer|DeepONet|MultiFeedForward")
    hidden_layers: list[int] = Field(description="Layer widths, e.g. [64, 64, 64]")
    activation: str = Field(description="tanh|sin|gelu|relu|swish|softplus")
    rationale: str = Field(description="Why this architecture suits the PDE characteristics")

class GeneratedCode(BaseModel):
    """Written to state['generated_code'] by the Code Agent."""
    model_module: str = Field(description="Contents of model.py")
    loss_module: str = Field(description="Contents of loss.py")
    data_module: str = Field(description="Contents of data.py")
    trainer_module: str = Field(description="Contents of train.py")
    full_script: str = Field(description="Concatenated executable script")

class ExecutionResult(BaseModel):
    """Written to state['execution_result'] by the Execution Agent."""
    success: bool
    stdout: str
    stderr: str
    traceback: Optional[str] = None

class EvaluationMetrics(BaseModel):
    """Written to state['evaluation_metrics'] by the Evaluator Agent."""
    convergence_steps: Optional[int] = None
    final_mse: Optional[float] = None
    param_count: Optional[int] = None
    gradient_health: str = Field(description="healthy|vanishing|exploding")
    quality_score: float = Field(ge=0.0, le=1.0, description="Composite quality 0–1")
    accept: bool = Field(description="True if quality_score >= threshold")
    feedback: str = Field(description="Refinement instruction if accept=False, else 'PASSED'")

---

4 — PDE Mathematical Constraints

The Code Agent MUST verify these constraints before finalizing code:

4.1 — Operator Well-Formedness

• All differential operators must match the PDE's spatial dimension (e.g., laplacian requires ≥2D input)
• Time derivatives MUST use grad(output, input, ...create_graph=True) for higher-order autograd compatibility
• The PDE residual function must return a tensor of shape (N, 1) where N = number of collocation points

4.2 — Boundary Condition Validation

BC Type	PINA Implementation	Constraint
Dirichlet	Condition(location=..., equation=FixedValue(value))	Value must be scalar or callable
Neumann	Condition(location=..., equation=Equation(neumann_fn))	Gradient direction must match boundary normal
Periodic	Custom Equation enforcing u(x_min) == u(x_max)	Both boundary locations must be specified
Robin	Custom Equation with α·u + β·∂u/∂n = g	Coefficients α, β must be non-zero

4.3 — Loss Weighting Strategy

# Default loss weights — adjust based on Feedback Agent diagnostics
LOSS_WEIGHTS = {
    "pde_residual": 1.0,      # Primary physics loss
    "boundary": 10.0,          # Boundary compliance (higher weight)
    "initial": 10.0,           # Initial condition compliance
    "data": 1.0,               # Sparse data fitting (if available)
}

Rule: If the Feedback Agent detects edge violations (high boundary MSE), increase boundary weight by 5× in the next iteration.

4.4 — Collocation Point Sampling

• Interior points: Latin Hypercube Sampling (LHS) via CartesianDomain.sample(n, mode='lhs')
• Boundary points: Uniform grid on domain boundaries, minimum 200 points per boundary face
• Temporal points: Uniform or LHS in the time domain

---

5 — Activation Function Selection Guide

PDE Characteristic	Recommended Activation	Reason
Smooth solutions (heat, diffusion)	tanh	Smooth, bounded, excellent gradient flow
Oscillatory / periodic solutions	sin (SIREN)	Native periodicity captures wave structure
Sharp gradients / shocks (Burgers')	softplus or gelu	Smooth approximation to ReLU, avoids dead neurons
Multi-scale phenomena	sin + tanh hybrid	Multi-frequency encoding
General purpose	gelu	Best general gradient properties

---

6 — Code Quality Standards

1. Type hints: All function signatures MUST have type annotations
2. Docstrings: Every class and public function MUST have a docstring explaining inputs, outputs, and physics meaning
3. Device handling: Always use device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
4. Reproducibility: Set torch.manual_seed(42) and pl.seed_everything(42) at script start
5. Float precision: Default to torch.float64 for PDE problems (higher precision needed for gradient computation)
6. No hardcoded paths: All file paths must be configurable via arguments or environment variables