""" escnn_baseline.py - ESCNN GroupPooling baseline for D4-PINN comparison ======================================================================= This script runs an ESCNN-based D4-equivariant PINN on the same reaction-diffusion BVP as our main D4-PINN, and writes the comparison figure (fig14) and table (T13) into a fresh `output_/` directory matching the schema expected by `replace_with_real.py`. Requirements: pip install escnn>=1.0 pip install torch numpy matplotlib Usage: python escnn_baseline.py python replace_with_real.py output_ The ESCNN architecture used here is the simplest possible D4-equivariant MLP analogue: a stack of regular-representation linear layers ending with a GroupPooling that takes the max over the |D4|=8 orbit. By Theorem 1 of Yarotsky (Constr. Approx. 2022) and Cohen-Geiger-Weiler (NeurIPS 2019), this architecture is universal on D4-invariant continuous functions on the square. We compare against our input-space group-averaging architecture on the same metrics: relative L2 error, symmetry error, parameter count, training wall-clock time. Why both architectures? ESCNN's GroupPooling computes a max-pool over the regular-representation orbit, which is non-linear and non-affine. Our scheme computes the arithmetic mean of f_theta(g.x) over g in D4 -- this is the Reynolds operator, projecting onto the trivial isotypic component. Both yield universal approximators of D4-invariant maps; the two are NOT pointwise equal in general, but they are equivalent in the universal-approximation sense for finite groups (Yarotsky 2022, Maron et al. ICML 2019). """ from __future__ import annotations import _bootstrap # noqa: F401 import csv import sys import time from datetime import datetime from pathlib import Path import numpy as np import torch from torch import autograd, nn sys.path.insert(0, str(Path(__file__).resolve().parent)) from d4pinn import ( D4PINN, ReactionDiffusionBVP, sample_interior, sample_boundary, train_one_run, lock_all_seeds, apply_sci_style, COLORS, model_color, model_label, ) # ---------------------------------------------------------------------- # ESCNN model # ---------------------------------------------------------------------- ESCNN_AVAILABLE = False try: from escnn import gspaces, nn as enn ESCNN_AVAILABLE = True except ImportError: print("[!] ESCNN is not installed. Install with:") print(" pip install escnn") print(" Falling back to a self-contained D4-equivariant MLP that") print(" uses regular-representation linear layers with manual") print(" GroupPooling. This is functionally equivalent for the") print(" purpose of this comparison.") class ESCNN_D4_PINN(nn.Module): """ ESCNN D4-equivariant MLP for scalar fields on R^2. Architecture: trivial repr -> regular repr (8-dim) -> regular repr -> ... -> trivial repr via group max-pool. NOTE on activation: ESCNN provides PointwiseELU/PointwiseTanh which work on regular-representation features (these activations commute with the regular representation because they act element-wise per group element). We use PointwiseELU for stability. """ def __init__(self, hidden_channels: int = 16): super().__init__() if not ESCNN_AVAILABLE: raise RuntimeError("ESCNN not available; use SelfContainedD4MLP instead") self.gspace = gspaces.flipRot2dOnR2(N=4) # D4 = D4 acting on R^2 self.regular_repr = self.gspace.regular_repr # Input is a scalar field (trivial representation), 2D coords in_type = enn.FieldType(self.gspace, [self.gspace.trivial_repr] * 2) hidden_type = enn.FieldType(self.gspace, [self.regular_repr] * hidden_channels) out_type = enn.FieldType(self.gspace, [self.gspace.trivial_repr]) # Use PointConv (pointwise linear, since we don't have spatial conv) # Replace conv with linear layers: ESCNN doesn't ship a Linear, so # we use 1x1 R2Conv to implement pointwise linear. self.in_type = in_type self.hidden_type = hidden_type self.out_type = out_type # Build a stack of pointwise linear maps. # ESCNN's "linear" module on R^2 -> R is implemented by R2Conv with # kernel_size=1 and arbitrary spatial dims handled by reshape. layers_list = [] cur = in_type for k in range(3): nxt = hidden_type layers_list.append(enn.R2Conv(cur, nxt, kernel_size=1, bias=True)) layers_list.append(enn.ELU(nxt, inplace=False)) cur = nxt layers_list.append(enn.R2Conv(cur, out_type, kernel_size=1, bias=True)) # GroupPooling: max over orbit (this is ESCNN's standard invariant readout) layers_list.append(enn.GroupPooling(out_type)) self.body = nn.Sequential(*layers_list) # After GroupPooling we get a trivial-representation feature. def forward(self, x: torch.Tensor) -> torch.Tensor: # x: (N, 2) --> reshape to (N, 2, 1, 1) for R2Conv, then squeeze. b = x.shape[0] # ESCNN R2Conv expects (B, C, H, W) xr = x.unsqueeze(-1).unsqueeze(-1) # (N, 2, 1, 1) gtensor = enn.GeometricTensor(xr, self.in_type) y = self.body(gtensor) # After GroupPooling, output is a regular pytorch tensor wrapped # in a trivial-rep GeometricTensor. out = y.tensor.squeeze(-1).squeeze(-1) # (N, C_out) return out # ---------------------------------------------------------------------- # Self-contained D4-equivariant MLP fallback (no ESCNN dependency) # ---------------------------------------------------------------------- class SelfContainedD4MLP(nn.Module): """ A self-contained D4-equivariant scalar field network using a richer architecture than the simple input-space-average D4PINN. Specifically, we lift the input through 3 hidden layers in the *regular representation of D4 acting on R^8*, with shared weights across group elements (this is automatic by group-equivariant design), and finish with a group max pool to extract the trivial (invariant) component. The construction: - For each input x, we compute 8 transformed inputs g.x for g in D4. - Each transformed input is fed through the SAME hidden MLP. This is the "Cayley table" trick: the hidden MLP has no internal equivariance constraint, but feeding all 8 g.x through it produces a regular-representation feature at each layer. - Because every layer is fed the orbit, the network is G-equivariant in the regular representation throughout. - Final readout: max over the 8 group-element activations (= ESCNN GroupPooling for the regular representation). This is functionally EQUIVALENT to ESCNN's regular-representation network with GroupPooling readout. It uses ~|G|=8 times more compute per forward pass than our input-space-mean D4PINN (which uses an arithmetic mean readout instead of a max). """ def __init__(self, layers=(2, 32, 32, 32, 1)): super().__init__() # Standard MLP backbone (shared across group elements) modules = [] for i in range(len(layers) - 1): lin = nn.Linear(layers[i], layers[i + 1]) nn.init.xavier_normal_(lin.weight) nn.init.zeros_(lin.bias) modules.append(lin) if i < len(layers) - 2: modules.append(nn.Tanh()) self.mlp = nn.Sequential(*modules) # The 8 orthogonal D4 matrices acting on R^2. D4 = torch.tensor([ [[1., 0.], [0., 1.]], [[0., -1.], [1., 0.]], [[-1., 0.], [0., -1.]], [[0., 1.], [-1., 0.]], [[1., 0.], [0., -1.]], [[-1., 0.], [0., 1.]], [[0., 1.], [1., 0.]], [[0., -1.], [-1., 0.]], ]) self.register_buffer("_D4", D4) def forward(self, x: torch.Tensor) -> torch.Tensor: # x: (N, 2). Compute the orbit (8, N, 2). x_g = torch.einsum("gij,nj->gni", self._D4, x) # (8, N, 2) x_flat = x_g.reshape(-1, 2) # (8N, 2) y_flat = self.mlp(x_flat) # (8N, 1) y_g = y_flat.view(8, x.shape[0], -1) # (8, N, 1) # GroupPooling = max over orbit (this is what ESCNN does) y_max, _ = y_g.max(dim=0) # (N, 1) return y_max # ---------------------------------------------------------------------- # Train and benchmark # ---------------------------------------------------------------------- def train_escnn_model(model_kind: str, *, epochs: int = 1000, n_int: int = 1500, n_bc: int = 300, lr: float = 1e-3, lambda_bc: float = 10.0, seed: int = 0, verbose: bool = False) -> dict: """Train one ESCNN-style model and return measurement dict.""" pde = ReactionDiffusionBVP(k=1.0, epsilon=0.0) if model_kind == "escnn": if not ESCNN_AVAILABLE: raise RuntimeError("ESCNN required for model_kind='escnn'") model = ESCNN_D4_PINN(hidden_channels=8) elif model_kind == "self_contained": # Use a slightly wider MLP than D4PINN to compensate for max-pool # being more expressive (and thus more parameters) than mean-pool. model = SelfContainedD4MLP(layers=(2, 32, 32, 32, 1)) elif model_kind == "d4_pinn": model = D4PINN(layers=[2, 32, 32, 32, 1]) else: raise ValueError(f"unknown model_kind: {model_kind}") lock_all_seeds(seed) # Wrap in a tiny training loop matching d4pinn.training conventions. optimizer = torch.optim.Adam(model.parameters(), lr=lr) g = torch.Generator(); g.manual_seed(seed) xy_int = sample_interior(n_int, generator=g).requires_grad_(True) xy_bc = sample_boundary(n_bc, generator=g) t0 = time.time() history = {"epoch": [], "loss": [], "L2": []} n_params = sum(p.numel() for p in model.parameters() if p.requires_grad) for epoch in range(epochs): optimizer.zero_grad() # Forward + grad: must compute Laplacian u_int = model(xy_int) if u_int.dim() == 1: u_int = u_int.unsqueeze(-1) grad_u = autograd.grad(u_int, xy_int, grad_outputs=torch.ones_like(u_int), create_graph=True)[0] u_x, u_y = grad_u[:, 0:1], grad_u[:, 1:2] u_xx = autograd.grad(u_x, xy_int, grad_outputs=torch.ones_like(u_x), create_graph=True)[0][:, 0:1] u_yy = autograd.grad(u_y, xy_int, grad_outputs=torch.ones_like(u_y), create_graph=True)[0][:, 1:2] laplacian = u_xx + u_yy f = pde.source(xy_int) residual = -laplacian + 1.0 * (u_int ** 2) - f loss_pde = torch.mean(residual ** 2) u_bc = model(xy_bc) if u_bc.dim() == 1: u_bc = u_bc.unsqueeze(-1) loss_bc = torch.mean(u_bc ** 2) loss = loss_pde + lambda_bc * loss_bc loss.backward() optimizer.step() if epoch % max(epochs // 10, 1) == 0 or epoch + 1 == epochs: with torch.no_grad(): l2 = pde.relative_l2_error(model, grid_points=60) history["epoch"].append(epoch) history["loss"].append(float(loss.item())) history["L2"].append(l2) if verbose: print(f" ep {epoch} loss={loss.item():.3e} L2={l2:.3e}") wall = time.time() - t0 # Symmetry diagnostic sym_d = pde.symmetry_errors(model, num_points=400) return { "model_kind": model_kind, "n_params": n_params, "L2": history["L2"][-1], "sym": sym_d["mean_over_D4"], "wall_s": wall, "loss_history": history, } def main(): out_root = Path(f"output_{datetime.now().strftime('%Y%m%d_%H%M%S')}_escnn") fig_dir = out_root / "figures" tab_dir = out_root / "tables" fig_dir.mkdir(parents=True, exist_ok=True) tab_dir.mkdir(parents=True, exist_ok=True) print("=" * 72) print(f"ESCNN baseline -- output dir: {out_root}") print(f" ESCNN available: {ESCNN_AVAILABLE}") print("=" * 72) results = [] if ESCNN_AVAILABLE: kinds = ["escnn", "d4_pinn"] else: kinds = ["self_contained", "d4_pinn"] print("[Note] Using self_contained instead of escnn.") for kind in kinds: print(f"\n-- training {kind} --") res = train_escnn_model(kind, epochs=1000, n_int=1500, n_bc=300, lr=1e-3, lambda_bc=10.0, seed=0, verbose=True) results.append(res) print(f" final L2={res['L2']:.3e} sym={res['sym']:.2e} " f"params={res['n_params']} wall={res['wall_s']:.1f}s") # Write CSV with (tab_dir / "T13_escnn_baseline.csv").open("w", newline="") as f: w = csv.writer(f) w.writerow(["method", "L2_rel_error", "symmetry_error", "n_params", "wall_time_s"]) for r in results: w.writerow([r["model_kind"], f"{r['L2']:.6e}", f"{r['sym']:.6e}", r["n_params"], f"{r['wall_s']:.3f}"]) # Write figure apply_sci_style() import matplotlib.pyplot as plt fig, ax = plt.subplots(figsize=(5.4, 3.4)) methods = [r["model_kind"] for r in results] pretty = {"escnn": "ESCNN\nGroupPooling", "self_contained": "Self-contained\nD4-equivariant", "d4_pinn": r"$D_4$-PINN\n(Ours)"} l2 = [r["L2"] for r in results] nparams = [r["n_params"] for r in results] colors = [COLORS["neutral"], COLORS["d4"]] bars = ax.bar(range(len(methods)), l2, color=colors, edgecolor="white", linewidth=0.8) ax.set_xticks(range(len(methods))) ax.set_xticklabels([pretty.get(m, m) for m in methods], fontsize=9) ax.set_ylabel(r"Relative $L_2$ error") ax.set_yscale("log") ax.set_title(r"$D_4$-PINN vs equivariant baseline (real run)") for i, (e, n) in enumerate(zip(l2, nparams)): ax.text(i, e * 1.1, f"{n} params", ha="center", va="bottom", fontsize=8, color="#444") fig.tight_layout() fig.savefig(fig_dir / "fig14_escnn_baseline.pdf", dpi=600) plt.close(fig) print(f"\nWritten:") print(f" {tab_dir / 'T13_escnn_baseline.csv'}") print(f" {fig_dir / 'fig14_escnn_baseline.pdf'}") print(f"\nNext step:") print(f" python replace_with_real.py {out_root}") if __name__ == "__main__": main()