diff --git a/.gitignore b/.gitignore index 7c368e9..abc1240 100644 --- a/.gitignore +++ b/.gitignore @@ -21,4 +21,7 @@ venv/ # Editor / IDE .cursor/ .vscode/ -.idea/ \ No newline at end of file +.idea/ + +# Local, machine-specific benchmark configs (real paths); not for the repo +*.local.yaml \ No newline at end of file diff --git a/CHANGELOG.md b/CHANGELOG.md index 57d4877..ce15a96 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -17,9 +17,52 @@ Versioning](https://semver.org/spec/v2.0.0.html). **`physicsnemo-cfd-create-custom-metric`** — each with a `SKILL.md` and `evals/` cases to guide adding a new `CFDModel`, `DatasetAdapter`, or metric. +- **DrivAerStar surface dataset adapter** (`datasets/adapters/drivaerstar.py`, + in-memory transforms with opt-in disk caching) and **`DatasetConfig.label`** + for decoupling report display names from adapter lookup / caching. +- **Probabilistic / uncertainty-quantification (UQ) benchmarking:** + additive UQ path across the evaluation stack. + - **`CFDModel` capability model:** `SUPPORTS_UQ` / `UQ_METHOD` + (`"none"` | `"analytic"` | `"sampling"`) class flags plus an optional + `decode_distribution` hook; deterministic wrappers are unchanged + (`SUPPORTS_UQ=False` -> UQ metrics report `NaN`). + - **`FieldDistribution`** payload (`datasets/schema.py`) carrying + `mean` + `std` / `epistemic_std` / `aleatoric_std` (and optional + `samples` / `quantiles`), with `build_predictive_distribution` and + `as_distribution` helpers; all in physical units. + - **Engine sampling loop** (`benchmarks/uq_inference.py`): `run.uq` + controls (`enabled`, `num_samples`, `retain_samples`, `device_metrics`), + streaming Welford aggregation of `N` stochastic passes, and an optional + `predict_ensemble(model_input, n)` fast-path for multi-member models. + - **Pooled + sample UQ metrics** (`metrics/builtin/uq.py`): `nlpd`, + `nlpd_epistemic`, `calibration_zrms`, `coverage_95`, `sharpness_std`, + `sharpness_epistemic_std`, `uncertainty_error_spearman` + (+ `_epistemic`), `sparsification_ause` (+ `_epistemic`), and `drag_uq` + (configurable `coeff` / `drag_direction`), implemented via + reducer/sample-metric protocols so statistics pool correctly across + cases and ranks. + - **`sparsification_plot`** report visual, and std / epistemic-std + companion arrays exported to inference / comparison meshes + (`output.std_mesh_field_names`, `output.epistemic_std_mesh_field_names`). + - **Example UQ model wrappers** (GeoTransolver family, one per + archetype): `geotransolver_gp_surface` (analytic GP field head), + `geotransolver_mc_dropout_surface` (Concrete-Dropout sampling), and + `geotransolver_ensemble_surface` (explicit multi-member ensemble), + alongside the deterministic `geotransolver_drivaerstar_surface`. Shared + inference plumbing in + `models/common_wrapper_utils/geotransolver_runtime.py`. + - **Example config:** `workflows/benchmarking/conf/config_uq_surface.yaml` + scores deterministic + analytic GP + MC-Dropout + ensemble on the same + cases. ### Changed +- **`physicsnemo-cfd-create-model-wrapper` skill:** documents the UQ + wrapper contract (capability flags, `FieldDistribution`, analytic vs. + sampling archetypes, the engine sampling loop, and UQ metrics/config), + with `ExampleAnalyticUQWrapper` / `ExampleSamplingUQWrapper` reference + implementations. + ### Deprecated ### Removed diff --git a/physicsnemo/cfd/evaluation/benchmarks/distributed_utils.py b/physicsnemo/cfd/evaluation/benchmarks/distributed_utils.py index 7af10d4..db50664 100644 --- a/physicsnemo/cfd/evaluation/benchmarks/distributed_utils.py +++ b/physicsnemo/cfd/evaluation/benchmarks/distributed_utils.py @@ -85,6 +85,12 @@ def merge_benchmark_result_shards(shards: list[dict[str, Any]]) -> dict[str, Any f"Distributed merge mismatch: expected {model!r}×{dataset!r}, " f"got {s['model']!r}×{s['dataset']!r}" ) + from physicsnemo.cfd.evaluation.benchmarks.uq_inference import ( + finalize_reducer_metrics, + finalize_sample_metrics, + is_uq_partial_key, + ) + per_case: list[dict[str, Any]] = [] for s in active: per_case.extend(s.get("per_case") or []) @@ -95,8 +101,22 @@ def merge_benchmark_result_shards(shards: list[dict[str, Any]]) -> dict[str, Any all_metric_values.setdefault(k, []).append(float(v)) metrics_summary: dict[str, float] = {} for mname, values in all_metric_values.items(): + # Reducer / sample statistics are finalized (pooled / collected), not mean-over-cases. + if is_uq_partial_key(mname): + continue valid = [v for v in values if v == v] metrics_summary[mname] = sum(valid) / len(valid) if valid else float("nan") + inference_domain = active[0].get("inference_domain") + metrics_summary.update(finalize_reducer_metrics(per_case, inference_domain)) + metrics_summary.update(finalize_sample_metrics(per_case, inference_domain)) + # Preserve NaN placeholders for configured-but-unavailable UQ metrics: each shard already + # finalized them (via the configured metric names in _run_single_model_dataset), so carry any + # summary key missing from the merged set forward as-is. setdefault never clobbers a value + # recomputed above from real partials. + for s in active: + for k, v in (s.get("metrics") or {}).items(): + if isinstance(v, (int, float)): + metrics_summary.setdefault(k, float(v)) case_ids_raw = [str(r["case_id"]) for r in per_case if r.get("case_id") is not None] case_ids_no_dup = sorted(set(case_ids_raw), key=natural_sort_key) return { @@ -105,6 +125,7 @@ def merge_benchmark_result_shards(shards: list[dict[str, Any]]) -> dict[str, Any "cases": case_ids_no_dup, "metrics": metrics_summary, "per_case": per_case, + "inference_domain": inference_domain, } diff --git a/physicsnemo/cfd/evaluation/benchmarks/engine.py b/physicsnemo/cfd/evaluation/benchmarks/engine.py index 4e3a844..901dfed 100644 --- a/physicsnemo/cfd/evaluation/benchmarks/engine.py +++ b/physicsnemo/cfd/evaluation/benchmarks/engine.py @@ -33,6 +33,15 @@ replace mesh or visualization workflows. The cache fingerprint includes ``run.seed`` (influences subsampling / ``randperm`` RNG in model wrappers). +When ``save_inference_mesh`` is enabled, per-case ``inference___.vt[p|u]`` +meshes are written for every scored case unless ``reports.visual_case_ids`` is set — in which case +only those ids get a mesh (all other cases are still scored for metrics). This keeps large validation +runs from dumping one VTP per case when only a couple are needed for inspection / report visuals. +The dataset label is embedded so multi-dataset sweeps that reuse case ids (e.g. per body-style +classes) do not overwrite each other. ``reports.save_comparison_meshes`` writes richer +``___comparison.vt[p|u]`` files (prediction + ground truth + std side by side) +and honours the same ``visual_case_ids`` gating. + When ``save_inference_mesh`` is enabled but exporting ``inference__.vt[p|u]`` fails, the full traceback is logged and persisted for audit: ``per_case[]`` keys and ``benchmark_artifacts.json`` (``inference_mesh_write_failures``, ``comparison_mesh_build_failures``, ``comparison_mesh_save_failures``) @@ -98,13 +107,32 @@ class BenchmarkPolicyError(RuntimeError): from physicsnemo.cfd.evaluation.common.inference_seed import seed_inference_rng from physicsnemo.cfd.evaluation.common.natural_sort import natural_sorted from physicsnemo.cfd.evaluation.datasets.progress import log_dataset -from physicsnemo.cfd.evaluation.datasets.schema import normalize_inference_domain_str +from physicsnemo.cfd.evaluation.datasets.schema import ( + FieldDistribution, + distribution_mean, + normalize_inference_domain_str, +) from physicsnemo.cfd.evaluation.models import get_model_wrapper from physicsnemo.cfd.evaluation.models.model_registry import ( get_inference_domain_for_model, ) from physicsnemo.cfd.evaluation.metrics import get_metric from physicsnemo.cfd.evaluation.metrics.mesh_bridge import build_comparison_mesh +from physicsnemo.cfd.postprocessing_tools.metric_registry import ( + is_reducer_metric, + is_sample_metric, +) +from physicsnemo.cfd.evaluation.benchmarks.uq_inference import ( + compute_sparsification_payload, + finalize_reducer_metrics, + finalize_sample_metrics, + is_uq_partial_key, + make_reducer_partial_key, + make_sample_partial_key, + run_sampling_inference, + select_inference_path, + strip_reducer_partials, +) # Recoverable VTK / NumPy / mesh_bridge failures. Avoid bare ``except Exception`` — # unexpected subclasses propagate so regressions are not mistaken for metric NaNs. @@ -208,6 +236,24 @@ def _retain_comparison_mesh_for_visual_context( return case_id in allow +def _sanitize_path_token(token: str) -> str: + """Make a string safe to embed in an output filename (drop path/whitespace chars).""" + return "".join(c if (c.isalnum() or c in "-.") else "_" for c in str(token)) + + +def _save_inference_mesh_for_case(reports: ReportsConfig | None, case_id: str) -> bool: + """Whether this case should write an ``inference__`` mesh. + + Saving every case's mesh is wasteful for large validation sets. When + ``reports.visual_case_ids`` is set, restrict inference-mesh writes to exactly those cases + (the ones you inspect / that the report visuals use); all other cases are scored for + metrics only. When it is ``None`` there is no restriction (write every case, back-compat). + """ + if reports is None or reports.visual_case_ids is None: + return True + return case_id in reports.visual_case_ids + + def _normalize_metrics_config( metrics: list[str | dict[str, Any]], ) -> list[tuple[str, dict]]: @@ -265,12 +311,14 @@ def _save_inference_mesh_if_requested( run_config: RunConfig, model_config: ModelConfig, output_config: OutputConfig, + reports: ReportsConfig | None, wrapper: Any, case: Any, case_id: str, predictions: dict[str, Any], output_dir: str, dataset_name: str, + dataset_label: str | None = None, ) -> str | None: """ Write ``inference__.vtp`` or ``.vtu`` when requested. @@ -307,13 +355,21 @@ def _save_inference_mesh_if_requested( """ if not run_config.save_inference_mesh: return None + if not _save_inference_mesh_for_case(reports, case_id): + return None import pyvista as pv m_dom = case.inference_domain - out_path = ( - Path(output_dir) - / f"inference_{model_config.name}_{case_id}{'.vtp' if m_dom == 'surface' else '.vtu'}" + # Include the dataset label so multi-dataset sweeps (e.g. per body-style classes that reuse the + # same numeric case ids) do not overwrite one another's meshes. Falls back to model+case only. + ext = ".vtp" if m_dom == "surface" else ".vtu" + label_tok = _sanitize_path_token(dataset_label) if dataset_label else "" + stem = ( + f"inference_{model_config.name}_{label_tok}_{case_id}" + if label_tok + else f"inference_{model_config.name}_{case_id}" ) + out_path = Path(output_dir) / f"{stem}{ext}" log_dataset( dataset_name, f"Writing inference mesh (predictions only) to {out_path}…", @@ -341,12 +397,31 @@ def _save_inference_mesh_if_requested( mesh = mesh.cast_to_unstructured_grid() names = output_config.volume_mesh_field_names + if m_dom == "surface": + std_names = output_config.std_mesh_field_names + epi_names = output_config.epistemic_std_mesh_field_names + else: + std_names = output_config.std_volume_mesh_field_names + epi_names = output_config.epistemic_std_volume_mesh_field_names + data_target = ( mesh.cell_data if wrapper.output_location == "cell" else mesh.point_data ) for canonical_key, mesh_name in names.items(): - if canonical_key in predictions: - data_target[mesh_name] = predictions[canonical_key] + if canonical_key not in predictions: + continue + value = predictions[canonical_key] + data_target[mesh_name] = distribution_mean(value) + # Attach uncertainty companions so exported meshes carry UQ for ParaView / visuals. + if isinstance(value, FieldDistribution): + if value.std is not None: + data_target[std_names.get(canonical_key) or f"{mesh_name}Std"] = ( + value.std + ) + if value.epistemic_std is not None: + data_target[ + epi_names.get(canonical_key) or f"{mesh_name}EpistemicStd" + ] = value.epistemic_std mesh.save(str(out_path)) log_dataset(dataset_name, f"Wrote inference mesh: {out_path}") except _MESH_IO_BRIDGE_ERRORS: @@ -412,6 +487,34 @@ def _call_metric( return fn(gt, predictions, **mkwargs) +def _call_reducer_partial( + metric: Any, + gt: dict, + predictions: dict, + *, + case: Any, + comparison_mesh: Any, + metric_dtype: str | None, + output: OutputConfig, + mkwargs: dict[str, Any], +) -> dict[str, float]: + """Invoke a reducer metric's ``partial`` with extended kwargs, falling back to the basics. + + Returns the per-case extensive sufficient statistics (additive across cases). + """ + extended = dict(mkwargs) + extended.update( + case=case, + comparison_mesh=comparison_mesh, + metric_dtype=metric_dtype, + output=output, + ) + try: + return metric.partial(gt, predictions, **extended) + except TypeError: + return metric.partial(gt, predictions, **mkwargs) + + def _run_single( model_config: ModelConfig, dataset_config: DatasetConfig, @@ -471,6 +574,10 @@ def _run_single( and ``mesh_ctx`` mapping case id -> comparison mesh for visuals. """ adapter_class = get_adapter(dataset_config.name) + # Adapter is resolved from ``name``; the result rows are keyed by ``display_name`` (== label or + # name) so the same adapter at different roots can appear as distinct dataset rows (e.g. one + # DrivAerStar adapter over estateback / fastback / notchback). Caching stays on ``name`` + root. + ds_label = dataset_config.display_name m_dom = _effective_inference_domain(model_config) d_dom = adapter_class.inference_domain_from_kwargs(dataset_config.kwargs) if m_dom != d_dom: @@ -486,7 +593,7 @@ def _run_single( return ( { "model": model_config.name, - "dataset": dataset_config.name, + "dataset": ds_label, "skipped": True, "skip_reason": reason, "cases": [], @@ -524,7 +631,7 @@ def _run_single( return ( { "model": model_config.name, - "dataset": dataset_config.name, + "dataset": ds_label, "cases": [], "metrics": {}, "per_case": [], @@ -559,6 +666,11 @@ def _run_single( output_dict=output_config_to_fingerprint_dict(output_config), metric_specs=metric_names, run_seed=run_config.seed, + run_uq={ + "enabled": run_config.uq.enabled, + "num_samples": run_config.uq.num_samples, + "retain_samples": run_config.uq.retain_samples, + }, ) log_dataset( dataset_config.name, @@ -636,20 +748,48 @@ def _run_single( case_cache[case_key] = case seed_inference_rng(run_config.seed, cid) model_input = wrapper.prepare_inputs(case) - raw = wrapper.predict(model_input) - predictions = wrapper.decode_outputs(raw, case, model_input) + # ``run.uq.enabled`` is the master switch: when off, EVERY wrapper (sampling AND analytic) + # takes the deterministic path — a single ``predict_deterministic`` + ``decode_outputs`` — + # so no distributions are emitted and no UQ metrics are produced, enabling apples-to-apples + # deterministic comparison runs. ``predict_deterministic`` (not ``predict``) is used so a + # stochastic sampler (e.g. MC-Dropout) returns a true point prediction here rather than one + # random draw. See :func:`select_inference_path`. + inference_path = select_inference_path( + supports_uq=bool(getattr(wrapper, "SUPPORTS_UQ", False)), + uq_method=getattr(wrapper, "UQ_METHOD", "none"), + uq_enabled=run_config.uq.enabled, + ) + if inference_path == "sampling": + # N stochastic passes; prepare_inputs already ran once (only the forward is repeated). + predictions = run_sampling_inference( + wrapper, + case, + model_input, + n=run_config.uq.num_samples, + run_seed=run_config.seed, + case_id=cid, + retain_samples=run_config.uq.retain_samples, + ) + elif inference_path == "analytic": + raw = wrapper.predict(model_input) + predictions = wrapper.decode_distribution(raw, case, model_input) + else: + raw = wrapper.predict_deterministic(model_input) + predictions = wrapper.decode_outputs(raw, case, model_input) gt = case.ground_truth or {} inference_mesh_err = _save_inference_mesh_if_requested( run_config=run_config, model_config=model_config, output_config=output_config, + reports=reports, wrapper=wrapper, case=case, case_id=cid, predictions=predictions, output_dir=output_dir, dataset_name=dataset_config.name, + dataset_label=dataset_config.display_name, ) comparison_mesh = None @@ -670,6 +810,42 @@ def _run_single( for mname, mkwargs in metric_names: metric_fn = get_metric(mname, domain=m_dom) try: + if is_sample_metric(metric_fn): + # Sample metric: store per-geometry scalars under a reserved key. Collected + # (not summed) across cases + ranks and finalized after the loop / merge. + partial_stats = _call_reducer_partial( + metric_fn, + gt, + predictions, + case=case, + comparison_mesh=comparison_mesh, + metric_dtype=metric_dtype, + output=output_config, + mkwargs=mkwargs, + ) + for pkey, pval in partial_stats.items(): + rk = make_sample_partial_key(mname, pkey) + case_metrics[rk] = float(pval) + all_metric_values.setdefault(rk, []).append(float(pval)) + continue + if is_reducer_metric(metric_fn): + # Pooled metric: store per-case additive sufficient statistics under a + # reserved key so they cache + merge like scalars; finalized after the loop. + partial_stats = _call_reducer_partial( + metric_fn, + gt, + predictions, + case=case, + comparison_mesh=comparison_mesh, + metric_dtype=metric_dtype, + output=output_config, + mkwargs=mkwargs, + ) + for pkey, pval in partial_stats.items(): + rk = make_reducer_partial_key(mname, pkey) + case_metrics[rk] = float(pval) + all_metric_values.setdefault(rk, []).append(float(pval)) + continue out = _call_metric( metric_fn, gt, @@ -711,11 +887,16 @@ def _run_single( if comparison_mesh is not None and metric_dtype is not None: row["metric_dtype"] = metric_dtype if reports: - if reports.save_comparison_meshes: + if reports.save_comparison_meshes and _save_inference_mesh_for_case( + reports, cid + ): sub = Path(output_dir) / reports.comparison_mesh_subdir sub.mkdir(parents=True, exist_ok=True) ext = ".vtp" if case.inference_domain == "surface" else ".vtu" - cmp_p = sub / f"{cid}_comparison{ext}" + # Disambiguate by model + dataset label: the comparison mesh carries this model's + # predictions, and case ids repeat across body-style classes. + _lbl = _sanitize_path_token(dataset_config.display_name) + cmp_p = sub / f"{model_config.name}_{_lbl}_{cid}_comparison{ext}" try: comparison_mesh.save(str(cmp_p)) row["comparison_mesh_path"] = str(cmp_p.resolve()) @@ -747,19 +928,36 @@ def _run_single( f"Could not write metrics cache for case {cid!r}: {ex}", ) - # Aggregate (mean over cases) + # Aggregate (mean over cases) for pointwise metrics; reducer / sample partials (reserved + # keys) are finalized separately below (pooled over points, resp. collected over geometries). metrics_summary = {} for mname, values in all_metric_values.items(): + if is_uq_partial_key(mname): + continue valid = [v for v in values if v == v] # filter nan metrics_summary[mname] = sum(valid) / len(valid) if valid else float("nan") + # Pooled reducer + sample-wise metrics. In distributed runs these are recomputed after the + # merge on rank 0 (merge_benchmark_result_shards) from the merged per-case partials; this + # local pass keeps single-process runs correct. Pass the configured metric names so a + # deterministic wrapper (no UQ partials) still reports configured UQ metrics as NaN rather + # than omitting them (consistent schema; fail_on_any_metric_nan can flag them). + configured_metric_names = [mname for mname, _ in metric_names] + metrics_summary.update( + finalize_reducer_metrics(per_case, m_dom, configured_metric_names) + ) + metrics_summary.update( + finalize_sample_metrics(per_case, m_dom, configured_metric_names) + ) + return ( { "model": model_config.name, - "dataset": dataset_config.name, + "dataset": ds_label, "cases": cases, "metrics": metrics_summary, "per_case": per_case, + "inference_domain": m_dom, }, mesh_ctx, ) @@ -899,6 +1097,24 @@ def run_benchmark( dm, results, meshes_by_run ) + # Reducer sufficient statistics (reserved ``_uq::`` keys) have been folded into each run's + # ``metrics`` summary; drop them from the reported per-case rows so outputs show real values. + # Before stripping, harvest the sample-metric per-geometry curves for the sparsification visual. + # These are kept in a side list aligned with ``results`` (not on the result dicts) so the numpy + # curve arrays never reach the JSON/CSV/HTML report serializers. + sparsification_by_run: list[dict[str, Any]] = [] + for r in results: + if r.get("skipped"): + sparsification_by_run.append({}) + continue + sparsification_by_run.append( + compute_sparsification_payload( + r.get("per_case") or [], r.get("inference_domain") + ) + ) + for r in results: + strip_reducer_partials(r.get("per_case") or []) + if is_rank0 and config.benchmark.reproducibility.save_artifacts: artifacts = Path(output_dir) / "benchmark_artifacts.json" skipped = [r for r in results if r.get("skipped")] @@ -950,7 +1166,10 @@ def run_benchmark( config, results, output_dir, - context={"comparison_meshes_by_run": meshes_by_run}, + context={ + "comparison_meshes_by_run": meshes_by_run, + "uq_sparsification_by_run": sparsification_by_run, + }, ) _enforce_benchmark_policy(config, results) @@ -1025,6 +1244,12 @@ def _config_to_dict(c: Config) -> dict: "enabled": c.run.metrics_cache.enabled, "path": c.run.metrics_cache.path, }, + "uq": { + "enabled": c.run.uq.enabled, + "num_samples": c.run.uq.num_samples, + "retain_samples": c.run.uq.retain_samples, + "device_metrics": c.run.uq.device_metrics, + }, "distributed": c.run.distributed, "fail_on_all_skipped": c.run.fail_on_all_skipped, "fail_on_any_metric_nan": c.run.fail_on_any_metric_nan, @@ -1038,6 +1263,7 @@ def _config_to_dict(c: Config) -> dict: }, "dataset": { "name": c.dataset.name, + "label": c.dataset.label, "root": c.dataset.root, "case_ids": c.dataset.case_ids, "kwargs": c.dataset.kwargs, @@ -1047,6 +1273,12 @@ def _config_to_dict(c: Config) -> dict: "volume_mesh_field_names": c.output.volume_mesh_field_names, "ground_truth_mesh_field_names": c.output.ground_truth_mesh_field_names, "ground_truth_volume_mesh_field_names": c.output.ground_truth_volume_mesh_field_names, + # UQ uncertainty array-name maps (drive the std / epistemic-std companions attached to + # comparison meshes and read back by drag_uq); serialized for reproducibility. + "std_mesh_field_names": c.output.std_mesh_field_names, + "epistemic_std_mesh_field_names": c.output.epistemic_std_mesh_field_names, + "std_volume_mesh_field_names": c.output.std_volume_mesh_field_names, + "epistemic_std_volume_mesh_field_names": c.output.epistemic_std_volume_mesh_field_names, "streamlines_vector_canonical": c.output.streamlines_vector_canonical, }, "metrics": c.metrics, diff --git a/physicsnemo/cfd/evaluation/benchmarks/metrics_cache.py b/physicsnemo/cfd/evaluation/benchmarks/metrics_cache.py index 26251c8..2ae3c9f 100644 --- a/physicsnemo/cfd/evaluation/benchmarks/metrics_cache.py +++ b/physicsnemo/cfd/evaluation/benchmarks/metrics_cache.py @@ -92,6 +92,7 @@ def metrics_cache_fingerprint( output_dict: dict[str, Any], metric_specs: list[tuple[str, dict[str, Any]]], run_seed: int = 42, + run_uq: dict[str, Any] | None = None, ) -> str: """ Build a SHA-256 fingerprint for cache lookup and invalidation. @@ -170,7 +171,7 @@ def metrics_cache_fingerprint( }, "output": _fingerprint_jsonify(output_dict), "metrics": specs_serializable, - "run": {"seed": int(run_seed)}, + "run": {"seed": int(run_seed), "uq": _fingerprint_jsonify(run_uq or {})}, } return hashlib.sha256(_stable_json(payload).encode("utf-8")).hexdigest() diff --git a/physicsnemo/cfd/evaluation/benchmarks/report.py b/physicsnemo/cfd/evaluation/benchmarks/report.py index 1e0db6c..fe1ef67 100644 --- a/physicsnemo/cfd/evaluation/benchmarks/report.py +++ b/physicsnemo/cfd/evaluation/benchmarks/report.py @@ -31,7 +31,12 @@ def _sanitize_for_strict_json(obj: Any) -> Any: - """Recursively replace NaN/Inf floats with None for RFC 8259–valid JSON.""" + """Recursively coerce a result payload into RFC 8259–valid JSON. + + Replaces NaN/Inf floats with ``None`` and defensively converts NumPy scalars/arrays + (which ``json.dump`` cannot serialize) into plain Python types so a stray array in a + metric payload can never crash the whole report. + """ if isinstance(obj, float): return None if (math.isnan(obj) or math.isinf(obj)) else obj if isinstance(obj, dict): @@ -40,6 +45,15 @@ def _sanitize_for_strict_json(obj: Any) -> Any: return [_sanitize_for_strict_json(x) for x in obj] if isinstance(obj, tuple): return tuple(_sanitize_for_strict_json(x) for x in obj) + # NumPy types are not JSON-serializable; coerce to native Python and recurse. + if ( + hasattr(obj, "item") + and not isinstance(obj, type) + and getattr(obj, "ndim", None) == 0 + ): + return _sanitize_for_strict_json(obj.item()) + if hasattr(obj, "tolist") and getattr(obj, "ndim", None) is not None: + return _sanitize_for_strict_json(obj.tolist()) return obj diff --git a/physicsnemo/cfd/evaluation/benchmarks/uq_inference.py b/physicsnemo/cfd/evaluation/benchmarks/uq_inference.py new file mode 100644 index 0000000..6e094dd --- /dev/null +++ b/physicsnemo/cfd/evaluation/benchmarks/uq_inference.py @@ -0,0 +1,435 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Engine-side UQ plumbing: the sampling inference loop and pooled reducer finalization. + +1. :func:`run_sampling_inference` drives ``N`` stochastic passes for a + ``UQ_METHOD="sampling"`` wrapper (MC-Dropout, ensembles) and + streams them into a :class:`~physicsnemo.cfd.evaluation.datasets.schema.FieldDistribution` + per field via Welford aggregation (no ``N`` full fields held in memory). ``prepare_inputs`` + is *not* repeated — only the forward pass (and per-pass decode). + +2. The pooled **reducer** metrics emit per-case additive sufficient + statistics. The engine stores each such statistic in the per-case ``metrics`` dict under a + reserved ``_uq::`` prefix so it flows through the existing per-case cache and distributed + merge (both handle per-case scalars). :func:`finalize_reducer_metrics` sums those across all + cases and calls each metric's ``finalize`` once to get the pooled value(s). +""" + +from __future__ import annotations + +from typing import Any + +import numpy as np + +from physicsnemo.cfd.evaluation.common.inference_seed import seed_inference_rng +from physicsnemo.cfd.evaluation.datasets.schema import ( + FieldDistribution, + build_predictive_distribution, +) + +# -------------------------------------------------------------------------------------------- +# Reducer partial statistics: reserved per-case keys +# -------------------------------------------------------------------------------------------- + +#: Per-case ``metrics`` keys carrying a reducer metric's extensive sufficient statistics are +#: namespaced with this prefix so they (a) cache like any per-case scalar, (b) survive the +#: distributed merge, and (c) are excluded from the deterministic mean-over-cases aggregation +#: and from the final per-case report rows. +REDUCER_PARTIAL_PREFIX = "_uq::" + + +def make_reducer_partial_key(metric_name: str, partial_key: str) -> str: + """Reserved per-case key for one reducer statistic (``_uq::{metric}::{stat}``).""" + return f"{REDUCER_PARTIAL_PREFIX}{metric_name}::{partial_key}" + + +def is_reducer_partial_key(key: str) -> bool: + """True for keys produced by :func:`make_reducer_partial_key`.""" + return key.startswith(REDUCER_PARTIAL_PREFIX) + + +def _split_reducer_partial_key(key: str) -> tuple[str, str]: + """Split ``_uq::{metric}::{stat}`` into ``(metric_name, stat_key)``.""" + body = key[len(REDUCER_PARTIAL_PREFIX) :] + metric_name, _, partial_key = body.partition("::") + return metric_name, partial_key + + +def finalize_reducer_metrics( + per_case_rows: list[dict[str, Any]], + domain: str | None, + configured_metric_names: list[str] | None = None, +) -> dict[str, float]: + """Sum reducer sufficient statistics over cases and finalize to pooled metric value(s). + + Reads the reserved ``_uq::`` keys off each row's ``metrics`` dict, sums them per + ``(metric, stat)``, resolves each metric from the registry, and calls ``finalize``. Dict + returns expand to ``f"{metric}_{subkey}"`` (matching the engine's pointwise expansion). + Returns the mapping of finalized metric key -> value to merge into the run summary. + + ``configured_metric_names`` (the metrics requested in the run config) drives NaN placeholders: + any configured reducer metric that produced **no** partials this run — e.g. a deterministic + wrapper emits no ``_uq::`` keys — is still finalized via ``finalize({})`` so the row reports + ``nlpd``/``coverage``/… as ``NaN`` rather than silently omitting them (consistent report + schema; ``fail_on_any_metric_nan`` can then flag an unavailable configured metric). + """ + # Resolve from the lightweight registry (no torch-backed metrics-package import needed; + # the metric instances were registered there when the builtin metrics loaded). + from physicsnemo.cfd.postprocessing_tools.metric_registry import ( + get_metric, + is_reducer_metric, + ) + + summed: dict[str, dict[str, float]] = {} + for row in per_case_rows: + for key, val in (row.get("metrics") or {}).items(): + if not is_reducer_partial_key(key): + continue + metric_name, partial_key = _split_reducer_partial_key(key) + if not partial_key: + continue + summed.setdefault(metric_name, {}) + summed[metric_name][partial_key] = summed[metric_name].get( + partial_key, 0.0 + ) + float(val) + + # Finalize every metric that produced partials, plus any configured reducer metric that did + # not (with empty stats -> NaN) so deterministic rows keep a consistent schema. + stats_by_metric: dict[str, dict[str, float]] = dict(summed) + for name in configured_metric_names or (): + stats_by_metric.setdefault(name, {}) + + out: dict[str, float] = {} + for metric_name, stats in stats_by_metric.items(): + try: + metric = get_metric(metric_name, domain=domain) + except KeyError: + continue + if not is_reducer_metric(metric): + continue + result = metric.finalize(stats) + if isinstance(result, dict): + for sub, v in result.items(): + out[f"{metric_name}_{sub}" if sub else metric_name] = float(v) + else: + out[metric_name] = float(result) + return out + + +# -------------------------------------------------------------------------------------------- +# Sample metric partials: reserved per-case keys (collected, not summed) +# -------------------------------------------------------------------------------------------- + +#: Per-case keys carrying a :class:`SampleMetric`'s per-geometry scalars. Distinct from the +#: reducer prefix so the sum-based finalize ignores them and the collect-based finalize picks +#: them up. ``"_uqs::".startswith("_uq::")`` is False, so the two namespaces are disjoint. +SAMPLE_PARTIAL_PREFIX = "_uqs::" + + +def make_sample_partial_key(metric_name: str, partial_key: str) -> str: + """Reserved per-case key for one sample-metric scalar (``_uqs::{metric}::{stat}``).""" + return f"{SAMPLE_PARTIAL_PREFIX}{metric_name}::{partial_key}" + + +def is_sample_partial_key(key: str) -> bool: + """True for keys produced by :func:`make_sample_partial_key`.""" + return key.startswith(SAMPLE_PARTIAL_PREFIX) + + +def _split_sample_partial_key(key: str) -> tuple[str, str]: + """Split ``_uqs::{metric}::{stat}`` into ``(metric_name, stat_key)``.""" + body = key[len(SAMPLE_PARTIAL_PREFIX) :] + metric_name, _, partial_key = body.partition("::") + return metric_name, partial_key + + +def finalize_sample_metrics( + per_case_rows: list[dict[str, Any]], + domain: str | None, + configured_metric_names: list[str] | None = None, +) -> dict[str, float]: + """Collect sample-metric per-geometry scalars over cases and finalize (e.g. sample-wise AUSE). + + Unlike :func:`finalize_reducer_metrics`, per-case values are **collected into lists** (not + summed): each ``_uqs::`` key contributes one value per case, appended in ``per_case`` order. + ``finalize_samples`` maps the per-key lists to the final value(s); dict returns expand to + ``f"{metric}_{subkey}"``. + + ``configured_metric_names`` drives NaN placeholders the same way as + :func:`finalize_reducer_metrics`: a configured sample metric with no partials this run (e.g. + a deterministic wrapper) is finalized with an empty mapping so AUSE / drag_uq report ``NaN`` + instead of being dropped from the row. + """ + from physicsnemo.cfd.postprocessing_tools.metric_registry import ( + get_metric, + is_sample_metric, + ) + + collected: dict[str, dict[str, list[float]]] = {} + for row in per_case_rows: + for key, val in (row.get("metrics") or {}).items(): + if not is_sample_partial_key(key): + continue + metric_name, partial_key = _split_sample_partial_key(key) + if not partial_key: + continue + collected.setdefault(metric_name, {}).setdefault(partial_key, []).append( + float(val) + ) + + stats_by_metric: dict[str, dict[str, list[float]]] = dict(collected) + for name in configured_metric_names or (): + stats_by_metric.setdefault(name, {}) + + out: dict[str, float] = {} + for metric_name, stats in stats_by_metric.items(): + try: + metric = get_metric(metric_name, domain=domain) + except KeyError: + continue + if not is_sample_metric(metric): + continue + result = metric.finalize_samples(stats) + if isinstance(result, dict): + for sub, v in result.items(): + out[f"{metric_name}_{sub}" if sub else metric_name] = float(v) + else: + out[metric_name] = float(result) + return out + + +def compute_sparsification_payload( + per_case_rows: list[dict[str, Any]], domain: str | None +) -> dict[str, dict[str, Any]]: + """Build the per-metric sparsification curves consumed by the ``sparsification_plot`` visual. + + Collects each sample metric's per-geometry scalars (the same reserved ``_uqs::`` keys that + :func:`finalize_sample_metrics` uses), and for every sample metric exposing a ``curves`` method + calls it once to produce ``{series_name: {fractions, by_uncertainty, oracle, full, ause, n}}``. + Returns ``{metric_name: {series_name: curve_dict}}`` (empty when no sample metrics ran). Must be + called *before* :func:`strip_reducer_partials` removes the reserved keys. + """ + from physicsnemo.cfd.postprocessing_tools.metric_registry import ( + get_metric, + is_sample_metric, + ) + + collected: dict[str, dict[str, list[float]]] = {} + for row in per_case_rows: + for key, val in (row.get("metrics") or {}).items(): + if not is_sample_partial_key(key): + continue + metric_name, partial_key = _split_sample_partial_key(key) + if not partial_key: + continue + collected.setdefault(metric_name, {}).setdefault(partial_key, []).append( + float(val) + ) + + out: dict[str, dict[str, Any]] = {} + for metric_name, stats in collected.items(): + try: + metric = get_metric(metric_name, domain=domain) + except KeyError: + continue + if not is_sample_metric(metric): + continue + curves_fn = getattr(metric, "curves", None) + if not callable(curves_fn): + continue + series = curves_fn(stats) + if series: + out[metric_name] = series + return out + + +def is_uq_partial_key(key: str) -> bool: + """True for any reserved UQ per-case key (reducer sufficient stat or sample-metric scalar).""" + return is_reducer_partial_key(key) or is_sample_partial_key(key) + + +def select_inference_path( + *, supports_uq: bool, uq_method: str, uq_enabled: bool +) -> str: + """Engine per-case dispatch: ``"sampling"`` | ``"analytic"`` | ``"deterministic"``. + + ``uq_enabled`` (``run.uq.enabled``) is the master switch: when it is off, EVERY wrapper takes + the deterministic path regardless of ``SUPPORTS_UQ`` / ``UQ_METHOD`` — so an analytic GP head + is not executed as a distribution and produces no UQ metrics, matching the documented behavior + and enabling apples-to-apples deterministic comparison runs. + """ + if uq_enabled and supports_uq and uq_method == "sampling": + return "sampling" + if uq_enabled and supports_uq and uq_method == "analytic": + return "analytic" + return "deterministic" + + +def strip_reducer_partials(per_case_rows: list[dict[str, Any]]) -> None: + """Remove reserved ``_uq::`` / ``_uqs::`` statistics from per-case ``metrics`` in place. + + Called once after aggregation / merge so the reported per-case rows show only real metric + values, not the internal statistics (already folded into the run summary by + :func:`finalize_reducer_metrics` / :func:`finalize_sample_metrics`). + """ + for row in per_case_rows: + metrics = row.get("metrics") + if not isinstance(metrics, dict): + continue + for key in [k for k in metrics if is_uq_partial_key(k)]: + del metrics[key] + + +# -------------------------------------------------------------------------------------------- +# Sampling inference loop (Welford streaming) +# -------------------------------------------------------------------------------------------- + + +class _Welford: + """Streaming mean / variance across passes for one field (arbitrary array shape). + + Tracks the across-pass (epistemic) statistics of the per-pass point predictions, plus an + optional running sum of per-pass aleatoric variances (when a pass is itself a + distribution), so the total predictive variance can be combined via the law of total + variance: ``total_var = mean_i(sigma_i^2) + var_i(mu_i)``. + """ + + def __init__(self) -> None: + self.n = 0 + self.mean: np.ndarray | None = None + self._m2: np.ndarray | None = None + self._aleatoric_var_sum: np.ndarray | None = None + self._has_aleatoric = False + + def update(self, x: np.ndarray, aleatoric_var: np.ndarray | None = None) -> None: + x = np.asarray(x, dtype=np.float64) + self.n += 1 + if self.mean is None: + self.mean = x.copy() + self._m2 = np.zeros_like(x) + else: + delta = x - self.mean + self.mean += delta / self.n + self._m2 += delta * (x - self.mean) + if aleatoric_var is not None: + av = np.asarray(aleatoric_var, dtype=np.float64) + self._has_aleatoric = True + if self._aleatoric_var_sum is None: + self._aleatoric_var_sum = av.copy() + else: + self._aleatoric_var_sum += av + + def result( + self, + ) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray | None] | None: + """Return ``(mean, total_std, epistemic_std, aleatoric_std_or_None)`` (physical units).""" + if self.n == 0 or self.mean is None: + return None + # Population variance across passes = epistemic uncertainty of the prediction. + epi_var = ( + self._m2 / self.n if self._m2 is not None else np.zeros_like(self.mean) + ) + epi_var = np.clip(epi_var, 0.0, None) + epistemic_std = np.sqrt(epi_var) + if self._has_aleatoric and self._aleatoric_var_sum is not None: + aleatoric_var = self._aleatoric_var_sum / self.n + aleatoric_std = np.sqrt(np.clip(aleatoric_var, 0.0, None)) + total_std = np.sqrt(np.clip(epi_var + aleatoric_var, 0.0, None)) + else: + aleatoric_std = None + total_std = epistemic_std + return self.mean, total_std, epistemic_std, aleatoric_std + + +def _decode_pass_to_arrays( + wrapper: Any, raw: Any, case: Any, model_input: Any +) -> dict[str, tuple[np.ndarray, np.ndarray | None]]: + """Decode one pass to ``{field: (mean_array, aleatoric_var_or_None)}`` (physical units).""" + preds = wrapper.decode_outputs(raw, case, model_input) + out: dict[str, tuple[np.ndarray, np.ndarray | None]] = {} + for key, value in preds.items(): + if isinstance(value, FieldDistribution): + mean = np.asarray(value.mean, dtype=np.float64) + ale_var = None + if value.std is not None: + ale_var = np.asarray(value.std, dtype=np.float64) ** 2 + out[key] = (mean, ale_var) + else: + out[key] = (np.asarray(value, dtype=np.float64), None) + return out + + +def run_sampling_inference( + wrapper: Any, + case: Any, + model_input: Any, + *, + n: int, + run_seed: int, + case_id: str, + retain_samples: bool = False, +) -> dict[str, FieldDistribution]: + """Drive ``n`` stochastic passes and aggregate to a distribution per field. + + Uses ``wrapper.predict_ensemble(model_input, n)`` when it returns an iterable (streamed one + output at a time, e.g. a generator over ensemble members), else calls + ``wrapper.predict(model_input)`` ``n`` times, reseeding per pass from + ``(run_seed, case_id, pass_index)`` so each pass has a distinct, reproducible RNG state. + """ + if n < 1: + raise ValueError(f"num_samples must be >= 1 for sampling inference, got {n}") + + ensemble = wrapper.predict_ensemble(model_input, n) + + accumulators: dict[str, _Welford] = {} + samples: dict[str, list[np.ndarray]] = {} + + def _consume(raw: Any) -> None: + for key, (mean, ale_var) in _decode_pass_to_arrays( + wrapper, raw, case, model_input + ).items(): + accumulators.setdefault(key, _Welford()).update(mean, ale_var) + if retain_samples: + samples.setdefault(key, []).append(mean) + + if ensemble is not None: + for raw in ensemble: + _consume(raw) + else: + for i in range(n): + seed_inference_rng(run_seed, f"{case_id}#pass{i}") + _consume(wrapper.predict(model_input)) + + distributions: dict[str, FieldDistribution] = {} + for key, acc in accumulators.items(): + res = acc.result() + if res is None: + continue + mean, total_std, epistemic_std, aleatoric_std = res + stacked = ( + np.stack(samples[key], axis=0) + if retain_samples and key in samples + else None + ) + distributions[key] = build_predictive_distribution( + mean=mean, + std=total_std, + epistemic_std=epistemic_std, + aleatoric_std=aleatoric_std, + samples=stacked, + ) + return distributions diff --git a/physicsnemo/cfd/evaluation/config.py b/physicsnemo/cfd/evaluation/config.py index 3b3f064..e137ad7 100644 --- a/physicsnemo/cfd/evaluation/config.py +++ b/physicsnemo/cfd/evaluation/config.py @@ -92,6 +92,36 @@ class MetricsCacheConfig: path: str = "" +@dataclass +class UQConfig: + """Uncertainty-quantification controls for probabilistic benchmarking. + + Attributes + ---------- + enabled : bool + Master switch. When ``False`` (default), sampling wrappers run their single + deterministic ``predict`` path and UQ metrics report ``NaN`` — behavior is identical + to the pre-UQ engine. + num_samples : int + Number of stochastic passes / ensemble members the engine drives for + ``UQ_METHOD="sampling"`` wrappers (MC-Dropout, ensembles). + Ignored by ``UQ_METHOD="analytic"`` wrappers (GP), which emit the distribution in one + pass. Enters the metrics-cache fingerprint so cached sampling results are invalidated + when it changes. + retain_samples : bool + Keep per-point per-pass samples on the :class:`FieldDistribution` (for sampling-based + metrics such as CRPS). Default ``False`` — only streaming mean/variance are kept. + device_metrics : bool + Reserved for the on-device pooled-metric fast-path. Currently a + no-op placeholder; pooled UQ metrics run on host. + """ + + enabled: bool = False + num_samples: int = 32 + retain_samples: bool = False + device_metrics: bool = False + + @dataclass class RunConfig: """Top-level run controls (device, output dir, seed, sharding, fail-on policies).""" @@ -104,6 +134,8 @@ class RunConfig: #: If False, inference CLI skips writing ``inference__.vtp|vtu`` (comparison mesh / visuals unchanged). save_inference_mesh: bool = True metrics_cache: MetricsCacheConfig = field(default_factory=MetricsCacheConfig) + #: Uncertainty-quantification controls (probabilistic benchmarking). Additive; default off. + uq: UQConfig = field(default_factory=UQConfig) #: When True (default) and launched multi-process (e.g. ``torchrun``), shard cases across ranks via #: ``DistributedManager`` and merge before reports. When False, each rank runs the full case list (debug only). distributed: bool = True @@ -169,6 +201,16 @@ class DatasetConfig: #: If set, only these case IDs are run; if ``None``, the adapter lists all cases under :attr:`root`. case_ids: list[str] | None = None kwargs: dict[str, Any] = field(default_factory=dict) + #: Optional display label for benchmark result rows. Defaults to :attr:`name`. Set this to give + #: the SAME adapter multiple distinct rows in one matrix run (e.g. one DrivAerStar adapter at + #: three different class roots labeled ``estateback`` / ``fastback`` / ``notchback``). The + #: adapter is always resolved from :attr:`name`; only the reported/aggregated dataset key changes. + label: str | None = None + + @property + def display_name(self) -> str: + """Label used as the ``dataset`` key in results (falls back to the adapter :attr:`name`).""" + return self.label or self.name @dataclass @@ -260,6 +302,15 @@ class OutputConfig: ground_truth_volume_mesh_field_names: dict[str, str] = field( default_factory=lambda: dict(DEFAULT_GROUND_TRUTH_VOLUME_MESH_FIELD_NAMES) ) + #: Optional VTK array names for the **total predictive std** companion of each canonical field + #: (surface / volume). When a wrapper returns a :class:`FieldDistribution`, the engine attaches + #: these to exported inference / comparison meshes for ParaView + spatial-UQ visuals. Empty + #: (default) → names are auto-derived by suffixing the prediction name with ``"Std"``. + std_mesh_field_names: dict[str, str] = field(default_factory=dict) + std_volume_mesh_field_names: dict[str, str] = field(default_factory=dict) + #: As above for the **epistemic std** companion (auto-derived suffix ``"EpistemicStd"``). + epistemic_std_mesh_field_names: dict[str, str] = field(default_factory=dict) + epistemic_std_volume_mesh_field_names: dict[str, str] = field(default_factory=dict) #: Canonical key for ``streamlines_comparison`` report visual (volume vector field). streamlines_vector_canonical: str = "velocity" #: If True and surface fields are on points (e.g. MeshGraphNet / FiGNet), kNN-IDW them to cell @@ -293,6 +344,11 @@ def from_dict(cls, data: dict[str, Any]) -> "Config": mc_raw = {} elif not isinstance(mc_raw, dict): raise TypeError("run.metrics_cache must be a mapping if provided") + uq_raw = run_raw.pop("uq", None) + if uq_raw is None: + uq_raw = {} + elif not isinstance(uq_raw, dict): + raise TypeError("run.uq must be a mapping if provided") run = RunConfig( device=str(run_raw.get("device", "cuda:0")), output_dir=str(run_raw.get("output_dir", "benchmark_results")), @@ -303,6 +359,12 @@ def from_dict(cls, data: dict[str, Any]) -> "Config": enabled=_parse_bool(mc_raw.get("enabled"), default=False), path=str(mc_raw.get("path") or ""), ), + uq=UQConfig( + enabled=_parse_bool(uq_raw.get("enabled"), default=False), + num_samples=int(uq_raw.get("num_samples", 32)), + retain_samples=_parse_bool(uq_raw.get("retain_samples"), default=False), + device_metrics=_parse_bool(uq_raw.get("device_metrics"), default=False), + ), distributed=bool(run_raw.get("distributed", True)), fail_on_all_skipped=bool(run_raw.get("fail_on_all_skipped", False)), fail_on_any_metric_nan=bool(run_raw.get("fail_on_any_metric_nan", False)), @@ -329,6 +391,16 @@ def from_dict(cls, data: dict[str, Any]) -> "Config": volume_mesh_field_names=vol_fn, ground_truth_mesh_field_names=gt_mesh, ground_truth_volume_mesh_field_names=gt_vol, + std_mesh_field_names=dict(out.get("std_mesh_field_names") or {}), + std_volume_mesh_field_names=dict( + out.get("std_volume_mesh_field_names") or {} + ), + epistemic_std_mesh_field_names=dict( + out.get("epistemic_std_mesh_field_names") or {} + ), + epistemic_std_volume_mesh_field_names=dict( + out.get("epistemic_std_volume_mesh_field_names") or {} + ), streamlines_vector_canonical=str( out.get("streamlines_vector_canonical") or "velocity" ), diff --git a/physicsnemo/cfd/evaluation/datasets/__init__.py b/physicsnemo/cfd/evaluation/datasets/__init__.py index 9f408d7..fae9810 100644 --- a/physicsnemo/cfd/evaluation/datasets/__init__.py +++ b/physicsnemo/cfd/evaluation/datasets/__init__.py @@ -18,7 +18,11 @@ from physicsnemo.cfd.evaluation.datasets.schema import ( CanonicalCase, + FieldDistribution, + as_distribution, build_predictions_dict, + build_predictive_distribution, + distribution_mean, ) from physicsnemo.cfd.evaluation.datasets.adapter_registry import ( get_adapter, @@ -31,6 +35,10 @@ __all__ = [ "CanonicalCase", "build_predictions_dict", + "FieldDistribution", + "build_predictive_distribution", + "as_distribution", + "distribution_mean", "DatasetAdapter", "register_adapter", "get_adapter", diff --git a/physicsnemo/cfd/evaluation/datasets/adapters/__init__.py b/physicsnemo/cfd/evaluation/datasets/adapters/__init__.py index d3331d5..835ee75 100644 --- a/physicsnemo/cfd/evaluation/datasets/adapters/__init__.py +++ b/physicsnemo/cfd/evaluation/datasets/adapters/__init__.py @@ -18,9 +18,11 @@ from physicsnemo.cfd.evaluation.datasets.adapter_registry import register_adapter from physicsnemo.cfd.evaluation.datasets.adapters.drivaerml import DrivAerMLAdapter +from physicsnemo.cfd.evaluation.datasets.adapters.drivaerstar import DrivAerStarAdapter from physicsnemo.cfd.evaluation.datasets.adapters.ahmed import AhmedAdapter register_adapter("drivaerml", DrivAerMLAdapter) +register_adapter("drivaerstar", DrivAerStarAdapter) register_adapter("ahmed", AhmedAdapter) -__all__ = ["DrivAerMLAdapter", "AhmedAdapter"] +__all__ = ["DrivAerMLAdapter", "DrivAerStarAdapter", "AhmedAdapter"] diff --git a/physicsnemo/cfd/evaluation/datasets/adapters/drivaerstar.py b/physicsnemo/cfd/evaluation/datasets/adapters/drivaerstar.py new file mode 100644 index 0000000..3b5ffa4 --- /dev/null +++ b/physicsnemo/cfd/evaluation/datasets/adapters/drivaerstar.py @@ -0,0 +1,401 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""DrivAerStar dataset adapter: flat directory of legacy surface ``.vtk`` files. + +`DrivAerStar `_ +is an industrial-grade automotive CFD dataset whose surface meshes ship as legacy +``.vtk`` PolyData in a flat directory (``.vtk``), with field names and a wall +shear stress sign convention that differ from DrivAerML (which the shipped GeoTransolver +/ Transolver checkpoints were trained on). This adapter bridges those differences so the +same model wrappers and metrics work unchanged: + +1. Rename ``Pressure`` → ``pMeanTrim`` (DrivAerML convention). +2. Combine the three WSS scalar components into one ``(N, 3)`` vector and flip its sign + to match DrivAerML (``flip_wss_sign``, default on). +3. Drop explicit ``Normals`` / ``Area`` arrays so downstream force integration and + rendering recompute them from mesh topology (DrivAerML convention). + +These transforms are cheap generic array ops, so by default they run **in memory** on each pass: +the prepared surface mesh is handed to the wrappers via +:attr:`~physicsnemo.cfd.evaluation.datasets.schema.CanonicalCase.reference_geometry` (predictions + +ground truth) and :attr:`~physicsnemo.cfd.evaluation.datasets.schema.CanonicalCase.geometry` (the +SDF / geometry branch — the DrivAerStar surface *is* its geometry, so no STL file is materialized). +Nothing is written to disk, avoiding a full-dataset duplicate and any stale-cache class of bug. +``mesh_path`` points at the source ``.vtk`` (used only for run-index / directory context; the +forward pass reads the in-memory meshes, not this file). + +Set ``cache_prepared: true`` to additionally **persist** a canonical ``.vtp`` (+ triangulated +``.stl``) under ``///`` for external inspection (ParaView, etc.). +Those writes are guarded by a sidecar signature (``.prepared.json``) recording the +transform kwargs (``flip_wss_sign``, field names, ``remove_normals_area``, ...) and the source +file's identity (mtime/size), so a changed transform rebuilds rather than reusing a stale artifact. + +.. note:: + The datapipe wrappers (GeoTransolver, Transolver) parse an integer run index from the + case id via ``run_id_from_case_id``, so case ids must be integer-like (a decimal string or + ``run_``). DrivAerStar VTK stems are integers, so this holds out of the box. Rename files + or pass ``glob_pattern`` accordingly if yours are not. +""" + +import json +from pathlib import Path +from typing import Any + +import numpy as np +import pyvista as pv + +from physicsnemo.cfd.evaluation.common.natural_sort import natural_sorted +from physicsnemo.cfd.evaluation.datasets.adapter_registry import DatasetAdapter +from physicsnemo.cfd.evaluation.datasets.progress import log_dataset +from physicsnemo.cfd.evaluation.datasets.schema import ( + CanonicalCase, + InferenceDomain, + coerce_inference_domain_or_default, +) +from physicsnemo.cfd.evaluation.datasets.vtk_ground_truth import ( + extract_pressure_wss_from_mesh, +) + +#: DrivAerStar source array names (legacy ``.vtk`` cell data). +DEFAULT_SOURCE_PRESSURE_NAME = "Pressure" +DEFAULT_SOURCE_WSS_COMPONENT_NAMES: tuple[str, str, str] = ( + "WallShearStressi", + "WallShearStressj", + "WallShearStressk", +) + +#: Canonical (DrivAerML) VTK array names written into the prepared VTP. Chosen so the +#: default ``output.ground_truth_mesh_field_names`` (which target DrivAerML names) work +#: without per-config overrides. +DEFAULT_PRESSURE_OUT_NAME = "pMeanTrim" +DEFAULT_SHEAR_OUT_NAME = "wallShearStressMeanTrim" + + +def _run_id_from_case_id(case_id: str) -> int: + """Parse an integer run index from ``run_`` or a decimal id (mirrors the datapipe + wrappers' ``run_id_from_case_id`` so the STL name resolves to the same tag).""" + s = str(case_id).strip() + if not s: + raise ValueError("case_id is empty") + suffix = s[4:] if s.startswith("run_") else s + return int(suffix) + + +class DrivAerStarAdapter(DatasetAdapter): + """Adapter for DrivAerStar surface ``.vtk`` files in a flat directory. + + Case ids are the VTK file stems. Ground truth is exposed under the canonical keys + ``pressure`` (N,) and ``shear_stress`` (N, 3). + + Optional kwargs (from ``dataset.kwargs`` in config): + + - ``inference_domain``: ``"surface"`` (default). Volume is not supported (DrivAerStar + surface meshes only). + - ``glob_pattern``: source-file glob relative to ``root`` (default ``"*.vtk"``). + Supports ``**`` for nested layouts (case id remains the file stem). + - ``pressure_field_name``: source pressure array name (default ``"Pressure"``). + - ``wss_component_names``: three source WSS scalar names (default + ``("WallShearStressi", "WallShearStressj", "WallShearStressk")``). + - ``flip_wss_sign``: wall-shear-stress **sign convention** knob. When ``True`` (default) the + combined WSS is negated so the DrivAerStar ground truth matches the DrivAerML sign + convention — use this for DrivAerML-convention checkpoints. Set it ``False`` for checkpoints + trained directly on the **native** DrivAerStar WSS sign (e.g. the GeoTransolver + ``transformer_models`` checkpoints scored in the UQ example config, which is why that config + sets ``flip_wss_sign: false``). This MUST match how your checkpoints were trained; a mismatch + silently flips WSS — and therefore drag and lift — for every case. + - ``remove_normals_area``: drop explicit ``Normals`` / ``Area`` arrays (default ``True``). + - ``gt_data_type``: ``auto`` / ``cell`` / ``point`` passed to GT extraction + (default ``"cell"``; DrivAerStar fields are cell-centered). + - ``pressure_out_name`` / ``shear_out_name``: canonical VTK array names for the prepared + pressure / WSS arrays (defaults ``"pMeanTrim"`` / ``"wallShearStressMeanTrim"``). + - ``cache_prepared``: also persist the prepared ``.vtp`` (+ triangulated ``.stl``) to disk for + external inspection (default ``False`` — preparation is in-memory only). + - ``prepared_subdir``: cache directory name under ``root`` when ``cache_prepared`` (default + ``"_prepared"``). + - ``force_reprepare``: rebuild the on-disk cache even if present, when ``cache_prepared`` + (default ``False``). + """ + + @classmethod + def inference_domain(cls) -> InferenceDomain: + """Class-level default inference domain (DrivAerStar is surface-only).""" + return "surface" + + @classmethod + def inference_domain_from_kwargs( + cls, kwargs: dict[str, Any] | None + ) -> InferenceDomain: + """Resolve inference domain from ``dataset.kwargs`` (surface only).""" + kw = kwargs or {} + domain = coerce_inference_domain_or_default( + kw.get("inference_domain"), + default="surface", + parameter="dataset.kwargs.inference_domain", + ) + if domain != "surface": + raise NotImplementedError( + "DrivAerStarAdapter supports surface inference only; got " + f"inference_domain={domain!r}." + ) + return domain + + def __init__(self, root: str, **kwargs: Any) -> None: + self.root = Path(root) + if not self.root.exists(): + raise FileNotFoundError(f"DrivAerStar root not found: {self.root}") + # Validate / normalize domain (raises for volume). + self.inference_domain_from_kwargs(kwargs) + + self._glob_pattern: str = kwargs.get("glob_pattern", "*.vtk") + self._pressure_field_name: str = kwargs.get( + "pressure_field_name", DEFAULT_SOURCE_PRESSURE_NAME + ) + wss = kwargs.get("wss_component_names", DEFAULT_SOURCE_WSS_COMPONENT_NAMES) + self._wss_component_names: tuple[str, ...] = tuple(wss) + if len(self._wss_component_names) != 3: + raise ValueError( + "wss_component_names must have exactly 3 entries; got " + f"{self._wss_component_names!r}" + ) + + self._flip_wss_sign: bool = bool(kwargs.get("flip_wss_sign", True)) + self._remove_normals_area: bool = bool(kwargs.get("remove_normals_area", True)) + self._gt_data_type: str = kwargs.get("gt_data_type", "cell") + self._cache_prepared: bool = bool(kwargs.get("cache_prepared", False)) + self._prepared_subdir: str = kwargs.get("prepared_subdir", "_prepared") + self._pressure_out_name: str = kwargs.get( + "pressure_out_name", DEFAULT_PRESSURE_OUT_NAME + ) + self._shear_out_name: str = kwargs.get("shear_out_name", DEFAULT_SHEAR_OUT_NAME) + self._force_reprepare: bool = bool(kwargs.get("force_reprepare", False)) + + self._prepared_root = self.root / self._prepared_subdir + + def _source_path(self, case_id: str) -> Path: + """Resolve the source ``.vtk`` for a case id (stem), honoring ``glob_pattern`` nesting.""" + suffix = Path(self._glob_pattern).suffix or ".vtk" + direct = self.root / f"{case_id}{suffix}" + if direct.exists(): + return direct + for candidate in self.root.glob(self._glob_pattern): + if candidate.stem == case_id: + return candidate + raise FileNotFoundError( + f"Source mesh for case {case_id!r} not found under {self.root} " + f"(glob_pattern={self._glob_pattern!r})" + ) + + def _stl_tag(self, case_id: str) -> str: + """STL tag matching the wrappers' ``run_id_from_case_id`` when derivable.""" + try: + return str(_run_id_from_case_id(case_id)) + except ValueError: + # Non-integer id: the datapipe wrappers can't use this case, but keep a stable + # STL name so GT-only / other consumers still work (one STL per case dir). + return case_id + + def list_cases(self) -> list[str]: + """Return case ids: stems of source ``.vtk`` files under ``root`` (excluding cache).""" + prepared = self._prepared_root.resolve() + case_ids: list[str] = [] + for p in self.root.glob(self._glob_pattern): + if not p.is_file(): + continue + if prepared in p.resolve().parents: + continue + case_ids.append(p.stem) + return natural_sorted(case_ids) + + def _transform_source_mesh(self, case_id: str) -> tuple[pv.PolyData, Path]: + """Read the raw ``.vtk`` and apply the in-memory canonicalization transforms. + + Shared core of both the default (in-memory) and the opt-in cached paths: rename pressure, + combine + optionally sign-flip WSS, and drop ``Normals`` / ``Area``. No disk writes. Returns + the prepared surface mesh and the source path. + """ + src_path = self._source_path(case_id) + mesh = pv.read(str(src_path)) + if not isinstance(mesh, pv.PolyData): + mesh = mesh.extract_surface() + self._rename_pressure(mesh, case_id) + self._combine_wss(mesh, case_id) + if self._remove_normals_area: + self._drop_normals_area(mesh) + return mesh, src_path + + #: Bump when the preparation transform changes in a way that invalidates cached artifacts. + _PREPARE_SIGNATURE_VERSION = 1 + + def _transform_signature(self, src_path: Path) -> dict[str, Any]: + """Signature of everything that affects the prepared artifacts (transform kwargs + source id). + + A cached VTP/STL is only reused when this matches the sidecar written when it was prepared, + so changing e.g. ``flip_wss_sign`` or a field name — or editing the source ``.vtk`` — + forces a rebuild instead of a stale cache hit. + """ + st = src_path.stat() + return { + "version": self._PREPARE_SIGNATURE_VERSION, + "flip_wss_sign": self._flip_wss_sign, + "remove_normals_area": self._remove_normals_area, + "pressure_field_name": self._pressure_field_name, + "wss_component_names": list(self._wss_component_names), + "pressure_out_name": self._pressure_out_name, + "shear_out_name": self._shear_out_name, + "source_name": src_path.name, + "source_mtime_ns": st.st_mtime_ns, + "source_size": st.st_size, + } + + @staticmethod + def _read_signature(sidecar_path: Path) -> dict[str, Any] | None: + """Load a prepared-case sidecar; treat a missing/corrupt file as no signature.""" + try: + with sidecar_path.open("r", encoding="utf-8") as fh: + data = json.load(fh) + return data if isinstance(data, dict) else None + except (OSError, ValueError): + return None + + def _write_prepared_cache( + self, case_id: str, mesh: pv.PolyData, src_path: Path + ) -> str: + """Persist the prepared ``.vtp`` (+ triangulated ``.stl``) for external inspection. + + Only called when ``cache_prepared`` is set. Writes are guarded by the sidecar signature so a + changed transform (or edited source) rebuilds rather than reusing a stale artifact. Returns + the prepared VTP path (used as the case ``mesh_path`` in cached mode). + """ + case_dir = self._prepared_root / case_id + case_dir.mkdir(parents=True, exist_ok=True) + vtp_path = case_dir / f"{case_id}.vtp" + stl_path = case_dir / f"drivaer_{self._stl_tag(case_id)}.stl" + sidecar_path = case_dir / f"{case_id}.prepared.json" + + signature = self._transform_signature(src_path) + cache_valid = ( + not self._force_reprepare + and vtp_path.exists() + and stl_path.exists() + and self._read_signature(sidecar_path) == signature + ) + if not cache_valid: + log_dataset("drivaerstar", f"Caching prepared VTP + STL for {case_id!r}") + mesh.save(str(vtp_path)) + mesh.triangulate().save(str(stl_path)) + # Write the sidecar only after successful saves so a crash mid-write does not leave a + # signature that would validate partial/absent artifacts. + with sidecar_path.open("w", encoding="utf-8") as fh: + json.dump(signature, fh, indent=2, sort_keys=True) + return str(vtp_path) + + @staticmethod + def _available_arrays(mesh: pv.PolyData) -> str: + """Human-readable listing of cell/point array names for error messages.""" + return ( + f"cell_data={sorted(mesh.cell_data.keys())}, " + f"point_data={sorted(mesh.point_data.keys())}" + ) + + def _rename_pressure(self, mesh: pv.PolyData, case_id: str) -> None: + for data in (mesh.cell_data, mesh.point_data): + if self._pressure_field_name in data: + data[self._pressure_out_name] = data.pop(self._pressure_field_name) + return + raise ValueError( + f"DrivAerStar case {case_id!r}: pressure array " + f"{self._pressure_field_name!r} not found in the source mesh " + f"(check ``pressure_field_name``). Available arrays: " + f"{self._available_arrays(mesh)}." + ) + + def _combine_wss(self, mesh: pv.PolyData, case_id: str) -> None: + for data in (mesh.cell_data, mesh.point_data): + if all(k in data for k in self._wss_component_names): + wss = np.stack( + [np.asarray(data.pop(k)) for k in self._wss_component_names], + axis=1, + ).astype(np.float32) + if self._flip_wss_sign: + wss = -wss + data[self._shear_out_name] = wss + return + raise ValueError( + f"DrivAerStar case {case_id!r}: wall-shear-stress components " + f"{list(self._wss_component_names)!r} not all found in a single data group " + f"(check ``wss_component_names``). Available arrays: " + f"{self._available_arrays(mesh)}." + ) + + @staticmethod + def _drop_normals_area(mesh: pv.PolyData) -> None: + for key in ("Normals", "Area"): + if key in mesh.cell_data: + del mesh.cell_data[key] + if key in mesh.point_data: + del mesh.point_data[key] + + def load_case(self, case_id: str) -> CanonicalCase: + """Transform the case in memory (optionally caching) and load it into the canonical schema. + + The prepared surface mesh is returned as both ``reference_geometry`` (predictions + GT) and + ``geometry`` (the wrappers' SDF branch — surface *is* geometry for DrivAerStar), so the + forward pass needs no on-disk VTP/STL. ``mesh_path`` is the source ``.vtk`` unless + ``cache_prepared`` persisted a ``.vtp`` (then it points there). + """ + log_dataset( + "drivaerstar", + f"load_case({case_id!r}): root={self.root}", + ) + mesh, src_path = self._transform_source_mesh(case_id) + if self._cache_prepared: + mesh_path = self._write_prepared_cache(case_id, mesh, src_path) + else: + mesh_path = str(src_path) + + gt_dict, gt_loc = extract_pressure_wss_from_mesh( + mesh, + data_type=self._gt_data_type, + pressure_names=(self._pressure_out_name,), + shear_names=(self._shear_out_name,), + ) + ground_truth = gt_dict if gt_dict else None + mesh_type = gt_loc if gt_loc is not None else "cell" + + meta: dict[str, Any] = { + "dataset": "drivaerstar", + "case": case_id, + "branch": "surface", + "source_vtk": src_path.name, + "prepared_cached": self._cache_prepared, + } + if ground_truth: + meta["ground_truth_location"] = gt_loc + meta["ground_truth_fields"] = list(ground_truth.keys()) + + return CanonicalCase( + case_id=case_id, + mesh_path=mesh_path, + mesh_type=mesh_type, + ground_truth=ground_truth, + metadata=meta, + inference_domain="surface", + reference_geometry=mesh, + # DrivAerStar surface *is* the geometry: hand the same mesh to the SDF/geometry branch + # so the wrappers derive stl_coordinates/faces/centers in memory (no STL file needed). + geometry=mesh, + ) diff --git a/physicsnemo/cfd/evaluation/datasets/schema.py b/physicsnemo/cfd/evaluation/datasets/schema.py index 151c282..e7a3594 100644 --- a/physicsnemo/cfd/evaluation/datasets/schema.py +++ b/physicsnemo/cfd/evaluation/datasets/schema.py @@ -23,6 +23,12 @@ import numpy as np +#: Array payloads on :class:`FieldDistribution` may be NumPy arrays or (for the on-device +#: metric path) framework tensors such as ``torch.Tensor``. We +#: intentionally do not import ``torch`` here to keep the ``datasets`` layer dependency-light; +#: metrics convert to NumPy at their boundary. +ArrayLike = Any + # Surface vs volume inference (which manifold the case uses). Combined surface+volume is deferred. InferenceDomain = Literal["surface", "volume"] @@ -75,6 +81,13 @@ class CanonicalCase: #: may skip a redundant ``pv.read(case.mesh_path)`` for the same case. Adapters are not #: required to populate this field. reference_geometry: Any | None = None + #: Optional in-memory geometry mesh (:class:`pyvista.PolyData`) for the wrappers' SDF / + #: geometry-embedding branch, replacing an on-disk STL. When set, the datapipe I/O derives + #: ``stl_coordinates`` / ``stl_faces`` / ``stl_centers`` from it instead of globbing an STL + #: file next to ``mesh_path``. Adapters where the surface *is* the geometry (e.g. DrivAerStar) + #: can set this to avoid materializing an STL; adapters with a distinct geometry file leave it + #: ``None`` to keep the file-based lookup. + geometry: Any | None = None def __post_init__(self) -> None: if self.mesh_type not in ("point", "cell", "unknown"): @@ -119,3 +132,99 @@ def build_predictions_dict( if v is not None: out[k] = np.asarray(v, dtype=np.float32) return out + + +@dataclass +class FieldDistribution: + """Per-point/-cell predictive distribution for one canonical field. + + All arrays are in physical (de-normalized) units — the same space as the + deterministic predictions and ground truth. UQ wrappers denormalize ``mean`` and + the std/variance channels before returning, so metrics never touch normalization + stats. ``mean`` / ``std`` may be NumPy arrays or (for the on-device metric path) + tensors such as ``torch.Tensor``; UQ metrics convert to NumPy at their + boundary. + + The ``mean`` / ``std`` pair is the Gaussian summary every UQ method can provide + and that the Gaussian metrics (NLPD, zRMS, coverage) consume. ``samples`` and + ``quantiles`` / ``quantile_levels`` are optional non-Gaussian escape hatches for + methods whose predictive law is not Gaussian (quantile regression, conformal, + evidential, generative); distribution-agnostic metrics (CRPS, quantile coverage) use + them when present. + """ + + mean: ArrayLike # (N,) or (N, C) + std: ArrayLike | None = None # total predictive std (epistemic + aleatoric) + epistemic_std: ArrayLike | None = None + aleatoric_std: ArrayLike | None = None + samples: ArrayLike | None = None # optional (S, N[, C]) for CRPS / non-Gaussian + quantiles: ArrayLike | None = ( + None # optional (Q, N[, C]) values at ``quantile_levels`` + ) + quantile_levels: ArrayLike | None = ( + None # optional (Q,) in (0, 1), for interval methods + ) + + +def build_predictive_distribution( + *, + mean: ArrayLike, + std: ArrayLike | None = None, + epistemic_std: ArrayLike | None = None, + aleatoric_std: ArrayLike | None = None, + samples: ArrayLike | None = None, + quantiles: ArrayLike | None = None, + quantile_levels: ArrayLike | None = None, +) -> FieldDistribution: + """Build a :class:`FieldDistribution`, mirroring :func:`build_predictions_dict`. + + Use this in UQ wrappers' :meth:`~physicsnemo.cfd.evaluation.models.model_registry.CFDModel.decode_distribution` + (prefer keyword arguments) so the predictive-distribution payload stays consistent. + NumPy inputs are coerced to ``float32``; framework tensors (e.g. ``torch.Tensor``) are + passed through untouched so the on-device metric path can keep them on device. + """ + + def _coerce(x: ArrayLike | None) -> ArrayLike | None: + if x is None: + return None + if isinstance(x, np.ndarray): + return x.astype(np.float32, copy=False) + return x # tensor or other array-like: leave as-is + + return FieldDistribution( + mean=_coerce(mean), + std=_coerce(std), + epistemic_std=_coerce(epistemic_std), + aleatoric_std=_coerce(aleatoric_std), + samples=_coerce(samples), + quantiles=_coerce(quantiles), + quantile_levels=_coerce(quantile_levels), + ) + + +def as_distribution(predictions: dict[str, Any], key: str) -> FieldDistribution | None: + """Return the predictive distribution for ``key`` from a predictions dict. + + - A :class:`FieldDistribution` value is returned as-is. + - A plain array value is wrapped as a **degenerate** distribution (``std=None``) so + UQ metrics can treat deterministic and probabilistic wrappers uniformly. + - A missing / ``None`` value returns ``None``. + """ + value = predictions.get(key) + if value is None: + return None + if isinstance(value, FieldDistribution): + return value + return FieldDistribution(mean=value) + + +def distribution_mean(value: Any) -> Any: + """Unwrap the point estimate from a predictions value. + + Returns ``value.mean`` for a :class:`FieldDistribution`, else ``value`` unchanged. + Deterministic metrics use this so a probabilistic wrapper's :class:`FieldDistribution` + scores through the existing (mean-based) L2 / force / physics metrics. + """ + if isinstance(value, FieldDistribution): + return value.mean + return value diff --git a/physicsnemo/cfd/evaluation/metrics/builtin/__init__.py b/physicsnemo/cfd/evaluation/metrics/builtin/__init__.py index 7725bab..41c2ded 100644 --- a/physicsnemo/cfd/evaluation/metrics/builtin/__init__.py +++ b/physicsnemo/cfd/evaluation/metrics/builtin/__init__.py @@ -19,6 +19,7 @@ from physicsnemo.cfd.evaluation.metrics.builtin.forces import register_force_metrics from physicsnemo.cfd.evaluation.metrics.builtin.l2 import register_l2_metrics from physicsnemo.cfd.evaluation.metrics.builtin.physics import register_physics_metrics +from physicsnemo.cfd.evaluation.metrics.builtin.uq import register_uq_metrics def register_all_builtin_metrics() -> None: @@ -26,6 +27,7 @@ def register_all_builtin_metrics() -> None: register_l2_metrics() register_force_metrics() register_physics_metrics() + register_uq_metrics() register_all_builtin_metrics() diff --git a/physicsnemo/cfd/evaluation/metrics/builtin/uq.py b/physicsnemo/cfd/evaluation/metrics/builtin/uq.py new file mode 100644 index 0000000..545b195 --- /dev/null +++ b/physicsnemo/cfd/evaluation/metrics/builtin/uq.py @@ -0,0 +1,741 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Pooled (reducer) uncertainty-quantification metrics. + +Every metric here is a :class:`~physicsnemo.cfd.postprocessing_tools.metric_registry.ReducerMetric`: +``partial`` returns **per-case extensive sufficient statistics** (sums & counts, additive +across cases and distributed ranks) and ``finalize`` maps the globally-summed statistics to +the final per-channel value(s). This global pooling over all points — rather than a mean of +per-case values — is the statistically correct way to estimate calibration/coverage (points +per case differ). + +**Space.** All statistics are computed in **physical (de-normalized) units** — the space the +wrappers denormalize ``mean`` *and* ``std`` into. NLPD carries an additive per-channel +``log(scale)`` offset vs normalized space but its cross-method ranking is unchanged. + +Channels are inferred from the prediction arrays: scalar fields (e.g. ``pressure``) give one +channel; 3-vector fields (e.g. ``shear_stress``, ``velocity``) give ``_x`` / ``_y`` / ``_z`` +channels. A method with no total ``std`` (deterministic wrapper) or no ``epistemic_std`` +contributes nothing, so the corresponding metric finalizes to ``NaN``. +""" + +from __future__ import annotations + +import math +from typing import Any, Callable, Iterator + +import numpy as np +from scipy.stats import spearmanr + +from physicsnemo.cfd.evaluation.datasets.schema import as_distribution +from physicsnemo.cfd.postprocessing_tools.metric_registry import register_metric + +_LOG_2PI = math.log(2.0 * math.pi) +#: Numerical floor on variance for the log / division terms. +_VAR_FLOOR = 1.0e-12 + +#: ``np.trapz`` was deprecated in NumPy 2.0 and removed in 2.4; ``np.trapezoid`` is its drop-in +#: replacement. Resolve once so AUSE works across the pinned range of NumPy versions. +_trapezoid = getattr(np, "trapezoid", None) or np.trapz # type: ignore[attr-defined] + +#: Component suffixes for 3-vector fields. +_VEC3_SUFFIXES = ("x", "y", "z") + + +def _to_numpy_f64(x: Any) -> np.ndarray | None: + """Convert a NumPy array or framework tensor (e.g. ``torch.Tensor``) to ``float64`` NumPy.""" + if x is None: + return None + if isinstance(x, np.ndarray): + return x.astype(np.float64, copy=False) + # Duck-typed torch/cupy tensor: detach + move to host without importing torch here. + detach = getattr(x, "detach", None) + if callable(detach): + x = detach() + cpu = getattr(x, "cpu", None) + if callable(cpu): + x = cpu() + return np.asarray(x, dtype=np.float64) + + +def _iter_channels( + gt: dict[str, Any], + predictions: dict[str, Any], + *, + need_epistemic: bool, +) -> Iterator[tuple[str, np.ndarray, np.ndarray, np.ndarray, np.ndarray | None]]: + """Yield ``(channel_name, y, mu, sigma, sigma_epi)`` per scalar channel (physical units). + + Iterates every prediction key that resolves to a :class:`FieldDistribution` with a total + ``std`` (and, when ``need_epistemic``, an ``epistemic_std``) and has matching ground truth. + Vector fields are split into per-component channels. Arrays are 1-D ``float64`` over points. + """ + if not gt: + return + for key in predictions: + if key not in gt or gt[key] is None: + continue + dist = as_distribution(predictions, key) + if dist is None or dist.std is None: + continue + epi = dist.epistemic_std + if need_epistemic and epi is None: + continue + + y = _to_numpy_f64(gt[key]) + mu = _to_numpy_f64(dist.mean) + sig = _to_numpy_f64(dist.std) + epi_np = _to_numpy_f64(epi) + if y is None or mu is None or sig is None: + continue + + # Normalize to (N, C): scalars -> (N, 1); vectors already (N, 3). + def _as_2d(a: np.ndarray) -> np.ndarray | None: + if a.ndim == 1: + return a.reshape(-1, 1) + if a.ndim == 2: + return a + return None + + y2, mu2, sig2 = _as_2d(y), _as_2d(mu), _as_2d(sig) + epi2 = _as_2d(epi_np) if epi_np is not None else None + if y2 is None or mu2 is None or sig2 is None: + continue + if not (y2.shape == mu2.shape == sig2.shape): + continue + if epi2 is not None and epi2.shape != y2.shape: + epi2 = None + if need_epistemic: + continue + + n_comp = y2.shape[1] + for c in range(n_comp): + if n_comp == 1: + chan = key + elif n_comp == 3: + chan = f"{key}_{_VEC3_SUFFIXES[c]}" + else: + chan = f"{key}_{c}" + yield ( + chan, + y2[:, c], + mu2[:, c], + sig2[:, c], + epi2[:, c] if epi2 is not None else None, + ) + + +class _PooledUQMetric: + """Reducer metric over per-point Gaussian channel statistics. + + ``per_point`` returns a dict of **extensive** contributions (sums) for one channel, plus + ``"n"`` (point count). ``finalize_channel`` maps a channel's globally-summed stats to its + scalar value. Statistics from all channels are namespaced ``"{channel}::{stat}"`` so the + engine can sum them across cases with no schema change. + """ + + def __init__( + self, + per_point: Callable[..., dict[str, float]], + finalize_channel: Callable[[dict[str, float]], float], + *, + need_epistemic: bool = False, + ) -> None: + self._per_point = per_point + self._finalize_channel = finalize_channel + self._need_epistemic = need_epistemic + + def partial(self, gt: Any, predictions: Any, **_: Any) -> dict[str, float]: + stats: dict[str, float] = {} + for chan, y, mu, sig, epi in _iter_channels( + gt or {}, predictions or {}, need_epistemic=self._need_epistemic + ): + contrib = self._per_point(y, mu, sig, epi) + for stat, val in contrib.items(): + stats[f"{chan}::{stat}"] = stats.get(f"{chan}::{stat}", 0.0) + float( + val + ) + return stats + + def finalize(self, summed: dict[str, float]) -> float | dict[str, float]: + # Regroup "{channel}::{stat}" -> {channel: {stat: value}}. + by_channel: dict[str, dict[str, float]] = {} + for key, val in summed.items(): + if "::" not in key: + continue + chan, stat = key.split("::", 1) + by_channel.setdefault(chan, {})[stat] = val + if not by_channel: + # No contributions (e.g. deterministic wrapper): return the SAME headline sub-key a + # populated finalize would (``{metric}_mean``) so the report schema is consistent and + # ``fail_on_any_metric_nan`` can flag the unavailable configured metric. + return {"mean": float("nan")} + out: dict[str, float] = {} + values: list[float] = [] + for chan in sorted(by_channel): + v = self._finalize_channel(by_channel[chan]) + out[chan] = float(v) + if v == v: # skip NaN in the headline mean + values.append(float(v)) + out["mean"] = float(np.mean(values)) if values else float("nan") + return out + + +# --- per-point contribution / finalize functions ------------------------------------------- + + +def _nlpd_contrib_total(y, mu, sig, epi) -> dict[str, float]: + err = y - mu + var = np.clip(sig**2, _VAR_FLOOR, None) + nlpd = 0.5 * (_LOG_2PI + np.log(var) + err**2 / var) + return {"sum": float(nlpd.sum()), "n": float(y.size)} + + +def _nlpd_contrib_epistemic(y, mu, sig, epi) -> dict[str, float]: + err = y - mu + var = np.clip(epi**2, _VAR_FLOOR, None) + nlpd = 0.5 * (_LOG_2PI + np.log(var) + err**2 / var) + return {"sum": float(nlpd.sum()), "n": float(y.size)} + + +def _zrms_contrib(y, mu, sig, epi) -> dict[str, float]: + err = y - mu + var = np.clip(sig**2, _VAR_FLOOR, None) + return {"sum_z2": float((err**2 / var).sum()), "n": float(y.size)} + + +def _coverage95_contrib(y, mu, sig, epi) -> dict[str, float]: + err = y - mu + within = (np.abs(err) <= 1.96 * sig).astype(np.float64) + return {"within": float(within.sum()), "n": float(y.size)} + + +def _sharpness_total_contrib(y, mu, sig, epi) -> dict[str, float]: + return {"sum_sig": float(sig.sum()), "n": float(y.size)} + + +def _sharpness_epistemic_contrib(y, mu, sig, epi) -> dict[str, float]: + return {"sum_sig": float(epi.sum()), "n": float(y.size)} + + +def _mean_ratio(sum_key: str) -> Callable[[dict[str, float]], float]: + def _f(d: dict[str, float]) -> float: + n = d.get("n", 0.0) + return d.get(sum_key, 0.0) / n if n > 0 else float("nan") + + return _f + + +def _zrms_finalize(d: dict[str, float]) -> float: + n = d.get("n", 0.0) + return math.sqrt(d.get("sum_z2", 0.0) / n) if n > 0 else float("nan") + + +# --- pointwise per-point ranking quality: Spearman(|error|, uncertainty) --------------------- +# +# "Does per-point uncertainty rank per-point error *within* a geometry?" is a rank-correlation +# question, so the right per-case diagnostic is Spearman's rho — NOT a within-case AUSE +# (sparsification only makes sense *across* geometries; see ``_SampleAUSE`` below). Returned per +# channel (+ headline ``mean``) and averaged over cases by the engine. + + +def _spearman(a: np.ndarray, b: np.ndarray) -> float: + """Tie-aware Spearman rank correlation of ``a`` and ``b`` (via ``scipy.stats.spearmanr``). + + ``scipy`` handles ties with *average* ranks (ordinal ranks are wrong when values tie — and + ties are common here: zero-spread regions, quantized uncertainty, repeated ensemble draws). + Returns ``NaN`` for fewer than two points, mismatched sizes, or a constant input (rank + correlation is mathematically undefined when either variable has zero rank variance). + """ + a = np.asarray(a, dtype=np.float64) + b = np.asarray(b, dtype=np.float64) + if a.size < 2 or a.size != b.size: + return float("nan") + # Undefined for a constant input; short-circuit to avoid SciPy's "input is constant" warning. + if np.ptp(a) == 0.0 or np.ptp(b) == 0.0: + return float("nan") + return float(spearmanr(a, b).correlation) + + +class _UncertaintyErrorSpearman: + """Per-case Spearman(|error|, uncertainty) over :class:`FieldDistribution` predictions. + + Pointwise (per-case) metric averaged over cases by the engine. ``rank_epistemic`` correlates + the per-point absolute error with the epistemic std (the high-contrast OOD signal); otherwise + with the total predictive std. rho→1 means the uncertainty perfectly orders the error. + """ + + def __init__(self, *, rank_epistemic: bool) -> None: + self._rank_epistemic = rank_epistemic + + def __call__(self, gt: Any, predictions: Any, **_: Any) -> dict[str, float]: + out: dict[str, float] = {} + values: list[float] = [] + for chan, y, mu, sig, epi in _iter_channels( + gt or {}, predictions or {}, need_epistemic=self._rank_epistemic + ): + unc = epi if self._rank_epistemic else sig + if unc is None: + continue + rho = _spearman(np.abs(y - mu), unc) + out[chan] = float(rho) + if rho == rho: # skip NaN in the headline mean + values.append(float(rho)) + if not out: + return {"mean": float("nan")} + out["mean"] = float(np.mean(values)) if values else float("nan") + return out + + +# --- sample-wise sparsification / AUSE (SAMPLE metric: one number per geometry) -------------- +# +# AUSE answers the active-learning question at the *geometry* level: rank the dataset's +# geometries by their aggregate predicted uncertainty, drop the most-uncertain first, and check +# that the error of what remains falls (toward the oracle that drops true-worst-error first). +# This needs one (uncertainty, error) pair *per geometry* and a global sort over geometries — not +# an additive sufficient statistic — so it is a :class:`SampleMetric`: ``partial`` emits the +# per-case scalars, and ``finalize_samples`` collects them across all cases (and ranks) and +# computes AUSE once. Direct port of ``plot_field_gp_sparsification._sparsification`` / ``_ause``. + + +def _sparsification_curve( + err_per_sample: np.ndarray, rank_signal: np.ndarray +) -> tuple[np.ndarray, np.ndarray, np.ndarray, float, float] | None: + """Full sparsification curve for one (error, uncertainty) set of per-geometry scalars. + + Returns ``(fractions, rmse_by_uncertainty, rmse_oracle, full_rmse, ause)`` where each curve + gives the RMSE of the *retained* geometries after removing the ``fractions`` most-uncertain + (resp. true-worst-error) ones, or ``None`` when fewer than two geometries are present. + ``err_per_sample`` is one non-negative error scalar per geometry; ``rank_signal`` is the + per-geometry uncertainty used to order the uncertainty-driven removals. AUSE is the + normalized area between the uncertainty and oracle curves (lower is better; 0 == oracle). + Direct port of ``plot_field_gp_sparsification._sparsification`` / ``_ause``. + """ + err2 = np.asarray(err_per_sample, dtype=np.float64) ** 2 + n = err2.size + if n < 2: + return None + fr = np.arange(n, dtype=np.float64) / n + + def _suffix_rmse(order: np.ndarray) -> np.ndarray: + vals = err2[order] # ordered so the first entries are removed first + suffix_sum = np.cumsum(vals[::-1])[::-1] # sum over the retained tail + counts = n - np.arange(n) + return np.sqrt(suffix_sum / counts) + + by_unc = _suffix_rmse(np.argsort(-np.asarray(rank_signal, dtype=np.float64))) + oracle = _suffix_rmse(np.argsort(-err2)) + full = float(np.sqrt(err2.mean())) + ause = ( + float(_trapezoid((by_unc - oracle) / full, fr)) if full > 0.0 else float("nan") + ) + return fr, by_unc, oracle, full, ause + + +def _ause_curve_area(err_per_sample: np.ndarray, rank_signal: np.ndarray) -> float: + """Scalar AUSE only (see :func:`_sparsification_curve`); ``NaN`` for <2 geometries.""" + curve = _sparsification_curve(err_per_sample, rank_signal) + return float("nan") if curve is None else curve[4] + + +def _curve_payload( + err_per_sample: np.ndarray, rank_signal: np.ndarray +) -> dict[str, Any] | None: + """Serializable curve dict for the sparsification visual, or ``None`` for <2 geometries.""" + curve = _sparsification_curve(err_per_sample, rank_signal) + if curve is None: + return None + fr, by_unc, oracle, full, ause = curve + return { + "fractions": fr, + "by_uncertainty": by_unc, + "oracle": oracle, + "full": full, + "ause": ause, + "n": int(fr.size), + } + + +class _SampleAUSE: + """Sample-wise AUSE. ``partial`` → per-geometry (uncertainty, error); ``finalize_samples`` → AUSE. + + Per channel, ``partial`` reduces one geometry to its RMS error (``err``) and its mean + predicted std (``unc``; epistemic when ``rank_epistemic``). ``finalize_samples`` gathers those + across all geometries and computes the AUSE of ranking geometries by ``unc``, plus a + ``_spearman`` companion (rank correlation of per-geometry error vs. uncertainty) that + captures monotonic *trend alignment* independent of error scale — the standard complement to + AUSE when comparing methods whose absolute errors differ. + """ + + def __init__(self, *, rank_epistemic: bool) -> None: + self._rank_epistemic = rank_epistemic + + def partial(self, gt: Any, predictions: Any, **_: Any) -> dict[str, float]: + stats: dict[str, float] = {} + for chan, y, mu, sig, epi in _iter_channels( + gt or {}, predictions or {}, need_epistemic=self._rank_epistemic + ): + unc = epi if self._rank_epistemic else sig + if unc is None: + continue + err = float(np.sqrt(np.mean((y - mu) ** 2))) # per-geometry RMS error + stats[f"{chan}::err"] = err + stats[f"{chan}::unc"] = float(np.mean(unc)) # per-geometry mean uncertainty + return stats + + @staticmethod + def _regroup( + collected: dict[str, list[float]], + ) -> dict[str, dict[str, list[float]]]: + """Regroup ``{channel}::{err|unc}`` -> ``{channel: {"err": [...], "unc": [...]}}``.""" + by_channel: dict[str, dict[str, list[float]]] = {} + for key, vals in collected.items(): + if "::" not in key: + continue + chan, stat = key.rsplit("::", 1) + by_channel.setdefault(chan, {})[stat] = vals + return by_channel + + def finalize_samples( + self, collected: dict[str, list[float]] + ) -> float | dict[str, float]: + by_channel = self._regroup(collected) + if not by_channel: + # No per-geometry scalars (e.g. deterministic wrapper): return the same headline + # sub-keys a populated finalize would, as NaN, for a consistent report schema. + return {"mean": float("nan"), "mean_spearman": float("nan")} + out: dict[str, float] = {} + values: list[float] = [] + rhos: list[float] = [] + for chan in sorted(by_channel): + err = np.asarray(by_channel[chan].get("err", []), dtype=np.float64) + unc = np.asarray(by_channel[chan].get("unc", []), dtype=np.float64) + if err.size != unc.size or err.size < 2: + out[chan] = float("nan") + out[f"{chan}_spearman"] = float("nan") + continue + v = _ause_curve_area(err, unc) + out[chan] = float(v) + if v == v: # skip NaN in the headline mean + values.append(float(v)) + # Trend-alignment companion to AUSE: does higher per-geometry uncertainty rank the + # higher-error geometries (rho -> 1)? AUSE measures the *area* gap to the oracle; + # Spearman measures monotonic ordering, independent of error scale. + rho = _spearman(err, unc) + out[f"{chan}_spearman"] = float(rho) + if rho == rho: + rhos.append(float(rho)) + out["mean"] = float(np.mean(values)) if values else float("nan") + out["mean_spearman"] = float(np.mean(rhos)) if rhos else float("nan") + return out + + def curves(self, collected: dict[str, list[float]]) -> dict[str, dict[str, Any]]: + """Per-channel sparsification curves for the plot (one series per field channel).""" + by_channel = self._regroup(collected) + out: dict[str, dict[str, Any]] = {} + for chan in sorted(by_channel): + err = np.asarray(by_channel[chan].get("err", []), dtype=np.float64) + unc = np.asarray(by_channel[chan].get("unc", []), dtype=np.float64) + if err.size != unc.size: + continue + payload = _curve_payload(err, unc) + if payload is not None: + out[chan] = payload + return out + + +# --- sample-wise drag sparsification (SAMPLE metric): linear UQ propagation into Cd ---------- +# +# Drag is a *linear* functional of the surface fields, so the predicted-drag mean is the surface +# integral of the predicted mean field and the drag *variance* is the area/normal-weighted sum of +# the per-cell field variances. This propagates the per-cell MARGINAL std the engine produces for +# EVERY UQ method (an analytic GP posterior's marginals OR the sampling across-pass spread), so the +# same caveat applies to all of them — it is NOT specific to analytic models: treating cells as +# independent is exact for spatially-uncorrelated uncertainty but a lower bound for +# spatially-correlated (i.e. epistemic) uncertainty, which dominates a coherent surface integral. +# This is the decision-relevant panel: does the field UQ flag the geometries whose drag we predict +# worst? Direct port of ``field_gp_utils.compute_drag_uq_stats``, but reading physical-unit +# fields straight off the benchmark comparison mesh (so no normalization factors are needed). + + +def _drag_mean_and_std( + normals: np.ndarray, + area: np.ndarray, + direction: np.ndarray, + p: np.ndarray, + wss: np.ndarray, + p_std: np.ndarray, + wss_std: np.ndarray, + coeff: float, +) -> tuple[float, float]: + """Integrated drag coefficient and its linear-propagated std for one geometry (physical units). + + ``normals`` (N,3) cell unit normals, ``area`` (N,) cell areas, ``direction`` (3,) force + direction, ``p`` (N,) pressure, ``wss`` (N,3) wall shear stress, ``p_std`` (N,) / ``wss_std`` + (N,3) the per-cell field std to propagate. Mirrors ``compute_force_coefficients`` for the mean + and the per-channel drag weights ``w_p = coeff*(n·f)*a`` / ``w_tau = -coeff*a*f`` for the + variance ``Var[Cd] = Σ w_p² p_std² + Σ w_tau² wss_std²``. + """ + n_dot_f = normals @ direction # (N,) + c_p = coeff * float(np.sum(n_dot_f * area * p)) + c_f = -coeff * float(np.sum((wss @ direction) * area)) + drag = c_p + c_f + w_p = coeff * n_dot_f * area # (N,) + w_tau = -coeff * area[:, None] * direction[None, :] # (N, 3) + var = float(np.sum(w_p**2 * p_std**2)) + float(np.sum(w_tau**2 * wss_std**2)) + return drag, math.sqrt(max(var, 0.0)) + + +def _cell_array(mesh: Any, name: str | None) -> np.ndarray | None: + """Return ``mesh.cell_data[name]`` as ``float64`` (or ``None`` when absent).""" + if name is None: + return None + try: + cd = mesh.cell_data + if name in cd: + return np.asarray(cd[name], dtype=np.float64) + except (AttributeError, KeyError, TypeError): + return None + return None + + +class _SampleDragUQ: + """Sample-wise drag sparsification via linear UQ propagation into the ``Cd`` integral (surface). + + ``partial`` reduces one geometry to ``(drag_abs_err, drag_epistemic_std, drag_total_std)`` by + integrating the comparison mesh's predicted / GT pressure + WSS and propagating the per-cell + predictive std into the drag integral. ``finalize_samples`` computes the AUSE of ranking + geometries by drag epistemic-std and by drag total-std against ``|Cd_pred - Cd_true|``, and adds + ``_spearman`` (rank correlation of drag error vs. drag uncertainty) as the + scale-independent trend-alignment companion to the drag AUSE. + Requires the comparison mesh to carry cell-dof std companions. The benchmark attaches those + for any distribution-valued prediction using the *same* names as :func:`uq_std_field_names` + (configured ``output.std_mesh_field_names`` / ``epistemic_std_mesh_field_names`` when given, + else the auto-derived ``Std`` / ``EpistemicStd``), so drag_uq resolves them with + that shared helper and works with the default output config. It is a no-op for deterministic + wrappers or when std fields are unavailable. + + ``coeff`` (Cd prefactor ``2/(A·ρ·U²)``; only sets the overall scale, which cancels in AUSE) and + ``drag_direction`` are configurable exactly like the deterministic ``drag`` metric + (:func:`~physicsnemo.cfd.evaluation.metrics.builtin.forces.drag_error`): the constructor sets + the registration-time defaults and ``partial`` accepts per-run overrides from the metric's + config ``kwargs`` (e.g. ``- {name: drag_uq, coeff: 2.5, drag_direction: [1, 0, 0]}``). + """ + + def __init__( + self, + *, + coeff: float = 1.0, + drag_direction: tuple[float, float, float] = (1.0, 0.0, 0.0), + ) -> None: + self._coeff = float(coeff) + self._dir = tuple(float(x) for x in drag_direction) + + def partial( + self, + gt: Any, + predictions: Any, + *, + case: Any = None, + comparison_mesh: Any = None, + metric_dtype: str | None = None, + output: Any = None, + coeff: float | None = None, + drag_direction: list[float] | None = None, + **_: Any, + ) -> dict[str, float]: + predictions = predictions or {} + # Per-run overrides from config kwargs, else the registration-time defaults. + used_coeff = self._coeff if coeff is None else float(coeff) + direction = np.asarray( + self._dir if drag_direction is None else drag_direction, dtype=np.float64 + ) + pdist = as_distribution(predictions, "pressure") + wdist = as_distribution(predictions, "shear_stress") + # Deterministic wrapper (no std) or missing fields: nothing to propagate. + if ( + output is None + or pdist is None + or wdist is None + or pdist.std is None + or wdist.std is None + ): + return {} + from physicsnemo.cfd.evaluation.metrics.mesh_bridge import ( + resolve_comparison_mesh_for_metric, + uq_std_field_names, + ) + + mesh, dtype = resolve_comparison_mesh_for_metric( + predictions, + case=case, + comparison_mesh=comparison_mesh, + metric_dtype=metric_dtype, + output=output, + ) + # Drag integration here uses explicit *cell* normals/areas; require the cell dof (the + # config's surface_interpolate_point_to_cell_for_metrics guarantees it). + if mesh is None or dtype != "cell": + return {} + + prp = output.mesh_field_names.get("pressure") + prw = output.mesh_field_names.get("shear_stress") + gtp = output.ground_truth_mesh_field_names.get("pressure") + gtw = output.ground_truth_mesh_field_names.get("shear_stress") + # Uncertainty array names: mirror the mesh-attachment fallback so drag_uq works with the + # documented *default* output config (no explicit std_mesh_field_names). When pressure/WSS + # are not mapped to the mesh at all, ``_cell_array(None)`` below simply yields None. + std_p, epi_p = ( + uq_std_field_names( + prp, + output.std_mesh_field_names.get("pressure"), + output.epistemic_std_mesh_field_names.get("pressure"), + ) + if prp is not None + else (None, None) + ) + std_w, epi_w = ( + uq_std_field_names( + prw, + output.std_mesh_field_names.get("shear_stress"), + output.epistemic_std_mesh_field_names.get("shear_stress"), + ) + if prw is not None + else (None, None) + ) + + # Explicit per-cell normals + areas (physically correct surface integral; consistent + # weights for mean and variance). + geom = mesh.compute_normals( + cell_normals=True, point_normals=False, inplace=False + ) + geom = geom.compute_cell_sizes(length=False, area=True, volume=False) + normals = np.asarray(geom.cell_data["Normals"], dtype=np.float64) + area = np.asarray(geom.cell_data["Area"], dtype=np.float64).reshape(-1) + + p_pred = _cell_array(geom, prp) + wss_pred = _cell_array(geom, prw) + p_true = _cell_array(geom, gtp) + wss_true = _cell_array(geom, gtw) + p_std = _cell_array(geom, std_p) + wss_std = _cell_array(geom, std_w) + p_epi = _cell_array(geom, epi_p) + wss_epi = _cell_array(geom, epi_w) + if any(a is None for a in (p_pred, wss_pred, p_true, wss_true, p_std, wss_std)): + return {} + + drag_pred, drag_total_std = _drag_mean_and_std( + normals, area, direction, p_pred, wss_pred, p_std, wss_std, used_coeff + ) + drag_true, _ = _drag_mean_and_std( + normals, + area, + direction, + p_true, + wss_true, + np.zeros_like(p_true), + np.zeros_like(wss_true), + used_coeff, + ) + stats = { + "abs_err": abs(drag_pred - drag_true), + "total_std": drag_total_std, + } + if p_epi is not None and wss_epi is not None: + _, drag_epi_std = _drag_mean_and_std( + normals, area, direction, p_pred, wss_pred, p_epi, wss_epi, used_coeff + ) + stats["epi_std"] = drag_epi_std + return stats + + def finalize_samples( + self, collected: dict[str, list[float]] + ) -> float | dict[str, float]: + err = np.asarray(collected.get("abs_err", []), dtype=np.float64) + out: dict[str, float] = {} + for sub, key in (("epistemic", "epi_std"), ("total", "total_std")): + unc = np.asarray(collected.get(key, []), dtype=np.float64) + if err.size == unc.size and err.size >= 2: + out[sub] = _ause_curve_area(err, unc) + # Trend alignment: rank correlation of per-geometry drag error vs drag uncertainty. + out[f"{sub}_spearman"] = _spearman(err, unc) + else: + out[sub] = float("nan") + out[f"{sub}_spearman"] = float("nan") + return out + + def curves(self, collected: dict[str, list[float]]) -> dict[str, dict[str, Any]]: + """Drag sparsification curves ranked by epistemic and total drag std.""" + err = np.asarray(collected.get("abs_err", []), dtype=np.float64) + out: dict[str, dict[str, Any]] = {} + for sub, key in (("epistemic", "epi_std"), ("total", "total_std")): + unc = np.asarray(collected.get(key, []), dtype=np.float64) + if err.size != unc.size: + continue + payload = _curve_payload(err, unc) + if payload is not None: + out[sub] = payload + return out + + +def _make_uq_metrics() -> dict[str, _PooledUQMetric]: + """Instantiate the pooled UQ reducer metrics (shared across surface/volume domains).""" + return { + "nlpd": _PooledUQMetric(_nlpd_contrib_total, _mean_ratio("sum")), + "nlpd_epistemic": _PooledUQMetric( + _nlpd_contrib_epistemic, _mean_ratio("sum"), need_epistemic=True + ), + "calibration_zrms": _PooledUQMetric(_zrms_contrib, _zrms_finalize), + "coverage_95": _PooledUQMetric(_coverage95_contrib, _mean_ratio("within")), + "sharpness_std": _PooledUQMetric( + _sharpness_total_contrib, _mean_ratio("sum_sig") + ), + "sharpness_epistemic_std": _PooledUQMetric( + _sharpness_epistemic_contrib, _mean_ratio("sum_sig"), need_epistemic=True + ), + } + + +def _make_pointwise_uq_metrics() -> dict[str, _UncertaintyErrorSpearman]: + """Instantiate the pointwise (per-case) UQ metrics: Spearman(|error|, uncertainty).""" + return { + "uncertainty_error_spearman": _UncertaintyErrorSpearman(rank_epistemic=False), + "uncertainty_error_spearman_epistemic": _UncertaintyErrorSpearman( + rank_epistemic=True + ), + } + + +def _make_sample_uq_metrics() -> dict[str, Any]: + """Instantiate the sample-wise (per-geometry) UQ metrics: field + drag sparsification AUSE.""" + return { + "sparsification_ause": _SampleAUSE(rank_epistemic=False), + "sparsification_ause_epistemic": _SampleAUSE(rank_epistemic=True), + } + + +def register_uq_metrics() -> None: + """Register the pooled (reducer), pointwise, and sample-wise UQ metrics for both domains.""" + for domain in ("surface", "volume"): + for name, metric in _make_uq_metrics().items(): + register_metric(name, metric, domain=domain) + for name, metric in _make_pointwise_uq_metrics().items(): + register_metric(name, metric, domain=domain) + for name, metric in _make_sample_uq_metrics().items(): + register_metric(name, metric, domain=domain) + # Drag UQ integrates a surface force, so it is registered for the surface domain only. + register_metric("drag_uq", _SampleDragUQ(), domain="surface") diff --git a/physicsnemo/cfd/evaluation/metrics/mesh_bridge.py b/physicsnemo/cfd/evaluation/metrics/mesh_bridge.py index ad43d1f..137068c 100644 --- a/physicsnemo/cfd/evaluation/metrics/mesh_bridge.py +++ b/physicsnemo/cfd/evaluation/metrics/mesh_bridge.py @@ -24,12 +24,76 @@ import pyvista as pv from physicsnemo.cfd.evaluation.config import OutputConfig -from physicsnemo.cfd.evaluation.datasets.schema import CanonicalCase +from physicsnemo.cfd.evaluation.datasets.schema import ( + CanonicalCase, + FieldDistribution, + distribution_mean, +) from physicsnemo.cfd.postprocessing_tools.interpolation.interpolate_mesh_to_pc import ( interpolate_point_data_to_cell_centers, ) +def _unwrap_prediction_means(predictions: dict[str, Any] | None) -> dict[str, Any]: + """Return a predictions view with each :class:`FieldDistribution` replaced by its ``mean``. + + Deterministic array values pass through unchanged. Used so dof inference and mesh field + assignment (which expect plain arrays) work for probabilistic wrappers too. + """ + if not predictions: + return {} + return {k: distribution_mean(v) for k, v in predictions.items()} + + +def uq_std_field_names( + pred_name: str, + std_name: str | None = None, + epi_std_name: str | None = None, +) -> tuple[str, str]: + """Resolve the ``(std, epistemic_std)`` VTK array names for a prediction field. + + Single source of truth for the auto-derived uncertainty array names so the mesh-attachment + path (:func:`_attach_uq_std_companions`) and any metric that reads those arrays back off the + comparison mesh (e.g. ``drag_uq``) stay in agreement: use the configured name when given, else + suffix the prediction field name with ``"Std"`` / ``"EpistemicStd"``. + """ + return (std_name or f"{pred_name}Std", epi_std_name or f"{pred_name}EpistemicStd") + + +def _attach_uq_std_companions( + mesh: pv.DataSet, + preference: str, + predictions: dict[str, Any], + pairs: tuple[tuple[str, str], ...], + pred_map: dict[str, str], + std_map: dict[str, str], + epi_std_map: dict[str, str], +) -> list[str]: + """Attach ``*_std`` / ``*_epistemic_std`` arrays for any prediction that is a distribution. + + Names come from the configured std maps, else are auto-derived by suffixing the prediction + field name with ``"Std"`` / ``"EpistemicStd"``. Missing std channels are skipped silently. + Returns the list of attached array names so callers can interpolate them to cell dofs + alongside the mean fields (keeping std on the same dof as its prediction). + """ + attached: list[str] = [] + for canonical, _ in pairs: + value = predictions.get(canonical) + if not isinstance(value, FieldDistribution) or canonical not in pred_map: + continue + pred_name = pred_map[canonical] + std_name, epi_std_name = uq_std_field_names( + pred_name, std_map.get(canonical), epi_std_map.get(canonical) + ) + if value.std is not None: + _assign_field(mesh, preference, std_name, value.std) + attached.append(std_name) + if value.epistemic_std is not None: + _assign_field(mesh, preference, epi_std_name, value.epistemic_std) + attached.append(epi_std_name) + return attached + + def _infer_surface_preference( mesh: pv.PolyData, gt: dict[str, Any], @@ -210,20 +274,27 @@ def build_comparison_mesh( # Fresh read; no second consumer, no need to copy. mesh = pv.read(case.mesh_path) gt = case.ground_truth or {} + # Probabilistic wrappers return FieldDistribution; unwrap to the mean for dof inference and + # field assignment. Std channels are attached separately below. + pred_means = _unwrap_prediction_means(predictions) if case.inference_domain == "surface": gt_map = output.ground_truth_mesh_field_names pred_map = output.mesh_field_names + std_map = output.std_mesh_field_names + epi_std_map = output.epistemic_std_mesh_field_names pairs = (("pressure", "pressure"), ("shear_stress", "shear_stress")) if not isinstance(mesh, pv.PolyData): mesh = mesh.extract_surface() # Do not call point_data_to_cell_data here: it can change n_cells vs the mesh the # adapter used when extracting GT, causing "expected N cell values, got ..." mismatches. - preference = _infer_surface_preference(mesh, gt, case.mesh_type, predictions) + preference = _infer_surface_preference(mesh, gt, case.mesh_type, pred_means) elif case.inference_domain == "volume": - preference = _infer_volume_preference(mesh, gt, predictions, case.mesh_type) + preference = _infer_volume_preference(mesh, gt, pred_means, case.mesh_type) gt_map = output.ground_truth_volume_mesh_field_names pred_map = output.volume_mesh_field_names + std_map = output.std_volume_mesh_field_names + epi_std_map = output.epistemic_std_volume_mesh_field_names pairs = ( ("pressure", "pressure"), ("velocity", "velocity"), @@ -237,10 +308,14 @@ def build_comparison_mesh( _assign_field(mesh, preference, gt_map[canonical], gt[canonical]) if ( canonical in pred_map - and canonical in predictions - and predictions[canonical] is not None + and canonical in pred_means + and pred_means[canonical] is not None ): - _assign_field(mesh, preference, pred_map[canonical], predictions[canonical]) + _assign_field(mesh, preference, pred_map[canonical], pred_means[canonical]) + + std_names = _attach_uq_std_companions( + mesh, preference, predictions or {}, pairs, pred_map, std_map, epi_std_map + ) if ( case.inference_domain == "surface" @@ -253,6 +328,9 @@ def build_comparison_mesh( names.append(gt_map[canonical]) if canonical in pred_map: names.append(pred_map[canonical]) + # Interpolate the UQ std companions on the same dof as their mean fields so metrics that + # read std off cells (e.g. drag_uq) and the saved comparison mesh stay dof-consistent. + names.extend(std_names) names = list(dict.fromkeys(names)) interpolate_point_data_to_cell_centers( mesh, diff --git a/physicsnemo/cfd/evaluation/metrics/registry.py b/physicsnemo/cfd/evaluation/metrics/registry.py index ecf6e03..c8b8fdc 100644 --- a/physicsnemo/cfd/evaluation/metrics/registry.py +++ b/physicsnemo/cfd/evaluation/metrics/registry.py @@ -17,10 +17,23 @@ """Re-export bench metric registry for ``physicsnemo.cfd.evaluation``.""" from physicsnemo.cfd.postprocessing_tools.metric_registry import ( # noqa: F401 + ReducerMetric, + SampleMetric, get_metric, + is_reducer_metric, + is_sample_metric, list_metrics, register_metric, unregister_metric, ) -__all__ = ["register_metric", "unregister_metric", "get_metric", "list_metrics"] +__all__ = [ + "register_metric", + "unregister_metric", + "get_metric", + "list_metrics", + "ReducerMetric", + "is_reducer_metric", + "SampleMetric", + "is_sample_metric", +] diff --git a/physicsnemo/cfd/evaluation/models/common_wrapper_utils/geotransolver_runtime.py b/physicsnemo/cfd/evaluation/models/common_wrapper_utils/geotransolver_runtime.py new file mode 100644 index 0000000..2d79f11 --- /dev/null +++ b/physicsnemo/cfd/evaluation/models/common_wrapper_utils/geotransolver_runtime.py @@ -0,0 +1,609 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Shared GeoTransolver + ``TransolverDataPipe`` runtime helpers (composition, not inheritance). + +The GeoTransolver-family wrappers are **independent** :class:`CFDModel` subclasses (no shared base +beyond ``CFDModel``): the DrivAerML baseline, the ``transformer_models`` deterministic wrapper, the +GP-head wrapper, the MC-Dropout wrapper, and the ensemble wrapper. They share their plumbing through +the free functions here rather than class inheritance: + +* :func:`build_geotransolver_backbone` — construct the backbone and load its checkpoint. +* :func:`parse_runtime_kwargs` — pull the common inference kwargs into a :class:`GeoTransolverRuntimeConfig`. +* :func:`build_transolver_batch` — read a case, (re)build the datapipe, and produce the model batch. +* :func:`geotransolver_forward` / :func:`unscale_targets` — the blocked forward pass + target unscaling. +* :func:`decode_surface_predictions` / :func:`decode_volume_predictions` — physical-unit decode. + +The ``REDIMENSIONALIZE_OUTPUTS`` decision (whether to re-dimensionalize by ``ρu²`` / ``u`` / ``u·L``) +is owned by each wrapper and passed explicitly into the decode helpers. +""" + +from __future__ import annotations + +from contextlib import nullcontext +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +import numpy as np +import torch + +from physicsnemo.cfd.evaluation.common.checkpoint_compat import ( + parse_checkpoint_epoch, + trusted_torch_load_context, +) +from physicsnemo.cfd.evaluation.config import _parse_bool +from physicsnemo.cfd.evaluation.datasets.schema import ( + CanonicalCase, + build_predictions_dict, +) +from physicsnemo.cfd.evaluation.inference.progress import log_inference +from physicsnemo.cfd.evaluation.models.common_wrapper_utils.vtk_datapipe_io import ( + build_surface_data_dict, + build_volume_data_dict, + run_id_from_case_id, +) +from physicsnemo.cfd.evaluation.models.inference_autocast import cuda_bf16_autocast +from physicsnemo.cfd.evaluation.models.model_registry import Predictions, RawOutput + +# Optional physicsnemo imports (required for real inference). +try: + from physicsnemo.datapipes.cae.transolver_datapipe import TransolverDataPipe + from physicsnemo.distributed import DistributedManager + from physicsnemo.experimental.models.geotransolver import GeoTransolver + from physicsnemo.utils import load_model_weights + + _PHYSICSNEMO_AVAILABLE = True +except ImportError: + _PHYSICSNEMO_AVAILABLE = False + + +# Default GeoTransolver surface config (from geotransolver_surface + model/geotransolver.yaml). +DEFAULT_GEOTRANSOLVER_KW = dict( + functional_dim=6, + global_dim=2, + geometry_dim=3, + out_dim=4, + n_layers=20, + n_hidden=256, + dropout=0.0, + n_head=8, + act="gelu", + mlp_ratio=2, + slice_num=128, + use_te=False, + plus=False, + include_local_features=True, + radii=[0.01, 0.05, 0.25, 1.0, 2.5, 5.0], + neighbors_in_radius=[4, 8, 16, 64, 128, 256], + n_hidden_local=32, +) + +# Volume training defaults (geotransolver_volume.yaml). +DEFAULT_GEOTRANSOLVER_VOLUME_KW = { + **DEFAULT_GEOTRANSOLVER_KW, + "functional_dim": 7, + "out_dim": 5, +} + +#: ``model.kwargs`` keys forwarded to ``TransolverDataPipe`` (overrides of the static defaults). +DATAPIPE_KEYS = frozenset( + { + "include_normals", + "include_sdf", + "translational_invariance", + "scale_invariance", + "reference_scale", + "broadcast_global_features", + "include_geometry", + "return_mesh_features", + } +) + + +def geotransolver_available() -> bool: + """True when physicsnemo (GeoTransolver, TransolverDataPipe, load_model_weights) is importable.""" + return _PHYSICSNEMO_AVAILABLE + + +def global_fx_to_bnc(fx: torch.Tensor) -> torch.Tensor: + """GeoTransolver requires ``global_embedding`` shape (B, N_g, C_g). + + With ``broadcast_global_features=False``, ``TransolverDataPipe`` may stack ``fx`` as + (1, 1, C) before ``__call__`` adds another batch dimension, yielding (1, 1, 1, C). + Squeeze singleton middle dims until 3D. + """ + out = fx + while out.ndim > 3: + for d in range(1, out.ndim - 1): + if out.shape[d] == 1: + out = out.squeeze(d) + break + else: + raise ValueError( + "Cannot reshape global ``fx`` to 3D for GeoTransolver; " + f"shape={tuple(fx.shape)}" + ) + return out + + +def _surface_datapipe_static_kw() -> dict[str, Any]: + return dict( + include_normals=True, + include_sdf=False, + broadcast_global_features=False, + include_geometry=True, + translational_invariance=True, + scale_invariance=True, + reference_scale=[12.0, 4.5, 3.25], + return_mesh_features=True, + ) + + +def _volume_datapipe_static_kw() -> dict[str, Any]: + """Defaults aligned with ``data/core.yaml`` + ``geotransolver_volume.yaml``.""" + return dict( + include_normals=True, + include_sdf=True, + translational_invariance=True, + scale_invariance=True, + reference_scale=[12.0, 4.5, 3.25], + broadcast_global_features=False, + include_geometry=True, + return_mesh_features=False, + ) + + +@dataclass +class GeoTransolverRuntimeConfig: + """Common inference knobs parsed once in ``load`` and reused by every runtime helper.""" + + device: str = "cuda:0" + air_density: float = 1.205 + stream_velocity: float = 30.0 + batch_resolution: int = 2048 + cuda_bf16_autocast: bool = False + geometry_sampling: int = 300_000 + datapipe_user_kw: dict[str, Any] = field(default_factory=dict) + + +def split_datapipe_kwargs(kw: dict[str, Any]) -> dict[str, Any]: + """Pop ``TransolverDataPipe`` overrides out of a mutable ``model.kwargs`` dict.""" + return {k: kw.pop(k) for k in list(kw.keys()) if k in DATAPIPE_KEYS} + + +def parse_runtime_kwargs(kw: dict[str, Any], device: str) -> GeoTransolverRuntimeConfig: + """Parse the shared runtime kwargs into a :class:`GeoTransolverRuntimeConfig` (mutates ``kw``). + + Pops ``cuda_bf16_autocast``, the datapipe override keys, and the ignored ``resolution`` from + ``kw`` (benchmark inference always uses the full mesh, ``resolution=None``); the remaining + entries in ``kw`` are read non-destructively so a caller can inspect them afterwards. + """ + cfg = GeoTransolverRuntimeConfig( + device=device, + air_density=float(kw.get("air_density", 1.205)), + stream_velocity=float(kw.get("stream_velocity", 30.0)), + batch_resolution=int(kw.get("batch_resolution", 2048)), + cuda_bf16_autocast=_parse_bool( + kw.pop("cuda_bf16_autocast", None), default=False + ), + geometry_sampling=int(kw.get("geometry_sampling", 300_000)), + ) + # Benchmark inference uses all mesh points; ``resolution`` in model kwargs is ignored. + kw.pop("resolution", None) + cfg.datapipe_user_kw = split_datapipe_kwargs(kw) + return cfg + + +def ensure_distributed_initialized() -> None: + """Initialize ``DistributedManager`` once (the weight loaders rely on it).""" + if not DistributedManager.is_initialized(): + DistributedManager.initialize() + + +def resolve_checkpoint_file(checkpoint_path: str) -> tuple[Path, int]: + """Resolve a checkpoint knob to ``(directory, epoch)`` from a **specific, existing file**. + + The config must point ``checkpoint`` at a concrete checkpoint file whose name encodes the + epoch (``.0..mdlus`` or ``checkpoint.0..pt``) rather than a directory. + This makes the loaded epoch a single source of truth (no silent "latest epoch" fallback and + no drift between the backbone and a separately-loaded head). + + The file **must exist**: a mistyped or missing checkpoint raises :class:`FileNotFoundError` + here rather than letting the loader log-and-skip and silently return a randomly-initialized + model. Validation order is name-shape first (so a malformed config gives a clear message even + if the file happens to be absent), then existence. + """ + ckpt = Path(checkpoint_path) + if not checkpoint_path: + raise ValueError( + "checkpoint path is empty; point it at a specific checkpoint file." + ) + if ckpt.is_dir(): + raise ValueError( + f"checkpoint {checkpoint_path!r} is a directory; point it at a specific checkpoint " + "file whose name encodes the epoch (e.g. ``GeoTransolver.0.30.mdlus`` or " + "``checkpoint.0.100.pt``) so the loaded epoch is unambiguous." + ) + epoch = parse_checkpoint_epoch(checkpoint_path) + if epoch is None: + raise ValueError( + f"cannot parse an epoch from checkpoint file name {ckpt.name!r}; expected " + "``.0..mdlus`` or ``checkpoint.0..pt``." + ) + if not ckpt.is_file(): + raise FileNotFoundError( + f"checkpoint file {checkpoint_path!r} does not exist. A missing/mistyped checkpoint " + "would otherwise be skipped by the loader, leaving a randomly-initialized model." + ) + return ckpt.parent, epoch + + +def build_geotransolver_backbone( + *, + checkpoint_path: str, + device: str, + inference_mode: str, + concrete_dropout: bool = False, +) -> "GeoTransolver": + """Construct a ``GeoTransolver`` for ``surface``/``volume`` and load its checkpoint onto ``device``. + + ``checkpoint_path`` must be a specific checkpoint file (see :func:`resolve_checkpoint_file`); + a directory or an epoch-less name is rejected rather than silently loading the latest epoch. + + Set ``concrete_dropout=True`` to build the backbone with learned per-layer + :class:`~physicsnemo.nn.ConcreteDropout` layers, matching a checkpoint trained with + ``model.concrete_dropout=true`` (required for the MC-Dropout wrapper: without it the loader + would reject the extra ``p_logit`` parameters). The returned model is in ``eval()`` mode; + callers that want stochastic passes must re-enable the dropout layers. + + The exact file named by ``checkpoint_path`` is loaded (via + :func:`physicsnemo.utils.load_model_weights`); a missing file raises rather than silently + leaving the freshly-constructed random weights in place. + """ + model_kw = dict( + DEFAULT_GEOTRANSOLVER_VOLUME_KW + if inference_mode == "volume" + else DEFAULT_GEOTRANSOLVER_KW + ) + if concrete_dropout: + model_kw["concrete_dropout"] = True + # Strict: rejects a directory / epoch-less name and raises FileNotFoundError for a missing file. + resolve_checkpoint_file(checkpoint_path) + + ensure_distributed_initialized() + dev = torch.device(device) + model = GeoTransolver(**model_kw) + # Load the exact named file (not a dir + reconstructed name), so the checkpoint the config + # points at is the one that is loaded, and a missing file raises here. + with trusted_torch_load_context(): + load_model_weights(model, checkpoint_path, device=dev) + model = model.to(dev) + model.eval() + return model + + +def _move_reference_scale_to_device(dp: "TransolverDataPipe", device: str) -> None: + """``reference_scale`` is often created on CPU; mesh tensors live on ``device``.""" + if dp.config.scale_invariance and dp.config.reference_scale is not None: + dp.config.reference_scale = dp.config.reference_scale.to(torch.device(device)) + + +def build_surface_datapipe( + *, + geometry_sampling: int, + surface_factors: Any, + device: str, + user_kw: dict[str, Any], +) -> "TransolverDataPipe": + """Build a surface ``TransolverDataPipe`` (static defaults merged with ``user_kw``).""" + merged = {**_surface_datapipe_static_kw(), **user_kw} + merged["geometry_sampling"] = geometry_sampling + dp = TransolverDataPipe( + input_path=None, + model_type="surface", + resolution=None, + surface_factors=surface_factors, + volume_factors=None, + scaling_type="mean_std_scaling", + **merged, + ) + _move_reference_scale_to_device(dp, device) + return dp + + +def build_volume_datapipe( + *, geometry_sampling: int, volume_factors: Any, device: str, user_kw: dict[str, Any] +) -> "TransolverDataPipe": + """Build a volume ``TransolverDataPipe`` (static defaults merged with ``user_kw``).""" + merged = {**_volume_datapipe_static_kw(), **user_kw} + merged["geometry_sampling"] = geometry_sampling + dp = TransolverDataPipe( + input_path=None, + model_type="volume", + resolution=None, + surface_factors=None, + volume_factors=volume_factors, + scaling_type="mean_std_scaling", + **merged, + ) + _move_reference_scale_to_device(dp, device) + return dp + + +def _volume_length_scale(data_dict: dict[str, Any]) -> float: + """STL bounding-box max extent (matches DoMINO ``length_scale``), used to unscale νₜ by ``u·L``.""" + stl = data_dict["stl_coordinates"] + return float((stl.amax(dim=0) - stl.amin(dim=0)).max().item()) + + +@dataclass +class TransolverBatch: + """Result of :func:`build_transolver_batch`: the model batch plus the (cached) datapipe.""" + + batch: dict[str, Any] + datapipe: "TransolverDataPipe" + geometry_effective: int + volume_length_scale: float | None = None + + +def build_transolver_batch( + *, + case: CanonicalCase, + inference_mode: str, + cfg: GeoTransolverRuntimeConfig, + surface_factors: Any, + volume_factors: Any, + datapipe: "TransolverDataPipe | None", + geometry_effective: int | None, +) -> TransolverBatch: + """Read a case, (re)build the datapipe when the effective sampling changed, and run it. + + ``datapipe`` / ``geometry_effective`` are the caller's cached state; pass the returned + :class:`TransolverBatch` fields back next time to reuse the datapipe across cases of equal size. + """ + log_inference( + "geotransolver", + f"Reading case inputs (case {case.case_id}): mesh {case.mesh_path}, " + f"run dir {Path(case.mesh_path).parent}", + ) + run_dir = Path(case.mesh_path).parent + run_idx = run_id_from_case_id(case.case_id) + device = torch.device(cfg.device) + length_scale: float | None = None + # In-memory geometry (surface == geometry for e.g. DrivAerStar) bypasses the STL file glob. + geometry_mesh = getattr(case, "geometry", None) + + if inference_mode == "volume": + data_dict = build_volume_data_dict( + run_dir=run_dir, + vtu_path=case.mesh_path, + device=device, + air_density=cfg.air_density, + stream_velocity=cfg.stream_velocity, + run_idx=run_idx, + reference_mesh=case.reference_geometry, + geometry_mesh=geometry_mesh, + ) + length_scale = _volume_length_scale(data_dict) + n_stl = int(data_dict["stl_coordinates"].shape[0]) + safe_geo = max(1, min(cfg.geometry_sampling, n_stl)) + if datapipe is None or geometry_effective != safe_geo: + datapipe = build_volume_datapipe( + geometry_sampling=safe_geo, + volume_factors=volume_factors, + device=cfg.device, + user_kw=cfg.datapipe_user_kw, + ) + geometry_effective = safe_geo + else: + data_dict = build_surface_data_dict( + run_dir=run_dir, + vtp_path=case.mesh_path, + device=device, + air_density=cfg.air_density, + stream_velocity=cfg.stream_velocity, + run_idx=run_idx, + reference_mesh=case.reference_geometry, + geometry_mesh=geometry_mesh, + ) + n_stl = int(data_dict["stl_coordinates"].shape[0]) + n_surf = int(data_dict["surface_mesh_centers"].shape[0]) + safe_geo = max(1, min(cfg.geometry_sampling, n_stl, n_surf)) + if datapipe is None or geometry_effective != safe_geo: + datapipe = build_surface_datapipe( + geometry_sampling=safe_geo, + surface_factors=surface_factors, + device=cfg.device, + user_kw=cfg.datapipe_user_kw, + ) + geometry_effective = safe_geo + + batch = datapipe(data_dict) + return TransolverBatch( + batch=batch, + datapipe=datapipe, + geometry_effective=int(safe_geo), + volume_length_scale=length_scale, + ) + + +def make_forward_permutation(batch: dict[str, Any]) -> torch.Tensor: + """A random point permutation for :func:`geotransolver_forward` (one per case). + + Multi-pass UQ wrappers (ensemble / MC-dropout) should build this **once per case** and pass it + to every pass/member so the across-pass spread reflects genuine model/dropout differences, not + the random block-partition grouping. Note the permutation always covers *all* ``n`` points + (every point is predicted exactly once, order restored), so it never subsamples the cloud — it + only controls which points share a forward block. + """ + n = batch["embeddings"].shape[1] + return torch.randperm(n, device=batch["embeddings"].device) + + +def geotransolver_forward( + *, + model: "GeoTransolver", + batch: dict[str, Any], + batch_resolution: int, + cuda_bf16_autocast_enabled: bool, + device: str, + perm: torch.Tensor | None = None, +) -> torch.Tensor: + """Blocked GeoTransolver forward pass; returns model-space predictions (N, C), order restored. + + ``perm`` optionally fixes the point permutation used to form forward blocks (see + :func:`make_forward_permutation`); when ``None`` a fresh random permutation is drawn. Either + way every point is predicted exactly once and the output is returned in the original order. + """ + fx_bn_c = global_fx_to_bnc(batch["fx"]) + n = batch["embeddings"].shape[1] + batch_res = min(batch_resolution, n) + indices = ( + perm + if perm is not None + else torch.randperm(n, device=batch["embeddings"].device) + ) + index_blocks = torch.split(indices, batch_res) + preds_list: list[torch.Tensor] = [] + use_full_fx = "geometry" in batch + + ac_ctx = cuda_bf16_autocast(device) if cuda_bf16_autocast_enabled else nullcontext() + with torch.no_grad(): + with ac_ctx: + for index_block in index_blocks: + local_embeddings = batch["embeddings"][:, index_block] + local_fx = fx_bn_c if use_full_fx else fx_bn_c[:, index_block] + local_positions = local_embeddings[:, :, :3] + geometry_kw = batch["geometry"] if "geometry" in batch else None + outputs = model( + local_embedding=local_embeddings, + local_positions=local_positions, + global_embedding=local_fx, + geometry=geometry_kw, + ) + preds_list.append(outputs) + predictions = torch.cat(preds_list, dim=1) + inverse_indices = torch.empty_like(indices) + inverse_indices[indices] = torch.arange(n, device=indices.device) + predictions = predictions[:, inverse_indices] + return predictions.squeeze(0) + + +def unscale_targets( + *, + datapipe: "TransolverDataPipe", + predictions: torch.Tensor, + batch: dict[str, Any], + inference_mode: str, +) -> torch.Tensor: + """Invert the datapipe target scaling (``unscale_model_targets``) for surface/volume.""" + return datapipe.unscale_model_targets( + predictions, + air_density=batch.get("air_density"), + stream_velocity=batch.get("stream_velocity"), + factor_type="volume" if inference_mode == "volume" else "surface", + ) + + +def decode_surface_predictions( + raw_output: RawOutput, + *, + redimensionalize: bool, + air_density: float, + stream_velocity: float, +) -> Predictions: + """Map unscaled surface targets to canonical ``pressure`` / ``shear_stress`` (physical units). + + ``redimensionalize`` multiplies by the dynamic pressure ``ρu²`` for non-dimensional-target + (DrivAerML) checkpoints; it is the identity for physical-target (``transformer_models``) ones. + """ + pred = raw_output + if pred.dim() == 3: + pred = pred.squeeze(0) + log_inference("geotransolver", "Decoding outputs (pressure + WSS to numpy)…") + dynamic_pressure = air_density * (stream_velocity**2) if redimensionalize else 1.0 + pressure = (pred[:, 0] * dynamic_pressure).cpu().numpy().astype(np.float32) + wss = (pred[:, 1:4] * dynamic_pressure).cpu().numpy().astype(np.float32) + return build_predictions_dict(pressure=pressure, shear_stress=wss) + + +def decode_volume_predictions( + raw_output: RawOutput, + *, + redimensionalize: bool, + air_density: float, + stream_velocity: float, + length_scale: float | None, +) -> Predictions: + """Map unscaled volume targets to canonical ``velocity`` / ``pressure`` / ``turbulent_viscosity``. + + ``redimensionalize`` applies the ``u`` (velocity), ``ρu²`` (pressure) and ``u·L`` (νₜ) scales + for non-dimensional-target checkpoints; identity for physical-target ones. + """ + pred = raw_output + if pred.dim() == 3: + pred = pred.squeeze(0) + log_inference( + "geotransolver", + "Decoding outputs (velocity + pressure + nut → canonical volume keys)…", + ) + if redimensionalize: + if length_scale is None: + raise RuntimeError( + "Volume decode requires the STL length scale; prepare_inputs must run " + "before decode_outputs." + ) + u = float(stream_velocity) + rho = float(air_density) + dynamic_pressure = rho * (u**2) + nut_scale = u * length_scale + else: + u = 1.0 + dynamic_pressure = 1.0 + nut_scale = 1.0 + velocity = (pred[:, 0:3] * u).cpu().numpy().astype(np.float32) + pressure = (pred[:, 3] * dynamic_pressure).cpu().numpy().astype(np.float32) + turbulent_viscosity = (pred[:, 4] * nut_scale).cpu().numpy().astype(np.float32) + return build_predictions_dict( + velocity=velocity, + pressure=pressure, + turbulent_viscosity=turbulent_viscosity, + ) + + +__all__ = [ + "DEFAULT_GEOTRANSOLVER_KW", + "DEFAULT_GEOTRANSOLVER_VOLUME_KW", + "DATAPIPE_KEYS", + "GeoTransolverRuntimeConfig", + "TransolverBatch", + "geotransolver_available", + "global_fx_to_bnc", + "split_datapipe_kwargs", + "parse_runtime_kwargs", + "ensure_distributed_initialized", + "resolve_checkpoint_file", + "build_geotransolver_backbone", + "build_surface_datapipe", + "build_volume_datapipe", + "build_transolver_batch", + "make_forward_permutation", + "geotransolver_forward", + "unscale_targets", + "decode_surface_predictions", + "decode_volume_predictions", +] diff --git a/physicsnemo/cfd/evaluation/models/common_wrapper_utils/vtk_datapipe_io.py b/physicsnemo/cfd/evaluation/models/common_wrapper_utils/vtk_datapipe_io.py index 5de8712..b5584fc 100644 --- a/physicsnemo/cfd/evaluation/models/common_wrapper_utils/vtk_datapipe_io.py +++ b/physicsnemo/cfd/evaluation/models/common_wrapper_utils/vtk_datapipe_io.py @@ -34,10 +34,16 @@ ) -def read_stl_geometry(stl_path: str, device: torch.device) -> dict[str, torch.Tensor]: - """Read STL and return stl_coordinates, stl_faces, stl_centers for SDF/center of mass.""" - mesh_raw = pv.read(stl_path) - mesh = triangulate_surface_mesh(mesh_raw) +def stl_geometry_from_mesh( + mesh: pv.DataSet, device: torch.device +) -> dict[str, torch.Tensor]: + """Build stl_coordinates / stl_faces / stl_centers from an in-memory geometry mesh. + + Same payload as :func:`read_stl_geometry` (points, triangle connectivity, cell centers for + SDF / center-of-mass) but from a mesh already in memory — no file read. Used when a dataset's + surface *is* its geometry (e.g. DrivAerStar), so the datapipe geometry branch needs no STL file. + """ + mesh = triangulate_surface_mesh(mesh) stl_coordinates = torch.from_numpy(np.asarray(mesh.points)).to( device=device, dtype=torch.float32 ) @@ -53,6 +59,11 @@ def read_stl_geometry(stl_path: str, device: torch.device) -> dict[str, torch.Te } +def read_stl_geometry(stl_path: str, device: torch.device) -> dict[str, torch.Tensor]: + """Read STL and return stl_coordinates, stl_faces, stl_centers for SDF/center of mass.""" + return stl_geometry_from_mesh(pv.read(stl_path), device) + + def read_surface_from_vtp( vtp_path: str, device: torch.device, @@ -159,6 +170,19 @@ def _find_stl_in_dir(run_dir: Path, run_idx: int) -> Path: raise FileNotFoundError(f"No STL file found in {run_dir} for run_idx {run_idx}") +def _resolve_stl_geometry( + *, + run_dir: Path, + run_idx: int, + device: torch.device, + geometry_mesh: pv.DataSet | None, +) -> dict[str, torch.Tensor]: + """Geometry tensors from an in-memory mesh when given, else from the STL file in ``run_dir``.""" + if geometry_mesh is not None: + return stl_geometry_from_mesh(geometry_mesh, device) + return read_stl_geometry(str(_find_stl_in_dir(run_dir, run_idx)), device) + + def build_volume_data_dict( run_dir: Path, vtu_path: str, @@ -169,16 +193,19 @@ def build_volume_data_dict( n_output_fields: int = 5, *, reference_mesh: pv.DataSet | None = None, + geometry_mesh: pv.DataSet | None = None, ) -> dict[str, torch.Tensor]: """Build data dict for volume inference: STL + VTU + flow params (DrivAer-style run dir). - STL resolution matches :func:`build_surface_data_dict` — ``drivaer_.stl``, + When ``geometry_mesh`` is given the geometry tensors come from it (no STL file); otherwise the + STL is resolved from ``run_dir`` — ``drivaer_.stl``, ``drivaer__single_solid.stl``, then ``*_single_solid.stl``, then any ``*.stl``. Model evaluation locations are mesh points (see :func:`read_volume_from_vtu`). """ - stl_path = _find_stl_in_dir(run_dir, run_idx) - data_dict = read_stl_geometry(str(stl_path), device) + data_dict = _resolve_stl_geometry( + run_dir=run_dir, run_idx=run_idx, device=device, geometry_mesh=geometry_mesh + ) data_dict.update( read_volume_from_vtu( vtu_path, @@ -205,10 +232,16 @@ def build_surface_data_dict( run_idx: int = 1, *, reference_mesh: pv.DataSet | None = None, + geometry_mesh: pv.DataSet | None = None, ) -> dict[str, torch.Tensor]: - """Build data dict for surface inference: STL + VTP + flow params. Finds STL in run_dir.""" - stl_path = _find_stl_in_dir(run_dir, run_idx) - data_dict = read_stl_geometry(str(stl_path), device) + """Build data dict for surface inference: STL + VTP + flow params. + + When ``geometry_mesh`` is given the geometry tensors come from it (no STL file); otherwise the + STL is resolved from ``run_dir`` (see :func:`_find_stl_in_dir`). + """ + data_dict = _resolve_stl_geometry( + run_dir=run_dir, run_idx=run_idx, device=device, geometry_mesh=geometry_mesh + ) data_dict.update(read_surface_from_vtp(vtp_path, device, mesh=reference_mesh)) data_dict["air_density"] = torch.tensor( [air_density], device=device, dtype=torch.float32 diff --git a/physicsnemo/cfd/evaluation/models/model_registry.py b/physicsnemo/cfd/evaluation/models/model_registry.py index bae1701..9bba0f7 100644 --- a/physicsnemo/cfd/evaluation/models/model_registry.py +++ b/physicsnemo/cfd/evaluation/models/model_registry.py @@ -17,10 +17,11 @@ """CFDModel base class and registry for model wrappers.""" from abc import ABC, abstractmethod -from typing import Any, ClassVar, Literal, Optional, Type +from typing import Any, ClassVar, Iterable, Literal, Optional, Type from physicsnemo.cfd.evaluation.datasets.schema import ( CanonicalCase, + FieldDistribution, InferenceDomain, normalize_inference_domain_str, ) @@ -52,6 +53,18 @@ class CFDModel(ABC): INFERENCE_DOMAIN: ClassVar[InferenceDomain | None] = "surface" REQUIRES_REMOTE_ASSETS: ClassVar[bool] = True + #: Whether this wrapper produces a predictive *distribution* (uncertainty), not just a + #: point estimate. Deterministic wrappers leave this ``False``; UQ metrics then report + #: ``NaN`` for them (consistent with the engine's recoverable-metric behavior). + SUPPORTS_UQ: ClassVar[bool] = False + #: How the predictive distribution is produced: + #: ``"analytic"`` — one forward pass emits the distribution/params (GP, mean-variance, + #: evidential); the wrapper overrides :meth:`decode_distribution`. + #: ``"sampling"`` — the distribution is built from statistics over ``N`` stochastic + #: passes / ensemble members; the engine drives the passes and aggregates. + #: ``"none"`` — deterministic. + UQ_METHOD: ClassVar[Literal["none", "analytic", "sampling"]] = "none" + @classmethod def inference_domain_from_kwargs( cls, kwargs: dict[str, Any] @@ -93,6 +106,21 @@ def predict(self, model_input: ModelInput) -> RawOutput: """Run forward pass; return raw model output.""" ... + def predict_deterministic(self, model_input: ModelInput) -> RawOutput: + """Single **deterministic** forward pass (used when ``run.uq.enabled`` is off). + + The engine calls this (not :meth:`predict`) on the deterministic path so that turning UQ + off yields a true point prediction for *every* wrapper. The default simply delegates to + :meth:`predict`, which is correct for deterministic and analytic wrappers (their + :meth:`predict` is already deterministic). + + **Sampling** wrappers whose :meth:`predict` is stochastic (e.g. MC-Dropout keeps dropout + masks active) MUST override this to remove the stochasticity — e.g. disable dropout for one + pass, or return a single ensemble member — otherwise ``run.uq.enabled=false`` would still + return a random draw rather than a deterministic prediction. + """ + return self.predict(model_input) + @abstractmethod def decode_outputs( self, @@ -107,6 +135,46 @@ def decode_outputs( """ ... + def decode_distribution( + self, + raw_output: RawOutput, + case: CanonicalCase, + model_input: Optional[ModelInput] = None, + ) -> dict[str, "FieldDistribution"]: + """Map raw output to a per-field predictive distribution (physical units). + + **Analytic** UQ wrappers (``UQ_METHOD="analytic"``, e.g. a GP head) override this to + return :class:`~physicsnemo.cfd.evaluation.datasets.schema.FieldDistribution` with a + real ``std`` / ``epistemic_std``. The default wraps :meth:`decode_outputs` as + **degenerate** distributions (``std=None``) so callers can uniformly request a + distribution from any wrapper. + """ + preds = self.decode_outputs(raw_output, case, model_input) + return {k: FieldDistribution(mean=v) for k, v in preds.items()} + + def predict_ensemble( + self, model_input: ModelInput, n: int + ) -> Optional[Iterable[RawOutput]]: + """Optional multi-pass path for ``UQ_METHOD="sampling"`` wrappers. + + Return an **iterable of raw outputs** (stochastic passes or ensemble members) — ideally a + lazy generator — that the engine iterates, folding each into a streaming Welford + mean/variance + (:func:`~physicsnemo.cfd.evaluation.benchmarks.uq_inference.run_sampling_inference`). A + generator keeps only **one** raw output resident on the compute device at a time, so device + memory stays O(field), not O(n × field); a plain ``list`` materializes all outputs at once + and can OOM for large ``n`` / full-mesh fields, so prefer a generator. + + ``n`` is the requested budget (``run.uq.num_samples``). Honor it: yield ``n`` stochastic + passes for a per-model sampler (e.g. MC-Dropout), or ``min(n, member_count)`` members for a + fixed-size ensemble (which cannot fabricate more distinct members than it holds). + + Return ``None`` (the default) to have the engine instead call :meth:`predict` ``n`` times + (reseeding per pass) and stream those — appropriate for a stochastic :meth:`predict` where a + batched path offers no benefit. + """ + return None + def register_model(name: str, wrapper_class: Type[CFDModel]) -> None: """Register a model wrapper by name.""" diff --git a/physicsnemo/cfd/evaluation/models/wrappers/__init__.py b/physicsnemo/cfd/evaluation/models/wrappers/__init__.py index a627f73..71d013c 100644 --- a/physicsnemo/cfd/evaluation/models/wrappers/__init__.py +++ b/physicsnemo/cfd/evaluation/models/wrappers/__init__.py @@ -22,6 +22,18 @@ from physicsnemo.cfd.evaluation.models.wrappers.geotransolver import ( GeoTransolverWrapper, ) +from physicsnemo.cfd.evaluation.models.wrappers.geotransolver_drivaerstar import ( + GeoTransolverDrivAerStarWrapper, +) +from physicsnemo.cfd.evaluation.models.wrappers.geotransolver_gp import ( + GeoTransolverGPDrivAerStarWrapper, +) +from physicsnemo.cfd.evaluation.models.wrappers.ensemble_drivaerstar import ( + GeoTransolverEnsembleDrivAerStarWrapper, +) +from physicsnemo.cfd.evaluation.models.wrappers.mc_dropout import ( + GeoTransolverMCDropoutDrivAerStarWrapper, +) from physicsnemo.cfd.evaluation.models.wrappers.surface_baseline import ( SurfaceBaselineWrapper, ) @@ -35,6 +47,14 @@ register_model("xmgn_surface", XMGNWrapper) register_model("geotransolver_surface", GeoTransolverWrapper) register_model("geotransolver_volume", GeoTransolverWrapper) +register_model("geotransolver_drivaerstar_surface", GeoTransolverDrivAerStarWrapper) +register_model("geotransolver_gp_surface", GeoTransolverGPDrivAerStarWrapper) +register_model( + "geotransolver_mc_dropout_surface", GeoTransolverMCDropoutDrivAerStarWrapper +) +register_model( + "geotransolver_ensemble_surface", GeoTransolverEnsembleDrivAerStarWrapper +) register_model("transolver_surface", TransolverWrapper) register_model("transolver_volume", TransolverWrapper) register_model("domino_surface", DominoWrapper) @@ -46,6 +66,10 @@ "FIGNetWrapper", "XMGNWrapper", "GeoTransolverWrapper", + "GeoTransolverDrivAerStarWrapper", + "GeoTransolverGPDrivAerStarWrapper", + "GeoTransolverEnsembleDrivAerStarWrapper", + "GeoTransolverMCDropoutDrivAerStarWrapper", "TransolverWrapper", "DominoWrapper", "SurfaceBaselineWrapper", diff --git a/physicsnemo/cfd/evaluation/models/wrappers/ensemble_drivaerstar/__init__.py b/physicsnemo/cfd/evaluation/models/wrappers/ensemble_drivaerstar/__init__.py new file mode 100644 index 0000000..fc9dc8a --- /dev/null +++ b/physicsnemo/cfd/evaluation/models/wrappers/ensemble_drivaerstar/__init__.py @@ -0,0 +1,21 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from physicsnemo.cfd.evaluation.models.wrappers.ensemble_drivaerstar.wrapper import ( + GeoTransolverEnsembleDrivAerStarWrapper, +) + +__all__ = ["GeoTransolverEnsembleDrivAerStarWrapper"] diff --git a/physicsnemo/cfd/evaluation/models/wrappers/ensemble_drivaerstar/wrapper.py b/physicsnemo/cfd/evaluation/models/wrappers/ensemble_drivaerstar/wrapper.py new file mode 100644 index 0000000..8fec309 --- /dev/null +++ b/physicsnemo/cfd/evaluation/models/wrappers/ensemble_drivaerstar/wrapper.py @@ -0,0 +1,312 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Deep-/snapshot-ensemble sampling UQ wrapper for GeoTransolver (surface). + +An ensemble aggregates the predictions of several independently-saved model checkpoints (a sibling +sampling UQ method to :mod:`..mc_dropout`, which instead resamples dropout masks on one model). The +engine drives it through the SAME ``UQ_METHOD="sampling"`` path — but here each "pass" is a genuine +model, so the across-member spread +is a meaningful epistemic uncertainty (subject to the caveat below). + +.. note:: + As configured for the benchmark, the members are the **last K checkpoints of a single training + run** (a *snapshot ensemble*), NOT K independently-initialized runs (a *deep ensemble*). Snapshot + members are correlated, so the ensemble tends to be **under-dispersed** (over-confident) relative + to a true deep ensemble — but it needs no extra training and is a concrete, honest multi-model + example. Point ``checkpoint`` at checkpoints from separate runs (via ``member_checkpoints``) for a + true deep ensemble. + +Completely independent :class:`CFDModel` subclass (no inheritance from the other GeoTransolver +wrappers). It composes the shared GeoTransolver + ``TransolverDataPipe`` plumbing in +:mod:`physicsnemo.cfd.evaluation.models.common_wrapper_utils.geotransolver_runtime`, holding one +backbone per member. :meth:`predict_ensemble` runs up to ``run.uq.num_samples`` members on the +shared per-case batch and **yields** their raw outputs one at a time (a generator, so device memory +stays O(field) rather than O(K × field)); the engine's +:func:`~physicsnemo.cfd.evaluation.benchmarks.uq_inference.run_sampling_inference` folds each into a +streaming Welford mean + across-member (epistemic) std. All members share one per-case block +partition so the spread is model-only. ``run.uq.num_samples`` is honored as a budget: +``min(num_samples, member_count)`` members run (a fixed-size ensemble cannot fabricate more distinct +members than it holds; when ``num_samples`` exceeds the member count all members run and a warning is +logged). + +It decorates ``transformer_models`` (physical-target) checkpoints, so predictions are +re-standardized only — no dynamic-pressure re-dimensionalization +(:attr:`REDIMENSIONALIZE_OUTPUTS` = ``False``). + +Model kwargs (``model.kwargs``): + +- ``member_checkpoints`` (list[str], **required**): the explicit member checkpoint files, one per + ensemble member (any number ``K``). Each must be a specific **model-weights** file whose name + encodes an epoch (``.0..mdlus`` or a bare state-dict ``.pt``) — NOT the + ``checkpoint.0..pt`` training-state file; the exact file is loaded and a missing one + raises. List members from a single run (snapshot ensemble) or from separate runs (true deep + ensemble) — same code path either way. +- plus the shared GeoTransolver runtime kwargs (``batch_resolution``, ``geometry_sampling``, + ``cuda_bf16_autocast``; datapipe overrides; etc.); ``stats_path`` is the shared normalization. + +The top-level ``checkpoint`` field is still required by the benchmark (it anchors the run's asset +identity / cache fingerprint); point it at any one of the members (e.g. the final epoch). +""" + +import logging +from pathlib import Path +from typing import Any, ClassVar, Iterator, Literal, Optional + +from physicsnemo.cfd.evaluation.datasets.schema import ( + CanonicalCase, + InferenceDomain, + coerce_inference_domain_or_default, +) +from physicsnemo.cfd.evaluation.common.io import ( + load_transolver_surface_factors, + load_transolver_volume_factors, +) +from physicsnemo.cfd.evaluation.inference.progress import log_inference +from physicsnemo.cfd.evaluation.models.common_wrapper_utils.geotransolver_runtime import ( + GeoTransolverRuntimeConfig, + build_geotransolver_backbone, + build_transolver_batch, + decode_surface_predictions, + decode_volume_predictions, + geotransolver_available, + geotransolver_forward, + make_forward_permutation, + parse_runtime_kwargs, + unscale_targets, +) +from physicsnemo.cfd.evaluation.models.model_registry import ( + CFDModel, + ModelInput, + OutputLocation, + Predictions, + RawOutput, +) + +_LOG = logging.getLogger(__name__) + + +class GeoTransolverEnsembleDrivAerStarWrapper(CFDModel): + """K-member (snapshot/deep) ensemble over GeoTransolver checkpoints (surface). + + Decorates ``transformer_models`` (physical-target) checkpoints, so predictions are + re-standardized only (no dynamic-pressure re-dimensionalization): + :attr:`REDIMENSIONALIZE_OUTPUTS` = ``False``. + """ + + INFERENCE_DOMAIN: ClassVar[InferenceDomain | None] = None + OUTPUT_LOCATION: ClassVar[OutputLocation] = "cell" + REDIMENSIONALIZE_OUTPUTS: ClassVar[bool] = False + # UQ contract with the engine: "I produce uncertainty via repeated passes." predict_ensemble() + # returns one pass per member, and the engine aggregates the spread into mean + epistemic std. + SUPPORTS_UQ: ClassVar[bool] = True + UQ_METHOD: ClassVar[str] = "sampling" + + @property + def output_location(self) -> OutputLocation: + """See :attr:`CFDModel.output_location` (GeoTransolver predicts at cell centers).""" + return self.OUTPUT_LOCATION + + @classmethod + def inference_domain_from_kwargs( + cls, kwargs: dict[str, Any] + ) -> InferenceDomain | None: + """Align benchmark routing with :meth:`load` (default surface when omitted).""" + return coerce_inference_domain_or_default( + kwargs.get("inference_domain"), + default="surface", + parameter="model.kwargs.inference_domain", + ) + + def __init__(self) -> None: + self._models: list[Any] = [] + self._datapipe: Any = None + self._datapipe_geometry_effective: Optional[int] = None + self._surface_factors: Any = None + self._volume_factors: Any = None + self._inference_mode: Literal["surface", "volume"] = "surface" + self._volume_length_scale: Optional[float] = None + self._cfg: GeoTransolverRuntimeConfig = GeoTransolverRuntimeConfig() + + def load( + self, + checkpoint_path: str, + stats_path: str, + device: str, + **kwargs: Any, + ) -> "GeoTransolverEnsembleDrivAerStarWrapper": + """Build one GeoTransolver backbone per explicitly-listed member checkpoint.""" + if not geotransolver_available(): + raise RuntimeError( + "GeoTransolverEnsembleDrivAerStarWrapper requires physicsnemo (GeoTransolver, " + "TransolverDataPipe, load_model_weights)." + ) + kw = dict(kwargs) + member_checkpoints = kw.pop("member_checkpoints", None) + self._inference_mode = coerce_inference_domain_or_default( + kw.pop("inference_domain", None), + default="surface", + parameter="model.kwargs.inference_domain", + ) + self._cfg = parse_runtime_kwargs(kw, device) + + if self._inference_mode == "volume": + log_inference("ensemble", f"Loading volume normalization from {stats_path}") + self._volume_factors = load_transolver_volume_factors(stats_path, device) + if self._volume_factors is None: + raise FileNotFoundError( + "Volume inference requires ``global_stats.json`` or " + f"``volume_fields_normalization.npz`` (looked under {stats_path!r})." + ) + self._surface_factors = None + else: + log_inference( + "ensemble", f"Loading surface normalization from {stats_path}" + ) + self._surface_factors = load_transolver_surface_factors(stats_path, device) + self._volume_factors = None + + if not member_checkpoints: + raise ValueError( + "GeoTransolverEnsembleDrivAerStarWrapper requires `member_checkpoints`: an explicit list of " + "checkpoint files (one per member). Auto-discovery from a directory is not " + "supported — list every member path in model.kwargs.member_checkpoints." + ) + members = [str(p) for p in member_checkpoints] + + self._datapipe = None + self._datapipe_geometry_effective = None + # One backbone per member. Members are ~O(100 MB) each; a handful fit comfortably on device. + # Each is loaded from its own checkpoint file (strict epoch-in-name resolution). + self._models = [ + build_geotransolver_backbone( + checkpoint_path=member, + device=device, + inference_mode=self._inference_mode, + ) + for member in members + ] + _LOG.info( + "GeoTransolverEnsembleDrivAerStarWrapper: loaded %d members from %s", + len(self._models), + [Path(m).name for m in members], + ) + return self + + def prepare_inputs(self, case: CanonicalCase) -> ModelInput: + """Build the surface/volume batch once per case (shared by every member).""" + if not self._models: + raise RuntimeError( + "GeoTransolverEnsembleDrivAerStarWrapper: call load() first" + ) + result = build_transolver_batch( + case=case, + inference_mode=self._inference_mode, + cfg=self._cfg, + surface_factors=self._surface_factors, + volume_factors=self._volume_factors, + datapipe=self._datapipe, + geometry_effective=self._datapipe_geometry_effective, + ) + self._datapipe = result.datapipe + self._datapipe_geometry_effective = result.geometry_effective + if result.volume_length_scale is not None: + self._volume_length_scale = result.volume_length_scale + return {"batch": result.batch, "datapipe": result.datapipe} + + def _forward_member( + self, model: Any, model_input: ModelInput, perm: Any = None + ) -> RawOutput: + """Run one member's forward + target unscaling on the shared per-case batch. + + ``perm`` fixes the forward block-partition so every member sees the same point grouping; + the across-member spread then reflects genuine model differences, not partition noise. + """ + raw = geotransolver_forward( + model=model, + batch=model_input["batch"], + batch_resolution=self._cfg.batch_resolution, + cuda_bf16_autocast_enabled=self._cfg.cuda_bf16_autocast, + device=self._cfg.device, + perm=perm, + ) + return unscale_targets( + datapipe=model_input["datapipe"], + predictions=raw, + batch=model_input["batch"], + inference_mode=self._inference_mode, + ) + + def predict_ensemble( + self, model_input: ModelInput, n: int + ) -> Optional[Iterator[RawOutput]]: + """Yield up to ``n`` member raw outputs (lazy generator), sharing one per-case permutation. + + ``n`` (``run.uq.num_samples``) is honored as a budget: this yields ``min(n, member_count)`` + members. A fixed-size ensemble cannot fabricate more distinct members than it holds, so when + ``n`` exceeds the member count all members are used and a warning is logged (the two sampling + rows may then run different pass counts — e.g. 32 MC-Dropout passes vs 5 ensemble members). + This is a **generator**, so only one member's output is resident at a time — the engine's + Welford accumulator consumes each before the next is produced (O(field) device memory, not + O(K × field)). All members use the same fixed block-partition permutation so the spread is + model-only. + """ + if not self._models or self._datapipe is None: + raise RuntimeError( + "GeoTransolverEnsembleDrivAerStarWrapper: call load() first" + ) + k = min(int(n), len(self._models)) + if int(n) > len(self._models): + _LOG.warning( + "num_samples=%d exceeds ensemble member count %d; using all %d members.", + int(n), + len(self._models), + len(self._models), + ) + perm = make_forward_permutation(model_input["batch"]) + return ( + self._forward_member(model, model_input, perm) for model in self._models[:k] + ) + + def predict(self, model_input: ModelInput) -> RawOutput: + """Deterministic fallback (first member). The engine uses :meth:`predict_ensemble` for UQ.""" + if not self._models or self._datapipe is None: + raise RuntimeError( + "GeoTransolverEnsembleDrivAerStarWrapper: call load() first" + ) + return self._forward_member(self._models[0], model_input) + + def decode_outputs( + self, + raw_output: RawOutput, + case: CanonicalCase, + model_input: Optional[ModelInput] = None, + ) -> Predictions: + """Emit canonical surface/volume keys in physical units (no ``ρu²`` re-dimensionalization).""" + if self._inference_mode == "volume": + return decode_volume_predictions( + raw_output, + redimensionalize=self.REDIMENSIONALIZE_OUTPUTS, + air_density=self._cfg.air_density, + stream_velocity=self._cfg.stream_velocity, + length_scale=self._volume_length_scale, + ) + return decode_surface_predictions( + raw_output, + redimensionalize=self.REDIMENSIONALIZE_OUTPUTS, + air_density=self._cfg.air_density, + stream_velocity=self._cfg.stream_velocity, + ) diff --git a/physicsnemo/cfd/evaluation/models/wrappers/geotransolver_drivaerstar/__init__.py b/physicsnemo/cfd/evaluation/models/wrappers/geotransolver_drivaerstar/__init__.py new file mode 100644 index 0000000..d502538 --- /dev/null +++ b/physicsnemo/cfd/evaluation/models/wrappers/geotransolver_drivaerstar/__init__.py @@ -0,0 +1,21 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from physicsnemo.cfd.evaluation.models.wrappers.geotransolver_drivaerstar.wrapper import ( + GeoTransolverDrivAerStarWrapper, +) + +__all__ = ["GeoTransolverDrivAerStarWrapper"] diff --git a/physicsnemo/cfd/evaluation/models/wrappers/geotransolver_drivaerstar/wrapper.py b/physicsnemo/cfd/evaluation/models/wrappers/geotransolver_drivaerstar/wrapper.py new file mode 100644 index 0000000..89145ff --- /dev/null +++ b/physicsnemo/cfd/evaluation/models/wrappers/geotransolver_drivaerstar/wrapper.py @@ -0,0 +1,211 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""GeoTransolver wrapper for the DrivAerStar (``transformer_models``) convention. + +This module is **purely additive**: it does not modify the original +:class:`~physicsnemo.cfd.evaluation.models.wrappers.geotransolver.wrapper.GeoTransolverWrapper` +(the DrivAerML / DoMINO non-dimensional path). It only adds a sibling wrapper for checkpoints +trained with the ``examples/cfd/external_aerodynamics/transformer_models`` pipeline, whose +targets are standardized **physical** fields. + +Both wrappers are independent :class:`CFDModel` subclasses and compose the shared plumbing in +:mod:`physicsnemo.cfd.evaluation.models.common_wrapper_utils.geotransolver_runtime`. +""" + +from typing import Any, ClassVar, Literal, Optional + +from physicsnemo.cfd.evaluation.common.io import ( + load_transolver_surface_factors, + load_transolver_volume_factors, +) +from physicsnemo.cfd.evaluation.datasets.schema import ( + CanonicalCase, + InferenceDomain, + coerce_inference_domain_or_default, +) +from physicsnemo.cfd.evaluation.inference.progress import log_inference +from physicsnemo.cfd.evaluation.models.common_wrapper_utils.geotransolver_runtime import ( + GeoTransolverRuntimeConfig, + build_geotransolver_backbone, + build_transolver_batch, + decode_surface_predictions, + decode_volume_predictions, + geotransolver_available, + geotransolver_forward, + parse_runtime_kwargs, + unscale_targets, +) +from physicsnemo.cfd.evaluation.models.model_registry import ( + CFDModel, + ModelInput, + OutputLocation, + Predictions, + RawOutput, +) + + +class GeoTransolverDrivAerStarWrapper(CFDModel): + """GeoTransolver on **DrivAerStar** (``examples/.../transformer_models`` convention, physical targets). + + Registered as ``geotransolver_drivaerstar_surface``. Identical mechanics to + :class:`~physicsnemo.cfd.evaluation.models.wrappers.geotransolver.wrapper.GeoTransolverWrapper`, + but the training targets are **standardized physical fields** (``preprocess.py``: + ``(x-mean)/std``), so ``TransolverDataPipe.unscale_model_targets`` already returns physical units + (``norm*std + mean``) and :meth:`decode_outputs` must **not** re-dimensionalize by ``ρu²`` — + matching ``field_gp_utils`` and the DrivAerStar ``.vtk`` ground truth. Completely independent of + ``GeoTransolverWrapper`` (both subclass ``CFDModel`` and compose the same runtime helpers); only + :attr:`REDIMENSIONALIZE_OUTPUTS` differs. + """ + + INFERENCE_DOMAIN: ClassVar[InferenceDomain | None] = None + OUTPUT_LOCATION: ClassVar[OutputLocation] = "cell" + REDIMENSIONALIZE_OUTPUTS: ClassVar[bool] = False + + @property + def output_location(self) -> OutputLocation: + """See :attr:`CFDModel.output_location` (GeoTransolver predicts at cell centers).""" + return self.OUTPUT_LOCATION + + @classmethod + def inference_domain_from_kwargs( + cls, kwargs: dict[str, Any] + ) -> InferenceDomain | None: + """Align benchmark routing with :meth:`load` (default surface when omitted).""" + return coerce_inference_domain_or_default( + kwargs.get("inference_domain"), + default="surface", + parameter="model.kwargs.inference_domain", + ) + + def __init__(self) -> None: + self._model: Any = None + self._datapipe: Any = None + self._datapipe_geometry_effective: Optional[int] = None + self._surface_factors: Any = None + self._volume_factors: Any = None + self._inference_mode: Literal["surface", "volume"] = "surface" + self._volume_length_scale: Optional[float] = None + self._cfg: GeoTransolverRuntimeConfig = GeoTransolverRuntimeConfig() + + def load( + self, + checkpoint_path: str, + stats_path: str, + device: str, + **kwargs: Any, + ) -> "GeoTransolverDrivAerStarWrapper": + """Load the GeoTransolver backbone + physical-target normalization factors onto ``device``.""" + if not geotransolver_available(): + raise RuntimeError( + "GeoTransolver wrapper requires physicsnemo (GeoTransolver, " + "TransolverDataPipe, load_model_weights)." + ) + kw = dict(kwargs) + self._inference_mode = coerce_inference_domain_or_default( + kw.pop("inference_domain", None), + default="surface", + parameter="model.kwargs.inference_domain", + ) + self._cfg = parse_runtime_kwargs(kw, device) + + if self._inference_mode == "volume": + log_inference( + "geotransolver", f"Loading volume normalization from {stats_path}" + ) + self._volume_factors = load_transolver_volume_factors(stats_path, device) + if self._volume_factors is None: + raise FileNotFoundError( + "Volume inference requires ``global_stats.json`` (with velocity, " + "pressure, turbulent_viscosity) or ``volume_fields_normalization.npz`` " + f"next to stats/checkpoint (looked under {stats_path!r})." + ) + self._surface_factors = None + else: + log_inference( + "geotransolver", f"Loading surface normalization from {stats_path}" + ) + self._surface_factors = load_transolver_surface_factors(stats_path, device) + self._volume_factors = None + + self._datapipe = None + self._datapipe_geometry_effective = None + self._model = build_geotransolver_backbone( + checkpoint_path=checkpoint_path, + device=device, + inference_mode=self._inference_mode, + ) + return self + + def prepare_inputs(self, case: CanonicalCase) -> ModelInput: + """Build the surface/volume data dict, lazily (re)create the datapipe, and run it.""" + if self._model is None: + raise RuntimeError("GeoTransolverDrivAerStarWrapper: call load() first") + result = build_transolver_batch( + case=case, + inference_mode=self._inference_mode, + cfg=self._cfg, + surface_factors=self._surface_factors, + volume_factors=self._volume_factors, + datapipe=self._datapipe, + geometry_effective=self._datapipe_geometry_effective, + ) + self._datapipe = result.datapipe + self._datapipe_geometry_effective = result.geometry_effective + if result.volume_length_scale is not None: + self._volume_length_scale = result.volume_length_scale + return {"batch": result.batch, "datapipe": result.datapipe} + + def predict(self, model_input: ModelInput) -> RawOutput: + """Run blocked GeoTransolver forward passes and return unscaled (physical) targets.""" + if self._model is None or self._datapipe is None: + raise RuntimeError("GeoTransolverDrivAerStarWrapper: call load() first") + log_inference("geotransolver", "Running forward pass (predicting fields)…") + raw = geotransolver_forward( + model=self._model, + batch=model_input["batch"], + batch_resolution=self._cfg.batch_resolution, + cuda_bf16_autocast_enabled=self._cfg.cuda_bf16_autocast, + device=self._cfg.device, + ) + return unscale_targets( + datapipe=model_input["datapipe"], + predictions=raw, + batch=model_input["batch"], + inference_mode=self._inference_mode, + ) + + def decode_outputs( + self, + raw_output: RawOutput, + case: CanonicalCase, + model_input: Optional[ModelInput] = None, + ) -> Predictions: + """Emit canonical surface/volume keys in physical units (no ``ρu²`` re-dimensionalization).""" + if self._inference_mode == "volume": + return decode_volume_predictions( + raw_output, + redimensionalize=self.REDIMENSIONALIZE_OUTPUTS, + air_density=self._cfg.air_density, + stream_velocity=self._cfg.stream_velocity, + length_scale=self._volume_length_scale, + ) + return decode_surface_predictions( + raw_output, + redimensionalize=self.REDIMENSIONALIZE_OUTPUTS, + air_density=self._cfg.air_density, + stream_velocity=self._cfg.stream_velocity, + ) diff --git a/physicsnemo/cfd/evaluation/models/wrappers/geotransolver_gp/__init__.py b/physicsnemo/cfd/evaluation/models/wrappers/geotransolver_gp/__init__.py new file mode 100644 index 0000000..f6e7820 --- /dev/null +++ b/physicsnemo/cfd/evaluation/models/wrappers/geotransolver_gp/__init__.py @@ -0,0 +1,21 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from physicsnemo.cfd.evaluation.models.wrappers.geotransolver_gp.wrapper import ( + GeoTransolverGPDrivAerStarWrapper, +) + +__all__ = ["GeoTransolverGPDrivAerStarWrapper"] diff --git a/physicsnemo/cfd/evaluation/models/wrappers/geotransolver_gp/wrapper.py b/physicsnemo/cfd/evaluation/models/wrappers/geotransolver_gp/wrapper.py new file mode 100644 index 0000000..38f28de --- /dev/null +++ b/physicsnemo/cfd/evaluation/models/wrappers/geotransolver_gp/wrapper.py @@ -0,0 +1,473 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""GeoTransolver + Gaussian-Process field-head wrapper (analytic surface UQ). + +This is the ``UQ_METHOD="analytic"`` archetype and a **completely independent** :class:`CFDModel` +subclass (no inheritance from the deterministic GeoTransolver wrappers). It reuses the shared +GeoTransolver + ``TransolverDataPipe`` plumbing in +:mod:`physicsnemo.cfd.evaluation.models.common_wrapper_utils.geotransolver_runtime` (backbone +construction, VTP → model inputs) and swaps the deterministic readout for a pointwise multitask GP +head (:class:`physicsnemo.experimental.uq.FieldGPHead`): the GP posterior mean is the per-point +surface field and the posterior variance is the per-point uncertainty, in **one forward pass**. It +is a lift-and-shift of the standalone GP inference path +(``examples/.../transformer_models/src/inference_field_gp_zarr._predict_chunked``). + +It targets ``transformer_models`` (physical-target) checkpoints, so predictions are re-standardized +only — no dynamic-pressure re-dimensionalization (:attr:`REDIMENSIONALIZE_OUTPUTS` = ``False``). + +``decode_distribution`` returns, per canonical field, a +:class:`~physicsnemo.cfd.evaluation.datasets.schema.FieldDistribution` with ``mean``, total +``std`` (epistemic + observation noise), ``epistemic_std``, and ``aleatoric_std`` — all in +**physical units** (both mean and the std channels are denormalized here), so the pooled UQ +metrics consume them directly. + +Surface only (the GP head predicts 4 surface tasks: pressure + 3 wall-shear components). + +This wrapper loads **two** checkpoints — the GeoTransolver backbone and the ``FieldGPHead`` — and +both are named explicitly: + +- ``checkpoint`` must point at a **specific backbone file** whose name encodes the epoch + (``GeoTransolver.0..mdlus``); the epoch is parsed from that name. Pointing at a directory + (or an epoch-less name) is rejected — this removes the old ``checkpoint_epoch`` kwarg and the risk + of the backbone drifting to a "latest" epoch. +- ``gp_head_checkpoint`` should point at the matching ``FieldGPHead.0..pt``. The **exact + file named here is loaded** (via :func:`physicsnemo.utils.load_model_weights`), so a missing or + mistyped path raises rather than silently leaving the head randomly initialized (no cross- + validation against the backbone — passing matching checkpoints is the caller's responsibility). + It is **optional**: when omitted the wrapper falls back to the backbone's sibling + ``FieldGPHead.0..pt`` (which must exist). Either way the exact head file it loads is logged. + +Model kwargs (``model.kwargs`` in config), matching the trained checkpoint's GP settings: + +- ``gp_head_checkpoint`` (path): the ``FieldGPHead.0..pt`` to load (see above). +- ``gp_feature_norm`` (``"none"`` | ``"l2"`` | ``"layernorm"`` | ``"l2_radial"``): must match training. +- ``gp_lengthscale_range`` / ``gp_lengthscale_prior`` / ``gp_outputscale_prior``: GP kernel config. +- ``gp_n_inducing`` (default 256), ``gp_mlp_hidden`` (optional DKL MLP), ``num_tasks`` (default 4). +- ``gp_spectral_norm_coeff`` (float, default 0.0) / ``gp_dkl_residual`` (bool, default True): + DUE-style bi-Lipschitz DKL. When ``gp_spectral_norm_coeff > 0`` the head uses a spectral- + normalised + residual feature extractor (van Amersfoort et al., 2021) instead of the plain DKL + MLP; both must match training. The defaults reproduce the plain MLP (non-DUE) path unchanged. +- ``gp_inference_chunk_size`` (default 51200): points per backbone+GP chunk (drawn from a random + permutation, then inverted — see the standalone script for why contiguous chunks degrade preds). +- ``cuda_bf16_autocast`` (bool, default False): run forwards under bf16 autocast. +""" + +from contextlib import nullcontext +from pathlib import Path +from typing import Any, ClassVar, Optional + +import numpy as np +import torch + +from physicsnemo.cfd.evaluation.common.checkpoint_compat import ( + trusted_torch_load_context, +) +from physicsnemo.cfd.evaluation.common.io import load_transolver_surface_factors +from physicsnemo.cfd.evaluation.datasets.schema import ( + CanonicalCase, + FieldDistribution, + InferenceDomain, + build_predictions_dict, + build_predictive_distribution, + coerce_inference_domain_or_default, +) +from physicsnemo.cfd.evaluation.inference.progress import log_inference +from physicsnemo.cfd.evaluation.models.common_wrapper_utils.geotransolver_runtime import ( + GeoTransolverRuntimeConfig, + build_geotransolver_backbone, + build_transolver_batch, + geotransolver_available, + global_fx_to_bnc, + parse_runtime_kwargs, + resolve_checkpoint_file, +) +from physicsnemo.cfd.evaluation.models.inference_autocast import cuda_bf16_autocast +from physicsnemo.cfd.evaluation.models.model_registry import ( + CFDModel, + ModelInput, + OutputLocation, + Predictions, + RawOutput, +) + +try: + from physicsnemo.experimental.uq import FieldGPHead + from physicsnemo.utils import load_model_weights + + _GP_AVAILABLE = True +except ImportError: + _GP_AVAILABLE = False + +#: Surface field channels predicted by the GP head: pressure, wall-shear (x, y, z). +NUM_SURFACE_TASKS = 4 + + +class GeoTransolverGPDrivAerStarWrapper(CFDModel): + """GeoTransolver backbone + ``FieldGPHead`` for analytic per-point surface UQ. + + Trained on ``transformer_models`` (physical-target) checkpoints, so predictions are + re-standardized only (no dynamic-pressure re-dimensionalization): + :attr:`REDIMENSIONALIZE_OUTPUTS` = ``False``. Completely independent of the deterministic + GeoTransolver wrappers (subclasses ``CFDModel`` directly, composing the shared runtime helpers). + """ + + REDIMENSIONALIZE_OUTPUTS: ClassVar[bool] = False + INFERENCE_DOMAIN: ClassVar[InferenceDomain | None] = "surface" + OUTPUT_LOCATION: ClassVar[OutputLocation] = "cell" + SUPPORTS_UQ: ClassVar[bool] = True + UQ_METHOD: ClassVar[str] = "analytic" + + @property + def output_location(self) -> OutputLocation: + """See :attr:`CFDModel.output_location` (GeoTransolver predicts at cell centers).""" + return self.OUTPUT_LOCATION + + @classmethod + def inference_domain_from_kwargs( + cls, kwargs: dict[str, Any] + ) -> InferenceDomain | None: + """Surface only; reject volume explicitly (the GP head is a 4-task surface head).""" + dom = coerce_inference_domain_or_default( + kwargs.get("inference_domain"), + default="surface", + parameter="model.kwargs.inference_domain", + ) + if dom != "surface": + raise NotImplementedError( + "GeoTransolverGPDrivAerStarWrapper supports surface inference only; got " + f"inference_domain={dom!r}." + ) + return dom + + def __init__(self) -> None: + self._model: Any = None + self._datapipe: Any = None + self._datapipe_geometry_effective: Optional[int] = None + self._surface_factors: Any = None + self._cfg: GeoTransolverRuntimeConfig = GeoTransolverRuntimeConfig() + self._head: Optional[FieldGPHead] = None + self._gp_kw: dict[str, Any] = {} + self._checkpoint_epoch: Optional[int] = None + self._head_checkpoint: Optional[str] = None + self._gp_chunk_size: int = 51200 + self._field_mean_t: Optional[torch.Tensor] = None + self._field_std_t: Optional[torch.Tensor] = None + + def load( + self, + checkpoint_path: str, + stats_path: str, + device: str, + **kwargs: Any, + ) -> "GeoTransolverGPDrivAerStarWrapper": + """Build the GeoTransolver backbone + surface factors, and prepare the GP head. + + The backbone is constructed via the shared runtime helpers; the ``FieldGPHead`` is built + lazily on the first :meth:`predict` (its input dim is probed from the backbone features). + """ + if not _GP_AVAILABLE or not geotransolver_available(): + raise RuntimeError( + "GeoTransolverGPDrivAerStarWrapper requires physicsnemo with GeoTransolver + " + "physicsnemo.experimental.uq.FieldGPHead." + ) + kw = dict(kwargs) + # Surface-only; validate and force the routing hint before building the datapipe. + self.inference_domain_from_kwargs(kw) + kw.pop("inference_domain", None) + + # GP-head hyperparameters (must match the trained checkpoint). Pulled off here so they + # are not consumed by the shared runtime-kwargs parsing. + self._gp_chunk_size = int(kw.pop("gp_inference_chunk_size", 51200)) + # Explicit GP-head checkpoint file (optional). When omitted we fall back to the backbone's + # sibling ``FieldGPHead.0..pt``. No cross-checking against the backbone — passing a + # matching pair is the caller's responsibility. + head_ckpt_arg = kw.pop("gp_head_checkpoint", None) + self._gp_kw = { + "num_tasks": int(kw.pop("num_tasks", NUM_SURFACE_TASKS)), + "n_inducing": int(kw.pop("gp_n_inducing", 256)), + "mlp_hidden": ( + list(kw.pop("gp_mlp_hidden")) + if kw.get("gp_mlp_hidden") is not None + else kw.pop("gp_mlp_hidden", None) + ), + "lengthscale_range": tuple(kw.pop("gp_lengthscale_range", (0.01, 10.0))), + "lengthscale_prior": ( + tuple(kw.pop("gp_lengthscale_prior")) + if kw.get("gp_lengthscale_prior") is not None + else kw.pop("gp_lengthscale_prior", None) + ), + "outputscale_prior": ( + tuple(kw.pop("gp_outputscale_prior")) + if kw.get("gp_outputscale_prior") is not None + else kw.pop("gp_outputscale_prior", None) + ), + "feature_norm": str(kw.pop("gp_feature_norm", "none")), + # DUE-style bi-Lipschitz DKL extractor (van Amersfoort et al., 2021). When + # ``gp_spectral_norm_coeff > 0`` the head builds a spectral-normalised + residual + # feature extractor instead of the plain DKL MLP, so this MUST match training or the + # ``FieldGPHead`` state dict will not reconstruct. Defaults (0.0 / True) keep the plain + # MLP path, so existing (non-DUE) GP checkpoints are unaffected. + "spectral_norm_coeff": float(kw.pop("gp_spectral_norm_coeff", 0.0)), + "dkl_residual": bool(kw.pop("gp_dkl_residual", True)), + } + + self._cfg = parse_runtime_kwargs(kw, device) + log_inference( + "geotransolver_gp", f"Loading surface normalization from {stats_path}" + ) + self._surface_factors = load_transolver_surface_factors(stats_path, device) + if self._surface_factors is None: + raise FileNotFoundError( + "GeoTransolverGPDrivAerStarWrapper requires surface normalization " + "(``global_stats.json`` or ``surface_fields_normalization.npz``) at " + f"{stats_path!r}." + ) + # Cache flat (4,) standardization vectors for physical de-normalization of the GP outputs. + self._field_mean_t = self._surface_factors["mean"].reshape(-1).to(device) + self._field_std_t = self._surface_factors["std"].reshape(-1).to(device) + + self._datapipe = None + self._datapipe_geometry_effective = None + self._model = build_geotransolver_backbone( + checkpoint_path=checkpoint_path, + device=device, + inference_mode="surface", + ) + + # The head is loaded from the EXACT ``gp_head_checkpoint`` file when given (so the file the + # user names is the one read), else from the backbone's sibling ``FieldGPHead.0..pt``. + # Either way the file must exist (no cross-validation of contents — passing matching + # checkpoints is the caller's responsibility). + if head_ckpt_arg is not None: + # resolve_checkpoint_file validates existence + the epoch-in-name shape. + _, self._checkpoint_epoch = resolve_checkpoint_file(str(head_ckpt_arg)) + self._head_checkpoint = str(Path(head_ckpt_arg)) + else: + ckpt_dir, self._checkpoint_epoch = resolve_checkpoint_file(checkpoint_path) + self._head_checkpoint = str( + Path(ckpt_dir) / f"FieldGPHead.0.{self._checkpoint_epoch}.pt" + ) + if not Path(self._head_checkpoint).is_file(): + raise FileNotFoundError( + "GP head checkpoint not found next to the backbone: " + f"{self._head_checkpoint!r}. Pass model.kwargs.gp_head_checkpoint explicitly." + ) + log_inference( + "geotransolver_gp", + f"GP checkpoints -> backbone: {checkpoint_path} | head: {self._head_checkpoint}", + ) + return self + + def prepare_inputs(self, case: CanonicalCase) -> ModelInput: + """Build the surface data dict, lazily (re)create the datapipe, and run it (surface only).""" + if self._model is None: + raise RuntimeError("GeoTransolverGPDrivAerStarWrapper: call load() first") + result = build_transolver_batch( + case=case, + inference_mode="surface", + cfg=self._cfg, + surface_factors=self._surface_factors, + volume_factors=None, + datapipe=self._datapipe, + geometry_effective=self._datapipe_geometry_effective, + ) + self._datapipe = result.datapipe + self._datapipe_geometry_effective = result.geometry_effective + return {"batch": result.batch, "datapipe": result.datapipe} + + def _build_and_load_head(self, batch: dict) -> None: + """Probe the backbone feature dim on ``batch``, build the GP head, load ONLY its weights. + + The backbone was already loaded from its own ``checkpoint`` in :meth:`load`; here we load + just the ``FieldGPHead`` from the EXACT ``self._head_checkpoint`` file (the one named by + ``gp_head_checkpoint``, or the validated backbone sibling). Loading only the head from its + own file avoids reloading — or silently overwriting — the backbone, and a missing file + raises rather than leaving the head randomly initialized. + """ + dev = torch.device(self._cfg.device) + feature_dim = self._probe_feature_dim(batch) + self._head = FieldGPHead( + input_dim=feature_dim, + n_train=1, # unused at inference + **self._gp_kw, + ).to(dev) + with trusted_torch_load_context(): + load_model_weights(self._head, self._head_checkpoint, device=dev) + self._model.eval() + self._head.eval() + self._head.likelihood.eval() + log_inference( + "geotransolver_gp", + f"Loaded FieldGPHead from {self._head_checkpoint} " + f"(epoch {self._checkpoint_epoch}); feature_dim={feature_dim}, " + f"gp_dim={getattr(self._head, 'gp_input_dim', '?')}", + ) + + def _autocast_ctx(self): + return ( + cuda_bf16_autocast(self._cfg.device) + if self._cfg.cuda_bf16_autocast + else nullcontext() + ) + + @torch.no_grad() + def _probe_feature_dim(self, batch: dict) -> int: + n = min(batch["embeddings"].shape[1], 1024) + fx = global_fx_to_bnc(batch["fx"]) + local_emb = batch["embeddings"][:, :n] + geometry = batch.get("geometry") + with self._autocast_ctx(): + _, point_features = self._model( + global_embedding=fx, + local_embedding=local_emb, + geometry=geometry, + local_positions=local_emb[:, :, :3], + return_point_features=True, + ) + return int(point_features.shape[-1]) + + @torch.no_grad() + def _predict_chunked_gp( + self, batch: dict + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Backbone + GP over all points, chunked on a random permutation (order restored). + + Returns ``(mean, total_std, epistemic_std)`` each ``(N, T)`` on device, in the + normalized target space the GP was trained in. + """ + fx = global_fx_to_bnc(batch["fx"]) + geometry = batch.get("geometry") + embeddings = batch["embeddings"] + n = embeddings.shape[1] + chunk = self._gp_chunk_size + if chunk is None or chunk <= 0 or chunk > n: + chunk = n + + perm = torch.randperm(n, device=embeddings.device) + mean_blocks: list[torch.Tensor] = [] + std_blocks: list[torch.Tensor] = [] + epi_blocks: list[torch.Tensor] = [] + with self._autocast_ctx(): + for idx_block in torch.split(perm, chunk): + local_emb = embeddings[:, idx_block] + _, point_features = self._model( + global_embedding=fx, + local_embedding=local_emb, + geometry=geometry, + local_positions=local_emb[:, :, :3], + return_point_features=True, + ) + pred = self._head.predict(point_features) + mean_blocks.append(pred.mean) + std_blocks.append(pred.variance.clamp_min(0).sqrt()) + epi_blocks.append(pred.epistemic_variance.clamp_min(0).sqrt()) + + mean_perm = torch.cat(mean_blocks, dim=1) + std_perm = torch.cat(std_blocks, dim=1) + epi_perm = torch.cat(epi_blocks, dim=1) + inverse = torch.empty(n, dtype=torch.long, device=mean_perm.device) + inverse[perm] = torch.arange(n, device=mean_perm.device) + mean = mean_perm[:, inverse].squeeze(0) + std = std_perm[:, inverse].squeeze(0) + epi = epi_perm[:, inverse].squeeze(0) + return mean, std, epi + + def predict(self, model_input: ModelInput) -> RawOutput: + """One forward pass: backbone features → GP posterior (mean, total std, epistemic std).""" + if self._model is None: + raise RuntimeError("GeoTransolverGPDrivAerStarWrapper: call load() first") + batch = model_input["batch"] + if self._head is None: + self._build_and_load_head(batch) + log_inference("geotransolver_gp", "Running GP forward pass (mean + variance)…") + mean, std, epi = self._predict_chunked_gp(batch) + return {"mean": mean, "total_std": std, "epistemic_std": epi} + + def _to_physical( + self, mean: torch.Tensor, std: torch.Tensor, epi: torch.Tensor + ) -> dict[str, np.ndarray]: + """De-normalize (mean, total std, epistemic std) from GP target space to physical units. + + For the affine standardization ``phys = (norm * field_std + field_mean) * q`` (with + ``q`` the dynamic pressure ``rho * u^2``), the mean maps fully while std/variance scale + by ``|field_std * q|`` only (the additive ``field_mean`` and the affine offset drop out). + ``q`` collapses to 1 when the checkpoint's targets are already dimensional (the + ``transformer_models`` GP checkpoints; see ``REDIMENSIONALIZE_OUTPUTS``), so physical + fields are just ``norm*std + mean`` — matching ``field_gp_utils.compute_drag_uq_stats``. + """ + q = ( + self._cfg.air_density * (self._cfg.stream_velocity**2) + if self.REDIMENSIONALIZE_OUTPUTS + else 1.0 + ) + fm = self._field_mean_t + fs = self._field_std_t + scale = fs * q # (T,) + mean_phys = (mean * fs + fm) * q # (N, T) + std_phys = std * scale + epi_phys = epi * scale + # Aleatoric (observation-noise) std from total vs epistemic variance. + ale_phys = (std_phys**2 - epi_phys**2).clamp_min(0).sqrt() + return { + "mean": mean_phys.detach().cpu().numpy().astype(np.float32), + "std": std_phys.detach().cpu().numpy().astype(np.float32), + "epistemic_std": epi_phys.detach().cpu().numpy().astype(np.float32), + "aleatoric_std": ale_phys.detach().cpu().numpy().astype(np.float32), + } + + def decode_distribution( + self, + raw_output: RawOutput, + case: CanonicalCase, + model_input: Optional[ModelInput] = None, + ) -> dict[str, FieldDistribution]: + """Map GP posterior to per-field :class:`FieldDistribution`s in physical units.""" + phys = self._to_physical( + raw_output["mean"], raw_output["total_std"], raw_output["epistemic_std"] + ) + log_inference( + "geotransolver_gp", + "Decoding GP distribution (pressure + WSS, physical units)…", + ) + out: dict[str, FieldDistribution] = {} + out["pressure"] = build_predictive_distribution( + mean=phys["mean"][:, 0], + std=phys["std"][:, 0], + epistemic_std=phys["epistemic_std"][:, 0], + aleatoric_std=phys["aleatoric_std"][:, 0], + ) + out["shear_stress"] = build_predictive_distribution( + mean=phys["mean"][:, 1:4], + std=phys["std"][:, 1:4], + epistemic_std=phys["epistemic_std"][:, 1:4], + aleatoric_std=phys["aleatoric_std"][:, 1:4], + ) + return out + + def decode_outputs( + self, + raw_output: RawOutput, + case: CanonicalCase, + model_input: Optional[ModelInput] = None, + ) -> Predictions: + """Point-estimate (GP mean) predictions in physical units (for non-UQ metric paths).""" + phys = self._to_physical( + raw_output["mean"], raw_output["total_std"], raw_output["epistemic_std"] + ) + return build_predictions_dict( + pressure=phys["mean"][:, 0], shear_stress=phys["mean"][:, 1:4] + ) diff --git a/physicsnemo/cfd/evaluation/models/wrappers/mc_dropout/__init__.py b/physicsnemo/cfd/evaluation/models/wrappers/mc_dropout/__init__.py new file mode 100644 index 0000000..4da746c --- /dev/null +++ b/physicsnemo/cfd/evaluation/models/wrappers/mc_dropout/__init__.py @@ -0,0 +1,21 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from physicsnemo.cfd.evaluation.models.wrappers.mc_dropout.wrapper import ( + GeoTransolverMCDropoutDrivAerStarWrapper, +) + +__all__ = ["GeoTransolverMCDropoutDrivAerStarWrapper"] diff --git a/physicsnemo/cfd/evaluation/models/wrappers/mc_dropout/wrapper.py b/physicsnemo/cfd/evaluation/models/wrappers/mc_dropout/wrapper.py new file mode 100644 index 0000000..a955a4c --- /dev/null +++ b/physicsnemo/cfd/evaluation/models/wrappers/mc_dropout/wrapper.py @@ -0,0 +1,322 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""**MC-Dropout** sampling UQ over a Concrete-Dropout GeoTransolver (surface/volume). + +This is the MC-Dropout sampling baseline (``geotransolver_mc_dropout_surface``). The backbone was trained with +learned per-layer :class:`~physicsnemo.nn.ConcreteDropout` (``model.concrete_dropout=true``, +Gal-Hron-Kendall 2017), so its dropout rates are calibrated by the training objective rather than +hand-picked. UQ is obtained at inference exactly as in the training repo's ``setup_mc_dropout`` / +``mc_dropout_inference_loop``: put the network in ``eval()`` but keep the Concrete-Dropout layers +**stochastic** (``.train()``), then average over ``N`` forward passes. The engine drives the ``N`` +passes (``run.uq.num_samples``) and aggregates the across-pass spread into a +:class:`~physicsnemo.cfd.evaluation.datasets.schema.FieldDistribution` (mean + epistemic std) via +streaming Welford — same ``UQ_METHOD="sampling"`` path as the ensemble wrapper. + +Completely independent :class:`CFDModel` subclass (no inheritance from the deterministic +GeoTransolver wrappers). It composes the shared GeoTransolver + ``TransolverDataPipe`` plumbing in +:mod:`physicsnemo.cfd.evaluation.models.common_wrapper_utils.geotransolver_runtime`; the *only* +difference from a deterministic forward is that the Concrete-Dropout submodules stay stochastic, so +:meth:`predict` returns a different draw each call. + +It decorates ``transformer_models`` (physical-target) checkpoints, so predictions are +re-standardized only — no dynamic-pressure re-dimensionalization +(:attr:`REDIMENSIONALIZE_OUTPUTS` = ``False``). + +Model kwargs (``model.kwargs``): + +- the shared GeoTransolver runtime kwargs (``batch_resolution``, ``geometry_sampling``, + ``cuda_bf16_autocast``; datapipe overrides; etc.). ``checkpoint`` must point at a + Concrete-Dropout checkpoint file (trained with ``concrete_dropout=true``); ``stats_path`` is the + shared surface/volume normalization. + +The number of stochastic passes is **not** a model kwarg — it is the benchmark-wide +``run.uq.num_samples`` so every sampling method (this and the ensemble) is compared at the same +budget. +""" + +import logging +from typing import Any, ClassVar, Literal, Optional + +from physicsnemo.cfd.evaluation.datasets.schema import ( + CanonicalCase, + InferenceDomain, + coerce_inference_domain_or_default, +) +from physicsnemo.cfd.evaluation.common.io import ( + load_transolver_surface_factors, + load_transolver_volume_factors, +) +from physicsnemo.cfd.evaluation.inference.progress import log_inference +from physicsnemo.cfd.evaluation.models.common_wrapper_utils.geotransolver_runtime import ( + GeoTransolverRuntimeConfig, + build_geotransolver_backbone, + build_transolver_batch, + decode_surface_predictions, + decode_volume_predictions, + geotransolver_available, + geotransolver_forward, + make_forward_permutation, + parse_runtime_kwargs, + unscale_targets, +) +from physicsnemo.cfd.evaluation.models.model_registry import ( + CFDModel, + ModelInput, + OutputLocation, + Predictions, + RawOutput, +) + +# ConcreteDropout is only needed to (a) find the dropout submodules and (b) read their learned +# rates for a sanity log. Guard the import like geotransolver_runtime guards physicsnemo. +try: + from physicsnemo.nn import ConcreteDropout, get_concrete_dropout_rates + + _CONCRETE_DROPOUT_AVAILABLE = True +except ImportError: + _CONCRETE_DROPOUT_AVAILABLE = False + +_LOG = logging.getLogger(__name__) + + +class GeoTransolverMCDropoutDrivAerStarWrapper(CFDModel): + """MC-Dropout sampling over a Concrete-Dropout GeoTransolver (surface/volume). + + Decorates a ``transformer_models`` (physical-target) checkpoint, so predictions are + re-standardized only (no dynamic-pressure re-dimensionalization): + :attr:`REDIMENSIONALIZE_OUTPUTS` = ``False``. + """ + + INFERENCE_DOMAIN: ClassVar[InferenceDomain | None] = None + OUTPUT_LOCATION: ClassVar[OutputLocation] = "cell" + REDIMENSIONALIZE_OUTPUTS: ClassVar[bool] = False + # Contract with the engine: "I produce uncertainty, get it by calling predict() many times." + # The engine loops predict() run.uq.num_samples times per case and turns the spread of the + # stochastic passes into a mean + epistemic std (streaming Welford). + SUPPORTS_UQ: ClassVar[bool] = True + UQ_METHOD: ClassVar[str] = "sampling" + + @property + def output_location(self) -> OutputLocation: + """See :attr:`CFDModel.output_location` (GeoTransolver predicts at cell centers).""" + return self.OUTPUT_LOCATION + + @classmethod + def inference_domain_from_kwargs( + cls, kwargs: dict[str, Any] + ) -> InferenceDomain | None: + """Align benchmark routing with :meth:`load` (default surface when omitted).""" + return coerce_inference_domain_or_default( + kwargs.get("inference_domain"), + default="surface", + parameter="model.kwargs.inference_domain", + ) + + def __init__(self) -> None: + self._model: Any = None + self._datapipe: Any = None + self._datapipe_geometry_effective: Optional[int] = None + self._surface_factors: Any = None + self._volume_factors: Any = None + self._inference_mode: Literal["surface", "volume"] = "surface" + self._volume_length_scale: Optional[float] = None + self._cfg: GeoTransolverRuntimeConfig = GeoTransolverRuntimeConfig() + + def load( + self, + checkpoint_path: str, + stats_path: str, + device: str, + **kwargs: Any, + ) -> "GeoTransolverMCDropoutDrivAerStarWrapper": + """Load the Concrete-Dropout backbone and keep its dropout layers stochastic for MC passes.""" + if not geotransolver_available(): + raise RuntimeError( + "GeoTransolverMCDropoutDrivAerStarWrapper requires physicsnemo (GeoTransolver, " + "TransolverDataPipe, load_model_weights)." + ) + if not _CONCRETE_DROPOUT_AVAILABLE: + raise RuntimeError( + "GeoTransolverMCDropoutDrivAerStarWrapper requires physicsnemo.nn.ConcreteDropout " + "(train the backbone with model.concrete_dropout=true)." + ) + kw = dict(kwargs) + self._inference_mode = coerce_inference_domain_or_default( + kw.pop("inference_domain", None), + default="surface", + parameter="model.kwargs.inference_domain", + ) + self._cfg = parse_runtime_kwargs(kw, device) + + if self._inference_mode == "volume": + log_inference( + "mc_dropout", f"Loading volume normalization from {stats_path}" + ) + self._volume_factors = load_transolver_volume_factors(stats_path, device) + if self._volume_factors is None: + raise FileNotFoundError( + "Volume inference requires ``global_stats.json`` or " + f"``volume_fields_normalization.npz`` (looked under {stats_path!r})." + ) + self._surface_factors = None + else: + log_inference( + "mc_dropout", f"Loading surface normalization from {stats_path}" + ) + self._surface_factors = load_transolver_surface_factors(stats_path, device) + self._volume_factors = None + + self._datapipe = None + self._datapipe_geometry_effective = None + # Build WITH concrete_dropout=True so the checkpoint's learned ConcreteDropout parameters + # (p_logit per layer) load cleanly; a plain backbone would reject those extra keys. + self._model = build_geotransolver_backbone( + checkpoint_path=checkpoint_path, + device=device, + inference_mode=self._inference_mode, + concrete_dropout=True, + ) + self._enable_mc_dropout() + return self + + def _enable_mc_dropout(self) -> None: + """Keep the network in eval() but flip ConcreteDropout layers to train() (stochastic). + + This is the MC-Dropout trick: everything deterministic (norms, etc.) stays in eval, but the + Concrete-Dropout masks are re-sampled every forward pass, so N passes give N draws. Mirrors + the training repo's ``setup_mc_dropout``. + """ + self._model.eval() + n_dropout = 0 + for module in self._model.modules(): + if isinstance(module, ConcreteDropout): + module.train() + n_dropout += 1 + if n_dropout == 0: + raise RuntimeError( + "No ConcreteDropout layers found in the loaded checkpoint. Was it trained with " + "model.concrete_dropout=true? MC-Dropout has no source of stochasticity otherwise." + ) + rates = get_concrete_dropout_rates(self._model) + if rates: + vals = list(rates.values()) + _LOG.info( + "MC-Dropout enabled over %d ConcreteDropout layers; learned rates " + "min=%.4f max=%.4f mean=%.4f", + n_dropout, + min(vals), + max(vals), + sum(vals) / len(vals), + ) + else: + _LOG.info("MC-Dropout enabled over %d ConcreteDropout layers.", n_dropout) + + def prepare_inputs(self, case: CanonicalCase) -> ModelInput: + """Build the surface/volume data dict, lazily (re)create the datapipe, and run it.""" + if self._model is None: + raise RuntimeError( + "GeoTransolverMCDropoutDrivAerStarWrapper: call load() first" + ) + result = build_transolver_batch( + case=case, + inference_mode=self._inference_mode, + cfg=self._cfg, + surface_factors=self._surface_factors, + volume_factors=self._volume_factors, + datapipe=self._datapipe, + geometry_effective=self._datapipe_geometry_effective, + ) + self._datapipe = result.datapipe + self._datapipe_geometry_effective = result.geometry_effective + if result.volume_length_scale is not None: + self._volume_length_scale = result.volume_length_scale + # Fix ONE forward block-partition per case (reused across all N passes below) so the + # across-pass spread reflects only the resampled dropout masks, not random point grouping. + perm = make_forward_permutation(result.batch) + return {"batch": result.batch, "datapipe": result.datapipe, "perm": perm} + + def predict(self, model_input: ModelInput) -> RawOutput: + """One stochastic MC-Dropout pass (Concrete-Dropout layers resample their masks). + + No weights are modified: the stochasticity comes entirely from the dropout masks, which are + redrawn from the global torch RNG (re-seeded per pass by the engine's sampling loop, so + passes differ from each other yet are reproducible run-to-run). The forward block-partition + is held fixed across passes (from :meth:`prepare_inputs`) so only the dropout varies. + """ + if self._model is None or self._datapipe is None: + raise RuntimeError( + "GeoTransolverMCDropoutDrivAerStarWrapper: call load() first" + ) + raw = geotransolver_forward( + model=self._model, + batch=model_input["batch"], + batch_resolution=self._cfg.batch_resolution, + cuda_bf16_autocast_enabled=self._cfg.cuda_bf16_autocast, + device=self._cfg.device, + perm=model_input.get("perm"), + ) + return unscale_targets( + datapipe=model_input["datapipe"], + predictions=raw, + batch=model_input["batch"], + inference_mode=self._inference_mode, + ) + + def predict_deterministic(self, model_input: ModelInput) -> RawOutput: + """One dropout-free forward pass (used when ``run.uq.enabled`` is off). + + MC-Dropout's :meth:`predict` is stochastic (the ConcreteDropout layers stay in ``train()``), + so the base :meth:`CFDModel.predict_deterministic` (which delegates to :meth:`predict`) would + return a single random draw. Here we temporarily flip the ConcreteDropout modules back to + ``eval()`` — their deterministic expectation — run one pass, then restore their stochastic + state, so ``run.uq.enabled=false`` yields a genuine point prediction. + """ + if self._model is None: + raise RuntimeError( + "GeoTransolverMCDropoutDrivAerStarWrapper: call load() first" + ) + dropouts = [m for m in self._model.modules() if isinstance(m, ConcreteDropout)] + were_training = [m.training for m in dropouts] + for m in dropouts: + m.eval() + try: + return self.predict(model_input) + finally: + for m, was_training in zip(dropouts, were_training): + m.train(was_training) + + def decode_outputs( + self, + raw_output: RawOutput, + case: CanonicalCase, + model_input: Optional[ModelInput] = None, + ) -> Predictions: + """Emit canonical surface/volume keys in physical units (no ``ρu²`` re-dimensionalization).""" + if self._inference_mode == "volume": + return decode_volume_predictions( + raw_output, + redimensionalize=self.REDIMENSIONALIZE_OUTPUTS, + air_density=self._cfg.air_density, + stream_velocity=self._cfg.stream_velocity, + length_scale=self._volume_length_scale, + ) + return decode_surface_predictions( + raw_output, + redimensionalize=self.REDIMENSIONALIZE_OUTPUTS, + air_density=self._cfg.air_density, + stream_velocity=self._cfg.stream_velocity, + ) diff --git a/physicsnemo/cfd/evaluation/reports/builtin/__init__.py b/physicsnemo/cfd/evaluation/reports/builtin/__init__.py index 25524e8..2808f02 100644 --- a/physicsnemo/cfd/evaluation/reports/builtin/__init__.py +++ b/physicsnemo/cfd/evaluation/reports/builtin/__init__.py @@ -33,10 +33,13 @@ from physicsnemo.cfd.evaluation.reports.builtin.aggregate_volume import ( register_aggregate_volume, ) +from physicsnemo.cfd.evaluation.reports.builtin.sparsification_plot import ( + register_sparsification_visual, +) def register_all_builtin_visuals() -> None: - """Register every built-in visual (surface/volume comparisons, line plots, hexbin, streamlines, aggregate).""" + """Register every built-in visual (surface/volume comparisons, line plots, hexbin, streamlines, aggregate, sparsification).""" register_field_comparison_surface() register_plot_fields_volume() register_line_plot() @@ -44,6 +47,7 @@ def register_all_builtin_visuals() -> None: register_projections_hexbin() register_streamlines_visual() register_aggregate_volume() + register_sparsification_visual() register_all_builtin_visuals() diff --git a/physicsnemo/cfd/evaluation/reports/builtin/sparsification_plot.py b/physicsnemo/cfd/evaluation/reports/builtin/sparsification_plot.py new file mode 100644 index 0000000..5e6d90b --- /dev/null +++ b/physicsnemo/cfd/evaluation/reports/builtin/sparsification_plot.py @@ -0,0 +1,196 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Sample-level sparsification / AUSE visual for UQ benchmark runs. + +Sparsification answers the active-learning question at the *geometry* level: rank the +validation geometries by predicted uncertainty, drop the most-uncertain first, and check that +the error of what remains falls toward the **oracle** (which drops the true-worst-error first). +The gap between the two curves is the **AUSE** (Area Under the Sparsification Error curve; lower +is better, 0 == oracle). The flat **full-RMSE** line is the do-nothing / random baseline. + +This visual consumes the ``uq_sparsification`` payload the benchmark engine attaches to each +result (built by +:func:`~physicsnemo.cfd.evaluation.benchmarks.uq_inference.compute_sparsification_payload` from +the sample metrics ``sparsification_ause`` / ``sparsification_ause_epistemic`` and, for the +decision-relevant drag panel, ``drag_uq``). One figure is written per dataset with a panel per +(metric, series) and every model overlaid, so GP and MC-Dropout uncertainties can be compared +against each other and the oracle. Deterministic models produce no payload and are skipped. +""" + +from __future__ import annotations + +import math +from pathlib import Path +from typing import Any + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt + +from physicsnemo.cfd.evaluation.datasets.progress import log_dataset +from physicsnemo.cfd.evaluation.reports.registry import register_visual +from physicsnemo.cfd.evaluation.reports.visual_filenames import benchmark_visual_png + +_PALETTE = ( + "#1f77b4", + "#d62728", + "#2ca02c", + "#9467bd", + "#ff7f0e", + "#17becf", + "#8c564b", + "#e377c2", +) + + +def _panel_title(metric_name: str, series: str) -> str: + """Human-readable panel title from a ``(metric, series)`` key.""" + label = metric_name + if label.startswith("sparsification_ause"): + kind = "epistemic std" if label.endswith("_epistemic") else "total std" + return f"{series} \u2014 rank by {kind}" + if label == "drag_uq": + return f"Drag (Cd) \u2014 rank by {series} std" + return f"{metric_name}:{series}" + + +def _collect_panels( + runs: list[tuple[dict[str, Any], dict[str, Any]]], +) -> list[tuple[str, str]]: + """Ordered, de-duplicated ``(metric_name, series)`` panel keys across all (run, payload) pairs.""" + seen: dict[tuple[str, str], None] = {} + for _, payload in runs: + for metric_name in sorted(payload): + for series in payload[metric_name]: + seen.setdefault((metric_name, series), None) + return list(seen) + + +def _draw_panel( + ax: Any, + panel: tuple[str, str], + runs: list[tuple[dict[str, Any], dict[str, Any]]], + color_by_model: dict[str, str], +) -> bool: + """Overlay every model's curves for one panel; return ``True`` if anything was drawn.""" + metric_name, series = panel + drawn = False + for run, payload in runs: + curve = (payload.get(metric_name) or {}).get(series) + if not curve: + continue + model = run["model"] + color = color_by_model[model] + fr = curve["fractions"] + ax.plot( + fr, + curve["by_uncertainty"], + color=color, + lw=2.2, + label=f"{model} (AUSE={curve['ause']:.3f})", + ) + ax.plot(fr, curve["oracle"], color=color, ls="--", lw=1.5, alpha=0.8) + ax.axhline(curve["full"], color=color, ls=":", lw=1.1, alpha=0.6) + drawn = True + if drawn: + ax.set_title(_panel_title(metric_name, series), fontsize=11) + ax.set_xlabel("Fraction of most-uncertain geometries removed") + ax.set_ylabel("RMSE of retained geometries") + ax.set_xlim(0, 1) + ax.set_ylim(bottom=0) + ax.grid(True, alpha=0.25) + ax.legend(fontsize=8, frameon=False) + return drawn + + +def sparsification_plot( + config: Any, + results: list[dict[str, Any]], + output_dir: str, + *, + context: dict[str, Any] | None = None, + max_cols: int = 3, + dpi: int = 140, + **_: Any, +) -> None: + """Write one sparsification figure per dataset (panels = metric/series, models overlaid).""" + del config + # Curve payloads are passed in ``context`` aligned with ``results`` by index (kept off the + # result dicts so the numpy arrays never reach the JSON report). Fall back to a per-run key for + # back-compat / direct callers. + payloads = (context or {}).get("uq_sparsification_by_run") or [] + paired: list[tuple[dict[str, Any], dict[str, Any]]] = [] + for i, run in enumerate(results): + if run.get("skipped"): + continue + payload = payloads[i] if i < len(payloads) else None + if not payload: + payload = run.get("uq_sparsification") + if payload: + paired.append((run, payload)) + if not paired: + log_dataset( + "benchmark", + "Skip sparsification_plot: no run carries a uq_sparsification payload " + "(need a probabilistic model and the sparsification_ause / drag_uq metrics).", + ) + return + + out = Path(output_dir) / "visuals" + out.mkdir(parents=True, exist_ok=True) + + datasets: dict[str, list[tuple[dict[str, Any], dict[str, Any]]]] = {} + for run, payload in paired: + datasets.setdefault(run["dataset"], []).append((run, payload)) + + for dataset, runs in datasets.items(): + panels = _collect_panels(runs) + if not panels: + continue + models = sorted({run["model"] for run, _ in runs}) + color_by_model = {m: _PALETTE[i % len(_PALETTE)] for i, m in enumerate(models)} + n = len(panels) + ncols = min(max_cols, n) + nrows = math.ceil(n / ncols) + fig, axes = plt.subplots( + nrows, ncols, figsize=(6.2 * ncols, 4.8 * nrows), squeeze=False + ) + flat = [ax for row in axes for ax in row] + any_drawn = False + for ax, panel in zip(flat, panels): + any_drawn |= _draw_panel(ax, panel, runs, color_by_model) + for ax in flat[n:]: + ax.axis("off") + if not any_drawn: + plt.close(fig) + continue + fig.suptitle( + f"Sample-level sparsification \u2014 {dataset} " + "(solid: by uncertainty, dashed: oracle, dotted: full RMSE)", + fontsize=13, + ) + fig.tight_layout(rect=(0, 0, 1, 0.96)) + safe = benchmark_visual_png("sparsification", dataset) + fig.savefig(str(out / safe), bbox_inches="tight", dpi=dpi) + plt.close(fig) + log_dataset("benchmark", f"Wrote sparsification plot {out / safe}") + + +def register_sparsification_visual() -> None: + """Register the ``sparsification_plot`` visual.""" + register_visual("sparsification_plot", sparsification_plot) diff --git a/physicsnemo/cfd/postprocessing_tools/metric_registry.py b/physicsnemo/cfd/postprocessing_tools/metric_registry.py index 6c6a284..af24465 100644 --- a/physicsnemo/cfd/postprocessing_tools/metric_registry.py +++ b/physicsnemo/cfd/postprocessing_tools/metric_registry.py @@ -14,20 +14,93 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Named metric registry with optional domain scoping (surface / volume).""" +"""Named metric registry with optional domain scoping (surface / volume). + +Two metric kinds share this one registry: + +- **Pointwise** metrics are plain callables ``fn(gt, predictions, **extended) -> float | + dict[str, float]``. The engine collects one value per case and aggregates by mean over + cases. All deterministic (L2 / force / physics) metrics are this kind. +- **Pooled / reducer** metrics implement the :class:`ReducerMetric` protocol + (``partial`` + ``finalize``). The engine calls ``partial`` per case, **sums** the + returned extensive sufficient statistics across all cases (and across distributed + ranks), and calls ``finalize`` once per (model, dataset). This way, the pooling of + calibration/coverage metrics over all points is statistically correct. +- **Sample** metrics implement the :class:`SampleMetric` protocol (``partial`` + + ``finalize_samples``). Like reducers, ``partial`` runs per case, but the engine + **collects** (does not sum) the returned per-geometry scalars across all cases and + ranks, and ``finalize_samples`` maps that list to the final value(s). Used for metrics + that need a global sort over geometries — e.g. sample-wise sparsification / AUSE. +""" from __future__ import annotations -from typing import Any, Callable +from typing import Any, Callable, Protocol, runtime_checkable MetricFn = Callable[..., float | dict[str, float]] -_REGISTRY: dict[tuple[str, str | None], MetricFn] = {} + +@runtime_checkable +class ReducerMetric(Protocol): + """A pooled metric expressed as two pure functions over additive sufficient statistics.""" + + def partial(self, gt: Any, predictions: Any, **extended: Any) -> dict[str, float]: + """Per-case **extensive** sufficient statistics (sums & counts; additive across cases).""" + ... + + def finalize(self, summed: dict[str, float]) -> float | dict[str, float]: + """Map globally-summed statistics to the final metric value(s).""" + ... + + +@runtime_checkable +class SampleMetric(Protocol): + """A per-geometry metric whose per-case values are collected (not summed) then finalized.""" + + def partial(self, gt: Any, predictions: Any, **extended: Any) -> dict[str, float]: + """Per-case (per-geometry) scalar contributions (collected across cases, not summed).""" + ... + + def finalize_samples( + self, collected: dict[str, list[float]] + ) -> float | dict[str, float]: + """Map the per-key lists of per-geometry values to the final metric value(s).""" + ... + + +#: A registry entry is a pointwise callable, a reducer-metric, or a sample-metric instance. +MetricEntry = "MetricFn | ReducerMetric | SampleMetric" + +_REGISTRY: dict[tuple[str, str | None], Any] = {} -def register_metric(name: str, fn: MetricFn, *, domain: str | None = None) -> None: +def is_reducer_metric(obj: Any) -> bool: + """True if ``obj`` is a pooled reducer metric (has callable ``partial`` and ``finalize``). + + Checked structurally (not ``isinstance(..., ReducerMetric)``) so plain metric callables + — which are not classes with these methods — are never misclassified. + """ + return callable(getattr(obj, "partial", None)) and callable( + getattr(obj, "finalize", None) + ) + + +def is_sample_metric(obj: Any) -> bool: + """True if ``obj`` is a sample metric (has callable ``partial`` and ``finalize_samples``). + + Disjoint from :func:`is_reducer_metric` in practice: a sample metric exposes + ``finalize_samples`` (not ``finalize``), so the engine's collect-not-sum path is selected. + """ + return callable(getattr(obj, "partial", None)) and callable( + getattr(obj, "finalize_samples", None) + ) + + +def register_metric(name: str, fn: Any, *, domain: str | None = None) -> None: """Register a metric by name and optional domain. + ``fn`` is either a pointwise callable or a :class:`ReducerMetric` instance. + When ``domain`` is ``None`` the metric is domain-agnostic and acts as the fallback when no domain-specific variant exists. When set (e.g. ``"surface"`` or ``"volume"``), the metric is scoped to that domain and @@ -41,7 +114,7 @@ def unregister_metric(name: str, *, domain: str | None = None) -> None: _REGISTRY.pop((name, domain), None) -def get_metric(name: str, *, domain: str | None = None) -> MetricFn: +def get_metric(name: str, *, domain: str | None = None) -> Any: """Resolve a metric function by name, preferring a domain-specific variant. Lookup order: diff --git a/skills/physicsnemo-cfd-create-model-wrapper/SKILL.md b/skills/physicsnemo-cfd-create-model-wrapper/SKILL.md index 17c7720..694f017 100644 --- a/skills/physicsnemo-cfd-create-model-wrapper/SKILL.md +++ b/skills/physicsnemo-cfd-create-model-wrapper/SKILL.md @@ -97,6 +97,13 @@ 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` +(analytic), 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 @@ -320,6 +327,111 @@ register_model("my_model", MyModelWrapper) Then use `model.name: my_model` in any YAML config. +## Uncertainty quantification (UQ) wrappers + +If the model produces **uncertainty**, not just a point estimate, it +opts into the UQ path additively. Copy `ExampleAnalyticUQWrapper` or +`ExampleSamplingUQWrapper` from `references/example_wrapper.py`, and see +the shipped `workflows/benchmarking/conf/config_uq_surface.yaml` for a +complete four-row example config. + +### Capability flags (on `CFDModel`) + +```python +SUPPORTS_UQ: ClassVar[bool] = True +UQ_METHOD: ClassVar[str] = "analytic" # or "sampling"; default "none" +``` + +There are exactly **two archetypes**, split by *how the predictive +distribution is produced*: + +- **`UQ_METHOD="analytic"`** — the model emits the distribution (or its + parameters) in **one** forward pass: GP head, mean–variance / + heteroscedastic net, evidential, SNGP/DUQ, quantile regression. + Implement **`decode_distribution(raw, case, model_input=None) -> + dict[str, FieldDistribution]`**. The engine calls `predict` once, then + `decode_distribution`. +- **`UQ_METHOD="sampling"`** — UQ comes from **multiple evaluations**: + N stochastic passes of one model (MC-Dropout, weight samples) *or* one + pass each of K models (deep/snapshot ensemble). The **engine** drives + the passes and aggregates them (streaming Welford) — you do **not** + build the distribution. Just make `predict` produce a *different* draw + each call (keep dropout stochastic at inference; don't freeze the RNG — + the engine re-seeds per pass for reproducibility). Because `predict` is + stochastic here, also override **`predict_deterministic(model_input)`** + so `run.uq.enabled=false` gives a true point prediction (disable dropout + for one pass, or return a single member). Optionally implement + **`predict_ensemble(model_input, n) -> Iterable[RawOutput] | None`** to + yield the passes/members (prefer a lazy generator so only one output is + device-resident at a time). Honor `n`: yield `n` passes for a per-model + sampler, or `min(n, member_count)` members for a fixed-size ensemble + (it cannot fabricate more distinct members than it holds); return `None` + to fall back to N× `predict`. + +Deterministic wrappers keep `SUPPORTS_UQ=False`; UQ metrics simply report +`NaN` for them (a useful det-vs-UQ contrast in the same report). + +### `FieldDistribution` payload + +`decode_distribution` (analytic) and the engine's aggregator (sampling) +both yield a `FieldDistribution` per field. Build it with +`build_predictive_distribution(...)` (the UQ analogue of +`build_predictions_dict`): + +```python +FieldDistribution( + mean, # (N,) or (N, C) + std=..., # total predictive std (epistemic + aleatoric) + epistemic_std=..., # model/knowledge uncertainty (optional) + aleatoric_std=..., # data/noise uncertainty (optional) + # samples / quantiles / quantile_levels -> non-Gaussian escape hatches (CRPS, intervals) +) +``` + +- **Physical units, always.** Denormalize the `mean` **and** the std + channels before returning — metrics never touch normalization stats. + Means invert with `x*std + mean`; std/variance channels scale by `std` + **only** (the additive offset drops out). +- **Epistemic vs aleatoric.** Analytic: the wrapper splits them (a GP + gives posterior variance = epistemic, noise floor = aleatoric). + Sampling: the across-pass spread is **epistemic**; total std equals it + unless each pass is *itself* a distribution (ensemble of mean–variance + members, MC-Dropout + heteroscedastic head), in which case return + `(mean_i, aleatoric_var_i)` per pass and the engine combines them by + the law of total variance `total = mean_i(σ_i²) + var_i(μ_i)`. + +### `decode_outputs` is still required + +Keep returning the **point estimate** (the distribution mean) from +`decode_outputs`. The deterministic metrics (L2 / drag / lift) read it, +non-UQ runs use it, and — for `sampling` wrappers — the engine calls it +**once per pass** to get each pass's fields before aggregating. + +### Config + metrics + +Turn UQ on in the run config and add the pooled metrics you want: + +```yaml +run: + uq: + enabled: true + num_samples: 32 # passes for UQ_METHOD="sampling" (ignored by analytic/ensemble) +metrics: + - nlpd + - calibration_zrms + - coverage_95 + - sharpness_std + - uncertainty_error_spearman # + _epistemic variants + - sparsification_ause # + _epistemic + - { name: drag_uq, drag_direction: [1, 0, 0] } +``` + +Export std companions to inspected `.vtp`s via +`output.std_mesh_field_names` / `epistemic_std_mesh_field_names` (auto +-derived as `*Std` / `*EpistemicStd` if omitted). See +`workflows/benchmarking/conf/config_uq_surface.yaml` for a complete +four-row example (deterministic + analytic GP + MC-Dropout + ensemble). + ## Gotchas - **DistributedManager**: Model wrappers may call @@ -341,8 +453,9 @@ Then use `model.name: my_model` in any YAML config. ## Related resources -- `references/example_wrapper.py` — complete surface + volume `CFDModel` - templates to copy and adapt (bundled; available without the repo on - disk). +- `references/example_wrapper.py` — complete `CFDModel` templates to copy + and adapt (bundled; available without the repo on disk): deterministic + surface + volume, plus `ExampleAnalyticUQWrapper` and + `ExampleSamplingUQWrapper` for the two UQ archetypes. - `assets/global_stats.example.json` — sample mean/std stats layout for surface and volume. diff --git a/skills/physicsnemo-cfd-create-model-wrapper/references/example_wrapper.py b/skills/physicsnemo-cfd-create-model-wrapper/references/example_wrapper.py index 9ec30a4..2e26dc0 100644 --- a/skills/physicsnemo-cfd-create-model-wrapper/references/example_wrapper.py +++ b/skills/physicsnemo-cfd-create-model-wrapper/references/example_wrapper.py @@ -17,22 +17,28 @@ """Self-contained reference wrappers for the PhysicsNeMo CFD benchmarking workflow. These are complete, correct ``CFDModel`` implementations to **adapt** when writing a -new wrapper — one surface model and one volume model. They are built from the real -evaluation APIs and the patterns in the shipped baseline wrappers -(``physicsnemo/cfd/evaluation/models/wrappers/surface_baseline.py`` and -``volume_baseline.py``) plus the ``adding_a_new_model.ipynb`` tutorial, so they stay -usable as a template even when the full PhysicsNeMo source tree is not on disk. When the -repo is present, verify the imported names against it — these templates are hints, not -ground truth. +new wrapper. They cover: + * ``ExampleSurfaceWrapper`` / ``ExampleVolumeWrapper`` — deterministic models. + * ``ExampleAnalyticUQWrapper`` — a single-pass UQ model (GP head / mean-variance / + evidential): ``UQ_METHOD="analytic"``, returns a predictive distribution directly. + * ``ExampleSamplingUQWrapper`` — a multi-pass UQ model (MC-Dropout / deep ensemble): + ``UQ_METHOD="sampling"``, the engine drives the passes and aggregates them. + +They are built from the real evaluation APIs and the patterns in the shipped baseline and +UQ wrappers, so they stay usable as a template even when the full PhysicsNeMo source tree +is not on disk. When the repo is present, verify the imported names against it — these +templates are hints, not ground truth. Adapt, don't copy blindly. The four things every wrapper must mirror from the model's own training/inference code are flagged inline below: (1) NORMALIZATION scheme, (2) INPUT tier, (3) OUTPUT fields + channel order, (4) shapes. +For UQ wrappers there is a fifth: (5) the UQ CONTRACT (see the two UQ examples and the +"Uncertainty quantification" section of SKILL.md). """ from __future__ import annotations -from typing import Any, ClassVar, Optional +from typing import Any, ClassVar, Iterable, Optional import numpy as np import torch @@ -44,8 +50,10 @@ ) from physicsnemo.cfd.evaluation.datasets.schema import ( CanonicalCase, + FieldDistribution, InferenceDomain, build_predictions_dict, + build_predictive_distribution, ) from physicsnemo.cfd.evaluation.inference.progress import log_inference from physicsnemo.cfd.evaluation.models.model_registry import ( @@ -244,6 +252,263 @@ def _denormalize(self, tensor: torch.Tensor, field: str) -> torch.Tensor: return tensor * std + mean +class ExampleAnalyticUQWrapper(CFDModel): + """Single-pass UQ model (GP head / mean-variance / evidential): ``UQ_METHOD="analytic"``. + + The model itself emits the predictive distribution (or its parameters) in ONE forward + pass. The extra contract vs a deterministic wrapper is: + * declare ``SUPPORTS_UQ = True`` and ``UQ_METHOD = "analytic"``; + * implement ``decode_distribution`` to return a ``FieldDistribution`` per field. + ``decode_outputs`` is still provided (the distribution *mean*) so the deterministic + metrics (L2 / drag / lift) and non-UQ runs keep working unchanged. + """ + + INFERENCE_DOMAIN: ClassVar[InferenceDomain] = "surface" + OUTPUT_LOCATION: ClassVar[OutputLocation] = "cell" + # (5) UQ CONTRACT: one forward pass -> a distribution, materialized in decode_distribution. + SUPPORTS_UQ: ClassVar[bool] = True + UQ_METHOD: ClassVar[str] = "analytic" + + def __init__(self) -> None: + self._model: torch.nn.Module | None = None + self._stats: dict[str, Any] | None = None + self._device: str = "cpu" + + @property + def output_location(self) -> OutputLocation: + """Where predictions live (cell-centered surface, here).""" + return self.OUTPUT_LOCATION + + def load( + self, checkpoint_path: str, stats_path: str, device: str, **kwargs: Any + ) -> "ExampleAnalyticUQWrapper": + """Build the architecture + UQ head, load weights and normalization stats.""" + self._device = device + # Build the architecture + UQ head and load weights here (e.g. a GP head whose + # hyperparameters come from kwargs and MUST match training so the state dict loads). + self._stats = load_global_stats(stats_path, device) + log_inference("example_analytic_uq", f"Loaded from {checkpoint_path}") + return self + + def prepare_inputs(self, case: CanonicalCase) -> torch.Tensor: + """Read the case mesh and return the model-ready input tensor.""" + mesh = pv.read(case.mesh_path) + if not isinstance(mesh, pv.PolyData): + mesh = mesh.extract_surface() + coords = np.array(mesh.cell_centers().points, dtype=np.float32) + return torch.tensor(coords, device=self._device) + + def predict(self, model_input: torch.Tensor) -> dict[str, torch.Tensor]: + """One forward pass -> raw posterior summary (still normalized). + + Return whatever the head produces; here the Gaussian summary mean + total std + + the epistemic (model) component. A GP splits posterior variance (epistemic) from the + learned noise floor (aleatoric); a mean-variance net returns mean + a variance head. + """ + with torch.no_grad(): + n = model_input.shape[0] + # raw = self._model(model_input) # -> mean, variance, epistemic_variance, ... + raw = { + "mean": torch.zeros((n, 4), device=self._device), # p + WSS(3) + "total_std": torch.ones((n, 4), device=self._device), + "epistemic_std": torch.ones((n, 4), device=self._device), + } + return raw + + def _to_physical( + self, arr: torch.Tensor, field: str, *, is_std: bool + ) -> np.ndarray: + """Inverse mean-std. NOTE: means map with +mean; std/variance channels do NOT + add the offset (they scale by ``std`` only). Denormalize BOTH here so metrics never + touch normalization stats (all UQ metrics run in physical units).""" + if self._stats is None: + return arr.cpu().numpy() + mean = self._stats["mean"].get(field, 0.0) + std = self._stats["std"].get(field, 1.0) + out = arr * std if is_std else arr * std + mean + return out.cpu().numpy() + + def decode_distribution( + self, + raw_output: dict[str, torch.Tensor], + case: CanonicalCase, + model_input: Optional[torch.Tensor] = None, + ) -> dict[str, FieldDistribution]: + """Map the raw posterior to a ``FieldDistribution`` per canonical field (physical units). + + ``build_predictive_distribution`` mirrors ``build_predictions_dict`` but carries the + std channels. Provide ``std`` (total) and, when the method separates them, + ``epistemic_std`` / ``aleatoric_std`` — the pooled UQ metrics consume them directly. + """ + mean = raw_output["mean"] + tot = raw_output["total_std"] + epi = raw_output["epistemic_std"] + return { + "pressure": build_predictive_distribution( + mean=self._to_physical(mean[:, 0], "pressure", is_std=False), + std=self._to_physical(tot[:, 0], "pressure", is_std=True), + epistemic_std=self._to_physical(epi[:, 0], "pressure", is_std=True), + ), + "shear_stress": build_predictive_distribution( + mean=self._to_physical(mean[:, 1:4], "shear_stress", is_std=False), + std=self._to_physical(tot[:, 1:4], "shear_stress", is_std=True), + epistemic_std=self._to_physical( + epi[:, 1:4], "shear_stress", is_std=True + ), + ), + } + + def decode_outputs( + self, + raw_output: dict[str, torch.Tensor], + case: CanonicalCase, + model_input: Optional[torch.Tensor] = None, + ) -> dict[str, np.ndarray]: + """Point estimate (distribution mean) for the deterministic metric paths / non-UQ runs.""" + mean = raw_output["mean"] + return build_predictions_dict( + pressure=self._to_physical(mean[:, 0], "pressure", is_std=False), + shear_stress=self._to_physical(mean[:, 1:4], "shear_stress", is_std=False), + ) + + +class ExampleSamplingUQWrapper(CFDModel): + """Multi-pass UQ model (MC-Dropout / deep ensemble): ``UQ_METHOD="sampling"``. + + The wrapper only produces ONE (stochastic) pass at a time; the ENGINE drives ``N`` + passes and aggregates the spread into a ``FieldDistribution`` (streaming Welford). So the + contract is: + * declare ``SUPPORTS_UQ = True`` and ``UQ_METHOD = "sampling"``; + * make ``predict`` produce a *different* draw each call (dropout stays stochastic at + inference, or a weight sample is drawn) — the engine re-seeds per pass for + reproducibility, so do NOT freeze the RNG yourself; + * ``decode_outputs`` maps one raw pass to physical predictions (the engine calls it + once per pass, then aggregates); + * override ``predict_deterministic`` when ``predict`` is stochastic, so ``run.uq.enabled= + false`` yields a true point prediction (MC-Dropout disables dropout for one pass; an + ensemble returns a single member); + * OPTIONAL multi-pass path: implement ``predict_ensemble(model_input, n)`` to yield the + passes/members as an ``Iterable[RawOutput]`` (prefer a lazy generator so only one output + is device-resident at a time). Honor ``n``: yield ``n`` passes for a per-model sampler, or + ``min(n, member_count)`` members for a fixed-size ensemble. No ``decode_distribution`` is + needed — the engine builds the distribution. + + The number of passes is the benchmark-wide ``run.uq.num_samples`` (NOT a model kwarg), so + every sampling method is compared at the same budget. + """ + + INFERENCE_DOMAIN: ClassVar[InferenceDomain] = "surface" + OUTPUT_LOCATION: ClassVar[OutputLocation] = "cell" + # (5) UQ CONTRACT: engine calls predict() N times (or predict_ensemble once) and aggregates. + SUPPORTS_UQ: ClassVar[bool] = True + UQ_METHOD: ClassVar[str] = "sampling" + + def __init__(self) -> None: + self._models: list[torch.nn.Module] = ( + [] + ) # one for MC-Dropout, K for an ensemble + self._stats: dict[str, Any] | None = None + self._device: str = "cpu" + + @property + def output_location(self) -> OutputLocation: + """Where predictions live (cell-centered surface, here).""" + return self.OUTPUT_LOCATION + + def load( + self, checkpoint_path: str, stats_path: str, device: str, **kwargs: Any + ) -> "ExampleSamplingUQWrapper": + """Load the stochastic model(s) used to draw repeated inference passes.""" + self._device = device + # MC-Dropout: load ONE model and keep its dropout layers in train() mode (stochastic) + # while the rest is eval(). Ensemble: load one model per kwargs["member_checkpoints"]. + # model.load_state_dict(torch.load(checkpoint_path, weights_only=True)); model.eval() + # for m in model.modules(): + # if isinstance(m, torch.nn.Dropout): m.train() + self._stats = load_global_stats(stats_path, device) + log_inference("example_sampling_uq", f"Loaded from {checkpoint_path}") + return self + + def prepare_inputs(self, case: CanonicalCase) -> torch.Tensor: + """Read the case mesh and return the model-ready input tensor.""" + mesh = pv.read(case.mesh_path) + if not isinstance(mesh, pv.PolyData): + mesh = mesh.extract_surface() + coords = np.array(mesh.cell_centers().points, dtype=np.float32) + return torch.tensor(coords, device=self._device) + + def predict(self, model_input: torch.Tensor) -> dict[str, torch.Tensor]: + """ONE stochastic pass. Must differ call-to-call (dropout mask / weight sample).""" + with torch.no_grad(): + n = model_input.shape[0] + # raw = self._models[0](model_input) # dropout re-samples its mask each call + raw = { + "pressure": torch.randn(n, device=self._device), + "shear_stress": torch.randn((n, 3), device=self._device), + } + return raw + + def predict_deterministic( + self, model_input: torch.Tensor + ) -> dict[str, torch.Tensor]: + """One DETERMINISTIC pass for ``run.uq.enabled=false`` (no dropout / one member). + + The default ``CFDModel.predict_deterministic`` just calls ``predict``; override it whenever + ``predict`` is stochastic. MC-Dropout: flip the dropout modules to ``eval()`` for one pass + then restore them. Ensemble: return a single member. Otherwise turning UQ off would still + return one random draw instead of a point prediction. + """ + # MC-Dropout sketch: + # dropouts = [m for m in self._models[0].modules() if isinstance(m, torch.nn.Dropout)] + # were_training = [m.training for m in dropouts] + # for m in dropouts: m.eval() + # try: return self.predict(model_input) + # finally: + # for m, t in zip(dropouts, were_training): m.train(t) + return self.predict(model_input) + + def predict_ensemble( + self, model_input: torch.Tensor, n: int + ) -> Optional[Iterable[dict[str, torch.Tensor]]]: + """Optional multi-pass path: YIELD each pass/member (prefer a lazy generator). + + Return an ``Iterable[RawOutput]`` (ideally a generator, so only one output is device- + resident at a time — a ``list`` materializes all ``n`` at once and can OOM). Honor ``n``: + yield ``n`` stochastic passes for a per-model sampler, or ``min(n, member_count)`` members + for a fixed-size ensemble (it cannot fabricate more distinct members than it holds). Return + ``None`` to let the engine fall back to calling ``predict`` ``n`` times. + """ + if not self._models: + return None + k = min( + n, len(self._models) + ) # ensemble: honor the budget, capped at member count + return (self.predict(model_input) for _ in range(k)) + + def decode_outputs( + self, + raw_output: dict[str, torch.Tensor], + case: CanonicalCase, + model_input: Optional[torch.Tensor] = None, + ) -> dict[str, np.ndarray]: + """Denormalize ONE pass to physical predictions (the engine aggregates across passes).""" + p = raw_output["pressure"] * self._std("pressure") + self._mean("pressure") + w = raw_output["shear_stress"] * self._std("shear_stress") + self._mean( + "shear_stress" + ) + return build_predictions_dict( + pressure=p.cpu().numpy(), shear_stress=w.cpu().numpy().reshape(-1, 3) + ) + + def _mean(self, field: str) -> Any: + return 0.0 if self._stats is None else self._stats["mean"].get(field, 0.0) + + def _std(self, field: str) -> Any: + return 1.0 if self._stats is None else self._stats["std"].get(field, 1.0) + + # Registration: do this once at import time so the engine can resolve the name. register_model("example_surface", ExampleSurfaceWrapper) register_model("example_volume", ExampleVolumeWrapper) +register_model("example_analytic_uq", ExampleAnalyticUQWrapper) +register_model("example_sampling_uq", ExampleSamplingUQWrapper) diff --git a/test/ci_tests/test_drivaerstar_adapter.py b/test/ci_tests/test_drivaerstar_adapter.py new file mode 100644 index 0000000..cb09cdc --- /dev/null +++ b/test/ci_tests/test_drivaerstar_adapter.py @@ -0,0 +1,106 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""DrivAerStar adapter regression tests: the WSS sign-convention decision, pressure renaming, +normals/area removal, and the integrated ground-truth sign (drag/lift depend on it).""" + +from __future__ import annotations + +import numpy as np +import pyvista as pv +import pytest + +from physicsnemo.cfd.evaluation.datasets.adapters.drivaerstar import ( + DEFAULT_PRESSURE_OUT_NAME, + DEFAULT_SHEAR_OUT_NAME, + DrivAerStarAdapter, +) + + +def _plane_with_fields() -> pv.PolyData: + """A small plane carrying DrivAerStar-style cell arrays (constant WSS for sign checks).""" + mesh = pv.Plane(i_resolution=3, j_resolution=3) + n = mesh.n_cells + mesh.cell_data["Pressure"] = np.arange(n, dtype=np.float64) + mesh.cell_data["WallShearStressi"] = np.full(n, 1.0, dtype=np.float64) + mesh.cell_data["WallShearStressj"] = np.full(n, 2.0, dtype=np.float64) + mesh.cell_data["WallShearStressk"] = np.full(n, 3.0, dtype=np.float64) + mesh.cell_data["Normals"] = np.zeros((n, 3), dtype=np.float64) + mesh.cell_data["Area"] = np.ones(n, dtype=np.float64) + return mesh + + +def test_flip_wss_sign_true_negates_components(tmp_path) -> None: + """flip_wss_sign=True (default) negates the combined WSS to the DrivAerML convention.""" + adapter = DrivAerStarAdapter(root=str(tmp_path), flip_wss_sign=True) + mesh = _plane_with_fields() + adapter._combine_wss(mesh, "1") + wss = np.asarray(mesh.cell_data[DEFAULT_SHEAR_OUT_NAME]) + assert wss.shape[1] == 3 + assert np.allclose(wss[:, 0], -1.0) + assert np.allclose(wss[:, 1], -2.0) + assert np.allclose(wss[:, 2], -3.0) + # Source component arrays are consumed into the combined vector. + assert "WallShearStressi" not in mesh.cell_data + + +def test_flip_wss_sign_false_keeps_native_sign(tmp_path) -> None: + """flip_wss_sign=False keeps the native DrivAerStar sign (as the UQ example config sets).""" + adapter = DrivAerStarAdapter(root=str(tmp_path), flip_wss_sign=False) + mesh = _plane_with_fields() + adapter._combine_wss(mesh, "1") + wss = np.asarray(mesh.cell_data[DEFAULT_SHEAR_OUT_NAME]) + assert np.allclose(wss[:, 0], 1.0) + assert np.allclose(wss[:, 1], 2.0) + assert np.allclose(wss[:, 2], 3.0) + + +def test_combine_wss_missing_components_raises_loudly(tmp_path) -> None: + """A missing/misnamed WSS component fails loudly (not a silent no-op) during preparation.""" + adapter = DrivAerStarAdapter(root=str(tmp_path)) + with pytest.raises(ValueError): + adapter._combine_wss(pv.Plane(i_resolution=2, j_resolution=2), "1") + + +def test_rename_pressure_and_missing_raises(tmp_path) -> None: + """Pressure is renamed to the DrivAerML name; a missing pressure array raises loudly.""" + adapter = DrivAerStarAdapter(root=str(tmp_path)) + mesh = _plane_with_fields() + adapter._rename_pressure(mesh, "1") + assert DEFAULT_PRESSURE_OUT_NAME in mesh.cell_data + assert "Pressure" not in mesh.cell_data + with pytest.raises(ValueError): + adapter._rename_pressure(pv.Plane(i_resolution=2, j_resolution=2), "1") + + +def test_drop_normals_area_removes_explicit_arrays() -> None: + """Explicit Normals/Area arrays are dropped so forces recompute from mesh topology.""" + mesh = _plane_with_fields() + DrivAerStarAdapter._drop_normals_area(mesh) + assert "Normals" not in mesh.cell_data and "Area" not in mesh.cell_data + assert "Normals" not in mesh.point_data and "Area" not in mesh.point_data + + +def test_load_case_ground_truth_respects_wss_sign(tmp_path) -> None: + """End-to-end load_case exposes ground-truth WSS in the configured sign (drives drag/lift).""" + _plane_with_fields().save(str(tmp_path / "1.vtk")) + adapter = DrivAerStarAdapter(root=str(tmp_path), flip_wss_sign=True) + case = adapter.load_case("1") + assert case.ground_truth is not None + shear = np.asarray(case.ground_truth["shear_stress"]) + assert np.allclose(shear[:, 0], -1.0) + assert np.allclose(shear[:, 1], -2.0) + assert np.allclose(shear[:, 2], -3.0) diff --git a/test/ci_tests/test_l2_errors.py b/test/ci_tests/test_l2_errors.py index 9014deb..976b736 100644 --- a/test/ci_tests/test_l2_errors.py +++ b/test/ci_tests/test_l2_errors.py @@ -80,11 +80,13 @@ def test_dof_coordinates_cell_returns_cell_centers() -> None: def test_bounds_mask_none_short_circuits() -> None: + """A None bounds spec short-circuits to no mask.""" coords = np.zeros((5, 3)) assert _bounds_mask(coords, None) is None def test_bounds_mask_inclusive_aabb() -> None: + """The bounds mask is an inclusive axis-aligned box test.""" coords = np.array( [ [0.0, 0.0, 0.0], @@ -100,6 +102,7 @@ def test_bounds_mask_inclusive_aabb() -> None: def test_classify_fields_scalar_vs_vector() -> None: + """Field classification distinguishes scalar from vector arrays by component count.""" grid = _grid() grid.point_data["scalar"] = np.zeros(grid.n_points, dtype=np.float64) grid.point_data["vector"] = np.zeros((grid.n_points, 3), dtype=np.float64) diff --git a/test/ci_tests/test_uq_inference.py b/test/ci_tests/test_uq_inference.py new file mode 100644 index 0000000..8a132b1 --- /dev/null +++ b/test/ci_tests/test_uq_inference.py @@ -0,0 +1,295 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for the engine-side UQ plumbing (sampling loop + pooled reducer finalization).""" + +from __future__ import annotations + +import math + +import numpy as np + +# Registers built-in metrics (including the pooled UQ reducers) into the registry. +import physicsnemo.cfd.evaluation.metrics # noqa: F401 +from physicsnemo.cfd.evaluation.benchmarks.uq_inference import ( + finalize_reducer_metrics, + finalize_sample_metrics, + is_reducer_partial_key, + make_reducer_partial_key, + run_sampling_inference, + select_inference_path, + strip_reducer_partials, + _Welford, +) +from physicsnemo.cfd.evaluation.datasets.schema import FieldDistribution +from physicsnemo.cfd.postprocessing_tools.metric_registry import get_metric + + +def test_reducer_partial_key_round_trip_and_strip() -> None: + """Reducer partial keys round-trip and are stripped from reported per-case rows.""" + key = make_reducer_partial_key("nlpd", "pressure::sum") + assert key == "_uq::nlpd::pressure::sum" + assert is_reducer_partial_key(key) + rows = [{"case_id": "a", "metrics": {"l2_pressure": 0.1, key: 2.0}}] + strip_reducer_partials(rows) + assert rows[0]["metrics"] == {"l2_pressure": 0.1} + + +def test_welford_population_variance() -> None: + """``_Welford`` yields the streaming mean and population (epistemic) std across passes.""" + w = _Welford() + for x in (np.array([1.0, 2.0]), np.array([3.0, 2.0]), np.array([5.0, 2.0])): + w.update(x) + mean, total, epi, ale = w.result() + assert np.allclose(mean, [3.0, 2.0]) + assert np.allclose(epi, [np.sqrt(8 / 3), 0.0]) # population variance of {1,3,5} + assert ale is None and np.allclose(total, epi) + + +class _EnsembleWrapper: + SUPPORTS_UQ = True + UQ_METHOD = "sampling" + + def __init__(self, outs, hetero_std=None): + self._outs = outs + self._hetero_std = hetero_std + + def predict_ensemble(self, model_input, n): + return self._outs[:n] + + def predict(self, model_input): # pragma: no cover - ensemble path used + raise AssertionError("ensemble fast-path should be used") + + def decode_outputs(self, raw, case, model_input=None): + if self._hetero_std is not None: + return {"pressure": FieldDistribution(mean=raw, std=self._hetero_std)} + return {"pressure": raw} + + +def test_run_sampling_inference_epistemic_and_samples() -> None: + """Ensemble sampling gives the across-member mean, epistemic std, and retained samples.""" + outs = [np.array([1.0, 10.0]), np.array([3.0, 10.0]), np.array([5.0, 10.0])] + d = run_sampling_inference( + _EnsembleWrapper(outs), + None, + None, + n=3, + run_seed=0, + case_id="c", + retain_samples=True, + ) + fd = d["pressure"] + assert np.allclose(fd.mean, [3.0, 10.0]) + assert np.allclose(fd.epistemic_std, [np.sqrt(8 / 3), 0.0]) + assert np.allclose(fd.std, fd.epistemic_std) # no aleatoric component + assert fd.samples.shape == (3, 2) + + +def test_run_sampling_inference_zero_spread_gives_zero_epistemic() -> None: + """Identical passes -> zero epistemic std.""" + d = run_sampling_inference( + _EnsembleWrapper([np.array([2.0, 2.0])] * 4), + None, + None, + n=4, + run_seed=0, + case_id="c", + ) + assert np.allclose(d["pressure"].epistemic_std, 0.0) + + +def test_run_sampling_inference_law_of_total_variance() -> None: + """Per-pass aleatoric std combines with across-pass spread via the law of total variance.""" + outs = [np.array([1.0, 10.0]), np.array([3.0, 10.0]), np.array([5.0, 10.0])] + d = run_sampling_inference( + _EnsembleWrapper(outs, hetero_std=np.array([2.0, 2.0])), + None, + None, + n=3, + run_seed=0, + case_id="c", + ) + fd = d["pressure"] + # total_var = var_i(mu_i) + mean_i(sigma_i^2) = 8/3 + 4 + assert np.allclose(fd.aleatoric_std, [2.0, 2.0]) + assert np.allclose(fd.std[0], np.sqrt(8 / 3 + 4)) + + +def test_finalize_reducer_metrics_pools_over_cases() -> None: + """Reducer finalization pools sufficient statistics over cases (not a mean of case means).""" + rng = np.random.default_rng(0) + + def _case(n, sigma): + mu = rng.normal(size=n).astype(np.float32) + y = (mu + sigma * rng.normal(size=n)).astype(np.float32) + pred = { + "pressure": FieldDistribution( + mean=mu, + std=np.full(n, sigma, np.float32), + epistemic_std=np.full(n, sigma * 0.5, np.float32), + ) + } + gt = {"pressure": y} + metrics = {} + for mn in ("coverage_95", "calibration_zrms", "sharpness_std"): + m = get_metric(mn, domain="surface") + for pk, pv in m.partial(gt, pred).items(): + metrics[make_reducer_partial_key(mn, pk)] = pv + return {"case_id": str(n), "metrics": metrics} + + rows = [_case(300_000, 1.0), _case(100_000, 2.0)] + summary = finalize_reducer_metrics(rows, "surface") + assert abs(summary["coverage_95_pressure"] - 0.95) < 0.01 + assert abs(summary["calibration_zrms_pressure"] - 1.0) < 0.02 + # pooled sharpness = (300000*1 + 100000*2) / 400000 = 1.25 (not mean-of-means 1.5) + assert abs(summary["sharpness_std_pressure"] - 1.25) < 1e-3 + + +# -------------------------------------------------------------------------------------------- +# Configured-metric NaN placeholders: deterministic rows report NaN, not omitted keys +# -------------------------------------------------------------------------------------------- + + +def test_finalize_reducer_metrics_emits_nan_for_configured_but_absent() -> None: + """A deterministic row (no ``_uq::`` partials) still reports configured reducer metrics as NaN. + + The placeholder uses the SAME headline key a populated finalize would (``{metric}_mean``) so + deterministic and UQ rows share a schema. + """ + rows = [{"case_id": "a", "metrics": {"l2_pressure": 0.1}}] # no reducer partials + # Without the configured names the metric is simply absent (legacy behavior)... + assert "nlpd_mean" not in finalize_reducer_metrics(rows, "surface") + # ...with them it is finalized to NaN so the report schema stays consistent. + summary = finalize_reducer_metrics(rows, "surface", ["nlpd", "coverage_95"]) + assert "nlpd_mean" in summary and math.isnan(summary["nlpd_mean"]) + assert "coverage_95_mean" in summary and math.isnan(summary["coverage_95_mean"]) + # An unknown / non-reducer configured name is ignored (not every metric is a reducer). + assert "l2_mean" not in finalize_reducer_metrics(rows, "surface", ["l2", "nlpd"]) + + +def test_finalize_sample_metrics_emits_nan_for_configured_but_absent() -> None: + """A deterministic row (no ``_uqs::`` partials) still reports configured sample metrics as NaN.""" + rows = [{"case_id": "a", "metrics": {"l2_pressure": 0.1}}] + summary = finalize_sample_metrics( + rows, "surface", ["sparsification_ause", "drag_uq"] + ) + assert math.isnan(summary["sparsification_ause_mean"]) + # drag_uq returns a dict -> expands to _, all NaN. + assert math.isnan(summary["drag_uq_epistemic"]) + assert math.isnan(summary["drag_uq_total"]) + + +# -------------------------------------------------------------------------------------------- +# Master switch dispatch (run.uq.enabled) + streaming ensemble generator +# -------------------------------------------------------------------------------------------- + + +def test_select_inference_path_master_switch() -> None: + """``run.uq.enabled`` is the master switch: off -> every wrapper goes deterministic.""" + # UQ enabled: sampling / analytic wrappers take their UQ path. + assert ( + select_inference_path(supports_uq=True, uq_method="sampling", uq_enabled=True) + == "sampling" + ) + assert ( + select_inference_path(supports_uq=True, uq_method="analytic", uq_enabled=True) + == "analytic" + ) + # Master switch OFF: EVERY wrapper (incl. analytic GP) goes deterministic -> no UQ metrics. + assert ( + select_inference_path(supports_uq=True, uq_method="analytic", uq_enabled=False) + == "deterministic" + ) + assert ( + select_inference_path(supports_uq=True, uq_method="sampling", uq_enabled=False) + == "deterministic" + ) + # Non-UQ wrapper is always deterministic. + assert ( + select_inference_path(supports_uq=False, uq_method="none", uq_enabled=True) + == "deterministic" + ) + + +class _GeneratorEnsembleWrapper: + """Ensemble whose ``predict_ensemble`` is a *generator* (one output resident at a time).""" + + SUPPORTS_UQ = True + UQ_METHOD = "sampling" + + def __init__(self, outs): + self._outs = outs + self.live = 0 + self.max_live = 0 + + def predict_ensemble(self, model_input, n): + for o in self._outs: + self.live += 1 + self.max_live = max(self.max_live, self.live) + yield o + self.live -= 1 # engine consumed it before the next is produced + + def predict(self, model_input): # pragma: no cover + raise AssertionError("generator ensemble path should be used") + + def decode_outputs(self, raw, case, model_input=None): + return {"pressure": raw} + + +def test_run_sampling_inference_consumes_generator_streaming() -> None: + """A generator ``predict_ensemble`` is consumed lazily (never all members resident at once).""" + outs = [np.array([1.0, 10.0]), np.array([3.0, 10.0]), np.array([5.0, 10.0])] + w = _GeneratorEnsembleWrapper(outs) + d = run_sampling_inference(w, None, None, n=3, run_seed=0, case_id="c") + assert np.allclose(d["pressure"].mean, [3.0, 10.0]) + assert np.allclose(d["pressure"].epistemic_std, [np.sqrt(8 / 3), 0.0]) + assert w.max_live == 1 # streaming: only one member output alive at any time + + +class _BudgetEnsembleWrapper: + """Fixed-size ensemble that honors the budget: yields ``min(n, member_count)`` members.""" + + SUPPORTS_UQ = True + UQ_METHOD = "sampling" + + def __init__(self, members): + self._members = members + self.yielded = 0 + + def predict_ensemble(self, model_input, n): + k = min(int(n), len(self._members)) + for m in self._members[:k]: + self.yielded += 1 + yield m + + def predict(self, model_input): # pragma: no cover - ensemble path used + raise AssertionError("ensemble path should be used") + + def decode_outputs(self, raw, case, model_input=None): + return {"pressure": raw} + + +def test_ensemble_predict_ensemble_honors_num_samples_budget() -> None: + """The ensemble caps at its member count but otherwise uses only ``n`` members (honors budget).""" + members = [np.array([float(i)]) for i in range(5)] + # n < K: only the first n members contribute. + w_small = _BudgetEnsembleWrapper(members) + run_sampling_inference(w_small, None, None, n=2, run_seed=0, case_id="c") + assert w_small.yielded == 2 + # n > K: cannot fabricate members -> all K used. + w_big = _BudgetEnsembleWrapper(members) + run_sampling_inference(w_big, None, None, n=32, run_seed=0, case_id="c") + assert w_big.yielded == 5 diff --git a/test/ci_tests/test_uq_metrics.py b/test/ci_tests/test_uq_metrics.py new file mode 100644 index 0000000..febc8b0 --- /dev/null +++ b/test/ci_tests/test_uq_metrics.py @@ -0,0 +1,429 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for the UQ contract (``FieldDistribution``) and pooled UQ reducer metrics.""" + +from __future__ import annotations + +import math + +import numpy as np +import pytest + +from physicsnemo.cfd.evaluation.datasets.schema import ( + FieldDistribution, + as_distribution, + build_predictive_distribution, + distribution_mean, +) +from physicsnemo.cfd.postprocessing_tools.metric_registry import ( + is_reducer_metric, + is_sample_metric, +) +from physicsnemo.cfd.evaluation.metrics.builtin.uq import ( + _ause_curve_area, + _curve_payload, + _drag_mean_and_std, + _make_pointwise_uq_metrics, + _make_sample_uq_metrics, + _make_uq_metrics, + _SampleDragUQ, + _spearman, + _trapezoid, +) + + +def test_build_predictive_distribution_coerces_numpy_to_float32() -> None: + """``build_predictive_distribution`` coerces mean/std/epistemic arrays to float32.""" + fd = build_predictive_distribution( + mean=np.ones(4, dtype=np.float64), std=np.full(4, 2.0), epistemic_std=np.ones(4) + ) + assert fd.mean.dtype == np.float32 + assert fd.std.dtype == np.float32 + assert fd.epistemic_std.dtype == np.float32 + + +def test_as_distribution_passthrough_wrap_and_none() -> None: + """``as_distribution`` returns distributions as-is, wraps plain arrays, and yields None otherwise.""" + fd = FieldDistribution(mean=np.zeros(3), std=np.ones(3)) + preds = {"pressure": fd, "shear_stress": np.arange(6.0).reshape(3, 2), "z": None} + assert as_distribution(preds, "pressure") is fd + wrapped = as_distribution(preds, "shear_stress") + assert isinstance(wrapped, FieldDistribution) and wrapped.std is None + assert as_distribution(preds, "z") is None + assert as_distribution(preds, "missing") is None + + +def test_distribution_mean_unwraps() -> None: + """``distribution_mean`` returns the mean array for both distributions and plain arrays.""" + fd = FieldDistribution(mean=np.arange(3.0), std=np.ones(3)) + assert distribution_mean(fd) is fd.mean + arr = np.zeros(3) + assert distribution_mean(arr) is arr + + +def _run_reducer(metric, cases): + """Mimic the engine: sum per-case ``partial`` stats, then ``finalize`` once.""" + summed: dict[str, float] = {} + for gt, pred in cases: + for key, val in metric.partial(gt, pred).items(): + summed[key] = summed.get(key, 0.0) + val + return metric.finalize(summed) + + +def _gaussian_case(rng, n, sigma, epi=None): + mu = rng.normal(size=n).astype(np.float32) + y = (mu + sigma * rng.normal(size=n)).astype(np.float32) + fd = FieldDistribution( + mean=mu, + std=np.full(n, sigma, np.float32), + epistemic_std=None if epi is None else np.full(n, epi, np.float32), + ) + return {"pressure": y}, {"pressure": fd} + + +def test_uq_metrics_are_registered_as_reducers() -> None: + """The pooled UQ metrics are all reducer metrics.""" + metrics = _make_uq_metrics() + assert set(metrics) >= { + "nlpd", + "nlpd_epistemic", + "calibration_zrms", + "coverage_95", + "sharpness_std", + "sharpness_epistemic_std", + } + for m in metrics.values(): + assert is_reducer_metric(m) + + +def test_pooled_calibration_on_synthetic_gaussian() -> None: + """A perfectly calibrated Gaussian gives cov95≈0.95, zrms≈1, and the analytic NLPD.""" + rng = np.random.default_rng(0) + metrics = _make_uq_metrics() + n, sigma = 400_000, 1.0 + cases = [_gaussian_case(rng, n, sigma)] + assert abs(_run_reducer(metrics["coverage_95"], cases)["pressure"] - 0.95) < 0.01 + assert ( + abs(_run_reducer(metrics["calibration_zrms"], cases)["pressure"] - 1.0) < 0.02 + ) + nlpd = _run_reducer(metrics["nlpd"], cases)["pressure"] + expected = 0.5 * (math.log(2 * math.pi) + math.log(sigma**2) + 1.0) + assert abs(nlpd - expected) < 0.02 + + +def test_pooling_is_global_not_mean_of_case_means() -> None: + """Sharpness over two cases with different point counts pools by point, not by case.""" + rng = np.random.default_rng(1) + metrics = _make_uq_metrics() + n1, n2, s1, s2 = 300_000, 100_000, 1.0, 3.0 + cases = [_gaussian_case(rng, n1, s1), _gaussian_case(rng, n2, s2)] + pooled = _run_reducer(metrics["sharpness_std"], cases)["pressure"] + expected_pooled = (n1 * s1 + n2 * s2) / (n1 + n2) + mean_of_means = (s1 + s2) / 2.0 + assert abs(pooled - expected_pooled) < 1e-3 + assert abs(pooled - mean_of_means) > 0.05 + + +def test_deterministic_prediction_yields_nan() -> None: + """A deterministic prediction (no std) yields NaN NLPD under the headline ``mean`` sub-key.""" + metrics = _make_uq_metrics() + cases = [ + ( + {"pressure": np.zeros(16, np.float32)}, + {"pressure": FieldDistribution(mean=np.zeros(16, np.float32))}, + ) + ] + # No channels contribute -> finalize({}) returns the {"mean": NaN} headline (consistent schema). + assert math.isnan(_run_reducer(metrics["nlpd"], cases)["mean"]) + + +def test_epistemic_metric_requires_epistemic_std() -> None: + """Epistemic metrics need an epistemic_std: absent -> NaN headline; present -> finite.""" + rng = np.random.default_rng(2) + metrics = _make_uq_metrics() + # total std present but epistemic_std absent -> epistemic metric is NaN (headline "mean"). + cases_no_epi = [_gaussian_case(rng, 1000, 1.0)] + assert math.isnan(_run_reducer(metrics["nlpd_epistemic"], cases_no_epi)["mean"]) + # epistemic_std present -> finite + cases_epi = [_gaussian_case(rng, 1000, 1.0, epi=0.5)] + res = _run_reducer(metrics["sharpness_epistemic_std"], cases_epi) + assert abs(res["pressure"] - 0.5) < 1e-3 + + +def test_vector_field_splits_into_components() -> None: + """A 3-vector field expands into per-component channels plus the headline mean.""" + rng = np.random.default_rng(3) + metrics = _make_uq_metrics() + n, sigma, epi = 50_000, 1.0, 0.5 + mu = rng.normal(size=(n, 3)).astype(np.float32) + y = (mu + sigma * rng.normal(size=(n, 3))).astype(np.float32) + fd = FieldDistribution( + mean=mu, + std=np.full((n, 3), sigma, np.float32), + epistemic_std=np.full((n, 3), epi, np.float32), + ) + res = _run_reducer( + metrics["nlpd_epistemic"], [({"shear_stress": y}, {"shear_stress": fd})] + ) + assert set(res) == {"shear_stress_x", "shear_stress_y", "shear_stress_z", "mean"} + + +def test_partial_outputs_are_additive_scalars() -> None: + """Reducer ``partial`` returns a flat dict of floats so it caches like per-case scalars.""" + rng = np.random.default_rng(4) + metrics = _make_uq_metrics() + gt, pred = _gaussian_case(rng, 100, 1.0) + part = metrics["coverage_95"].partial(gt, pred) + assert part and all(isinstance(v, float) for v in part.values()) + + +# -------------------------------------------------------------------------------------------- +# Pointwise per-point ranking quality: Spearman(|error|, uncertainty) +# -------------------------------------------------------------------------------------------- + + +def test_spearman_high_when_uncertainty_tracks_error() -> None: + """Spearman ≈ 1 when the per-point std equals the noise scale that generated the error.""" + rng = np.random.default_rng(5) + spear = _make_pointwise_uq_metrics()["uncertainty_error_spearman"] + n = 20_000 + mu = rng.normal(size=n).astype(np.float32) + sig = rng.uniform(0.2, 3.0, size=n).astype(np.float32) + # Make |error| a monotonic function of sig so ranks align almost perfectly. + y = (mu + sig * 3.0).astype(np.float32) + informative = spear( + {"pressure": y}, {"pressure": FieldDistribution(mean=mu, std=sig)} + ) + assert informative["pressure"] > 0.99 + + random_sig = rng.uniform(0.2, 3.0, size=n).astype(np.float32) + uninformative = spear( + {"pressure": y}, {"pressure": FieldDistribution(mean=mu, std=random_sig)} + ) + assert abs(uninformative["pressure"]) < 0.1 + + +def test_spearman_tie_aware_and_constant_input() -> None: + """Ties use *average* ranks and a constant input is undefined (NaN), not spuriously 1.0.""" + from scipy.stats import spearmanr + + # Increasing error paired with completely CONSTANT uncertainty: rank correlation is undefined + # (zero rank variance). The old ordinal-rank implementation wrongly reported rho = 1.0. + err = np.array([1.0, 2.0, 3.0, 4.0, 5.0]) + const_unc = np.full(5, 0.7) + assert math.isnan(_spearman(err, const_unc)) + assert math.isnan(_spearman(const_unc, err)) + # Tie-aware: half the uncertainties tie but the trend is still monotone -> strongly positive + # (average ranks give ~0.949; ordinal ranks would spuriously read higher/inconsistent). + tied_unc = np.array([0.1, 0.1, 0.2, 0.2, 0.3]) + mono_err = np.array([1.0, 1.5, 2.0, 2.5, 3.0]) + assert _spearman(mono_err, tied_unc) > 0.9 + assert ( + abs( + _spearman(mono_err, tied_unc) + - float(spearmanr(mono_err, tied_unc).correlation) + ) + < 1e-12 + ) + # Matches scipy (which we delegate to) on a small tied example. + a = np.array([1.0, 2.0, 2.0, 3.0]) + b = np.array([10.0, 20.0, 20.0, 40.0]) + assert abs(_spearman(a, b) - float(spearmanr(a, b).correlation)) < 1e-12 + # Degenerate sizes -> NaN. + assert math.isnan(_spearman(np.array([1.0]), np.array([1.0]))) + assert math.isnan(_spearman(np.array([1.0, 2.0]), np.array([1.0]))) + + +def test_trapezoid_available_and_ause_uses_it() -> None: + """``_trapezoid`` resolves (np.trapz was removed in NumPy 2.4) and AUSE integrates cleanly.""" + assert callable(_trapezoid) + assert _trapezoid(np.array([0.0, 1.0, 0.0]), np.array([0.0, 0.5, 1.0])) == 0.5 + # Perfectly-ranked errors -> AUSE ~ 0 (no AttributeError from a missing np.trapz). + err = np.array([0.2, 0.5, 1.0, 2.0, 4.0]) + assert abs(_ause_curve_area(err, err)) < 1e-9 + + +def test_spearman_deterministic_and_epistemic_variants() -> None: + """Deterministic (no std) and epistemic-without-epistemic-std both give a NaN Spearman mean.""" + spear = _make_pointwise_uq_metrics()["uncertainty_error_spearman"] + spear_epi = _make_pointwise_uq_metrics()["uncertainty_error_spearman_epistemic"] + n = 4000 + mu = np.zeros(n, np.float32) + y = np.random.default_rng(6).normal(size=n).astype(np.float32) + assert math.isnan( + spear({"pressure": y}, {"pressure": FieldDistribution(mean=mu)})["mean"] + ) + fd_total = FieldDistribution(mean=mu, std=np.ones(n, np.float32)) + assert math.isnan(spear_epi({"pressure": y}, {"pressure": fd_total})["mean"]) + + +# -------------------------------------------------------------------------------------------- +# Sample-wise sparsification / AUSE (sample metric: one number per geometry) +# -------------------------------------------------------------------------------------------- + + +def _collect_samples(metric, cases): + """Mimic the engine: gather per-case ``partial`` scalars into lists, then finalize_samples.""" + collected: dict[str, list[float]] = {} + for gt, pred in cases: + for key, val in metric.partial(gt, pred).items(): + collected.setdefault(key, []).append(val) + return metric.finalize_samples(collected) + + +def test_sample_ause_is_registered_as_sample_metric() -> None: + """The sample-wise AUSE metrics are sample metrics (finalize_samples, not finalize).""" + for m in _make_sample_uq_metrics().values(): + assert is_sample_metric(m) + assert not is_reducer_metric( + m + ) # sample metrics use finalize_samples, not finalize + + +def test_sample_ause_informative_below_random_over_geometries() -> None: + """Rank geometries by mean uncertainty: informative uncertainty gives lower AUSE than random.""" + rng = np.random.default_rng(7) + ause = _make_sample_uq_metrics()["sparsification_ause"] + + def _geom(scale, rank_sig): + n = 2000 + mu = np.zeros(n, np.float32) + y = (scale * rng.normal(size=n)).astype(np.float32) # error grows with `scale` + fd = FieldDistribution(mean=mu, std=np.full(n, rank_sig, np.float32)) + return {"pressure": y}, {"pressure": fd} + + scales = [0.2, 0.5, 1.0, 2.0, 4.0, 8.0] + # Informative: per-geometry uncertainty == its error scale (ranks geometries correctly). + informative = _collect_samples(ause, [_geom(s, s) for s in scales]) + # Uninformative: uncertainty unrelated to error scale. + shuffled = list(scales) + rng.shuffle(shuffled) + random_rank = _collect_samples( + ause, [_geom(s, r) for s, r in zip(scales, shuffled)] + ) + assert informative["pressure"] <= random_rank["pressure"] + 1e-9 + assert abs(informative["pressure"]) < 1e-6 # perfect ranking -> AUSE ~ 0 + # Spearman companion: informative ranking -> rho ~ 1, and never worse than random ranking. + assert informative["pressure_spearman"] >= random_rank["pressure_spearman"] - 1e-9 + assert informative["pressure_spearman"] > 0.99 + + +def test_sample_ause_needs_two_geometries() -> None: + """AUSE needs >= 2 geometries; a single geometry yields NaN for that channel.""" + ause = _make_sample_uq_metrics()["sparsification_ause"] + n = 100 + fd = FieldDistribution(mean=np.zeros(n, np.float32), std=np.ones(n, np.float32)) + res = _collect_samples( + ause, [({"pressure": np.ones(n, np.float32)}, {"pressure": fd})] + ) + assert math.isnan(res["pressure"]) + + +def test_sample_ause_curves_shape() -> None: + """``curves`` returns a per-channel payload with aligned fraction/curve arrays and AUSE.""" + ause = _make_sample_uq_metrics()["sparsification_ause"] + collected = { + "pressure::err": [0.2, 0.5, 1.0, 2.0], + "pressure::unc": [0.2, 0.5, 1.0, 2.0], # perfect ranking + } + curves = ause.curves(collected) + assert "pressure" in curves + c = curves["pressure"] + assert c["n"] == 4 and len(c["fractions"]) == 4 + assert len(c["by_uncertainty"]) == 4 and len(c["oracle"]) == 4 + assert abs(c["ause"]) < 1e-9 # uncertainty == error order -> curve == oracle + + +# -------------------------------------------------------------------------------------------- +# Sample-wise drag sparsification via linear UQ propagation into Cd (drag_uq) +# -------------------------------------------------------------------------------------------- + + +def test_drag_mean_and_std_matches_manual_propagation() -> None: + """``_drag_mean_and_std`` reproduces the force integral and the diagonal variance sum.""" + rng = np.random.default_rng(11) + n = 500 + normals = rng.normal(size=(n, 3)) + normals /= np.linalg.norm(normals, axis=1, keepdims=True) + area = rng.uniform(0.1, 1.0, size=n) + direction = np.array([1.0, 0.0, 0.0]) + p = rng.normal(size=n) + wss = rng.normal(size=(n, 3)) + p_std = rng.uniform(0.01, 0.1, size=n) + wss_std = rng.uniform(0.01, 0.1, size=(n, 3)) + drag, std = _drag_mean_and_std( + normals, area, direction, p, wss, p_std, wss_std, 1.0 + ) + c_p = float(np.sum((normals @ direction) * area * p)) + c_f = -float(np.sum((wss @ direction) * area)) + assert abs(drag - (c_p + c_f)) < 1e-9 + w_p = (normals @ direction) * area + w_tau = -area[:, None] * direction[None, :] + var = float(np.sum(w_p**2 * p_std**2)) + float(np.sum(w_tau**2 * wss_std**2)) + assert abs(std - math.sqrt(var)) < 1e-9 + # zero std -> zero drag std (deterministic mean, no uncertainty). + _, zero_std = _drag_mean_and_std( + normals, area, direction, p, wss, np.zeros(n), np.zeros((n, 3)), 1.0 + ) + assert zero_std == 0.0 + + +def test_drag_uq_finalize_and_curves_rank_geometries() -> None: + """Informative drag std -> AUSE ~ 0; anti-correlated std -> larger AUSE; both series present.""" + drag_uq = _SampleDragUQ() + err = [0.2, 0.5, 1.0, 2.0, 4.0, 8.0] + collected = { + "abs_err": err, + "epi_std": err, # perfectly ranks the drag error + "total_std": err[::-1], # anti-correlated + } + fin = drag_uq.finalize_samples(collected) + assert set(fin) == { + "epistemic", + "total", + "epistemic_spearman", + "total_spearman", + } + assert abs(fin["epistemic"]) < 1e-9 + assert fin["total"] > fin["epistemic"] + # Trend-alignment companion: perfect rank -> rho=+1, anti-correlated -> rho=-1. + assert abs(fin["epistemic_spearman"] - 1.0) < 1e-9 + assert abs(fin["total_spearman"] + 1.0) < 1e-9 + curves = drag_uq.curves(collected) + assert set(curves) == {"epistemic", "total"} + assert curves["epistemic"]["n"] == 6 + + +def test_drag_uq_partial_skips_deterministic() -> None: + """No std on the prediction (deterministic wrapper) -> ``partial`` is a no-op.""" + drag_uq = _SampleDragUQ() + preds = { + "pressure": FieldDistribution(mean=np.zeros(8, np.float32)), + "shear_stress": FieldDistribution(mean=np.zeros((8, 3), np.float32)), + } + # Also confirms ``partial`` accepts the same per-run overrides as the deterministic drag metric. + assert ( + drag_uq.partial( + {}, preds, output=object(), coeff=2.5, drag_direction=[1.0, 0.0, 0.0] + ) + == {} + ) + + +def test_curve_payload_none_for_single_sample() -> None: + """The sparsification curve payload is None for a single geometry.""" + assert _curve_payload(np.array([1.0]), np.array([1.0])) is None diff --git a/test/ci_tests/test_uq_wrappers.py b/test/ci_tests/test_uq_wrappers.py new file mode 100644 index 0000000..f292b70 --- /dev/null +++ b/test/ci_tests/test_uq_wrappers.py @@ -0,0 +1,200 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""UQ wrapper plumbing: checkpoint routing, block-partition coverage, field-name fallback, +registry routing, and the UQ contract class attributes.""" + +from __future__ import annotations + +import pytest +import torch + +from physicsnemo.cfd.evaluation.metrics.mesh_bridge import uq_std_field_names +from physicsnemo.cfd.evaluation.models.common_wrapper_utils.geotransolver_runtime import ( + make_forward_permutation, + resolve_checkpoint_file, +) +from physicsnemo.cfd.evaluation.models.model_registry import CFDModel + + +# -------------------------------------------------------------------------------------------- +# Checkpoint routing: a specific, EXISTING file (epoch-in-name) -> (dir, epoch); everything else +# (directory / bad name / missing file) rejects rather than silently loading a random model. +# -------------------------------------------------------------------------------------------- + + +def test_resolve_checkpoint_file_parses_epoch(tmp_path) -> None: + """An existing checkpoint file name resolves to ``(directory, epoch)``.""" + ckpt = tmp_path / "GeoTransolver.0.30.mdlus" + ckpt.write_bytes(b"") + directory, epoch = resolve_checkpoint_file(str(ckpt)) + assert directory == tmp_path and epoch == 30 + pt = tmp_path / "checkpoint.0.100.pt" + pt.write_bytes(b"") + d2, e2 = resolve_checkpoint_file(str(pt)) + assert d2 == tmp_path and e2 == 100 + + +def test_resolve_checkpoint_file_missing_file_raises(tmp_path) -> None: + """A well-named but MISSING checkpoint raises FileNotFoundError (no random-model fallback).""" + with pytest.raises(FileNotFoundError): + resolve_checkpoint_file(str(tmp_path / "GeoTransolver.0.30.mdlus")) + + +def test_resolve_checkpoint_file_rejects_directory_and_epochless(tmp_path) -> None: + """A directory, empty path, or epoch-less file name is rejected (no silent latest-epoch).""" + with pytest.raises(ValueError): + resolve_checkpoint_file(str(tmp_path)) # directory -> ambiguous epoch + with pytest.raises(ValueError): + resolve_checkpoint_file("") # empty + with pytest.raises(ValueError): + resolve_checkpoint_file( + str(tmp_path / "GeoTransolver.mdlus") + ) # no epoch in name (name-shape checked before existence) + + +# -------------------------------------------------------------------------------------------- +# Deterministic-prediction hook: run.uq.enabled=false must give a POINT prediction even for a +# stochastic sampler (MC-Dropout keeps dropout active in predict()). +# -------------------------------------------------------------------------------------------- + + +class _EchoWrapper(CFDModel): + OUTPUT_LOCATION = "cell" + + @property + def output_location(self): + return self.OUTPUT_LOCATION + + def load(self, *a, **k): + return self + + def prepare_inputs(self, case): + return None + + def predict(self, model_input): + return model_input + + def decode_outputs(self, raw_output, case, model_input=None): + return {"pressure": raw_output} + + +class _StochasticDropoutWrapper(CFDModel): + OUTPUT_LOCATION = "cell" + SUPPORTS_UQ = True + UQ_METHOD = "sampling" + + def __init__(self): + self.drop = torch.nn.Dropout(p=0.5) + self.drop.train() # stochastic, as MC-Dropout keeps it + + @property + def output_location(self): + return self.OUTPUT_LOCATION + + def load(self, *a, **k): + return self + + def prepare_inputs(self, case): + return None + + def predict(self, model_input): + return self.drop(torch.ones(4096)) + + def decode_outputs(self, raw_output, case, model_input=None): + return {"pressure": raw_output} + + def predict_deterministic(self, model_input): + was_training = self.drop.training + self.drop.eval() + try: + return self.predict(model_input) + finally: + self.drop.train(was_training) + + +def test_predict_deterministic_default_delegates_to_predict() -> None: + """The base hook simply calls predict() (correct for deterministic / analytic wrappers).""" + w = _EchoWrapper() + assert w.predict_deterministic("x") == w.predict("x") == "x" + + +def test_predict_deterministic_override_removes_stochasticity() -> None: + """A sampling wrapper's override disables dropout for one pass, then restores stochastic state.""" + w = _StochasticDropoutWrapper() + det = w.predict_deterministic(None) + assert torch.all(det == 1.0) # dropout off -> identity, no zeros + assert w.drop.training is True # stochastic state restored for subsequent UQ passes + torch.manual_seed(0) + assert torch.any(w.predict(None) == 0.0) # a stochastic pass zeros some entries + + +# -------------------------------------------------------------------------------------------- +# Block-partition permutation covers the WHOLE point cloud (no subsampling of samples) +# -------------------------------------------------------------------------------------------- + + +def test_make_forward_permutation_covers_all_points() -> None: + """The per-case forward permutation is a full bijection over all points (no dropped indices).""" + n = 137 + batch = {"embeddings": torch.zeros(1, n, 3)} + perm = make_forward_permutation(batch) + assert perm.shape == (n,) + # Every point index appears exactly once -> the forward predicts all points, order restorable. + assert torch.equal(torch.sort(perm).values, torch.arange(n)) + + +# -------------------------------------------------------------------------------------------- +# Automatic uncertainty field-name fallback (shared by mesh attachment and drag_uq) +# -------------------------------------------------------------------------------------------- + + +def test_uq_std_field_names_fallback_and_override() -> None: + """Std field names auto-derive from the prediction name unless explicitly overridden.""" + # No config -> derive from the prediction name (what comparison meshes attach by default). + assert uq_std_field_names("pressure") == ("pressureStd", "pressureEpistemicStd") + # Configured names win. + assert uq_std_field_names("pressure", "pStd", "pEpi") == ("pStd", "pEpi") + # Partial override: only the total std configured, epistemic still auto-derived. + assert uq_std_field_names("wallShearStress", "wssStd", None) == ( + "wssStd", + "wallShearStressEpistemicStd", + ) + + +# -------------------------------------------------------------------------------------------- +# Registry routing + UQ contract attributes for the three example UQ wrappers +# -------------------------------------------------------------------------------------------- + + +def test_uq_wrappers_registered_with_expected_contract() -> None: + """The three UQ wrappers are registered and expose the expected SUPPORTS_UQ/UQ_METHOD contract.""" + # Importing the wrappers package registers every model wrapper. + import physicsnemo.cfd.evaluation.models.wrappers # noqa: F401 + from physicsnemo.cfd.evaluation.models.model_registry import get_model_wrapper + + gp = get_model_wrapper("geotransolver_gp_surface") + mc = get_model_wrapper("geotransolver_mc_dropout_surface") + ens = get_model_wrapper("geotransolver_ensemble_surface") + + # Analytic GP head: single-pass distribution. + assert gp.SUPPORTS_UQ is True and gp.UQ_METHOD == "analytic" + # Sampling wrappers: repeated-pass spread. + assert mc.SUPPORTS_UQ is True and mc.UQ_METHOD == "sampling" + assert ens.SUPPORTS_UQ is True and ens.UQ_METHOD == "sampling" + # All three decorate physical-target (transformer_models) checkpoints -> no ρu² re-dim. + for cls in (gp, mc, ens): + assert cls.REDIMENSIONALIZE_OUTPUTS is False diff --git a/workflows/benchmarking/conf/config_uq_surface.yaml b/workflows/benchmarking/conf/config_uq_surface.yaml new file mode 100644 index 0000000..775a289 --- /dev/null +++ b/workflows/benchmarking/conf/config_uq_surface.yaml @@ -0,0 +1,227 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# ============================================================================ +# EXAMPLE: probabilistic / uncertainty-quantification (UQ) surface benchmark. +# +# It scores four models on the SAME cases so you can read the pooled UQ metrics +# (NLPD, calibration z-RMS, 95% coverage, sharpness, sparsification/AUSE) +# alongside the usual accuracy metrics (L2 / drag / lift). The four rows cover +# every path in the UQ engine and act as the reference examples for the three +# UQ wrapper archetypes: +# +# * geotransolver_drivaerstar_surface -> deterministic (SUPPORTS_UQ=False). +# UQ metrics finalize to NaN; included for the det-vs-UQ contrast. +# * geotransolver_gp_surface -> UQ_METHOD="analytic" (GP field head): +# one forward pass returns the posterior; decode_distribution() emits +# mean + total/epistemic std. +# * geotransolver_mc_dropout_surface -> UQ_METHOD="sampling" (one model, N +# stochastic passes): the engine runs run.uq.num_samples passes and +# aggregates them (Welford). +# * geotransolver_ensemble_surface -> UQ_METHOD="sampling" (K models, one +# pass each): predict_ensemble() returns one output per member. +# +# ---------------------------------------------------------------------------- +# BEFORE RUNNING: edit the paths in the `_paths` block and `benchmark.datasets` +# below to point at your own checkpoints / stats / data. They are placeholders +# (``/path/to/...``); nothing here is machine-specific. +# +# Run (on a GPU node): +# python main.py --config-name=config_uq_surface +# torchrun --standalone --nnodes=1 --nproc_per_node=8 main.py --config-name=config_uq_surface +# ============================================================================ + +defaults: + - _self_ + +hydra: + run: + dir: ${run.output_dir} + output_subdir: hydra + job: + chdir: false + +case_id: null + +run: + device: "cuda:0" + output_dir: "benchmark_results_uq_surface" + seed: 42 + batch_size: 1 + # Write per-case .vtp with Pred/Std/EpistemicStd arrays for ParaView spatial-UQ inspection. + save_inference_mesh: true + metrics_cache: + enabled: true + path: ${run.output_dir}/metrics_cache + # ---- Uncertainty-quantification controls (additive; default off) ---- + uq: + enabled: true + # Stochastic passes per case for UQ_METHOD="sampling" wrappers. Ignored by analytic / + # deterministic rows, and by the ensemble (its pass count is its number of members). + num_samples: 32 + # Keep the raw per-pass samples on the FieldDistribution (memory-heavy; only for CRPS / debugging). + retain_samples: false + device_metrics: false + +# ---- Edit these to your environment -------------------------------------- +# All checkpoints below are assumed to share ONE surface-normalization file. The GP / MC-Dropout / +# ensemble rows all decorate the same GeoTransolver backbone family (physical, standardize-only +# targets), so they reuse the same stats. +_paths: + surface_stats: &surface_stats /path/to/surface_fields_normalization.npz + # Deterministic + ensemble members come from a standard training run. Point these at the + # MODEL-weights file physicsnemo writes per saved epoch (``GeoTransolver.0..mdlus``), + # NOT the ``checkpoint.0..pt`` training-state file — the wrapper loads the exact file + # named here (and raises if it is missing), so an epoch's `.mdlus` is what you want. + baseline_ckpt: &baseline_ckpt /path/to/run/checkpoints/GeoTransolver.0.100.mdlus + # A backbone trained with Concrete-Dropout (model.concrete_dropout=true) for the MC-Dropout row. + concrete_dropout_ckpt: &concrete_dropout_ckpt /path/to/run_concrete_dropout/checkpoints/GeoTransolver.0.100.mdlus + +benchmark: + mode: matrix + models: + # (1) Deterministic baseline. SUPPORTS_UQ=False -> UQ metrics report NaN. + - name: geotransolver_drivaerstar_surface + inference_domain: surface + checkpoint: *baseline_ckpt + stats_path: *surface_stats + kwargs: { batch_resolution: 60000, geometry_sampling: 300000 } + + # (2) Analytic GP field head. This wrapper loads TWO checkpoints, both named + # explicitly: `checkpoint` is the GeoTransolver backbone (epoch parsed from the file name) and + # `gp_head_checkpoint` is the matching FieldGPHead (read as given; pass matching files). The GP + # hyperparameters here MUST match the training run so the FieldGPHead reconstructs before the + # state dict loads. + + - name: geotransolver_gp_surface + inference_domain: surface + checkpoint: /path/to/run_field_gp_due/checkpoints_field_gp/GeoTransolver.0.100.mdlus + stats_path: *surface_stats + kwargs: + gp_head_checkpoint: /path/to/run_field_gp_due/checkpoints_field_gp/FieldGPHead.0.100.pt + gp_feature_norm: "l2_radial" + gp_n_inducing: 1024 + gp_mlp_hidden: [128, 128, 16] + gp_spectral_norm_coeff: 0.95 # > 0 -> DUE bi-Lipschitz extractor + gp_dkl_residual: true + gp_lengthscale_range: [0.01, 2.0] + gp_lengthscale_prior: [2.0, 3.0] + gp_outputscale_prior: [2.0, 0.5] + num_tasks: 4 + gp_inference_chunk_size: 51200 + geometry_sampling: 300000 + + # (3) MC-Dropout: run.uq.num_samples stochastic passes over the Concrete-Dropout backbone + # (the ConcreteDropout layers stay stochastic at inference). + - name: geotransolver_mc_dropout_surface + inference_domain: surface + checkpoint: *concrete_dropout_ckpt + stats_path: *surface_stats + kwargs: + batch_resolution: 60000 + geometry_sampling: 300000 + + # (4) K-member ensemble: one pass per member. List the member checkpoints explicitly (from a + # single run = snapshot ensemble, or from separate runs = true deep ensemble; same code path). + # `checkpoint` just anchors the cache fingerprint -> point it at any one member. + - name: geotransolver_ensemble_surface + inference_domain: surface + checkpoint: *baseline_ckpt + stats_path: *surface_stats + kwargs: + # Model-weights files (one per member); NOT the checkpoint.0..pt training state. + member_checkpoints: + - /path/to/run/checkpoints/GeoTransolver.0.96.mdlus + - /path/to/run/checkpoints/GeoTransolver.0.97.mdlus + - /path/to/run/checkpoints/GeoTransolver.0.98.mdlus + - /path/to/run/checkpoints/GeoTransolver.0.99.mdlus + - /path/to/run/checkpoints/GeoTransolver.0.100.mdlus + batch_resolution: 60000 + geometry_sampling: 300000 + + datasets: + - name: drivaerstar + root: /path/to/drivaerstar/surface_files/class_F/val + # flip_wss_sign controls the wall-shear-stress sign convention. The adapter default (True) + # negates DrivAerStar WSS to the DrivAerML convention; it is set False here because the + # GeoTransolver checkpoints in this example were trained on the NATIVE DrivAerStar WSS sign + # (the transformer_models pipeline). Match this to how YOUR checkpoints were trained — a + # mismatch silently flips WSS, and therefore drag and lift, for every case. + kwargs: { inference_domain: surface, flip_wss_sign: false } + + reproducibility: + log_env: false + save_artifacts: false + +output: + surface_interpolate_point_to_cell_for_metrics: true + surface_metrics_idw_k: 5 + mesh_field_names: + pressure: "pressure_pred" + shear_stress: "wall_shear_stress_pred" + ground_truth_mesh_field_names: + pressure: "pressure_true" + shear_stress: "wall_shear_stress_true" + # Uncertainty companions on exported meshes (auto-derived as *Std / *EpistemicStd if omitted). + std_mesh_field_names: + pressure: "pressure_std" + shear_stress: "wall_shear_stress_std" + epistemic_std_mesh_field_names: + pressure: "pressure_epistemic_std" + shear_stress: "wall_shear_stress_epistemic_std" + +metrics: + # Deterministic (scored on the distribution mean for UQ wrappers). + - l2_pressure + - l2_shear_stress + - l2_pressure_area_weighted + # coeff (Cd prefactor 2/(A*rho*U^2)) + drag_direction are optional; drag_direction is shown + # explicitly so the drag_uq panel below uses the SAME direction as the deterministic drag. + - { name: drag, drag_direction: [1, 0, 0] } + - lift + # Pooled UQ reducer metrics (NaN for the deterministic row; finite for GP / sampling rows). + - nlpd + - nlpd_epistemic + - calibration_zrms + - coverage_95 + - sharpness_std + - sharpness_epistemic_std + # Per-point ranking quality: Spearman(|error|, uncertainty) within each geometry (higher=better). + - uncertainty_error_spearman + - uncertainty_error_spearman_epistemic + # Sample-wise sparsification: rank geometries by aggregate uncertainty vs error (AUSE, lower=better). + - sparsification_ause + - sparsification_ause_epistemic + # Decision-relevant drag sparsification: rank geometries by uncertainty propagated into the Cd + # integral (drag std) vs |Cd_pred - Cd_true|. NOTE: this applies to BOTH analytic (GP) and + # sampling (MC-Dropout / ensemble) rows, because drag_uq propagates the per-cell MARGINAL std the + # engine produces for every method. That diagonal (per-cell independent) propagation is exact for + # spatially-uncorrelated uncertainty but a lower bound for spatially-correlated (i.e. epistemic) + # uncertainty — which dominates a coherent surface integral — so the drag-EPISTEMIC ranking is a + # lower bound; read it alongside the field-level epistemic metrics. + - { name: drag_uq, drag_direction: [1, 0, 0] } + +reports: + enabled: true + save_comparison_meshes: false + comparison_mesh_subdir: comparison_meshes + # Only these cases write an inference__.vtp (and drive report visuals); every other + # case is still scored for metrics. Keep this short so large runs don't dump one VTP per case. + visual_case_ids: ["00035", "00039"] + # Sample-level sparsification figure per dataset (field panels from sparsification_ause + a Cd + # panel from drag_uq), with the GP / MC-Dropout / ensemble models overlaid vs the oracle. + visuals: + - sparsification_plot