Physicsnemo Cfd Create Model Wrapper

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).

Uncertainty-quantification (UQ) models set two more class variables (SUPPORTS_UQ, UQ_METHOD) and implement extra hooks — decode_distribution (closed-form), or for sampling a stochastic predict plus predict_deterministic (and optionally predict_ensemble). This is fully additive: deterministic wrappers leave the defaults (SUPPORTS_UQ=False) and are unaffected. See the "Uncertainty quantification" section below.

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.