Skip to content

Normalization

Flow matching transports \(\mathcal{N}(0, I)\) to the data, so raw PDE fields with non-unit scale make the velocity regression badly conditioned. Always normalize.

Statistics are keyed by raw field name (source, solution, kappa, initial, final) rather than by role, which is why one normalizer stays correct across problem='forward' and problem='inverse'.

Warning

Never refit normalization on validation or test data — share the same train normalizer instance.

normalization

Field Normalization for PDE Datasets

Flow matching transports a standard Gaussian base distribution to the data distribution. When the data has a very different scale from N(0, I) the velocity targets x_1 - x_0 inherit the raw data magnitude, which makes the regression problem badly conditioned. Standardizing each PDE field to roughly zero mean and unit variance removes that mismatch.

Statistics must always be fitted on the training split and then reused verbatim for validation/test data, otherwise the evaluation leaks information about the held-out set:

train_ds = generator.generate(num_samples=1000, seed=0)
test_ds  = generator.generate(num_samples=200, seed=1)

normalizer = FieldNormalizer.from_dataset(train_ds)
train_ds.set_normalizer(normalizer)
test_ds.set_normalizer(normalizer)      # same statistics, not refitted

Metrics should be reported in physical units, so predictions are mapped back with denormalize() (or denormalize_channels() for targets that concatenate several fields) before computing errors.

FieldNormalizer

Per-field mean/std standardization keyed by PDE field name.

Fields are addressed by their raw name ('source', 'solution', 'kappa', 'initial', 'final') rather than by their role in the learning problem, so a single normalizer stays correct when the same data is used for both the forward and the inverse direction.

Parameters:

Name Type Description Default
stats Optional[Dict[str, Dict[str, float]]]

Mapping {field_name: {'mean': float, 'std': float}}.

 None
eps float

Floor applied to standard deviations to avoid division by zero.

 1e-08
Example

normalizer = FieldNormalizer({'solution': {'mean': 2.0, 'std': 4.0}}) z = normalizer.normalize('solution', torch.tensor([6.0])) z tensor([1.]) normalizer.denormalize('solution', z) tensor([6.])

Source code in flowpde/datasets/normalization.py 
 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
 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
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
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
class FieldNormalizer:
    """
    Per-field mean/std standardization keyed by PDE field name.

    Fields are addressed by their raw name (`'source'`, `'solution'`,
    `'kappa'`, `'initial'`, `'final'`) rather than by their role in the
    learning problem, so a single normalizer stays correct when the same data
    is used for both the forward and the inverse direction.

    Args:
        stats: Mapping `{field_name: {'mean': float, 'std': float}}`.
        eps: Floor applied to standard deviations to avoid division by zero.

    Example:
        >>> normalizer = FieldNormalizer({'solution': {'mean': 2.0, 'std': 4.0}})
        >>> z = normalizer.normalize('solution', torch.tensor([6.0]))
        >>> z
        tensor([1.])
        >>> normalizer.denormalize('solution', z)
        tensor([6.])
    """

    def __init__(
        self,
        stats: Optional[Dict[str, Dict[str, float]]] = None,
        eps: float = 1e-8,
    ):
        self.eps = eps
        self.stats: Dict[str, Dict[str, float]] = {}
        if stats:
            for name, field_stats in stats.items():
                self.add_field(
                    name,
                    mean=float(field_stats["mean"]),
                    std=float(field_stats["std"]),
                )

    # Construction

    @classmethod
    def from_dataset(
        cls,
        dataset: "Any",
        fields: Optional[Iterable[str]] = None,
        eps: float = 1e-8,
    ) -> "FieldNormalizer":
        """
        Build a normalizer from statistics a dataset already carries.

        Generators compute per-field mean/std at construction time and store
        them in `metadata['stats']`, so no second pass over the data is
        needed.

        Args:
            dataset: Dataset exposing `get_stats()` (`PDEDataset`,
                `DarcyDataset`).
            fields: Restrict to these field names.  Defaults to every field
                that has statistics, excluding `'obs_mask'`.
            eps: Floor applied to standard deviations.

        Returns:
            A fitted `FieldNormalizer`.
        """
        if not hasattr(dataset, "get_stats"):
            raise TypeError(
                f"{type(dataset).__name__} does not expose get_stats(); "
                "pass statistics to FieldNormalizer(...) directly."
            )

        stats = dataset.get_stats()
        if not stats:
            raise ValueError(
                "Dataset carries no normalization statistics. It was probably "
                "not built through a generator's wrap_dataset()."
            )

        selected = list(fields) if fields is not None else [
            name for name in stats if name != "obs_mask"
        ]

        missing = [name for name in selected if name not in stats]
        if missing:
            raise KeyError(
                f"Dataset has no statistics for field(s): {missing}. "
                f"Available: {sorted(stats)}"
            )

        return cls({name: stats[name] for name in selected}, eps=eps)

    @classmethod
    def from_tensors(
        cls,
        tensors: Dict[str, Tensor],
        eps: float = 1e-8,
    ) -> "FieldNormalizer":
        """Fit directly from raw tensors, one entry per field."""
        stats = {
            name: {"mean": tensor.mean().item(), "std": tensor.std().item()}
            for name, tensor in tensors.items()
        }
        return cls(stats, eps=eps)

    def add_field(self, name: str, mean: float, std: float) -> None:
        """Register statistics for a single field."""
        self.stats[name] = {"mean": float(mean), "std": max(float(std), self.eps)}

    # Application

    @property
    def fields(self) -> List[str]:
        """Names of all fields this normalizer knows about."""
        return sorted(self.stats)

    def has_field(self, name: str) -> bool:
        return name in self.stats

    def normalize(self, name: str, tensor: Tensor) -> Tensor:
        """
        Standardize a field to zero mean and unit variance.

        Fields with no registered statistics pass through unchanged, so
        auxiliary channels such as `obs_mask` stay binary.
        """
        if name not in self.stats:
            return tensor
        field_stats = self.stats[name]
        return (tensor - field_stats["mean"]) / field_stats["std"]

    def denormalize(self, name: str, tensor: Tensor) -> Tensor:
        """Map a standardized field back to physical units."""
        if name not in self.stats:
            return tensor
        field_stats = self.stats[name]
        return tensor * field_stats["std"] + field_stats["mean"]

    def denormalize_channels(
        self,
        names: Sequence[str],
        tensor: Tensor,
        channel_dim: int = 1,
    ) -> Tensor:
        """
        Denormalize a tensor whose channels concatenate several fields.

        Used for targets such as Darcy's `inverse_mode='both'`, where the
        target is `cat([kappa, source])`.

        Args:
            names: Field name per channel group, in channel order.
            tensor: Tensor with `len(names)` equal-sized channel groups.
            channel_dim: Dimension holding the channels (default: 1).

        Returns:
            Tensor of the same shape, in physical units.
        """
        if len(names) == 1:
            return self.denormalize(names[0], tensor)

        num_channels = tensor.shape[channel_dim]
        if num_channels % len(names) != 0:
            raise ValueError(
                f"Cannot split {num_channels} channels evenly across "
                f"{len(names)} fields ({list(names)})."
            )

        chunks = torch.chunk(tensor, len(names), dim=channel_dim)
        restored = [
            self.denormalize(name, chunk) for name, chunk in zip(names, chunks, strict=True)
        ]
        return torch.cat(restored, dim=channel_dim)

    # Serialization

    def state_dict(self) -> Dict[str, object]:
        """Serializable state, suitable for storing next to model weights."""
        return {"stats": {k: dict(v) for k, v in self.stats.items()}, "eps": self.eps}

    def load_state_dict(self, state: Dict[str, object]) -> "FieldNormalizer":
        """Restore statistics saved by `state_dict()`."""
        self.eps = float(state.get("eps", 1e-8))
        self.stats = {}
        for name, field_stats in state["stats"].items():
            self.add_field(name, field_stats["mean"], field_stats["std"])
        return self

    @classmethod
    def from_state_dict(cls, state: Dict[str, object]) -> "FieldNormalizer":
        """Rebuild a normalizer from a serialized state."""
        return cls().load_state_dict(state)

    def __repr__(self) -> str:
        entries = ", ".join(
            f"{name}(mean={s['mean']:.4g}, std={s['std']:.4g})"
            for name, s in sorted(self.stats.items())
        )
        return f"FieldNormalizer({entries})"
fields property

Names of all fields this normalizer knows about.

from_dataset(dataset, fields=None, eps=1e-08) classmethod

Build a normalizer from statistics a dataset already carries.

Generators compute per-field mean/std at construction time and store them in metadata['stats'], so no second pass over the data is needed.

Parameters:

Name Type Description Default
dataset 'Any'

Dataset exposing get_stats() (PDEDataset, DarcyDataset).

 required
fields Optional[Iterable[str]]

Restrict to these field names. Defaults to every field that has statistics, excluding 'obs_mask'.

 None
eps float

Floor applied to standard deviations.

 1e-08

Returns:

Type Description
'FieldNormalizer'

A fitted FieldNormalizer.
Source code in flowpde/datasets/normalization.py 
 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
@classmethod
def from_dataset(
    cls,
    dataset: "Any",
    fields: Optional[Iterable[str]] = None,
    eps: float = 1e-8,
) -> "FieldNormalizer":
    """
    Build a normalizer from statistics a dataset already carries.

    Generators compute per-field mean/std at construction time and store
    them in `metadata['stats']`, so no second pass over the data is
    needed.

    Args:
        dataset: Dataset exposing `get_stats()` (`PDEDataset`,
            `DarcyDataset`).
        fields: Restrict to these field names.  Defaults to every field
            that has statistics, excluding `'obs_mask'`.
        eps: Floor applied to standard deviations.

    Returns:
        A fitted `FieldNormalizer`.
    """
    if not hasattr(dataset, "get_stats"):
        raise TypeError(
            f"{type(dataset).__name__} does not expose get_stats(); "
            "pass statistics to FieldNormalizer(...) directly."
        )

    stats = dataset.get_stats()
    if not stats:
        raise ValueError(
            "Dataset carries no normalization statistics. It was probably "
            "not built through a generator's wrap_dataset()."
        )

    selected = list(fields) if fields is not None else [
        name for name in stats if name != "obs_mask"
    ]

    missing = [name for name in selected if name not in stats]
    if missing:
        raise KeyError(
            f"Dataset has no statistics for field(s): {missing}. "
            f"Available: {sorted(stats)}"
        )

    return cls({name: stats[name] for name in selected}, eps=eps)
from_tensors(tensors, eps=1e-08) classmethod

Fit directly from raw tensors, one entry per field.

Source code in flowpde/datasets/normalization.py 
125
126
127
128
129
130
131
132
133
134
135
136
@classmethod
def from_tensors(
    cls,
    tensors: Dict[str, Tensor],
    eps: float = 1e-8,
) -> "FieldNormalizer":
    """Fit directly from raw tensors, one entry per field."""
    stats = {
        name: {"mean": tensor.mean().item(), "std": tensor.std().item()}
        for name, tensor in tensors.items()
    }
    return cls(stats, eps=eps)
add_field(name, mean, std)

Register statistics for a single field.

Source code in flowpde/datasets/normalization.py 
138
139
140
def add_field(self, name: str, mean: float, std: float) -> None:
    """Register statistics for a single field."""
    self.stats[name] = {"mean": float(mean), "std": max(float(std), self.eps)}
normalize(name, tensor)

Standardize a field to zero mean and unit variance.

Fields with no registered statistics pass through unchanged, so auxiliary channels such as obs_mask stay binary.

Source code in flowpde/datasets/normalization.py 
152
153
154
155
156
157
158
159
160
161
162
def normalize(self, name: str, tensor: Tensor) -> Tensor:
    """
    Standardize a field to zero mean and unit variance.

    Fields with no registered statistics pass through unchanged, so
    auxiliary channels such as `obs_mask` stay binary.
    """
    if name not in self.stats:
        return tensor
    field_stats = self.stats[name]
    return (tensor - field_stats["mean"]) / field_stats["std"]
denormalize(name, tensor)

Map a standardized field back to physical units.

Source code in flowpde/datasets/normalization.py 
164
165
166
167
168
169
def denormalize(self, name: str, tensor: Tensor) -> Tensor:
    """Map a standardized field back to physical units."""
    if name not in self.stats:
        return tensor
    field_stats = self.stats[name]
    return tensor * field_stats["std"] + field_stats["mean"]
denormalize_channels(names, tensor, channel_dim=1)

Denormalize a tensor whose channels concatenate several fields.

Used for targets such as Darcy's inverse_mode='both', where the target is cat([kappa, source]).

Parameters:

Name Type Description Default
names Sequence[str]

Field name per channel group, in channel order.

 required
tensor Tensor

Tensor with len(names) equal-sized channel groups.

 required
channel_dim int

Dimension holding the channels (default: 1).

 1

Returns:

Type Description
Tensor

Tensor of the same shape, in physical units.
Source code in flowpde/datasets/normalization.py 
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
def denormalize_channels(
    self,
    names: Sequence[str],
    tensor: Tensor,
    channel_dim: int = 1,
) -> Tensor:
    """
    Denormalize a tensor whose channels concatenate several fields.

    Used for targets such as Darcy's `inverse_mode='both'`, where the
    target is `cat([kappa, source])`.

    Args:
        names: Field name per channel group, in channel order.
        tensor: Tensor with `len(names)` equal-sized channel groups.
        channel_dim: Dimension holding the channels (default: 1).

    Returns:
        Tensor of the same shape, in physical units.
    """
    if len(names) == 1:
        return self.denormalize(names[0], tensor)

    num_channels = tensor.shape[channel_dim]
    if num_channels % len(names) != 0:
        raise ValueError(
            f"Cannot split {num_channels} channels evenly across "
            f"{len(names)} fields ({list(names)})."
        )

    chunks = torch.chunk(tensor, len(names), dim=channel_dim)
    restored = [
        self.denormalize(name, chunk) for name, chunk in zip(names, chunks, strict=True)
    ]
    return torch.cat(restored, dim=channel_dim)
state_dict()

Serializable state, suitable for storing next to model weights.

Source code in flowpde/datasets/normalization.py 
209
210
211
def state_dict(self) -> Dict[str, object]:
    """Serializable state, suitable for storing next to model weights."""
    return {"stats": {k: dict(v) for k, v in self.stats.items()}, "eps": self.eps}
load_state_dict(state)

Restore statistics saved by state_dict().

Source code in flowpde/datasets/normalization.py 
213
214
215
216
217
218
219
def load_state_dict(self, state: Dict[str, object]) -> "FieldNormalizer":
    """Restore statistics saved by `state_dict()`."""
    self.eps = float(state.get("eps", 1e-8))
    self.stats = {}
    for name, field_stats in state["stats"].items():
        self.add_field(name, field_stats["mean"], field_stats["std"])
    return self
from_state_dict(state) classmethod

Rebuild a normalizer from a serialized state.

Source code in flowpde/datasets/normalization.py 
221
222
223
224
@classmethod
def from_state_dict(cls, state: Dict[str, object]) -> "FieldNormalizer":
    """Rebuild a normalizer from a serialized state."""
    return cls().load_state_dict(state)