Neural Implicit Flow (NIF): mesh-agnostic dimensionality reduction

• NIF is a mesh-agnostic dimensionality reduction paradigm for parametric spatiotemporal fields. Dimensionality reduction is often the first step in reduced-order modeling of large-scale dynamics.

• NIF represents fields at arbitrary spatial coordinates, including data from adaptive mesh refinement, without requiring interpolation onto a uniform grid. Details can be found in our paper.

• NIF is implemented in PyTorch. A ParameterNet maps time or physical parameters to the weights of a coordinate-dependent ShapeNet.

Read the PyTorch v2.0.0 documentation for the API reference, runnable examples, and all eight tutorials.

Highlights

Installation and requirements

Python 3.10+ and PyTorch 2.6+ are required; CUDA is optional. Install from a local checkout:

python -m pip install -e .
# Optional: support TensorBoard logging in the standalone training example.
python -m pip install -e '.[tensorboard]'

Choose a suitable PyTorch CUDA wheel for your machine when training on GPU. See native API and migration notes for optimizer conventions and data formats. The validation report and comparison records document the measured results and limitations.

Configuring and training NIF

The example below fits a one-dimensional traveling wave. The configuration dictionaries set the input/output dimensions, network widths, hidden layers, and activations. Input columns are ordered as [parameters, coordinates].

import torch
from nif import NIF
from nif.demo import TravelingWave
from nif.optimizers import AdaBeliefOptimizer

shape = dict(input_dim=1, output_dim=1, units=20, nlayers=2, activation="tanh")
parameter = dict(input_dim=1, latent_dim=2, units=10, nlayers=2, activation="tanh")
model = NIF(shape, parameter).to("cuda" if torch.cuda.is_available() else "cpu")
device = next(model.parameters()).device

data = torch.as_tensor(TravelingWave().data, dtype=torch.float32, device=device)
x, target = data[:, :2], data[:, 2:3]
optimizer = AdaBeliefOptimizer(model.parameters(), lr=1e-3)
for epoch in range(100):
    optimizer.zero_grad(set_to_none=True)
    loss = (model(x) - target).square().mean() + model.regularization_loss(x)
    loss.backward()
    optimizer.step()
print(float(loss.detach()))

The default model uses float32. On supported hardware, mixed_policy="mixed_bfloat16" or mixed_policy="mixed_float16" enables autocast while retaining float32 trainable parameters. See the API notes for details.

Distributed learning

Use DistributedDataParallel (DDP) to train one model across multiple GPUs, with one process per GPU. The following example shows the data sharding and synchronized training loop. Save it as train_wave_ddp.py in the repository root; it reuses TrainingObjective from the training example to evaluate MSE and regularization together inside the DDP forward pass.

import os
from pathlib import Path

import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader, DistributedSampler, TensorDataset
from nif import NIF
from nif.demo import TravelingWave
from examples.train_pytorch import TrainingObjective

def main():
    local_rank = int(os.environ["LOCAL_RANK"])
    use_cuda = torch.cuda.is_available()
    device = torch.device("cuda", local_rank) if use_cuda else torch.device("cpu")
    if use_cuda:
        torch.cuda.set_device(local_rank)
    dist.init_process_group(backend="nccl" if use_cuda else "gloo")
    try:
        torch.manual_seed(42)
        shape = dict(input_dim=1, output_dim=1, units=20, nlayers=2, activation="tanh")
        parameter = dict(input_dim=1, latent_dim=2, units=10, nlayers=2, activation="tanh")
        model = NIF(shape, parameter).to(device)
        objective = DDP(TrainingObjective(model), device_ids=[local_rank] if use_cuda else None)

        data = torch.as_tensor(TravelingWave().data, dtype=torch.float32)
        dataset = TensorDataset(data[:, :2], data[:, 2:3])
        if len(dataset) % dist.get_world_size():
            raise ValueError("Dataset size must be divisible by the number of processes")
        sampler = DistributedSampler(dataset, shuffle=True, seed=42)
        loader = DataLoader(dataset, batch_size=256, sampler=sampler)
        optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

        model.train()
        for epoch in range(20):
            sampler.set_epoch(epoch)  # Give every epoch a new shared shuffle.
            for features, targets in loader:
                features, targets = features.to(device), targets.to(device)
                optimizer.zero_grad(set_to_none=True)
                mse, loss = objective(features, targets)
                loss.backward()  # DDP averages gradients across processes.
                optimizer.step()

        if dist.get_rank() == 0:  # Write one copy of the underlying NIF weights.
            output = Path("pytorch_outputs/wave-ddp")
            output.mkdir(parents=True, exist_ok=True)
            torch.save(model.state_dict(), output / "model.pt")
            model.save_config(output / "config.json")
    finally:
        dist.destroy_process_group()

if __name__ == "__main__":
    main()

On Linux or WSL, launch this snippet with two CUDA GPUs and an NCCL-enabled PyTorch build, or with CPU/Gloo when CUDA is unavailable:

torchrun --standalone --nproc-per-node=2 train_wave_ddp.py

The complete example also includes a learning-rate schedule, globally reduced metrics, checkpoints, and TensorBoard logging:

torchrun --standalone --nproc-per-node=2 examples/train_pytorch.py --device cuda --epochs 20 --batch-size 256 --output-dir pytorch_outputs/wave-ddp --tensorboard

The batch size is per process: two processes with 256 rows each give a global batch of 512. Equal dataset shards also keep the final batch sizes equal, without padding or dropping samples. The complete example also runs on one device with ordinary python; it trains on the same traveling-wave data used above. The L-BFGS example below uses one device. See the training example validation and requirements for tested platforms and limits.

Fine-tuning with L-BFGS

After first-order training, refine the same model using a fixed dataset and a closure. This example continues from the model, x, and target above.

from nif.optimizers import ReferenceLBFGS

fine_tuner = ReferenceLBFGS(model.parameters(), max_iter=1000)
def closure():
    fine_tuner.zero_grad(set_to_none=True)
    loss = (model(x) - target).square().mean() + model.regularization_loss(x)
    loss.backward()
    return loss

fine_tuner.step(closure)
print("Final loss:", float(fine_tuner.result.objective_value))
print("Iterations:", int(fine_tuner.result.num_iterations))

Training schedules and monitoring

Add a StepLR schedule, TensorBoard scalars, and checkpoints to an explicit PyTorch training loop. This single-device example uses the model, x, and target from the first example and starts a new Adam optimizer. Install the optional TensorBoard dependency shown above.

from pathlib import Path
from torch.utils.tensorboard import SummaryWriter

run = Path("pytorch_outputs/wave-monitored")
run.mkdir(parents=True, exist_ok=True)
model.save_config(run / "config.json")
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.1)

model.train()
with SummaryWriter(log_dir=str(run / "tensorboard")) as writer:
    for epoch in range(20):
        learning_rate = optimizer.param_groups[0]["lr"]
        optimizer.zero_grad(set_to_none=True)
        mse = (model(x) - target).square().mean()
        loss = mse + model.regularization_loss(x)
        loss.backward()
        optimizer.step()

        writer.add_scalar("train/mse", mse.item(), epoch + 1)
        writer.add_scalar("train/objective", loss.item(), epoch + 1)
        writer.add_scalar("train/lr", learning_rate, epoch + 1)
        print(f"epoch={epoch + 1} mse={mse.item():.6g} lr={learning_rate:.3g}")
        scheduler.step()  # Decay the learning rate after the epoch's updates.
        torch.save({
            "epoch": epoch + 1,
            "model_state_dict": model.state_dict(),
            "optimizer_state_dict": optimizer.state_dict(),
            "scheduler_state_dict": scheduler.state_dict(),
        }, run / "checkpoint.pt")

    model.eval()
    with torch.no_grad():
        final_mse = (model(x) - target).square().mean().item()
    writer.add_scalar("eval/mse", final_mse, epoch + 1)

Open these logs with tensorboard --logdir pytorch_outputs/wave-monitored. The training scalars above describe each full-batch update; eval/mse measures the final weights. Both use the normalized data. The context manager flushes and closes the writer.

For minibatch training with the same schedule and monitoring, run the complete example on one CPU:

python examples/train_pytorch.py --device cpu --epochs 20 --batch-size 256 --output-dir pytorch_outputs/wave --tensorboard
tensorboard --logdir pytorch_outputs/wave

For distributed runs, the complete example averages metrics across processes and writes logs and checkpoints on rank zero.

Extracting subnetworks

Extract the parameter-to-latent, latent-to-weight, and coordinate-to-output subnetworks. These views share the model’s tensors and retain their gradient connections.

parameter_inputs, coordinates = x[:, :1], x[:, 1:2]
latent = model.model_p_to_lr()(parameter_inputs)
weights = model.model_lr_to_w()(latent)
u_pred = model.model_x_to_u_given_w()(coordinates, weights)
torch.testing.assert_close(u_pred, model(x))

Input-output Jacobians and Hessians

from nif.layers import JacobianLayer, HessianLayer

output, jacobian = JacobianLayer(model, y_index=[0], x_index=[1])(x)
output, jacobian, hessian = HessianLayer(model, y_index=[0], x_index=[1])(x)

The layers return the complete output and selected derivatives, and retain parameter gradients for derivative-based training. See Tutorial 4.

Data normalization

PointWiseData keeps parameter, coordinate, and target columns in order and supports min-max or standard normalization. For example, with one parameter, one coordinate, and one output:

from nif.data import PointWiseData

raw_data = TravelingWave().data_raw
normalized, mean, scale = PointWiseData.minmax_normalize(
    raw_data=raw_data, n_para=1, n_x=1, n_target=1
)

Large-scale record data

TFRDataset converts an NPZ array into record shards and reads them as PyTorch DataLoader batches. This small, self-contained IO example uses the model, optimizer, and device defined above:

from pathlib import Path
from tempfile import TemporaryDirectory
import numpy as np
from nif.data.tfr_dataset import TFRDataset

with TemporaryDirectory() as directory:
    root = Path(directory)
    np.savez(root / "wave.npz", data=TravelingWave().data)
    records = TFRDataset(n_feature=2, n_target=1)
    records.create_from_npz(
        num_pts_per_file=1000, npz_path=root / "wave.npz", npz_key="data",
        tfr_path=root / "record_shards", prefix="wave"
    )
    for shard in records.get_tfr_meta_dataset(root / "record_shards", epoch=1):
        for features, targets in records.gen_dataset_from_batch_file(shard, batch_size=128):
            features, targets = features.to(device), targets.to(device)
            optimizer.zero_grad(set_to_none=True)
            loss = (model(features) - targets).square().mean()
            loss = loss + model.regularization_loss(features)
            loss.backward()
            optimizer.step()

Tutorial 5 demonstrates the same workflow on a one-million-row dataset.

Save and load models

torch.save(model.state_dict(), "model.pt")
restored = NIF(shape, parameter).to(device)
restored.load_state_dict(torch.load("model.pt", map_location=device, weights_only=True))
torch.testing.assert_close(restored(x), model(x))
model.save_config("config.json")

Network pruning and quantization

Prune a copy of the trained float32 NIF from the first example, then run dynamic int8 inference on CPU. The pruning controller below acts on the final ParameterNet projection kernel, leaving its bias unchanged. This short example uses 20 full-batch updates to demonstrate the schedule; choose the pruning budget and target sparsity for your data.

import copy
from torch import nn
from nif.pruning import PolynomialMagnitudePruning

compressed = copy.deepcopy(model).cpu().float().train()
features, targets = x.detach().cpu().float(), target.detach().cpu().float()
projection = compressed.pnet_list[-1].linear
pruner = PolynomialMagnitudePruning(
    projection, name="weight", initial_sparsity=0.1, final_sparsity=0.9,
    begin_step=0, end_step=20, frequency=1,
)
pruning_optimizer = torch.optim.Adam(compressed.parameters(), lr=1e-4)
for epoch in range(20):
    pruning_optimizer.zero_grad(set_to_none=True)
    loss = (compressed(features) - targets).square().mean()
    loss = loss + compressed.regularization_loss(features)
    loss.backward()
    pruning_optimizer.step()
    pruner.on_epoch_end()  # Reapply the mask before checkpointing or inference.

pruner.strip()  # Keep sparse weights; remove pruning hooks before conversion.
compressed.eval()
print("Projection sparsity:", (projection.weight == 0).float().mean().item())
int8_model = torch.ao.quantization.quantize_dynamic(
    compressed, {nn.Linear}, dtype=torch.qint8, inplace=False
).eval()
with torch.no_grad():
    float_prediction = compressed(features)
    int8_prediction = int8_model(features)
    print("Pruned-model MAE:", (float_prediction - targets).abs().mean().item())
    print("Int8-model MAE:", (int8_prediction - targets).abs().mean().item())
    print("Quantization MAE:", (int8_prediction - float_prediction).abs().mean().item())

Dynamic quantization converts this model’s native nn.Linear layers in ParameterNet to packed int8 CPU kernels. Custom residual layers and the generated ShapeNet computation remain floating point. All three MAEs use normalized output units; compare them before choosing the compressed model.

Tutorial 7 shows the full multiscale pruning experiment and TorchScript export. Its SIREN projection uses HyperLinearForSIREN.w, which the tutorial converts to nn.Linear with the original weight orientation preserved before quantization.

Google Colab Tutorial

The notebooks below run the native PyTorch examples. Their scientific workloads are the default; set NIF_SMOKE=1 before running only for a short development check.

1. Hello world! A simple fitting on a 1D traveling wave

   • Learn the nif.NIF model, its configuration, and training loop.

   • Save and restore checkpoints; refine the fit with L-BFGS.

2. Tackling multi-scale data

   • Use nif.NIFMultiScale and compare fits to a high-frequency wave.

3. Learning a linear representation

   • Use nif.NIFMultiScaleLastLayerParameterized on a cylinder-flow dataset from an adaptive mesh.

4. Getting input-output derivatives

   • Use nif.layers.JacobianLayer and nif.layers.HessianLayer.

5. Large-scale data training

   • Create and stream record shards using nif.data.tfr_dataset.TFRDataset.

   • Train on batches from a one-million-row example dataset.

6. Revisit multi-scale data with regularization

   • Compare L1 and L2 regularization of ParameterNet weights and biases.

   • Inspect training fits and spurious signals on a refined space-time grid.

7. NIF compression

   • Explore magnitude pruning, dynamic int8 quantization, and exported inference.

8. Revisit multi-scale data with Sobolev training

   • Add spatial-derivative supervision with nif.layers.JacobianLayer and compare field and derivative errors.

   • Inspect how derivative information affects the refined spatial prediction.

Issues, bugs, requests, ideas

Use the issues tracker to report bugs.

How to cite

For the software, use GitHub’s citation metadata and the Zenodo software record. The latter DOI represents all archived software versions. Cite the version-specific DOI from the Zenodo record when referring to a particular release.

If you find NIF is helpful, please cite our JMLR paper:

@article{JMLR:v24:22-0365,
  author  = {Shaowu Pan and Steven L. Brunton and J. Nathan Kutz},
  title   = {Neural Implicit Flow: a mesh-agnostic dimensionality reduction paradigm of spatio-temporal data},
  journal = {Journal of Machine Learning Research},
  year    = {2023},
  volume  = {24},
  number  = {41},
  pages   = {1--60},
  url     = {http://jmlr.org/papers/v24/22-0365.html}
}

Contributors

• Shaowu Pan

• Yaser Afshar

License

LGPL-2.1 License