diff --git a/doc/source/modules/api.rst b/doc/source/modules/api.rst
index 2ab095714..de837d4ae 100644
--- a/doc/source/modules/api.rst
+++ b/doc/source/modules/api.rst
@@ -166,6 +166,22 @@ Codomain Attention Neural Operators (CODANO)
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+.. _readouts_api:
+
+Readouts
+--------
+
+Modules for mapping spatial field outputs ``(B, C, *spatial)`` to
+task-specific scalars or vectors (e.g. scalar quantities of interest).
+
+.. autosummary::
+ :toctree: generated
+ :template: class.rst
+
+ ResolutionInvariantReadout
+
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
.. raw:: html
diff --git a/examples/models/GALLERY_HEADER.rst b/examples/models/GALLERY_HEADER.rst
index 5fd51b91c..4264f1ef5 100644
--- a/examples/models/GALLERY_HEADER.rst
+++ b/examples/models/GALLERY_HEADER.rst
@@ -3,4 +3,7 @@
Models
------
-Examples showcasing our end-to-end NO architectures.
\ No newline at end of file
+Examples showcasing our end-to-end NO architectures.
+This gallery includes field-to-field PDE surrogates as well as field-to-scalar
+readout examples for quantities of interest and continuum-electrostatics
+energies.
diff --git a/examples/models/plot_FNO_darcy_scalar_qoi.py b/examples/models/plot_FNO_darcy_scalar_qoi.py
new file mode 100644
index 000000000..104e46ffd
--- /dev/null
+++ b/examples/models/plot_FNO_darcy_scalar_qoi.py
@@ -0,0 +1,196 @@
+"""
+FNO scalar QoI prediction with resolution-invariant readout
+===========================================================
+
+This example trains an FNO to predict a scalar quantity of interest (QoI)
+from Darcy-Flow fields. We define the target as:
+
+``y_scalar = mean(y_field)``
+
+and use a resolution-invariant readout to map ``(B, C, H, W) -> (B, 1)``.
+The model is trained at 16×16 and evaluated at both the training resolution
+and zero-shot at 32×32.
+
+By default the example uses a small in-memory synthetic dataset so it runs
+without downloading anything. To use the real Darcy-Flow dataset from Zenodo,
+set ``USE_DARCY_DATA = True`` below and note that output-field normalization is
+disabled (``encode_output=False``) because the target is a derived scalar QoI
+rather than the full output field.
+"""
+
+# %%
+# .. raw:: html
+#
+#
+#
+# Configuration
+# -------------
+# Set ``USE_DARCY_DATA = True`` to download and use the real Darcy-Flow
+# dataset from Zenodo instead of the synthetic stand-in.
+
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader, Dataset
+
+from neuralop import Trainer
+from neuralop.data.datasets import load_darcy_flow_small
+from neuralop.data.transforms.data_processors import DefaultDataProcessor
+from neuralop.models import FNO, ResolutionInvariantReadout
+from neuralop.training import AdamW
+
+USE_DARCY_DATA = False # set True to use the real Zenodo dataset
+N_EPOCHS = 3
+BATCH_SIZE = 32
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+
+# %%
+# .. raw:: html
+#
+#
+#
+# Dataset
+# -------
+# We use a tiny synthetic Darcy-like dataset when ``USE_DARCY_DATA`` is False,
+# or load the real small Darcy-Flow dataset otherwise.
+
+
+class ScalarQOILoss(nn.Module):
+ """Sum of squared errors against scalar QoI computed from y, ignoring extra kwargs."""
+
+ def forward(self, out, y, **kwargs):
+ y_scalar = y.reshape(y.shape[0], -1).mean(dim=1, keepdim=True)
+ return torch.sum((out - y_scalar) ** 2)
+
+
+class SyntheticDarcyLikeDataset(Dataset):
+ """Small deterministic dataset for smoke testing the QoI example."""
+
+ def __init__(self, n_samples: int, resolution: int):
+ self.n_samples = n_samples
+ coords = torch.linspace(0.0, 1.0, resolution)
+ yy, xx = torch.meshgrid(coords, coords, indexing="ij")
+ self.base_x = (xx + yy).unsqueeze(0)
+ self.base_y = (torch.sin(torch.pi * xx) * torch.cos(torch.pi * yy)).unsqueeze(0)
+
+ def __len__(self):
+ return self.n_samples
+
+ def __getitem__(self, idx):
+ scale = (idx + 1) / self.n_samples
+ x = self.base_x * scale
+ y = self.base_y * (1.0 + 0.5 * scale)
+ return {"x": x, "y": y}
+
+
+if USE_DARCY_DATA:
+ # Disable output normalization: target is a derived scalar, not the field.
+ train_loader, test_loaders, data_processor = load_darcy_flow_small(
+ n_train=256,
+ batch_size=BATCH_SIZE,
+ n_tests=[64, 64],
+ test_resolutions=[16, 32],
+ test_batch_sizes=[BATCH_SIZE, BATCH_SIZE],
+ encode_output=False,
+ )
+else:
+ train_loader = DataLoader(
+ SyntheticDarcyLikeDataset(n_samples=32, resolution=16),
+ batch_size=BATCH_SIZE,
+ shuffle=True,
+ )
+ test_loaders = {
+ 16: DataLoader(
+ SyntheticDarcyLikeDataset(n_samples=8, resolution=16), batch_size=BATCH_SIZE
+ ),
+ 32: DataLoader(
+ SyntheticDarcyLikeDataset(n_samples=8, resolution=32), batch_size=BATCH_SIZE
+ ),
+ }
+ data_processor = DefaultDataProcessor()
+
+data_processor = data_processor.to(device)
+
+
+# %%
+# .. raw:: html
+#
+#
+#
+# Model
+# -----
+# The FNO maps input fields to a latent field, then the
+# :class:`~neuralop.models.ResolutionInvariantReadout` pools the field to a
+# single scalar per sample.
+
+model = FNO(
+ n_modes=(8, 8),
+ in_channels=1,
+ out_channels=16,
+ hidden_channels=24,
+ projection_channel_ratio=2,
+ readout=ResolutionInvariantReadout(
+ in_channels=16,
+ out_dim=1,
+ reduce="mean",
+ head="mlp",
+ mlp_hidden_dim=32,
+ ),
+).to(device)
+
+
+# %%
+# .. raw:: html
+#
+#
+#
+# Training
+# --------
+
+optimizer = AdamW(model.parameters(), lr=5e-3, weight_decay=1e-4)
+scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=N_EPOCHS)
+qoi_loss = ScalarQOILoss()
+
+trainer = Trainer(
+ model=model,
+ n_epochs=N_EPOCHS,
+ device=device,
+ data_processor=data_processor,
+ wandb_log=False,
+ eval_interval=1,
+ use_distributed=False,
+ verbose=True,
+)
+
+trainer.train(
+ train_loader=train_loader,
+ test_loaders=test_loaders,
+ optimizer=optimizer,
+ scheduler=scheduler,
+ regularizer=False,
+ training_loss=qoi_loss,
+ eval_losses={"qoi_mse": qoi_loss},
+)
+
+
+# %%
+# .. raw:: html
+#
+#
+#
+# Zero-shot evaluation at multiple resolutions
+# --------------------------------------------
+# The readout pools the field regardless of spatial size, so the same trained
+# model can be evaluated at resolutions it was never trained on.
+
+model.eval()
+with torch.no_grad():
+ for resolution in sorted(test_loaders):
+ sample = next(iter(test_loaders[resolution]))
+ sample = data_processor.preprocess(sample)
+ out = model(sample["x"])
+ y_scalar = (
+ sample["y"].reshape(sample["y"].shape[0], -1).mean(dim=1, keepdim=True)
+ )
+ mse = torch.mean((out - y_scalar) ** 2).item()
+ print(f"Resolution {resolution}x{resolution} scalar QoI MSE: {mse:.6f}")
diff --git a/examples/models/plot_FNO_lpb_screened_energy.py b/examples/models/plot_FNO_lpb_screened_energy.py
new file mode 100644
index 000000000..29a740044
--- /dev/null
+++ b/examples/models/plot_FNO_lpb_screened_energy.py
@@ -0,0 +1,449 @@
+"""
+FNO for Linearized Poisson-Boltzmann Screened Electrostatic Energy
+==================================================================
+
+We train a 3-D Fourier Neural Operator (FNO) to predict a scalar screened
+electrostatic energy from volumetric fields relevant to continuum
+electrostatics: the charge density :math:`\\rho(\\mathbf{r})` and the inverse
+Debye length :math:`\\kappa`.
+
+This is a natural application for neural operators:
+
+1. The input is already a continuous 3-D field.
+2. The target comes from a PDE solve, not from an arbitrary tabular label.
+3. Resolution invariance matters because the same physical system may be
+ discretized on different voxel grids in real electrostatics workflows.
+
+To keep the example self-contained, ground-truth labels are generated by
+solving a synthetic *linearized Poisson-Boltzmann* problem in Fourier space on
+a periodic box with homogeneous dielectric:
+
+.. math::
+
+ \\left(-\\nabla^2 + \\kappa^2\\right)\\varphi(\\mathbf{r})
+ = \\rho(\\mathbf{r}) / \\varepsilon
+
+and then computing the screened electrostatic energy
+
+.. math::
+
+ E_\\mathrm{el} = \\tfrac{1}{2}\\int \\rho(\\mathbf{r})\\,\\varphi(\\mathbf{r})\\,d^3r.
+
+We train on a coarse grid (16³) and evaluate on the *same latent test systems*
+discretized at both 16³ and 24³, which gives a direct resolution-invariance
+check rather than comparing unrelated random samples.
+"""
+
+# %%
+# .. raw:: html
+#
+#
+#
+# Import dependencies
+# -------------------
+# We use PyTorch for synthetic-data generation and Fourier-space solves, and
+# ``neuralop`` for the FNO model, trainer, and resolution-invariant readout.
+
+import math
+import sys
+
+import matplotlib.pyplot as plt
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader, Dataset
+
+from neuralop import Trainer
+from neuralop.data.transforms.data_processors import DefaultDataProcessor
+from neuralop.models import FNO, ResolutionInvariantReadout
+from neuralop.training import AdamW
+from neuralop.utils import count_model_params
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+
+# %%
+# .. raw:: html
+#
+#
+#
+# Synthetic linearized-PB dataset
+# -------------------------------
+# Continuum electrostatics is widely used in protein chemistry and drug design
+# to estimate electrostatic contributions to solvation, pKa shifts, and binding.
+# The *linearized* Poisson-Boltzmann equation is a standard screened-electrostatics
+# model. In this example we use a deliberately simplified setting: homogeneous
+# dielectric, periodic boundary conditions, and a spatially constant screening
+# parameter for each sample. In that setting the Fourier-space solution is analytic:
+#
+# .. math::
+#
+# \hat{\varphi}(\mathbf{k}) =
+# \frac{\hat{\rho}(\mathbf{k})}
+# {\varepsilon (|\mathbf{k}|^2 + \kappa^2)}.
+#
+# We use this exact spectral solve to generate labels for random mixtures of
+# Gaussian partial charges. The charge density and screening field are the
+# model inputs, while the scalar screened electrostatic energy is the target.
+
+EPSILON = 80.0 # solvent dielectric
+BOX_SIZE = 8.0 # Angstrom
+RHO_SIGMA = 0.35 # Angstrom
+N_TRAIN = 16
+N_TEST_HI = 24
+N_TRAIN_SAMPLES = 192
+N_TEST_SAMPLES = 48
+N_CHARGES_RANGE = (3, 6)
+
+
+def _grid_coordinates(n: int, box_size: float) -> torch.Tensor:
+ """Return a centered (n, n, n, 3) coordinate grid on a periodic box."""
+ h = box_size / n
+ coords_1d = (torch.arange(n, dtype=torch.float32) - n / 2 + 0.5) * h
+ gz, gy, gx = torch.meshgrid(coords_1d, coords_1d, coords_1d, indexing="ij")
+ return torch.stack([gx, gy, gz], dim=-1)
+
+
+def _make_charge_density(
+ grid: torch.Tensor,
+ positions: torch.Tensor,
+ charges: torch.Tensor,
+ sigma: float,
+) -> torch.Tensor:
+ """Build a smooth charge density from Gaussian partial charges."""
+ delta = grid.unsqueeze(-2) - positions.view(1, 1, 1, -1, 3)
+ r2 = (delta**2).sum(dim=-1)
+ norm = 1.0 / ((2.0 * math.pi * sigma**2) ** 1.5)
+ gaussians = torch.exp(-r2 / (2.0 * sigma**2)) * norm
+ return (gaussians * charges.view(1, 1, 1, -1)).sum(dim=-1)
+
+
+def _spectral_k2(n: int, box_size: float) -> torch.Tensor:
+ """Return the squared wave-number tensor for a periodic cube."""
+ h = box_size / n
+ freqs = 2.0 * math.pi * torch.fft.fftfreq(n, d=h)
+ kz, ky, kx = torch.meshgrid(freqs, freqs, freqs, indexing="ij")
+ return kx**2 + ky**2 + kz**2
+
+
+def _solve_linearized_pb(
+ rho: torch.Tensor, kappa: float, epsilon: float, k2: torch.Tensor
+) -> torch.Tensor:
+ """Solve (-Δ + kappa^2) phi = rho / epsilon on a periodic cube."""
+ rho_hat = torch.fft.fftn(rho)
+ phi_hat = rho_hat / (epsilon * (k2 + kappa**2))
+ return torch.fft.ifftn(phi_hat).real
+
+
+def _electrostatic_energy(
+ rho: torch.Tensor, phi: torch.Tensor, box_size: float
+) -> float:
+ """Return 0.5 * integral rho(r) * phi(r) d^3r."""
+ return float(0.5 * (rho * phi).mean() * (box_size**3))
+
+
+def _sample_latent_systems(n_samples: int, seed: int) -> list:
+ """Sample latent charge configurations shared across resolutions."""
+ gen = torch.Generator()
+ gen.manual_seed(seed)
+
+ systems = []
+ for _ in range(n_samples):
+ n_charges = int(
+ torch.randint(
+ low=N_CHARGES_RANGE[0],
+ high=N_CHARGES_RANGE[1] + 1,
+ size=(1,),
+ generator=gen,
+ ).item()
+ )
+
+ # Keep charges away from the periodic boundary so the localized density
+ # looks molecule-like instead of wrapping around the box.
+ positions = torch.empty(n_charges, 3).uniform_(-2.0, 2.0, generator=gen)
+ charges = torch.empty(n_charges).uniform_(-1.0, 1.0, generator=gen)
+ charges -= charges.mean()
+ # Use a deliberately broad synthetic screening range to make the
+ # operator-learning task diverse; realistic biomolecular screening is
+ # often weaker than this.
+ kappa = float(torch.empty(1).uniform_(0.4, 1.2, generator=gen).item())
+ systems.append({"positions": positions, "charges": charges, "kappa": kappa})
+ return systems
+
+
+class LinearPBDataset(Dataset):
+ """Exact spectral linearized-PB labels on a fixed-resolution voxel grid."""
+
+ def __init__(self, systems: list, n_pts: int, box_size: float):
+ self.systems = systems
+ self.n_pts = n_pts
+ self.box_size = box_size
+ self.grid = _grid_coordinates(n_pts, box_size)
+ self.k2 = _spectral_k2(n_pts, box_size)
+
+ xs, ys, phis = [], [], []
+ for system in systems:
+ rho = _make_charge_density(
+ self.grid, system["positions"], system["charges"], sigma=RHO_SIGMA
+ )
+ phi = _solve_linearized_pb(
+ rho=rho,
+ kappa=system["kappa"],
+ epsilon=EPSILON,
+ k2=self.k2,
+ )
+ kappa_field = torch.full_like(rho, system["kappa"])
+ xs.append(torch.stack([rho, kappa_field], dim=0))
+ ys.append(_electrostatic_energy(rho, phi, box_size))
+ phis.append(phi)
+
+ self.x = torch.stack(xs)
+ self.y = torch.tensor(ys, dtype=torch.float32).unsqueeze(1)
+ self.phi = torch.stack(phis)
+
+ def __len__(self):
+ return len(self.x)
+
+ def __getitem__(self, idx):
+ return {"x": self.x[idx], "y": self.y[idx]}
+
+
+train_systems = _sample_latent_systems(N_TRAIN_SAMPLES, seed=0)
+test_systems = _sample_latent_systems(N_TEST_SAMPLES, seed=1)
+
+train_dataset = LinearPBDataset(train_systems, n_pts=N_TRAIN, box_size=BOX_SIZE)
+test_dataset_16 = LinearPBDataset(test_systems, n_pts=N_TRAIN, box_size=BOX_SIZE)
+test_dataset_24 = LinearPBDataset(test_systems, n_pts=N_TEST_HI, box_size=BOX_SIZE)
+
+train_loader = DataLoader(train_dataset, batch_size=16, shuffle=True)
+test_loaders = {
+ N_TRAIN: DataLoader(test_dataset_16, batch_size=16),
+ N_TEST_HI: DataLoader(test_dataset_24, batch_size=16),
+}
+data_processor = DefaultDataProcessor().to(device)
+
+
+# %%
+# .. raw:: html
+#
+#
+#
+# Visualizing the field solve
+# --------------------------
+# For one test system, we show the charge density and the corresponding exact
+# linearized-PB potential obtained by the Fourier-space solver. This makes the
+# operator-learning aspect explicit: the scalar energy label is derived from a
+# nontrivial field solve, not a hand-coded molecular descriptor.
+
+fig, axes = plt.subplots(2, 2, figsize=(7, 6))
+for row, dataset in enumerate([test_dataset_16, test_dataset_24]):
+ n = dataset.n_pts
+ sl = n // 2
+ rho_sl = dataset.x[0, 0, sl].numpy()
+ phi_sl = dataset.phi[0, sl].numpy()
+ im0 = axes[row, 0].imshow(rho_sl, cmap="RdBu_r")
+ axes[row, 0].set_title(f"{n}³ grid: rho slice")
+ plt.colorbar(im0, ax=axes[row, 0])
+ im1 = axes[row, 1].imshow(phi_sl, cmap="viridis")
+ axes[row, 1].set_title(f"{n}³ grid: exact phi slice")
+ plt.colorbar(im1, ax=axes[row, 1])
+ for ax in axes[row]:
+ ax.set_xticks([])
+ ax.set_yticks([])
+fig.suptitle("Same latent test system at two voxel resolutions", y=0.98)
+plt.tight_layout()
+fig.show()
+
+
+# %%
+# .. raw:: html
+#
+#
+#
+# Building the FNO model
+# ----------------------
+# The model consumes two 3-D input channels:
+#
+# * :math:`\\rho(\\mathbf{r})` — charge density
+# * :math:`\\kappa(\\mathbf{r})` — inverse Debye length
+# (constant within each synthetic sample)
+#
+# It returns a scalar screened electrostatic energy via
+# :class:`~neuralop.models.ResolutionInvariantReadout`.
+# The readout computes a grid-independent physical integral by applying
+# ``mean(field) × L³`` before the final head.
+
+FNO_OUT_CHANNELS = 8
+
+model = FNO(
+ n_modes=(4, 4, 4),
+ in_channels=2,
+ out_channels=FNO_OUT_CHANNELS,
+ hidden_channels=20,
+ n_layers=4,
+ projection_channel_ratio=2,
+ readout=ResolutionInvariantReadout(
+ in_channels=FNO_OUT_CHANNELS,
+ out_dim=1,
+ reduce="integral",
+ measure_per_dim=BOX_SIZE,
+ head="mlp",
+ mlp_hidden_dim=32,
+ ),
+).to(device)
+
+n_params = count_model_params(model)
+print(f"\nModel has {n_params} parameters.")
+sys.stdout.flush()
+
+
+# %%
+# .. raw:: html
+#
+#
+#
+# Loss function and training
+# --------------------------
+# The target is a scalar screened electrostatic energy, so we optimize MSE.
+
+
+class EnergyLoss(nn.Module):
+ """Sum of squared errors between predicted and exact linearized-PB screened energies."""
+
+ def forward(self, out, y, **kwargs):
+ return torch.sum((out - y) ** 2)
+
+
+energy_loss = EnergyLoss()
+
+N_EPOCHS = 12
+optimizer = AdamW(model.parameters(), lr=5e-3, weight_decay=1e-4)
+scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=N_EPOCHS)
+
+trainer = Trainer(
+ model=model,
+ n_epochs=N_EPOCHS,
+ device=device,
+ data_processor=data_processor,
+ wandb_log=False,
+ eval_interval=4,
+ use_distributed=False,
+ verbose=True,
+)
+
+trainer.train(
+ train_loader=train_loader,
+ test_loaders=test_loaders,
+ optimizer=optimizer,
+ scheduler=scheduler,
+ regularizer=False,
+ training_loss=energy_loss,
+ eval_losses={"energy_mse": energy_loss},
+)
+
+
+# %%
+# .. raw:: html
+#
+#
+#
+# Predicted vs true energy at both resolutions
+# -------------------------------------------
+# We evaluate the same underlying test systems on both the training grid (16³)
+# and an unseen finer grid (24³). This isolates resolution effects from sample
+# variation.
+
+model.eval()
+results = {}
+for res, loader in sorted(test_loaders.items()):
+ preds, trues = [], []
+ with torch.no_grad():
+ for batch in loader:
+ batch = data_processor.preprocess(batch)
+ out = model(batch["x"])
+ preds.append(out.cpu())
+ trues.append(batch["y"].cpu())
+ preds = torch.cat(preds).squeeze(1).numpy()
+ trues = torch.cat(trues).squeeze(1).numpy()
+ results[res] = (trues, preds)
+ h_val = BOX_SIZE / res
+ mse = float(((preds - trues) ** 2).mean())
+ print(f"Resolution {res}³ (h = {h_val:.3f} Å) MSE = {mse:.4f}")
+
+
+# %%
+# .. raw:: html
+#
+#
+#
+# Resolution-paired comparison
+# ----------------------------
+# Because both test sets come from the same latent systems, we can compare
+# coarse-grid and fine-grid predictions *sample by sample*.
+
+preds_16 = results[N_TRAIN][1]
+preds_24 = results[N_TEST_HI][1]
+
+fig, axes = plt.subplots(1, 2, figsize=(9, 4))
+
+for ax, (res, (true_vals, pred_vals)) in zip(axes, sorted(results.items())):
+ mse = float(((pred_vals - true_vals) ** 2).mean())
+ lims = [
+ min(true_vals.min(), pred_vals.min()) - 0.02,
+ max(true_vals.max(), pred_vals.max()) + 0.02,
+ ]
+ ax.scatter(true_vals, pred_vals, s=12, alpha=0.6)
+ ax.plot(lims, lims, "r--", lw=1, label="perfect")
+ ax.set_xlabel("Exact screened energy")
+ ax.set_ylabel("FNO prediction")
+ ax.set_title(f"{res}³ grid\nMSE = {mse:.4f}")
+ ax.legend(fontsize=8)
+
+plt.tight_layout()
+fig.show()
+
+
+# %%
+# .. raw:: html
+#
+#
+#
+# Direct cross-resolution consistency
+# -----------------------------------
+# Finally, we compare the model's prediction for each test system at 16³ and
+# 24³. Small scatter around the diagonal indicates that one trained model can
+# serve multiple electrostatics grid resolutions without re-fitting.
+
+fig = plt.figure(figsize=(4.5, 4))
+plt.scatter(preds_16, preds_24, s=14, alpha=0.7)
+lims = [
+ min(preds_16.min(), preds_24.min()) - 0.02,
+ max(preds_16.max(), preds_24.max()) + 0.02,
+]
+plt.plot(lims, lims, "r--", lw=1, label="perfect agreement")
+plt.xlabel("Prediction on 16³ grid")
+plt.ylabel("Prediction on 24³ grid")
+plt.title("Same latent systems across resolutions")
+plt.legend(fontsize=8)
+plt.tight_layout()
+fig.show()
+
+
+# %%
+# .. raw:: html
+#
+#
+#
+# Why this is a natural FNO example
+# ---------------------------------
+# This example highlights two reasons neural operators are attractive in
+# biomolecular electrostatics:
+#
+# 1. The labels come from a PDE solve rather than an arbitrary descriptor model.
+# 2. Resolution invariance is genuinely useful because the same physical charge
+# configuration can be evaluated on multiple voxel grids in continuum-electrostatics
+# workflows.
+#
+# In a production setting, one would replace this synthetic homogeneous,
+# periodic linearized-PB generator with outputs from a finite-difference PB
+# solver or APBS pipeline for proteins, ligands, or protein-ligand complexes,
+# while keeping the same ``FNO + ResolutionInvariantReadout`` structure.
diff --git a/neuralop/__init__.py b/neuralop/__init__.py
index 131b9b028..77b118767 100644
--- a/neuralop/__init__.py
+++ b/neuralop/__init__.py
@@ -1,6 +1,6 @@
__version__ = "2.0.0"
-from .models import TFNO, FNO, RNO, get_model
+from .models import TFNO, FNO, RNO, get_model, ResolutionInvariantReadout
from .data import datasets, transforms
from . import mpu
from .training import Trainer
diff --git a/neuralop/models/__init__.py b/neuralop/models/__init__.py
index f2ca3f9db..39ffec36d 100644
--- a/neuralop/models/__init__.py
+++ b/neuralop/models/__init__.py
@@ -14,3 +14,4 @@
from .rno import RNO
from .otno import OTNO
from .base_model import get_model
+from .readouts import ResolutionInvariantReadout
diff --git a/neuralop/models/fno.py b/neuralop/models/fno.py
index 45cac5446..09adf52b1 100644
--- a/neuralop/models/fno.py
+++ b/neuralop/models/fno.py
@@ -1,5 +1,5 @@
from functools import partialmethod
-from typing import Tuple, List, Union, Literal
+from typing import Optional, Tuple, List, Union, Literal
Number = Union[float, int]
@@ -43,6 +43,9 @@ class FNO(BaseModel, name="FNO"):
Number of channels in input function. Determined by the problem.
out_channels : int
Number of channels in output function. Determined by the problem.
+ When ``readout`` is set, this becomes the intermediate channel width
+ feeding the readout head; the model's actual output shape is then
+ ``(batch, readout.out_dim)`` rather than ``(batch, out_channels, *spatial)``.
hidden_channels : int
Width of the FNO (i.e. number of channels).
This significantly affects the number of parameters of the FNO.
@@ -142,6 +145,10 @@ class FNO(BaseModel, name="FNO"):
frequency and Nyquist frequency are real-valued before calling irfft. When False,
relies on cuFFT's irfftn to handle symmetry automatically, which may fail on
certain GPUs or input sizes, causing line artifacts. By default True.
+ readout : nn.Module, optional
+ Optional module applied after the projection layer to map field outputs
+ of shape ``(B, C, *spatial)`` to task-specific outputs (e.g., scalar or
+ vector quantities of interest). If ``None``, returns the projected field.
Examples
--------
@@ -183,7 +190,9 @@ def __init__(
use_channel_mlp: bool = True,
channel_mlp_dropout: float = 0,
channel_mlp_expansion: float = 0.5,
- channel_mlp_skip: Literal["linear", "identity", "soft-gating", None] = "soft-gating",
+ channel_mlp_skip: Literal[
+ "linear", "identity", "soft-gating", None
+ ] = "soft-gating",
fno_skip: Literal["linear", "identity", "soft-gating", None] = "linear",
resolution_scaling_factor: Union[Number, List[Number]] = None,
domain_padding: Union[Number, List[Number]] = None,
@@ -199,9 +208,27 @@ def __init__(
preactivation: bool = False,
conv_module: nn.Module = SpectralConv,
enforce_hermitian_symmetry: bool = True,
+ readout: Optional[nn.Module] = None,
):
if decomposition_kwargs is None:
decomposition_kwargs = {}
+ if complex_data and readout is not None:
+ raise ValueError(
+ "readout is not supported with complex_data=True. "
+ "The projection layer outputs a complex tensor, so the readout "
+ "would need to collapse it to real values — but the right approach "
+ "depends on your use case. Wrap your input or the projection output "
+ "explicitly before passing to a readout, for example: "
+ "x.real (real part only), x.abs() (modulus field then pool), or "
+ "x.mean(spatial_dims).abs() (pool then take modulus)."
+ )
+ if readout is not None and hasattr(readout, "in_channels"):
+ if readout.in_channels != out_channels:
+ raise ValueError(
+ f"readout.in_channels ({readout.in_channels}) must match "
+ f"out_channels ({out_channels}). The readout is applied after "
+ "the projection layer, which outputs out_channels channels."
+ )
super().__init__()
self.n_dim = len(n_modes)
@@ -233,6 +260,7 @@ def __init__(
self.preactivation = preactivation
self.complex_data = complex_data
self.fno_block_precision = fno_block_precision
+ self.readout = readout
## Positional embedding
if positional_embedding == "grid":
@@ -349,6 +377,14 @@ def forward(self, x, output_shape=None, **kwargs):
6. Projection of intermediate function representation to the output channels
+ 7. If ``readout`` is set, applies the readout module to map the field
+ ``(B, out_channels, *spatial)`` to task-specific output (e.g. a scalar
+ or vector quantity of interest). The output shape and meaning change
+ accordingly. Note that ``output_shape`` (step 4) changes the spatial
+ resolution of the field that enters the readout without changing the
+ physical domain, so ``measure_per_dim`` should always reflect the
+ physical extent of the domain.
+
Parameters
----------
x : tensor
@@ -393,6 +429,10 @@ def forward(self, x, output_shape=None, **kwargs):
x = self.projection(x)
+ # The readout consumes the projected field, so its input width matches out_channels.
+ if self.readout is not None:
+ x = self.readout(x)
+
return x
@property
diff --git a/neuralop/models/otno.py b/neuralop/models/otno.py
index 24d6e5c97..b012aa112 100644
--- a/neuralop/models/otno.py
+++ b/neuralop/models/otno.py
@@ -45,6 +45,13 @@ def __init__(
domain_padding=None,
**kwargs,
):
+ if kwargs.get("readout") is not None:
+ raise ValueError(
+ "readout is not supported by OTNO. OTNO.forward indexes into a "
+ "target mesh and returns output of shape (out_channels, n_target), "
+ "which is incompatible with ResolutionInvariantReadout "
+ "and similar heads that expect (B, C, *spatial) input."
+ )
super().__init__(
n_modes=n_modes,
hidden_channels=hidden_channels,
diff --git a/neuralop/models/readouts.py b/neuralop/models/readouts.py
new file mode 100644
index 000000000..a31e32703
--- /dev/null
+++ b/neuralop/models/readouts.py
@@ -0,0 +1,147 @@
+from typing import Literal, Optional, Sequence, Union
+
+import torch
+import torch.nn as nn
+
+
+class ResolutionInvariantReadout(nn.Module):
+ """Resolution-invariant readout for neural operator field outputs.
+
+ Maps tensors of shape ``(B, C, *spatial)`` to ``(B, out_dim)`` by
+ spatially pooling the field and applying a learned head. The pooling
+ mode is either a plain mean (``reduce="mean"``) or a physical integral
+ approximation (``reduce="integral"``) that scales the mean by the domain
+ volume so that the result is independent of grid resolution.
+
+ Parameters
+ ----------
+ in_channels : int
+ Number of channels in the input field ``(B, C, *spatial)``.
+ out_dim : int
+ Number of output dimensions.
+ reduce : {"mean", "integral"}, optional
+ Spatial reduction mode. ``"mean"`` computes the spatial mean.
+ ``"integral"`` multiplies the mean by the domain volume to
+ approximate the physical integral:
+ ``mean(field) * prod(measure_per_dim) ≈ ∫ f dⁿx``.
+ Default is ``"mean"``.
+ measure_per_dim : float or sequence of float, optional
+ Physical domain measure per spatial dimension, used only when
+ ``reduce="integral"``. If a scalar, the domain volume is
+ ``measure_per_dim ** n_spatial_dims``. If a sequence, its length
+ must match the number of spatial dimensions and the domain volume is
+ the product of the entries. Defaults to ``1.0``.
+
+ For non-cubic domains pass a sequence with one entry per dimension,
+ e.g. ``measure_per_dim=[Lx, Ly, Lz]`` for a box of side-lengths
+ *Lx*, *Ly*, *Lz* in Å (or any consistent length unit).
+ head : {"linear", "mlp"}, optional
+ Projection head applied after pooling. ``"linear"`` uses a single
+ ``nn.Linear``; ``"mlp"`` uses a two-layer MLP with an activation.
+ Default is ``"linear"``.
+ mlp_hidden_dim : int, optional
+ Hidden width of the MLP head. Defaults to ``in_channels`` when
+ ``head="mlp"``.
+ activation : nn.Module, optional
+ Activation function for the MLP head. Defaults to ``nn.GELU()``.
+ """
+
+ def __init__(
+ self,
+ in_channels: int,
+ out_dim: int,
+ reduce: Literal["mean", "integral"] = "mean",
+ measure_per_dim: Optional[Union[float, Sequence[float]]] = None,
+ head: Literal["linear", "mlp"] = "linear",
+ mlp_hidden_dim: Optional[int] = None,
+ activation: Optional[nn.Module] = None,
+ ):
+ super().__init__()
+
+ if reduce not in {"mean", "integral"}:
+ raise ValueError(f"reduce must be 'mean' or 'integral', got {reduce}.")
+ if head not in {"linear", "mlp"}:
+ raise ValueError(f"head must be 'linear' or 'mlp', got {head}.")
+
+ self.in_channels = in_channels
+ self.out_dim = out_dim
+ self.reduce = reduce
+
+ if measure_per_dim is None:
+ measure_per_dim = 1.0
+
+ # Store the raw measure values as a buffer so they move with the module
+ # and are saved in checkpoints. _scalar_mode selects between
+ # value**n_dims (scalar) and per-dimension products (sequence).
+ # Copy caller-provided sequences so later external mutation cannot
+ # invalidate the cached values or dimension checks.
+ if isinstance(measure_per_dim, (float, int)):
+ self.measure_per_dim = float(measure_per_dim)
+ self._scalar_mode = True
+ self.register_buffer(
+ "_measure_buffer",
+ torch.tensor(self.measure_per_dim, dtype=torch.float32),
+ )
+ else:
+ values = list(measure_per_dim)
+ self.measure_per_dim = values
+ self._scalar_mode = False
+ self.register_buffer(
+ "_measure_buffer",
+ torch.tensor(values, dtype=torch.float32),
+ )
+
+ if head == "linear":
+ self.head = nn.Linear(in_channels, out_dim)
+ else:
+ if mlp_hidden_dim is None:
+ mlp_hidden_dim = in_channels
+ if activation is None:
+ activation = nn.GELU()
+ self.head = nn.Sequential(
+ nn.Linear(in_channels, mlp_hidden_dim),
+ activation,
+ nn.Linear(mlp_hidden_dim, out_dim),
+ )
+
+ def extra_repr(self) -> str:
+ return (
+ f"in_channels={self.in_channels}, out_dim={self.out_dim}, "
+ f"reduce={self.reduce!r}, measure_per_dim={self.measure_per_dim}"
+ )
+
+ def _measure_product(
+ self, n_spatial_dims: int, device: torch.device, dtype: torch.dtype
+ ) -> torch.Tensor:
+ if self._scalar_mode:
+ value = self._measure_buffer.to(device=device, dtype=dtype)
+ return value.pow(n_spatial_dims)
+ else:
+ # Sequence case: validate length, then multiply the per-dimension measures.
+ if len(self.measure_per_dim) != n_spatial_dims:
+ raise ValueError(
+ f"measure_per_dim has length {len(self.measure_per_dim)},"
+ f" expected {n_spatial_dims}."
+ )
+ return self._measure_buffer.to(device=device, dtype=dtype).prod()
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ if x.ndim < 3:
+ raise ValueError(
+ f"Expected input with shape (B, C, *spatial), got {tuple(x.shape)}"
+ )
+ if x.is_complex():
+ raise ValueError(
+ "ResolutionInvariantReadout does not support complex inputs. "
+ "Convert to real first, e.g. x.real, x.abs(), or x.float()."
+ )
+
+ spatial_dims = tuple(range(2, x.ndim))
+ reduced = x.mean(dim=spatial_dims)
+
+ if self.reduce == "integral":
+ reduced = reduced * self._measure_product(
+ n_spatial_dims=len(spatial_dims), device=x.device, dtype=x.dtype
+ )
+
+ return self.head(reduced)
diff --git a/neuralop/models/tests/test_fno_readout.py b/neuralop/models/tests/test_fno_readout.py
new file mode 100644
index 000000000..9ae4cc129
--- /dev/null
+++ b/neuralop/models/tests/test_fno_readout.py
@@ -0,0 +1,277 @@
+import pytest
+import torch
+
+from neuralop.models import FNO, ResolutionInvariantReadout
+
+
+@pytest.mark.parametrize("n_dim", [1, 2, 3])
+@pytest.mark.parametrize("reduce", ["mean", "integral"])
+def test_fno_resolution_invariant_readout_shapes(reduce, n_dim):
+ device = "cuda" if torch.cuda.is_available() else "cpu"
+
+ out_dim = 7
+ n_modes = (8,) * n_dim
+ shape_1 = (12,) * n_dim
+ shape_2 = (16,) * n_dim
+ model = FNO(
+ n_modes=n_modes,
+ in_channels=2,
+ out_channels=4,
+ hidden_channels=12,
+ n_layers=2,
+ readout=ResolutionInvariantReadout(
+ in_channels=4,
+ out_dim=out_dim,
+ reduce=reduce,
+ measure_per_dim=[1.0] * n_dim,
+ ),
+ ).to(device)
+
+ x_1 = torch.randn(3, 2, *shape_1, device=device)
+ x_2 = torch.randn(3, 2, *shape_2, device=device)
+
+ y_1 = model(x_1)
+ y_2 = model(x_2)
+
+ assert y_1.shape == (3, out_dim)
+ assert y_2.shape == (3, out_dim)
+
+
+@pytest.mark.parametrize("n_dim", [1, 2, 3])
+def test_fno_resolution_invariant_readout_backward(n_dim):
+ device = "cuda" if torch.cuda.is_available() else "cpu"
+
+ n_modes = (8,) * n_dim
+ shape = (12,) * n_dim
+ model = FNO(
+ n_modes=n_modes,
+ in_channels=2,
+ out_channels=4,
+ hidden_channels=12,
+ n_layers=2,
+ readout=ResolutionInvariantReadout(
+ in_channels=4,
+ out_dim=5,
+ reduce="mean",
+ head="mlp",
+ mlp_hidden_dim=8,
+ ),
+ ).to(device)
+
+ x = torch.randn(4, 2, *shape, device=device)
+ target = torch.randn(4, 5, device=device)
+
+ out = model(x)
+ loss = (out - target).pow(2).mean()
+ loss.backward()
+
+ assert all(param.grad is not None for param in model.readout.head.parameters())
+
+
+def test_resolution_invariant_readout_non_unit_measure():
+ """integral reduce with non-unit measure_per_dim should scale pre-head pooling by domain volume.
+
+ The measure scales the spatially-pooled tensor before the linear head, so the
+ relationship is head(6 * pooled) not 6 * head(pooled) — these differ by the bias
+ term. We zero the bias to isolate the scaling behaviour.
+ """
+ device = "cuda" if torch.cuda.is_available() else "cpu"
+
+ readout_unit = ResolutionInvariantReadout(
+ in_channels=4, out_dim=3, reduce="integral", measure_per_dim=[1.0, 1.0]
+ ).to(device)
+ readout_scaled = ResolutionInvariantReadout(
+ in_channels=4, out_dim=3, reduce="integral", measure_per_dim=[2.0, 3.0]
+ ).to(device)
+
+ # Share weights and zero bias so head is purely linear: head(6x) == 6 * head(x)
+ readout_scaled.head.load_state_dict(readout_unit.head.state_dict())
+ with torch.no_grad():
+ readout_unit.head.bias.zero_()
+ readout_scaled.head.bias.zero_()
+
+ x = torch.randn(2, 4, 8, 8, device=device)
+ with torch.no_grad():
+ out_unit = readout_unit(x)
+ out_scaled = readout_scaled(x)
+
+ # domain volume = 2.0 * 3.0 = 6.0, so scaled output should be exactly 6x unit
+ torch.testing.assert_close(out_scaled, out_unit * 6.0)
+
+
+def test_resolution_invariant_readout_measure_length_mismatch():
+ """measure_per_dim with wrong length should raise ValueError on forward."""
+ readout = ResolutionInvariantReadout(
+ in_channels=4, out_dim=3, reduce="integral", measure_per_dim=[1.0, 1.0, 1.0]
+ )
+ x = torch.randn(2, 4, 8, 8) # 2 spatial dims, but measure_per_dim has 3
+ with pytest.raises(ValueError, match="measure_per_dim"):
+ readout(x)
+
+
+def test_resolution_invariant_readout_copies_measure_sequence():
+ measure_per_dim = [1.0, 2.0]
+ readout = ResolutionInvariantReadout(
+ in_channels=4, out_dim=3, reduce="integral", measure_per_dim=measure_per_dim
+ )
+
+ measure_per_dim.append(3.0)
+ x = torch.randn(2, 4, 8, 8)
+
+ out = readout(x)
+ assert out.shape == (2, 3)
+ assert readout.measure_per_dim == [1.0, 2.0]
+
+
+def test_resolution_invariant_readout_state_dict_keys_match():
+ scalar = ResolutionInvariantReadout(
+ in_channels=4, out_dim=3, reduce="integral", measure_per_dim=2.0
+ )
+ sequence = ResolutionInvariantReadout(
+ in_channels=4, out_dim=3, reduce="integral", measure_per_dim=[2.0, 2.0]
+ )
+
+ assert set(scalar.state_dict().keys()) == set(sequence.state_dict().keys())
+
+
+def test_resolution_invariant_readout_cross_load_scalar_sequence_raises():
+ scalar = ResolutionInvariantReadout(
+ in_channels=4, out_dim=3, reduce="integral", measure_per_dim=2.0
+ )
+ sequence = ResolutionInvariantReadout(
+ in_channels=4, out_dim=3, reduce="integral", measure_per_dim=[2.0, 2.0]
+ )
+ with pytest.raises(RuntimeError, match="size mismatch"):
+ sequence.load_state_dict(scalar.state_dict())
+
+
+def test_resolution_invariant_readout_mean_vs_integral_differ():
+ """mean and integral reduce should produce numerically different outputs."""
+ device = "cuda" if torch.cuda.is_available() else "cpu"
+
+ readout_mean = ResolutionInvariantReadout(
+ in_channels=4, out_dim=3, reduce="mean", measure_per_dim=[2.0, 2.0]
+ ).to(device)
+ readout_integral = ResolutionInvariantReadout(
+ in_channels=4, out_dim=3, reduce="integral", measure_per_dim=[2.0, 2.0]
+ ).to(device)
+
+ # Share weights so only the reduce mode differs
+ readout_integral.head.load_state_dict(readout_mean.head.state_dict())
+
+ x = torch.randn(2, 4, 8, 8, device=device)
+ with torch.no_grad():
+ out_mean = readout_mean(x)
+ out_integral = readout_integral(x)
+
+ assert not torch.allclose(
+ out_mean, out_integral
+ ), "mean and integral outputs should differ when measure_per_dim != 1"
+
+
+@pytest.mark.parametrize("n_dim", [1, 2, 3])
+@pytest.mark.parametrize("reduce", ["mean", "integral"])
+def test_resolution_invariant_readout_constant_field(n_dim, reduce):
+ """Readout of a constant field should be identical across resolutions.
+
+ A spatially-constant input has the same pooled value regardless of grid
+ size for both ``mean`` and ``integral`` reduction when the physical domain
+ measure is held fixed. This validates the invariance semantics directly,
+ not just output shapes.
+ """
+ torch.manual_seed(0)
+ in_channels = 4
+ out_dim = 3
+ readout = ResolutionInvariantReadout(
+ in_channels=in_channels,
+ out_dim=out_dim,
+ reduce=reduce,
+ measure_per_dim=([0.5] * n_dim if reduce == "integral" else [1.0] * n_dim),
+ )
+ readout.eval()
+
+ # Build a batch of constant fields: same value at every spatial location.
+ value = torch.randn(2, in_channels) # (B, C)
+
+ outputs = []
+ for size in (8, 16, 32):
+ spatial = (size,) * n_dim
+ # Expand value to (B, C, *spatial)
+ x = value.reshape(2, in_channels, *([1] * n_dim)).expand(
+ 2, in_channels, *spatial
+ )
+ with torch.no_grad():
+ outputs.append(readout(x))
+
+ for out in outputs[1:]:
+ torch.testing.assert_close(out, outputs[0])
+
+
+def test_fno_complex_data_with_readout_raises():
+ """FNO(complex_data=True, readout=...) should raise ValueError at construction."""
+ with pytest.raises(ValueError, match="complex_data"):
+ FNO(
+ n_modes=(4, 4),
+ in_channels=2,
+ out_channels=4,
+ hidden_channels=8,
+ n_layers=1,
+ complex_data=True,
+ readout=ResolutionInvariantReadout(in_channels=4, out_dim=1),
+ )
+
+
+@pytest.mark.parametrize("n_dim", [1, 2, 3])
+def test_fno_without_readout_preserves_field_output_shape(n_dim):
+ model = FNO(
+ n_modes=(8,) * n_dim,
+ in_channels=2,
+ out_channels=5,
+ hidden_channels=12,
+ n_layers=2,
+ readout=None,
+ )
+ x = torch.randn(3, 2, *([16] * n_dim))
+ y = model(x)
+ assert y.shape == (3, 5, *([16] * n_dim))
+
+
+def test_fno_readout_channel_mismatch_raises():
+ """FNO should raise ValueError when readout.in_channels != out_channels."""
+ with pytest.raises(ValueError, match="in_channels"):
+ FNO(
+ n_modes=(4, 4),
+ in_channels=2,
+ out_channels=4,
+ hidden_channels=8,
+ n_layers=1,
+ readout=ResolutionInvariantReadout(in_channels=8, out_dim=1),
+ )
+
+
+def test_resolution_invariant_readout_invalid_reduce_raises():
+ """Invalid reduce value should raise ValueError at construction."""
+ with pytest.raises(ValueError, match="reduce"):
+ ResolutionInvariantReadout(in_channels=4, out_dim=3, reduce="sum")
+
+
+def test_resolution_invariant_readout_invalid_head_raises():
+ """Invalid head value should raise ValueError at construction."""
+ with pytest.raises(ValueError, match="head"):
+ ResolutionInvariantReadout(in_channels=4, out_dim=3, head="transformer")
+
+
+def test_resolution_invariant_readout_low_rank_input_raises():
+ """Input with fewer than 3 dimensions should raise ValueError on forward."""
+ readout = ResolutionInvariantReadout(in_channels=4, out_dim=3)
+ x = torch.randn(2, 4) # missing spatial dim
+ with pytest.raises(ValueError, match=r"\(B, C, \*spatial\)"):
+ readout(x)
+
+
+def test_resolution_invariant_readout_complex_input_raises():
+ """Complex input should raise a clear ValueError on forward."""
+ readout = ResolutionInvariantReadout(in_channels=4, out_dim=3)
+ x = torch.randn(2, 4, 8, 8, dtype=torch.cfloat)
+ with pytest.raises(ValueError, match="complex"):
+ readout(x)
diff --git a/neuralop/models/tests/test_otno.py b/neuralop/models/tests/test_otno.py
index 2a795ad61..7c1bc3827 100644
--- a/neuralop/models/tests/test_otno.py
+++ b/neuralop/models/tests/test_otno.py
@@ -2,6 +2,7 @@
import pytest
from ..otno import OTNO
+from ..readouts import ResolutionInvariantReadout
# Fixed variables
in_channels = 9
@@ -106,3 +107,22 @@ def test_otno_output_shape(n_s_sqrt, n_target_points):
assert list(out.shape) == [batch_size, n_target_points]
assert out.isfinite().all()
+
+
+def test_otno_rejects_readout():
+ """OTNO should raise ValueError when readout is passed.
+
+ OTNO.forward removes the batch dimension and indexes into a target mesh
+ before projection, so its output is incompatible with readout heads that
+ expect (B, C, *spatial) input. Silently accepting readout would let users
+ think they are getting a scalar/vector QoI head when the head is never called.
+ """
+ with pytest.raises(ValueError, match="readout"):
+ OTNO(
+ n_modes=[8, 8],
+ hidden_channels=8,
+ in_channels=4,
+ out_channels=1,
+ n_layers=2,
+ readout=ResolutionInvariantReadout(in_channels=1, out_dim=3),
+ )