Create a Model Wrapper

Guide the user through adding a new CFD model to the benchmarking workflow by writing a CFDModel subclass.

First: gather whatever context already exists

Spend a little effort up front collecting any existing artifacts that reveal how the model actually behaves — but do not block on them. Look (without stopping to ask the user) for:

• the training/inference script and any preprocessing/normalization utilities it imports;
• the model's config files (YAML/JSON hyperparameters, channel lists, stats paths);
• the model class's docstrings and signatures (forward inputs/outputs, expected shapes).

Use these to pin down four things the wrapper must mirror exactly:

• Normalization — which scheme and which stats (see Normalization below).
• Inputs — what the model actually consumes (coordinates only, extra fields, geometry/STL).
• Output fields and order — which variables the forward pass returns and in what channel order.
• I/O shapes and dtypes — so prepare_inputs/decode_outputs match the trained graph.

If some or all of these aren't available, that's fine — proceed anyway: reconstruct from the checkpoint and stats file, pick the most likely option, and build the wrapper now. Do not stop and wait for answers before writing code. Just avoid silently guessing: a wrong normalization scheme or channel order produces plausible-looking but wrong predictions. So state each such assumption inline, and after delivering the wrapper, raise the still-uncertain choices as explicit open items for the user to confirm — e.g. normalization scheme, input tier, and output fields/channel order. This keeps you unblocked while making the risky decisions visible.

Reference files to read first

Start with the complete, ready-to-adapt templates bundled with this skill — they are always available even when the PhysicsNeMo-CFD source tree is not on disk:

• references/example_wrapper.py — full surface and volume CFDModel reference implementations (load, prepare_inputs, predict, decode_outputs, registration). Copy and adapt one of these.
• assets/global_stats.example.json — sample mean/std stats for both surface and volume.

When the PhysicsNeMo-CFD repo is present, also read these for the live interface (verify paths against the actual tree):

• physicsnemo/cfd/evaluation/models/model_registry.py — base class and registry
• physicsnemo/cfd/evaluation/datasets/schema.py — CanonicalCase, build_predictions_dict
• physicsnemo/cfd/evaluation/models/wrappers/surface_baseline.py — simplest concrete surface wrapper
• physicsnemo/cfd/evaluation/models/wrappers/volume_baseline.py — simplest concrete volume wrapper
• physicsnemo/cfd/evaluation/models/wrappers/__init__.py — how wrappers are registered
• physicsnemo/cfd/evaluation/common/io.py — mesh loading and normalization stats helpers
• workflows/benchmarking/notebooks/adding_a_new_model.ipynb — end-to-end tutorial

The CFDModel interface

Every wrapper must set two class variables and implement four methods:

The engine calls load once, then prepare_inputs → predict → decode_outputs(raw, case, model_input) per case (model_input is the dict from prepare_inputs; use when decode must mirror inference geometry).

Step 1: Write the wrapper class

Always generate a new, complete wrapper class for the requested model. Existing wrappers (e.g. surface_baseline.py) are references to read, not substitutes — when the user asks to write or create a wrapper, produce a full new class file even if similar ones already exist. Do not stop at "a wrapper already exists" or offer to reuse one in place of writing the requested one. Always tell the user how to register it: register_model(...) at import for a quick test, and an entry in wrappers/__init__.py to make it permanent (Step 6).

Anti-patterns (do not do these)

• Reusing an existing wrapper instead of writing the requested one — even if git status shows a similar file, write the new class.
• Assuming mean-std normalization — confirm the scheme; a wrong inverse gives plausible-but-wrong fields.
• Hardcoding pressure + shear_stress — pass only the fields the model predicts; add custom ones via **extra.
• Skipping __init__.py — registering only inline and never mentioning permanent registration.
• Echoing the interface table or full reference file back — wastes tokens; reference, don't repeat.

Copy references/example_wrapper.py and adapt it — it has full surface and volume implementations. Don't hand-write from scratch or paste the whole template back to the user. The class skeleton is just:

class MyModelWrapper(CFDModel):
    INFERENCE_DOMAIN: ClassVar[InferenceDomain] = "surface"  # or "volume"
    OUTPUT_LOCATION: ClassVar[OutputLocation] = "cell"        # or "point"

    @property
    def output_location(self): return self.OUTPUT_LOCATION
    def load(self, checkpoint_path, stats_path, device, **kwargs): ...   # weights + stats; return self
    def prepare_inputs(self, case): ...                                  # CanonicalCase -> model input
    def predict(self, model_input): ...                                  # forward pass -> raw output
    def decode_outputs(self, raw_output, case, model_input=None): ...    # denormalize -> build_predictions_dict(...)

Keep responses terse: state the few model-specific decisions (normalization scheme, input tier, output fields) and the file you wrote — don't echo the interface table or the full reference file back.

Key implementation considerations

Normalization (match the training script exactly): Most trained models normalize inputs/outputs, and decode_outputs must apply the inverse of whatever the model was trained with. First identify the scheme:

• Mean-std (z-score): x_norm = (x - mean) / std → inverse x = x_norm * std + mean. This is the repo's built-in format. Use load_global_stats(stats_path) from physicsnemo/cfd/evaluation/common/io.py; it reads mean/std_dev JSON and returns mean/std tensors.
• Min-max: x_norm = (x - min) / (max - min) → inverse x = x_norm * (max - min) + min. There is no built-in helper for this — store min/max (e.g. in your stats JSON) and apply the inverse yourself in decode_outputs. Do not feed a min-max file to load_global_stats; the keys won't match.

Confirm the scheme from the training/inference script or the stats file rather than assuming mean-std. Applying the wrong inverse yields wrong-but-plausible fields that still pass shape checks.

Inputs (handle the model's actual input tier): prepare_inputs receives a CanonicalCase. Pull what the model needs:

• Point cloud only: coordinates from case.mesh_path (vtp/vtu) via pv.read, as in the example — sufficient for many geometry-only models.
• Extra field inputs (e.g. inlet/freestream velocity, Reynolds number): read from case.metadata (or case.ground_truth for field arrays). Broadcast/concatenate them onto the per-point features as the training script did.
• Geometry/STL (e.g. SDF or BVH-based models): the STL/geometry path is typically on case.metadata; load it in prepare_inputs and build the geometric encoding the model expects.

Inspect case.metadata and case.ground_truth keys for a real case early — the dataset adapter decides what is available.

Outputs (predict only what the model produces, plus extras): build_predictions_dict takes pressure, shear_stress, velocity, turbulent_viscosity (all optional) and arbitrary **extra fields. So:

• A model that predicts only WSS magnitude, or no WSS at all, simply omits the missing keys — pass only what it produces.
• Extra/non-standard outputs (e.g. stagnation_pressure, temperature, mach) are passed as keyword args: build_predictions_dict(pressure=p, mach=m, temperature=t). Each becomes a prediction variable.
• For a custom field to appear in the written mesh and metrics, add a matching entry to output.mesh_field_names in the config (Step 4) and a corresponding metric if you want it scored.

Output shape: pressure must be (N,) float32. shear_stress must be (N, 3) float32 for surface. Volume fields: velocity is (N, 3), turbulent_viscosity is (N,). Custom scalar fields are (N,), vector fields (N, k).

Output location: If OUTPUT_LOCATION = "cell", return N = mesh.n_cells values. If "point", return N = mesh.n_points values.

Batching: For large meshes, prepare_inputs may need to subsample or batch. Use kwargs passed through load() (e.g., batch_resolution, geometry_sampling) to control this.

Step 2: Create checkpoint and stats files

Your model needs a checkpoint file and optionally a global_stats.json:

# Checkpoint: save your model's state dict
torch.save(model.state_dict(), "checkpoint.pt")

# Stats: JSON with mean/std_dev for denormalization
# Surface format:
{
    "mean": {"pressure": [0.0], "shear_stress": [0.0, 0.0, 0.0]},
    "std_dev": {"pressure": [1.0], "shear_stress": [1.0, 1.0, 1.0]}
}
# Volume format:
{
    "mean": {"pressure": [0.0], "velocity": [0.0, 0.0, 0.0], "turbulent_viscosity": [0.0]},
    "std_dev": {"pressure": [1.0], "velocity": [1.0, 1.0, 1.0], "turbulent_viscosity": [1.0]}
}

This mean/std_dev layout is what load_global_stats() expects (mean-std models). If your model was trained with min-max normalization, this helper does not apply — persist min/max per field in your own JSON and apply the inverse manually in decode_outputs (see Normalization above).

Step 3: Register and test

register_model("my_model", MyModelWrapper)

# Load a case from any registered dataset adapter
from physicsnemo.cfd.evaluation.datasets.adapters.drivaerml import DrivAerMLAdapter
adapter = DrivAerMLAdapter(root="/path/to/data", inference_domain="surface")
case = adapter.load_case(adapter.list_cases()[0])

# Run the full inference pipeline
wrapper = MyModelWrapper()
wrapper.load(checkpoint_path="checkpoint.pt", stats_path="global_stats.json", device="cuda:0")
model_input = wrapper.prepare_inputs(case)
raw_output = wrapper.predict(model_input)
predictions = wrapper.decode_outputs(raw_output, case, model_input)

assert "pressure" in predictions
assert predictions["pressure"].shape[0] > 0

Step 4: Run the full benchmark

from physicsnemo.cfd.evaluation.config import Config
from physicsnemo.cfd.evaluation.benchmarks.engine import run_benchmark

config = Config.from_dict({
    "run": {"device": "cuda:0", "output_dir": "results", "metrics_cache": {"enabled": False}},
    "benchmark": {
        "mode": "matrix",
        "models": [{
            "name": "my_model",
            "inference_domain": "surface",
            "checkpoint": "/path/to/checkpoint.pt",
            "stats_path": "/path/to/global_stats.json",
            "kwargs": {},
        }],
        "datasets": [{
            "name": "drivaerml",
            "root": "/path/to/drivaerml/data",
            "case_ids": ["run_1", "run_11"],
            "kwargs": {"align_ground_truth_to_model": True, "inference_domain": "surface"},
        }],
        "reproducibility": {"log_env": False, "save_artifacts": True},
    },
    "output": {"mesh_field_names": {"pressure": "pMeanTrimPred", "shear_stress": "wallShearStressMeanTrimPred"}},
    "metrics": ["l2_pressure", "l2_shear_stress", "l2_pressure_area_weighted", "drag", "lift"],
    "reports": {"enabled": False},
})
results = run_benchmark(config)

Results are written to benchmark_results.json (a JSON list of dicts, one per model×dataset combo).

Step 5: Visualize predictions

from physicsnemo.cfd.postprocessing_tools.visualization.utils import plot_fields, plot_field_comparisons

# Just the predicted fields (no GT comparison):
plotter = plot_fields(mesh, fields=["pMeanTrimPred"], view="xy", dtype="cell", window_size=[1800, 600])
plotter.screenshot("predicted_pressure.png")
plotter.close()

# Side-by-side with GT (GT | Pred | Error):
plotter = plot_field_comparisons(mesh, true_fields=["pMeanTrim"], pred_fields=["pMeanTrimPred"],
                                  view="xy", dtype="cell", window_size=[1800, 600])
plotter.screenshot("comparison.png")
plotter.close()

Step 6: Make permanent (optional)

Save the wrapper to physicsnemo/cfd/evaluation/models/wrappers/my_model.py and register in wrappers/__init__.py:

from physicsnemo.cfd.evaluation.models.wrappers.my_model import MyModelWrapper
register_model("my_model", MyModelWrapper)

Then use model.name: my_model in any YAML config.

Gotchas

• DistributedManager: Model wrappers may call DistributedManager.initialize(). In notebooks without torchrun, set env vars first: WORLD_SIZE=1, RANK=0, LOCAL_RANK=0, MASTER_ADDR=localhost, MASTER_PORT=12355.
• weights_only=True: Use this flag with torch.load() for safe deserialization (PyTorch 2.6+ default).
• Domain matching: The engine checks that model.INFERENCE_DOMAIN matches the dataset adapter's inference_domain_from_kwargs(). Mismatches are skipped in matrix mode or raise in single mode.
• GT alignment: When align_ground_truth_to_model: true in dataset kwargs, the engine converts GT data to match OUTPUT_LOCATION (point ↔ cell). This is automatic — the wrapper just needs correct class vars.
• Results JSON format: benchmark_results.json is a plain list[dict], not {"results": [...]}. Iterate directly: for combo in report:.

Related resources

• references/example_wrapper.py — complete surface + volume CFDModel templates to copy and adapt (bundled; available without the repo on disk).
• assets/global_stats.example.json — sample mean/std stats layout for surface and volume.