From 466c22cc9fa27a8d176049e76c863706b9b12f3e Mon Sep 17 00:00:00 2001 From: Kaustubh Tangsali Date: Thu, 9 Jul 2026 22:26:58 +0000 Subject: [PATCH 01/10] initial commit with uq benchmarking tooling --- .../benchmarks/distributed_utils.py | 13 + .../cfd/evaluation/benchmarks/engine.py | 198 ++++- .../evaluation/benchmarks/metrics_cache.py | 3 +- .../cfd/evaluation/benchmarks/report.py | 12 +- .../cfd/evaluation/benchmarks/uq_inference.py | 390 ++++++++++ physicsnemo/cfd/evaluation/config.py | 62 ++ .../cfd/evaluation/datasets/__init__.py | 8 + .../evaluation/datasets/adapters/__init__.py | 4 +- .../datasets/adapters/drivaerstar.py | 302 ++++++++ physicsnemo/cfd/evaluation/datasets/schema.py | 100 +++ .../evaluation/metrics/builtin/__init__.py | 2 + .../cfd/evaluation/metrics/builtin/uq.py | 690 ++++++++++++++++++ .../cfd/evaluation/metrics/mesh_bridge.py | 74 +- .../cfd/evaluation/metrics/registry.py | 15 +- .../geotransolver_runtime.py | 523 +++++++++++++ .../cfd/evaluation/models/model_registry.py | 50 ++ .../evaluation/models/wrappers/__init__.py | 15 + .../geotransolver_drivaerstar/__init__.py | 21 + .../geotransolver_drivaerstar/wrapper.py | 211 ++++++ .../wrappers/geotransolver_gp/__init__.py | 21 + .../wrappers/geotransolver_gp/wrapper.py | 424 +++++++++++ .../wrappers/mc_perturbation/__init__.py | 21 + .../wrappers/mc_perturbation/wrapper.py | 320 ++++++++ .../evaluation/reports/builtin/__init__.py | 6 +- .../reports/builtin/sparsification_plot.py | 196 +++++ .../postprocessing_tools/metric_registry.py | 83 ++- test/ci_tests/test_mc_perturbation_scope.py | 74 ++ test/ci_tests/test_uq_inference.py | 142 ++++ test/ci_tests/test_uq_metrics.py | 349 +++++++++ .../benchmarking/conf/config_uq_surface.yaml | 188 +++++ 30 files changed, 4494 insertions(+), 23 deletions(-) create mode 100644 physicsnemo/cfd/evaluation/benchmarks/uq_inference.py create mode 100644 physicsnemo/cfd/evaluation/datasets/adapters/drivaerstar.py create mode 100644 physicsnemo/cfd/evaluation/metrics/builtin/uq.py create mode 100644 physicsnemo/cfd/evaluation/models/common_wrapper_utils/geotransolver_runtime.py create mode 100644 physicsnemo/cfd/evaluation/models/wrappers/geotransolver_drivaerstar/__init__.py create mode 100644 physicsnemo/cfd/evaluation/models/wrappers/geotransolver_drivaerstar/wrapper.py create mode 100644 physicsnemo/cfd/evaluation/models/wrappers/geotransolver_gp/__init__.py create mode 100644 physicsnemo/cfd/evaluation/models/wrappers/geotransolver_gp/wrapper.py create mode 100644 physicsnemo/cfd/evaluation/models/wrappers/mc_perturbation/__init__.py create mode 100644 physicsnemo/cfd/evaluation/models/wrappers/mc_perturbation/wrapper.py create mode 100644 physicsnemo/cfd/evaluation/reports/builtin/sparsification_plot.py create mode 100644 test/ci_tests/test_mc_perturbation_scope.py create mode 100644 test/ci_tests/test_uq_inference.py create mode 100644 test/ci_tests/test_uq_metrics.py create mode 100644 workflows/benchmarking/conf/config_uq_surface.yaml diff --git a/physicsnemo/cfd/evaluation/benchmarks/distributed_utils.py b/physicsnemo/cfd/evaluation/benchmarks/distributed_utils.py index 7af10d4..91baac6 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,14 @@ 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)) 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 +117,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..e51f6b2 100644 --- a/physicsnemo/cfd/evaluation/benchmarks/engine.py +++ b/physicsnemo/cfd/evaluation/benchmarks/engine.py @@ -33,6 +33,11 @@ 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. + 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 +103,31 @@ 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, + 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 +231,21 @@ def _retain_comparison_mesh_for_visual_context( return case_id in allow +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,6 +303,7 @@ 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, @@ -307,6 +346,8 @@ 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 @@ -341,12 +382,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 +472,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, @@ -559,6 +647,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,14 +729,32 @@ 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) + supports_uq = bool(getattr(wrapper, "SUPPORTS_UQ", False)) + uq_method = getattr(wrapper, "UQ_METHOD", "none") + if supports_uq and uq_method == "sampling" and run_config.uq.enabled: + # N stochastic passes; prepare_inputs already ran once (see UQ design doc §6.2). + 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 supports_uq and uq_method == "analytic": + raw = wrapper.predict(model_input) + predictions = wrapper.decode_distribution(raw, case, model_input) + else: + raw = wrapper.predict(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, @@ -670,6 +781,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, @@ -747,12 +894,21 @@ 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. + metrics_summary.update(finalize_reducer_metrics(per_case, m_dom)) + metrics_summary.update(finalize_sample_metrics(per_case, m_dom)) + return ( { "model": model_config.name, @@ -760,6 +916,7 @@ def _run_single( "cases": cases, "metrics": metrics_summary, "per_case": per_case, + "inference_domain": m_dom, }, mesh_ctx, ) @@ -899,6 +1056,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 +1125,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 +1203,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, 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..2ee4be3 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,11 @@ 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..074c8f9 --- /dev/null +++ b/physicsnemo/cfd/evaluation/benchmarks/uq_inference.py @@ -0,0 +1,390 @@ +# 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 +) -> 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. + """ + # 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) + ) + + out: dict[str, float] = {} + for metric_name, stats in summed.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 +) -> 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}"``. + """ + 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, float] = {} + 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 + 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 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 available (single batched call), + 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..bdee31a 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 @@ -260,6 +292,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 +334,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 +349,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 +381,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..833be55 --- /dev/null +++ b/physicsnemo/cfd/evaluation/datasets/adapters/drivaerstar.py @@ -0,0 +1,302 @@ +# 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. Convert legacy ``.vtk`` → ``.vtp`` (XML PolyData) for the datapipe wrappers. +2. Rename ``Pressure`` → ``pMeanTrim`` (DrivAerML convention). +3. Combine the three WSS scalar components into one ``(N, 3)`` vector and flip its sign + to match DrivAerML (``flip_wss_sign``, default on). +4. Drop explicit ``Normals`` / ``Area`` arrays so downstream force integration and + rendering recompute them from mesh topology (DrivAerML convention). +5. Generate a triangulated STL geometry (DrivAerStar ships none — the surface *is* the + geometry) named so the wrappers' STL resolver finds it. + +Prepared VTP + STL are cached under ``///`` (one STL per +case dir, so the wrappers' STL lookup is unambiguous) and reused on subsequent runs. + +.. note:: + The datapipe wrappers (GeoTransolver, Transolver) parse an integer run index from the + case id via ``run_id_from_case_id`` to name the STL, 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. +""" + +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``: negate combined WSS to match DrivAerML (default ``True``). + - ``remove_normals_area``: drop explicit ``Normals`` / ``Area`` arrays (default ``True``). + - ``make_stl``: generate a triangulated STL per case (default ``True``). + - ``gt_data_type``: ``auto`` / ``cell`` / ``point`` passed to GT extraction + (default ``"cell"``; DrivAerStar fields are cell-centered). + - ``prepared_subdir``: cache directory name under ``root`` (default ``"_prepared"``). + - ``pressure_out_name`` / ``shear_out_name``: VTK array names written into the prepared + VTP (defaults ``"pMeanTrim"`` / ``"wallShearStressMeanTrim"``). + - ``force_reprepare``: rebuild cached VTP/STL even if present (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._make_stl: bool = bool(kwargs.get("make_stl", True)) + self._gt_data_type: str = kwargs.get("gt_data_type", "cell") + 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 _prepare_case(self, case_id: str) -> tuple[str, pv.PolyData]: + """Convert source ``.vtk`` to a canonical VTP (+ STL), caching under the prepared dir. + + Returns the prepared VTP path and the in-memory prepared mesh (for ``reference_geometry``). + """ + src_path = self._source_path(case_id) + 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" + + if self._force_reprepare or not vtp_path.exists(): + log_dataset("drivaerstar", f"Preparing VTP for {case_id!r} from {src_path}") + mesh = pv.read(str(src_path)) + if not isinstance(mesh, pv.PolyData): + mesh = mesh.extract_surface() + + self._rename_pressure(mesh) + self._combine_wss(mesh) + if self._remove_normals_area: + self._drop_normals_area(mesh) + + vtp_path.parent.mkdir(parents=True, exist_ok=True) + mesh.save(str(vtp_path)) + else: + mesh = pv.read(str(vtp_path)) + + if self._make_stl and (self._force_reprepare or not stl_path.exists()): + log_dataset("drivaerstar", f"Writing STL geometry for {case_id!r}") + geom = pv.read(str(src_path)) + if not isinstance(geom, pv.PolyData): + geom = geom.extract_surface() + geom.triangulate().save(str(stl_path)) + + return str(vtp_path), mesh + + def _rename_pressure(self, mesh: pv.PolyData) -> 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 + + def _combine_wss(self, mesh: pv.PolyData) -> 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 + + @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: + """Prepare (cache) the case and load it into the canonical schema.""" + log_dataset( + "drivaerstar", + f"load_case({case_id!r}): root={self.root}", + ) + vtp_path, mesh = self._prepare_case(case_id) + + 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": self._source_path(case_id).name, + "prepared_vtp": Path(vtp_path).name, + } + 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=vtp_path, + mesh_type=mesh_type, + ground_truth=ground_truth, + metadata=meta, + inference_domain="surface", + reference_geometry=mesh, + ) diff --git a/physicsnemo/cfd/evaluation/datasets/schema.py b/physicsnemo/cfd/evaluation/datasets/schema.py index 151c282..5803be6 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, see the UQ design doc §6.8) 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"] @@ -119,3 +125,97 @@ 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..b5677bb --- /dev/null +++ b/physicsnemo/cfd/evaluation/metrics/builtin/uq.py @@ -0,0 +1,690 @@ +# 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 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 + +#: 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: + return 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 _rankdata(x: np.ndarray) -> np.ndarray: + """Ordinal ranks (0..n-1) of ``x`` via argsort-of-argsort. + + Continuous std / error values effectively never tie, so ordinal (vs average-tie) ranks are + fine here and avoid a SciPy dependency. + """ + order = np.argsort(x, kind="stable") + ranks = np.empty(x.size, dtype=np.float64) + ranks[order] = np.arange(x.size, dtype=np.float64) + return ranks + + +def _spearman(a: np.ndarray, b: np.ndarray) -> float: + """Spearman rank correlation = Pearson correlation of the ranks of ``a`` and ``b``.""" + if a.size < 2: + return float("nan") + ra = _rankdata(a) + rb = _rankdata(b) + ra -= ra.mean() + rb -= rb.mean() + denom = math.sqrt(float((ra**2).sum()) * float((rb**2).sum())) + if denom <= 0.0: + return float("nan") + return float((ra * rb).sum() / denom) + + +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(np.trapz((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``. + """ + + 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: + return float("nan") + out: dict[str, float] = {} + values: 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") + continue + v = _ause_curve_area(err, unc) + 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 + + 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-point field variances (diagonal-posterior linear error propagation; a lower bound under +# spatial correlation, but exact-enough for *ranking* geometries — which is all sparsification +# needs). 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|``. + Requires the comparison mesh to carry cell-dof std companions (the benchmark attaches them when + ``output.std_mesh_field_names`` / ``epistemic_std_mesh_field_names`` are set), so 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, + ) + + 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") + std_p = output.std_mesh_field_names.get("pressure") + std_w = output.std_mesh_field_names.get("shear_stress") + epi_p = output.epistemic_std_mesh_field_names.get("pressure") + epi_w = output.epistemic_std_mesh_field_names.get("shear_stress") + + # 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) + else: + out[sub] = 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..335ae71 100644 --- a/physicsnemo/cfd/evaluation/metrics/mesh_bridge.py +++ b/physicsnemo/cfd/evaluation/metrics/mesh_bridge.py @@ -24,12 +24,60 @@ 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 _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] + if value.std is not None: + name = std_map.get(canonical) or f"{pred_name}Std" + _assign_field(mesh, preference, name, value.std) + attached.append(name) + if value.epistemic_std is not None: + name = epi_std_map.get(canonical) or f"{pred_name}EpistemicStd" + _assign_field(mesh, preference, name, value.epistemic_std) + attached.append(name) + return attached + + def _infer_surface_preference( mesh: pv.PolyData, gt: dict[str, Any], @@ -210,20 +258,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 +292,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 +312,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..3b185f3 --- /dev/null +++ b/physicsnemo/cfd/evaluation/models/common_wrapper_utils/geotransolver_runtime.py @@ -0,0 +1,523 @@ +# 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, and the MC weight-perturbation proxy. 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_checkpoint + + _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_checkpoint) 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 (``load_checkpoint`` relies on it).""" + if not DistributedManager.is_initialized(): + DistributedManager.initialize() + + +def build_geotransolver_backbone( + *, checkpoint_path: str, device: str, inference_mode: str +) -> "GeoTransolver": + """Construct a ``GeoTransolver`` for ``surface``/``volume`` and load its checkpoint onto ``device``.""" + model_kw = dict( + DEFAULT_GEOTRANSOLVER_VOLUME_KW + if inference_mode == "volume" + else DEFAULT_GEOTRANSOLVER_KW + ) + checkpoint_dir = Path(checkpoint_path) + if checkpoint_dir.is_file(): + checkpoint_dir = checkpoint_dir.parent + + ensure_distributed_initialized() + dev = torch.device(device) + model = GeoTransolver(**model_kw) + ckpt_args: dict[str, Any] = {"path": str(checkpoint_dir), "models": model} + epoch = parse_checkpoint_epoch(checkpoint_path) + if epoch is not None: + ckpt_args["epoch"] = epoch + with trusted_torch_load_context(): + _ = load_checkpoint(device=dev, **ckpt_args) + 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 + + 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, + ) + 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, + ) + 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 geotransolver_forward( + *, + model: "GeoTransolver", + batch: dict[str, Any], + batch_resolution: int, + cuda_bf16_autocast_enabled: bool, + device: str, +) -> torch.Tensor: + """Blocked GeoTransolver forward pass; returns model-space predictions (N, C), order restored.""" + fx_bn_c = global_fx_to_bnc(batch["fx"]) + n = batch["embeddings"].shape[1] + batch_res = min(batch_resolution, n) + indices = 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", + "build_geotransolver_backbone", + "build_surface_datapipe", + "build_volume_datapipe", + "build_transolver_batch", + "geotransolver_forward", + "unscale_targets", + "decode_surface_predictions", + "decode_volume_predictions", +] diff --git a/physicsnemo/cfd/evaluation/models/model_registry.py b/physicsnemo/cfd/evaluation/models/model_registry.py index bae1701..5d753a5 100644 --- a/physicsnemo/cfd/evaluation/models/model_registry.py +++ b/physicsnemo/cfd/evaluation/models/model_registry.py @@ -21,6 +21,7 @@ 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 (see the UQ design doc §4.2): + #: ``"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] @@ -107,6 +120,43 @@ 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[list[RawOutput]]: + """Optional fast-path for ``UQ_METHOD="sampling"`` wrappers. + + Return ``n`` raw outputs (stochastic passes or ensemble members) in one call, letting + the engine skip its per-pass Python loop. **Only implement this when all ``n`` raw + outputs fit simultaneously in the compute device's memory (typically GPU), since the + returned list holds them all at once** — for full-mesh CFD fields that is usually ``n×`` + the single-pass footprint and will OOM for large ``n``. + + Return ``None`` (the default, used by all shipped sampling wrappers) to have the engine + instead call :meth:`predict` ``n`` times and stream the passes through Welford + aggregation (:func:`~physicsnemo.cfd.evaluation.benchmarks.uq_inference.run_sampling_inference`): + only **one** raw output is resident on the GPU at a time, and the running mean/variance + accumulators are a handful of host (CPU) arrays of one field's size — i.e. memory is + O(1) in ``n``. Prefer this default unless you have measured that the batched path fits. + """ + 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..f9c22b6 100644 --- a/physicsnemo/cfd/evaluation/models/wrappers/__init__.py +++ b/physicsnemo/cfd/evaluation/models/wrappers/__init__.py @@ -22,6 +22,15 @@ 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.mc_perturbation import ( + MCPerturbationDrivAerStarWrapper, +) from physicsnemo.cfd.evaluation.models.wrappers.surface_baseline import ( SurfaceBaselineWrapper, ) @@ -35,6 +44,9 @@ 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("mc_perturbation_surface", MCPerturbationDrivAerStarWrapper) register_model("transolver_surface", TransolverWrapper) register_model("transolver_volume", TransolverWrapper) register_model("domino_surface", DominoWrapper) @@ -46,6 +58,9 @@ "FIGNetWrapper", "XMGNWrapper", "GeoTransolverWrapper", + "GeoTransolverDrivAerStarWrapper", + "GeoTransolverGPDrivAerStarWrapper", + "MCPerturbationDrivAerStarWrapper", "TransolverWrapper", "DominoWrapper", "SurfaceBaselineWrapper", 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..b38e34e --- /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_checkpoint)." + ) + 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..0ea34fd --- /dev/null +++ b/physicsnemo/cfd/evaluation/models/wrappers/geotransolver_gp/wrapper.py @@ -0,0 +1,424 @@ +# 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, §6.3), so the pooled UQ +metrics consume them directly. + +Surface only (the GP head predicts 4 surface tasks: pressure + 3 wall-shear components). + +Model kwargs (``model.kwargs`` in config), matching the trained checkpoint's GP settings: + +- ``checkpoint_epoch`` (int, required for a checkpoint *directory*): epoch tag of the + ``GeoTransolver.0..mdlus`` + ``FieldGPHead.0..pt`` pair to load. +- ``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_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, +) +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_checkpoint + + _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_dir: Optional[str] = None + self._checkpoint_epoch: Optional[int] = 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)) + ce = kw.pop("checkpoint_epoch", None) + self._checkpoint_epoch = int(ce) if ce is not None else 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")), + } + + 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", + ) + + ckpt_dir = Path(checkpoint_path) + if ckpt_dir.is_file(): + ckpt_dir = ckpt_dir.parent + self._checkpoint_dir = str(ckpt_dir) + 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 its weights.""" + 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(): + loaded_epoch = load_checkpoint( + path=self._checkpoint_dir, + models=[self._model, self._head], + device=dev, + epoch=self._checkpoint_epoch, + ) + self._model.eval() + self._head.eval() + self._head.likelihood.eval() + log_inference( + "geotransolver_gp", + f"Loaded FieldGPHead (epoch {loaded_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_perturbation/__init__.py b/physicsnemo/cfd/evaluation/models/wrappers/mc_perturbation/__init__.py new file mode 100644 index 0000000..6f4d436 --- /dev/null +++ b/physicsnemo/cfd/evaluation/models/wrappers/mc_perturbation/__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_perturbation.wrapper import ( + MCPerturbationDrivAerStarWrapper, +) + +__all__ = ["MCPerturbationDrivAerStarWrapper"] diff --git a/physicsnemo/cfd/evaluation/models/wrappers/mc_perturbation/wrapper.py b/physicsnemo/cfd/evaluation/models/wrappers/mc_perturbation/wrapper.py new file mode 100644 index 0000000..7f35321 --- /dev/null +++ b/physicsnemo/cfd/evaluation/models/wrappers/mc_perturbation/wrapper.py @@ -0,0 +1,320 @@ +# 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. + +"""Monte-Carlo **weight-perturbation** proxy for a sampling UQ method (surface). + +.. warning:: + This is a **throwaway plumbing proxy**, *not* a valid UQ method. It exists only to exercise + the ``UQ_METHOD="sampling"`` engine path (N ``predict`` calls, per-pass seeding, streaming + Welford aggregation, pooled metrics, overlaid reports) *before* a real MC-Dropout wrapper is + available — **no training and no new checkpoint required** (UQ design doc §8.2). The + uncertainties it produces are **not calibrated or physically meaningful**. Swap it for the + real ``mc_dropout_surface`` wrapper (same config) once that lands. + +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` for its +deterministic forward, and makes it stochastic in :meth:`predict`: before each pass it multiplies a +subset of weights by ``(1 + eps)`` with ``eps ~ N(0, perturb_std^2)`` (a zero-cost SWAG-flavored +pseudo-ensemble), runs the forward, then restores the weights. The engine drives the N passes and +aggregates them into a +:class:`~physicsnemo.cfd.evaluation.datasets.schema.FieldDistribution` (sample mean + +across-pass std = epistemic). + +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``): + +- ``perturb_std`` (float, default ``0.02``): stddev of the multiplicative Gaussian weight noise. + ``perturb_std → 0`` collapses to the deterministic model (epistemic std → 0), a useful sanity test. +- ``perturb_scope`` (str, default ``"last_n_layers=4"``): which weights to perturb — + ``"all"``, ``"last_n_layers=K"``, or ``"fraction=f"`` (last fraction ``f`` of parameter tensors). +- plus the shared GeoTransolver runtime kwargs (``cuda_bf16_autocast``; datapipe overrides; etc.); + ``checkpoint`` / ``stats_path`` are the trained baseline. +""" + +import logging +import re +from typing import Any, ClassVar, Literal, Optional + +import torch + +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, + parse_runtime_kwargs, + unscale_targets, +) +from physicsnemo.cfd.evaluation.models.model_registry import ( + CFDModel, + ModelInput, + OutputLocation, + Predictions, + RawOutput, +) + +_LOG = logging.getLogger(__name__) + + +def _parse_perturb_scope(scope: str, named_params: list[str]) -> set[str]: + """Resolve which parameter names to perturb from a ``perturb_scope`` string. + + Supports ``"all"``, ``"last_n_layers=K"`` (params whose max embedded layer index is within + the top ``K`` distinct indices), and ``"fraction=f"`` (last fraction of parameter tensors). + Falls back to the last 25% of tensors when a layer-based scope finds no numeric indices. + """ + scope = (scope or "").strip() + if scope in ("", "all"): + return set(named_params) + + if scope.startswith("fraction="): + try: + frac = float(scope.split("=", 1)[1]) + except ValueError: + frac = 0.25 + frac = min(max(frac, 0.0), 1.0) + k = max(1, int(round(frac * len(named_params)))) + return set(named_params[-k:]) + + if scope.startswith("last_n_layers="): + try: + k = int(scope.split("=", 1)[1]) + except ValueError: + k = 4 + + def _layer_index(name: str) -> int | None: + # Largest integer among dotted segments that are pure digits (the block index). + ids = [int(t) for t in re.split(r"[.\[\]]", name) if t.isdigit()] + return max(ids) if ids else None + + indexed = {n: _layer_index(n) for n in named_params} + distinct = sorted({i for i in indexed.values() if i is not None}) + if distinct: + keep = set(distinct[-k:]) + return {n for n, i in indexed.items() if i in keep} + # No numeric layer ids -> fall back to last 25% of tensors. + + k = max(1, int(round(0.25 * len(named_params)))) + return set(named_params[-k:]) + + +class MCPerturbationDrivAerStarWrapper(CFDModel): + """Weight-perturbation sampling proxy over the baseline GeoTransolver (surface). + + 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 + 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() + self._perturb_std: float = 0.02 + self._perturb_scope: str = "last_n_layers=4" + self._perturb_names: set[str] = set() + + def load( + self, + checkpoint_path: str, + stats_path: str, + device: str, + **kwargs: Any, + ) -> "MCPerturbationDrivAerStarWrapper": + """Load the baseline GeoTransolver, then pick the weights to perturb per ``perturb_scope``.""" + if not geotransolver_available(): + raise RuntimeError( + "MCPerturbationDrivAerStarWrapper requires physicsnemo (GeoTransolver, " + "TransolverDataPipe, load_checkpoint)." + ) + kw = dict(kwargs) + self._perturb_std = float(kw.pop("perturb_std", 0.02)) + self._perturb_scope = str(kw.pop("perturb_scope", "last_n_layers=4")) + 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_perturbation", 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_perturbation", 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, + ) + + _LOG.warning( + "MCPerturbationDrivAerStarWrapper is a THROWAWAY sampling PROXY (weight perturbation, " + "perturb_std=%s, perturb_scope=%r); its uncertainties are NOT calibrated or " + "physically meaningful. Replace with the real MC-Dropout wrapper for valid UQ.", + self._perturb_std, + self._perturb_scope, + ) + names = [n for n, _ in self._model.named_parameters()] + self._perturb_names = _parse_perturb_scope(self._perturb_scope, names) + _LOG.info( + "Perturbing %d / %d parameter tensors per pass.", + len(self._perturb_names), + len(names), + ) + 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("MCPerturbationDrivAerStarWrapper: 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 _deterministic_predict(self, model_input: ModelInput) -> RawOutput: + """The unperturbed forward + target unscaling (shared runtime helpers).""" + 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 predict(self, model_input: ModelInput) -> RawOutput: + """One stochastic pass: perturb selected weights, run the forward, restore weights. + + Uses the global torch RNG (seeded per pass by the engine's sampling loop) so passes are + distinct and reproducible. ``perturb_std == 0`` reproduces the deterministic forward. + """ + if self._model is None or self._datapipe is None: + raise RuntimeError("MCPerturbationDrivAerStarWrapper: call load() first") + + saved: dict[str, torch.Tensor] = {} + if self._perturb_std > 0.0 and self._perturb_names: + with torch.no_grad(): + for name, p in self._model.named_parameters(): + if name in self._perturb_names: + saved[name] = p.detach().clone() + eps = torch.randn_like(p) * self._perturb_std + p.mul_(1.0 + eps) + try: + return self._deterministic_predict(model_input) + finally: + if saved: + with torch.no_grad(): + for name, p in self._model.named_parameters(): + if name in saved: + p.copy_(saved[name]) + + 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..ff115a2 --- /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/test/ci_tests/test_mc_perturbation_scope.py b/test/ci_tests/test_mc_perturbation_scope.py new file mode 100644 index 0000000..488af93 --- /dev/null +++ b/test/ci_tests/test_mc_perturbation_scope.py @@ -0,0 +1,74 @@ +# 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 MC weight-perturbation proxy's ``perturb_scope`` parser. + +Requires torch (imported by the wrapper module); skipped in the CPU-only lint env, runs on the +GPU node. This only exercises the pure parameter-selection logic — no model / forward pass. +""" + +from __future__ import annotations + +import pytest + +pytest.importorskip("torch") + +from physicsnemo.cfd.evaluation.models.wrappers.mc_perturbation.wrapper import ( # noqa: E402 + _parse_perturb_scope, +) + +# A representative GeoTransolver-like parameter-name layout: 20 transformer blocks + heads. +_NAMES = ( + ["preprocess.weight", "preprocess.bias"] + + [f"blocks.{i}.attn.qkv.weight" for i in range(20)] + + [f"blocks.{i}.mlp.fc1.weight" for i in range(20)] + + ["head.weight", "head.bias"] +) + + +def test_scope_all_selects_everything() -> None: + assert _parse_perturb_scope("all", _NAMES) == set(_NAMES) + assert _parse_perturb_scope("", _NAMES) == set(_NAMES) + + +def test_scope_last_n_layers_selects_top_block_indices() -> None: + sel = _parse_perturb_scope("last_n_layers=2", _NAMES) + # Blocks 18 and 19 (the two highest indices) across both attn and mlp param groups. + assert sel == { + "blocks.18.attn.qkv.weight", + "blocks.19.attn.qkv.weight", + "blocks.18.mlp.fc1.weight", + "blocks.19.mlp.fc1.weight", + } + + +def test_scope_fraction_selects_last_tensors() -> None: + sel = _parse_perturb_scope("fraction=0.5", _NAMES) + assert len(sel) == round(0.5 * len(_NAMES)) + # Fraction takes the *tail* of the tensor list (includes the head params). + assert "head.weight" in sel and "preprocess.weight" not in sel + + +def test_scope_last_n_layers_falls_back_when_no_indices() -> None: + names = ["a.weight", "b.weight", "c.weight", "d.weight"] + sel = _parse_perturb_scope("last_n_layers=2", names) + # No numeric block indices -> fall back to last 25% of tensors (>=1). + assert sel == {"d.weight"} + + +def test_scope_bad_value_defaults_gracefully() -> None: + sel = _parse_perturb_scope("fraction=notanumber", _NAMES) + assert len(sel) == round(0.25 * len(_NAMES)) diff --git a/test/ci_tests/test_uq_inference.py b/test/ci_tests/test_uq_inference.py new file mode 100644 index 0000000..6f7e6fa --- /dev/null +++ b/test/ci_tests/test_uq_inference.py @@ -0,0 +1,142 @@ +# 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 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, + is_reducer_partial_key, + make_reducer_partial_key, + run_sampling_inference, + 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: + 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: + 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: + 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: + 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: + 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: + 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 diff --git a/test/ci_tests/test_uq_metrics.py b/test/ci_tests/test_uq_metrics.py new file mode 100644 index 0000000..56055be --- /dev/null +++ b/test/ci_tests/test_uq_metrics.py @@ -0,0 +1,349 @@ +# 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 ( + _curve_payload, + _drag_mean_and_std, + _make_pointwise_uq_metrics, + _make_sample_uq_metrics, + _make_uq_metrics, + _SampleDragUQ, +) + + +def test_build_predictive_distribution_coerces_numpy_to_float32() -> None: + 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: + 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: + 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: + 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: + metrics = _make_uq_metrics() + cases = [ + ( + {"pressure": np.zeros(16, np.float32)}, + {"pressure": FieldDistribution(mean=np.zeros(16, np.float32))}, + ) + ] + assert math.isnan(_run_reducer(metrics["nlpd"], cases)) + + +def test_epistemic_metric_requires_epistemic_std() -> None: + rng = np.random.default_rng(2) + metrics = _make_uq_metrics() + # total std present but epistemic_std absent -> epistemic metric is NaN + cases_no_epi = [_gaussian_case(rng, 1000, 1.0)] + assert math.isnan(_run_reducer(metrics["nlpd_epistemic"], cases_no_epi)) + # 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: + 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_deterministic_and_epistemic_variants() -> None: + 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: + 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 + + +def test_sample_ause_needs_two_geometries() -> None: + 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"} + assert abs(fin["epistemic"]) < 1e-9 + assert fin["total"] > fin["epistemic"] + 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: + assert _curve_payload(np.array([1.0]), np.array([1.0])) is None diff --git a/workflows/benchmarking/conf/config_uq_surface.yaml b/workflows/benchmarking/conf/config_uq_surface.yaml new file mode 100644 index 0000000..bbb603f --- /dev/null +++ b/workflows/benchmarking/conf/config_uq_surface.yaml @@ -0,0 +1,188 @@ +# 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. + +# Probabilistic (UQ) surface benchmark — compares deterministic, analytic-GP, and a +# sampling (weight-perturbation) proxy model on the *same* cases, scoring the pooled UQ +# metrics (NLPD, calibration zRMS, 95% coverage, sharpness) alongside the usual L2 / drag / lift. +# See docs/design/uq_metrics_integration.md. +# +# Run (on a GPU node): +# python main.py --config-name=config_uq_surface +# +# The three model rows exercise every path in the UQ engine: +# * geotransolver_drivaerstar_surface -> deterministic baseline. UQ metrics finalize to NaN +# (no std); included so the report shows the det-vs-UQ contrast. +# * geotransolver_gp_surface -> UQ_METHOD="analytic": one forward pass, GP posterior std. +# * mc_perturbation_surface -> UQ_METHOD="sampling": run.uq.num_samples stochastic passes, +# aggregated by the engine (Welford). THROWAWAY proxy — its +# uncertainties are NOT calibrated; swap for real MC-Dropout later. + +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). + num_samples: 32 + # Keep the raw per-pass samples on the FieldDistribution (memory-heavy; only for CRPS / debugging). + retain_samples: false + device_metrics: false + +# Shared anchors for the surface normalization (both checkpoints were trained with the same npz) +# and the deterministic baseline checkpoint (used by both the baseline and the perturbation proxy). +# Paths are the in-container mounts from interactive_job.sh: +# /workspace -> /lustre/fsw/portfolios/coreai/users/ktangsali +# /data -> /lustre/fsw/portfolios/coreai/projects/coreai_modulus_cae +_paths: + transformer_src: &transformer_src /workspace/gp_fields/physicsnemo/examples/cfd/external_aerodynamics/transformer_models + surface_stats: &surface_stats /workspace/gp_fields/physicsnemo/examples/cfd/external_aerodynamics/transformer_models/src/surface_fields_normalization.npz + baseline_ckpt: &baseline_ckpt /workspace/gp_fields/physicsnemo/examples/cfd/external_aerodynamics/transformer_models/runs/geotransolver/surface/baseline_fastback_bf/checkpoints/checkpoint.0.100.pt + +benchmark: + mode: matrix + # All three rows are transformer_models checkpoints (physical, standardize-only targets: + # physical = norm*std + mean). They use the DrivAerStar GeoTransolver wrappers, which do NOT + # re-dimensionalize by ρu² (that is only for the non-dimensional DrivAerML/HF + # `geotransolver_surface` wrapper). The deterministic, GP, and MC-perturbation wrappers are three + # completely independent CFDModel subclasses (no shared base), each with REDIMENSIONALIZE_OUTPUTS=False. + models: + # Deterministic baseline (same checkpoint the perturbation proxy decorates). Epoch 100 = final. + - name: geotransolver_drivaerstar_surface + inference_domain: surface + checkpoint: *baseline_ckpt + stats_path: *surface_stats + kwargs: { batch_resolution: 60000, geometry_sampling: 300000 } + + # Analytic GP field head. Directory holds GeoTransolver.0..mdlus + FieldGPHead.0..pt. + # checkpoint_epoch: 30 = the best-calibrated checkpoint chosen in the prior eval. The GP + # hyperparameters below MATCH the field_gp_v4_long_fastback_bf training run (feature_norm=l2_radial, + # lengthscale_range=[0.01,0.5], lengthscale_prior=[2.0,6.0]; the rest are the conf defaults). + - name: geotransolver_gp_surface + inference_domain: surface + checkpoint: /workspace/gp_fields/physicsnemo/examples/cfd/external_aerodynamics/transformer_models/runs/geotransolver/surface/field_gp_v4_long_fastback_bf/checkpoints_field_gp + stats_path: *surface_stats + kwargs: + checkpoint_epoch: 30 + gp_feature_norm: "l2_radial" + gp_n_inducing: 256 + gp_mlp_hidden: [128, 16] + gp_lengthscale_range: [0.01, 0.5] + gp_lengthscale_prior: [2.0, 6.0] + gp_outputscale_prior: [2.0, 0.5] + num_tasks: 4 + gp_inference_chunk_size: 51200 + geometry_sampling: 300000 + + # Sampling proxy (weight perturbation) over the deterministic baseline — no new training. + - name: mc_perturbation_surface + inference_domain: surface + checkpoint: *baseline_ckpt + stats_path: *surface_stats + kwargs: + perturb_std: 0.02 + perturb_scope: "last_n_layers=4" + batch_resolution: 60000 + geometry_sampling: 300000 + + datasets: + - name: drivaerstar + root: /data/datasets/drivaerstar/surface_files/class_F/val + # These checkpoints were TRAINED on DrivAerStar (class_F = Fastback) zarr, whose WSS keeps the + # raw .vtk sign — so do NOT flip to the DrivAerML convention. The adapter default + # flip_wss_sign=True targets DrivAerML-trained checkpoints; leaving it on here negates the GT + # WSS relative to the model and inflates WSS relative-L2 to ~2.0 (pred ≈ -true) while pressure + # (never flipped) stays correct. Pressure/drag match at flip_wss_sign=false. + 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; shown explicitly so the + # drag_uq panel below uses the SAME direction as the deterministic drag it is diagnosing. + - { 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 GP uncertainty propagated into the + # Cd integral (drag std) vs |Cd_pred - Cd_true| (AUSE_epistemic / AUSE_total, lower=better). + # Takes the same coeff / drag_direction kwargs as the deterministic `drag` metric above. + - { 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 + # Drag Cd panel from drag_uq), with the GP and MC-perturbation models overlaid vs the oracle. + visuals: + - sparsification_plot From 31b8711833ea3d00fa9c3b0a733b7535130a3dc2 Mon Sep 17 00:00:00 2001 From: Kaustubh Tangsali Date: Thu, 9 Jul 2026 23:36:50 +0000 Subject: [PATCH 02/10] update checkpoint name and avoid disk writes of transforms --- .../datasets/adapters/drivaerstar.py | 193 +++++++++++++----- physicsnemo/cfd/evaluation/datasets/schema.py | 7 + .../geotransolver_runtime.py | 50 ++++- .../common_wrapper_utils/vtk_datapipe_io.py | 53 ++++- .../wrappers/geotransolver_gp/wrapper.py | 20 +- .../benchmarking/conf/config_uq_surface.yaml | 9 +- 6 files changed, 251 insertions(+), 81 deletions(-) diff --git a/physicsnemo/cfd/evaluation/datasets/adapters/drivaerstar.py b/physicsnemo/cfd/evaluation/datasets/adapters/drivaerstar.py index 833be55..0c7ce65 100644 --- a/physicsnemo/cfd/evaluation/datasets/adapters/drivaerstar.py +++ b/physicsnemo/cfd/evaluation/datasets/adapters/drivaerstar.py @@ -23,25 +23,35 @@ / Transolver checkpoints were trained on). This adapter bridges those differences so the same model wrappers and metrics work unchanged: -1. Convert legacy ``.vtk`` → ``.vtp`` (XML PolyData) for the datapipe wrappers. -2. Rename ``Pressure`` → ``pMeanTrim`` (DrivAerML convention). -3. Combine the three WSS scalar components into one ``(N, 3)`` vector and flip its sign +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). -4. Drop explicit ``Normals`` / ``Area`` arrays so downstream force integration and +3. Drop explicit ``Normals`` / ``Area`` arrays so downstream force integration and rendering recompute them from mesh topology (DrivAerML convention). -5. Generate a triangulated STL geometry (DrivAerStar ships none — the surface *is* the - geometry) named so the wrappers' STL resolver finds it. -Prepared VTP + STL are cached under ``///`` (one STL per -case dir, so the wrappers' STL lookup is unambiguous) and reused on subsequent runs. +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`` to name the STL, 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. + 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 @@ -102,13 +112,16 @@ class DrivAerStarAdapter(DatasetAdapter): ``("WallShearStressi", "WallShearStressj", "WallShearStressk")``). - ``flip_wss_sign``: negate combined WSS to match DrivAerML (default ``True``). - ``remove_normals_area``: drop explicit ``Normals`` / ``Area`` arrays (default ``True``). - - ``make_stl``: generate a triangulated STL per case (default ``True``). - ``gt_data_type``: ``auto`` / ``cell`` / ``point`` passed to GT extraction (default ``"cell"``; DrivAerStar fields are cell-centered). - - ``prepared_subdir``: cache directory name under ``root`` (default ``"_prepared"``). - - ``pressure_out_name`` / ``shear_out_name``: VTK array names written into the prepared - VTP (defaults ``"pMeanTrim"`` / ``"wallShearStressMeanTrim"``). - - ``force_reprepare``: rebuild cached VTP/STL even if present (default ``False``). + - ``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 @@ -155,8 +168,8 @@ def __init__(self, root: str, **kwargs: Any) -> None: 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._make_stl: bool = bool(kwargs.get("make_stl", 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 @@ -201,49 +214,110 @@ def list_cases(self) -> list[str]: case_ids.append(p.stem) return natural_sorted(case_ids) - def _prepare_case(self, case_id: str) -> tuple[str, pv.PolyData]: - """Convert source ``.vtk`` to a canonical VTP (+ STL), caching under the prepared dir. + def _transform_source_mesh(self, case_id: str) -> tuple[pv.PolyData, Path]: + """Read the raw ``.vtk`` and apply the in-memory canonicalization transforms. - Returns the prepared VTP path and the in-memory prepared mesh (for ``reference_geometry``). + 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" - - if self._force_reprepare or not vtp_path.exists(): - log_dataset("drivaerstar", f"Preparing VTP for {case_id!r} from {src_path}") - mesh = pv.read(str(src_path)) - if not isinstance(mesh, pv.PolyData): - mesh = mesh.extract_surface() - - self._rename_pressure(mesh) - self._combine_wss(mesh) - if self._remove_normals_area: - self._drop_normals_area(mesh) - - vtp_path.parent.mkdir(parents=True, exist_ok=True) + 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)) - else: - mesh = pv.read(str(vtp_path)) - - if self._make_stl and (self._force_reprepare or not stl_path.exists()): - log_dataset("drivaerstar", f"Writing STL geometry for {case_id!r}") - geom = pv.read(str(src_path)) - if not isinstance(geom, pv.PolyData): - geom = geom.extract_surface() - geom.triangulate().save(str(stl_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) - return str(vtp_path), mesh + @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) -> None: + 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) -> None: + 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( @@ -254,6 +328,12 @@ def _combine_wss(self, mesh: pv.PolyData) -> None: 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: @@ -264,12 +344,22 @@ def _drop_normals_area(mesh: pv.PolyData) -> None: del mesh.point_data[key] def load_case(self, case_id: str) -> CanonicalCase: - """Prepare (cache) the case and load it into the canonical schema.""" + """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}", ) - vtp_path, mesh = self._prepare_case(case_id) + 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, @@ -284,8 +374,8 @@ def load_case(self, case_id: str) -> CanonicalCase: "dataset": "drivaerstar", "case": case_id, "branch": "surface", - "source_vtk": self._source_path(case_id).name, - "prepared_vtp": Path(vtp_path).name, + "source_vtk": src_path.name, + "prepared_cached": self._cache_prepared, } if ground_truth: meta["ground_truth_location"] = gt_loc @@ -293,10 +383,13 @@ def load_case(self, case_id: str) -> CanonicalCase: return CanonicalCase( case_id=case_id, - mesh_path=vtp_path, + 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 5803be6..1f7c14a 100644 --- a/physicsnemo/cfd/evaluation/datasets/schema.py +++ b/physicsnemo/cfd/evaluation/datasets/schema.py @@ -81,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"): diff --git a/physicsnemo/cfd/evaluation/models/common_wrapper_utils/geotransolver_runtime.py b/physicsnemo/cfd/evaluation/models/common_wrapper_utils/geotransolver_runtime.py index 3b185f3..8f7f521 100644 --- a/physicsnemo/cfd/evaluation/models/common_wrapper_utils/geotransolver_runtime.py +++ b/physicsnemo/cfd/evaluation/models/common_wrapper_utils/geotransolver_runtime.py @@ -214,28 +214,55 @@ def ensure_distributed_initialized() -> None: DistributedManager.initialize() +def resolve_checkpoint_file(checkpoint_path: str) -> tuple[Path, int]: + """Resolve a checkpoint knob to ``(directory, epoch)`` from a **specific file name**. + + 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). ``physicsnemo.load_checkpoint`` + itself takes the parent directory plus the epoch, which are both derived here. + """ + 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``." + ) + return ckpt.parent, epoch + + def build_geotransolver_backbone( *, checkpoint_path: str, device: str, inference_mode: str ) -> "GeoTransolver": - """Construct a ``GeoTransolver`` for ``surface``/``volume`` and load its checkpoint onto ``device``.""" + """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. + """ model_kw = dict( DEFAULT_GEOTRANSOLVER_VOLUME_KW if inference_mode == "volume" else DEFAULT_GEOTRANSOLVER_KW ) - checkpoint_dir = Path(checkpoint_path) - if checkpoint_dir.is_file(): - checkpoint_dir = checkpoint_dir.parent + checkpoint_dir, epoch = resolve_checkpoint_file(checkpoint_path) ensure_distributed_initialized() dev = torch.device(device) model = GeoTransolver(**model_kw) - ckpt_args: dict[str, Any] = {"path": str(checkpoint_dir), "models": model} - epoch = parse_checkpoint_epoch(checkpoint_path) - if epoch is not None: - ckpt_args["epoch"] = epoch with trusted_torch_load_context(): - _ = load_checkpoint(device=dev, **ckpt_args) + _ = load_checkpoint( + path=str(checkpoint_dir), models=model, epoch=epoch, device=dev + ) model = model.to(dev) model.eval() return model @@ -325,6 +352,8 @@ def build_transolver_batch( 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( @@ -335,6 +364,7 @@ def build_transolver_batch( 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]) @@ -356,6 +386,7 @@ def build_transolver_batch( 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]) @@ -512,6 +543,7 @@ def decode_volume_predictions( "split_datapipe_kwargs", "parse_runtime_kwargs", "ensure_distributed_initialized", + "resolve_checkpoint_file", "build_geotransolver_backbone", "build_surface_datapipe", "build_volume_datapipe", 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/wrappers/geotransolver_gp/wrapper.py b/physicsnemo/cfd/evaluation/models/wrappers/geotransolver_gp/wrapper.py index 0ea34fd..b59c89e 100644 --- a/physicsnemo/cfd/evaluation/models/wrappers/geotransolver_gp/wrapper.py +++ b/physicsnemo/cfd/evaluation/models/wrappers/geotransolver_gp/wrapper.py @@ -37,10 +37,14 @@ Surface only (the GP head predicts 4 surface tasks: pressure + 3 wall-shear components). +The ``checkpoint`` knob must point at a **specific checkpoint file** whose name encodes the epoch +(``GeoTransolver.0..mdlus``); the epoch is parsed from that name and the sibling +``FieldGPHead.0..pt`` is loaded from the same directory. Pointing at a directory (or an +epoch-less name) is rejected — this removes the old ``checkpoint_epoch`` kwarg and the risk of the +backbone and head drifting to different (or "latest") epochs. + Model kwargs (``model.kwargs`` in config), matching the trained checkpoint's GP settings: -- ``checkpoint_epoch`` (int, required for a checkpoint *directory*): epoch tag of the - ``GeoTransolver.0..mdlus`` + ``FieldGPHead.0..pt`` pair to load. - ``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). @@ -50,7 +54,6 @@ """ from contextlib import nullcontext -from pathlib import Path from typing import Any, ClassVar, Optional import numpy as np @@ -76,6 +79,7 @@ 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 ( @@ -174,8 +178,6 @@ def load( # 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)) - ce = kw.pop("checkpoint_epoch", None) - self._checkpoint_epoch = int(ce) if ce is not None else None self._gp_kw = { "num_tasks": int(kw.pop("num_tasks", NUM_SURFACE_TASKS)), "n_inducing": int(kw.pop("gp_n_inducing", 256)), @@ -221,10 +223,12 @@ def load( inference_mode="surface", ) - ckpt_dir = Path(checkpoint_path) - if ckpt_dir.is_file(): - ckpt_dir = ckpt_dir.parent + # Single source of truth for the epoch: parsed from the checkpoint file name (same file + # the backbone just loaded). The sibling ``FieldGPHead.0..pt`` is loaded from this + # directory at this epoch in ``_build_and_load_head`` — no drift, no "latest" fallback. + ckpt_dir, ckpt_epoch = resolve_checkpoint_file(checkpoint_path) self._checkpoint_dir = str(ckpt_dir) + self._checkpoint_epoch = ckpt_epoch return self def prepare_inputs(self, case: CanonicalCase) -> ModelInput: diff --git a/workflows/benchmarking/conf/config_uq_surface.yaml b/workflows/benchmarking/conf/config_uq_surface.yaml index bbb603f..4ca3065 100644 --- a/workflows/benchmarking/conf/config_uq_surface.yaml +++ b/workflows/benchmarking/conf/config_uq_surface.yaml @@ -86,16 +86,17 @@ benchmark: stats_path: *surface_stats kwargs: { batch_resolution: 60000, geometry_sampling: 300000 } - # Analytic GP field head. Directory holds GeoTransolver.0..mdlus + FieldGPHead.0..pt. - # checkpoint_epoch: 30 = the best-calibrated checkpoint chosen in the prior eval. The GP + # Analytic GP field head. Point `checkpoint` at the specific GeoTransolver backbone file for the + # epoch you want; the epoch (30 here = the best-calibrated checkpoint from the prior eval) is + # parsed from that name and the sibling FieldGPHead.0.30.pt is loaded from the same directory. + # (No `checkpoint_epoch` knob / directory scanning -> the backbone and head can't drift.) The GP # hyperparameters below MATCH the field_gp_v4_long_fastback_bf training run (feature_norm=l2_radial, # lengthscale_range=[0.01,0.5], lengthscale_prior=[2.0,6.0]; the rest are the conf defaults). - name: geotransolver_gp_surface inference_domain: surface - checkpoint: /workspace/gp_fields/physicsnemo/examples/cfd/external_aerodynamics/transformer_models/runs/geotransolver/surface/field_gp_v4_long_fastback_bf/checkpoints_field_gp + checkpoint: /workspace/gp_fields/physicsnemo/examples/cfd/external_aerodynamics/transformer_models/runs/geotransolver/surface/field_gp_v4_long_fastback_bf/checkpoints_field_gp/GeoTransolver.0.30.mdlus stats_path: *surface_stats kwargs: - checkpoint_epoch: 30 gp_feature_norm: "l2_radial" gp_n_inducing: 256 gp_mlp_hidden: [128, 16] From 9871f3a247d5e0b7d9456d12929894560ee37593 Mon Sep 17 00:00:00 2001 From: Kaustubh Tangsali Date: Fri, 10 Jul 2026 03:59:21 +0000 Subject: [PATCH 03/10] add docs --- .../wrappers/mc_perturbation/wrapper.py | 23 ++++++++++++++++++- 1 file changed, 22 insertions(+), 1 deletion(-) diff --git a/physicsnemo/cfd/evaluation/models/wrappers/mc_perturbation/wrapper.py b/physicsnemo/cfd/evaluation/models/wrappers/mc_perturbation/wrapper.py index 7f35321..5078d8f 100644 --- a/physicsnemo/cfd/evaluation/models/wrappers/mc_perturbation/wrapper.py +++ b/physicsnemo/cfd/evaluation/models/wrappers/mc_perturbation/wrapper.py @@ -139,6 +139,9 @@ class MCPerturbationDrivAerStarWrapper(CFDModel): INFERENCE_DOMAIN: ClassVar[InferenceDomain | None] = None OUTPUT_LOCATION: ClassVar[OutputLocation] = "cell" REDIMENSIONALIZE_OUTPUTS: ClassVar[bool] = False + # These two flags are the contract with the engine: "I produce uncertainty, and the way to + # get it is to call predict() many times". The engine then loops predict() run.uq.num_samples + # times per case and turns the spread of the passes into a mean + std (see predict()). SUPPORTS_UQ: ClassVar[bool] = True UQ_METHOD: ClassVar[str] = "sampling" @@ -214,6 +217,9 @@ def load( self._datapipe = None self._datapipe_geometry_effective = None + # Load the ordinary trained baseline once. We never retrain or hold an ensemble of + # checkpoints in memory -- the "ensemble" is faked at inference time by jittering this one + # model's weights (that is what makes this a zero-cost proxy). self._model = build_geotransolver_backbone( checkpoint_path=checkpoint_path, device=device, @@ -227,6 +233,9 @@ def load( self._perturb_std, self._perturb_scope, ) + # Decide *which* weight tensors get jittered (e.g. only the last few layers). Perturbing + # every weight usually destroys the prediction; perturbing near the output gives spread + # without wrecking accuracy. We resolve the name set once here and reuse it every pass. names = [n for n, _ in self._model.named_parameters()] self._perturb_names = _parse_perturb_scope(self._perturb_scope, names) _LOG.info( @@ -280,17 +289,29 @@ def predict(self, model_input: ModelInput) -> RawOutput: if self._model is None or self._datapipe is None: raise RuntimeError("MCPerturbationDrivAerStarWrapper: call load() first") + # This method IS one sample of the pseudo-ensemble. The engine calls it N times per case; + # each call does three steps: (1) nudge the weights, (2) run a normal forward, (3) put the + # weights back exactly as they were. Because the noise differs each call, the N predictions + # differ slightly, and their spread is what we report as (fake) uncertainty. + + # Step 1: back up each weight tensor we are about to touch, then multiply it by (1 + eps), + # eps ~ N(0, perturb_std^2). Multiplicative noise scales the jitter to each weight's own + # magnitude. randn_like draws from the global torch RNG, which the engine re-seeds per pass, + # so passes are different from each other but reproducible run-to-run. saved: dict[str, torch.Tensor] = {} if self._perturb_std > 0.0 and self._perturb_names: with torch.no_grad(): for name, p in self._model.named_parameters(): if name in self._perturb_names: - saved[name] = p.detach().clone() + saved[name] = p.detach().clone() # keep the original to restore later eps = torch.randn_like(p) * self._perturb_std p.mul_(1.0 + eps) try: + # Step 2: a plain deterministic forward -- but on the now-jittered weights. return self._deterministic_predict(model_input) finally: + # Step 3: ALWAYS restore the original weights (even if the forward raised), so the next + # pass starts from the clean baseline and perturbations never accumulate. if saved: with torch.no_grad(): for name, p in self._model.named_parameters(): From 241bfcb1fa4af28b2236d7857eb89af838ee6984 Mon Sep 17 00:00:00 2001 From: Kaustubh Tangsali Date: Sun, 12 Jul 2026 22:44:04 +0000 Subject: [PATCH 04/10] add a sample ensemble wrapper, add multi-class evaluation --- .../cfd/evaluation/benchmarks/engine.py | 10 +- physicsnemo/cfd/evaluation/config.py | 10 + .../evaluation/models/wrappers/__init__.py | 5 + .../wrappers/ensemble_drivaerstar/__init__.py | 21 ++ .../wrappers/ensemble_drivaerstar/wrapper.py | 275 ++++++++++++++++++ .../conf/config_uq_multiclass_surface.yaml | 200 +++++++++++++ 6 files changed, 518 insertions(+), 3 deletions(-) create mode 100644 physicsnemo/cfd/evaluation/models/wrappers/ensemble_drivaerstar/__init__.py create mode 100644 physicsnemo/cfd/evaluation/models/wrappers/ensemble_drivaerstar/wrapper.py create mode 100644 workflows/benchmarking/conf/config_uq_multiclass_surface.yaml diff --git a/physicsnemo/cfd/evaluation/benchmarks/engine.py b/physicsnemo/cfd/evaluation/benchmarks/engine.py index e51f6b2..1b6e477 100644 --- a/physicsnemo/cfd/evaluation/benchmarks/engine.py +++ b/physicsnemo/cfd/evaluation/benchmarks/engine.py @@ -559,6 +559,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: @@ -574,7 +578,7 @@ def _run_single( return ( { "model": model_config.name, - "dataset": dataset_config.name, + "dataset": ds_label, "skipped": True, "skip_reason": reason, "cases": [], @@ -612,7 +616,7 @@ def _run_single( return ( { "model": model_config.name, - "dataset": dataset_config.name, + "dataset": ds_label, "cases": [], "metrics": {}, "per_case": [], @@ -912,7 +916,7 @@ def _run_single( return ( { "model": model_config.name, - "dataset": dataset_config.name, + "dataset": ds_label, "cases": cases, "metrics": metrics_summary, "per_case": per_case, diff --git a/physicsnemo/cfd/evaluation/config.py b/physicsnemo/cfd/evaluation/config.py index bdee31a..e137ad7 100644 --- a/physicsnemo/cfd/evaluation/config.py +++ b/physicsnemo/cfd/evaluation/config.py @@ -201,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 diff --git a/physicsnemo/cfd/evaluation/models/wrappers/__init__.py b/physicsnemo/cfd/evaluation/models/wrappers/__init__.py index f9c22b6..8011b92 100644 --- a/physicsnemo/cfd/evaluation/models/wrappers/__init__.py +++ b/physicsnemo/cfd/evaluation/models/wrappers/__init__.py @@ -28,6 +28,9 @@ from physicsnemo.cfd.evaluation.models.wrappers.geotransolver_gp import ( GeoTransolverGPDrivAerStarWrapper, ) +from physicsnemo.cfd.evaluation.models.wrappers.ensemble_drivaerstar import ( + EnsembleDrivAerStarWrapper, +) from physicsnemo.cfd.evaluation.models.wrappers.mc_perturbation import ( MCPerturbationDrivAerStarWrapper, ) @@ -47,6 +50,7 @@ register_model("geotransolver_drivaerstar_surface", GeoTransolverDrivAerStarWrapper) register_model("geotransolver_gp_surface", GeoTransolverGPDrivAerStarWrapper) register_model("mc_perturbation_surface", MCPerturbationDrivAerStarWrapper) +register_model("ensemble_surface", EnsembleDrivAerStarWrapper) register_model("transolver_surface", TransolverWrapper) register_model("transolver_volume", TransolverWrapper) register_model("domino_surface", DominoWrapper) @@ -60,6 +64,7 @@ "GeoTransolverWrapper", "GeoTransolverDrivAerStarWrapper", "GeoTransolverGPDrivAerStarWrapper", + "EnsembleDrivAerStarWrapper", "MCPerturbationDrivAerStarWrapper", "TransolverWrapper", "DominoWrapper", 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..c21e05e --- /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 ( + EnsembleDrivAerStarWrapper, +) + +__all__ = ["EnsembleDrivAerStarWrapper"] 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..8db4f8c --- /dev/null +++ b/physicsnemo/cfd/evaluation/models/wrappers/ensemble_drivaerstar/wrapper.py @@ -0,0 +1,275 @@ +# 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 is the canonical "real" counterpart to the throwaway MC weight-perturbation proxy +(:mod:`..mc_perturbation`): instead of jittering one model's weights, it aggregates the predictions +of several independently-saved model checkpoints. 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 every member on the shared per-case batch and +returns their raw outputs; the engine's +:func:`~physicsnemo.cfd.evaluation.benchmarks.uq_inference.run_sampling_inference` aggregates them +(Welford) into mean + across-member (epistemic) std. Because ``predict_ensemble`` returns exactly +the member count, ``run.uq.num_samples`` is ignored for this row. + +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 checkpoint file whose name encodes an + epoch (``checkpoint.0..pt`` / ``.0..mdlus``). 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, 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, + parse_runtime_kwargs, + unscale_targets, +) +from physicsnemo.cfd.evaluation.models.model_registry import ( + CFDModel, + ModelInput, + OutputLocation, + Predictions, + RawOutput, +) + +_LOG = logging.getLogger(__name__) + + +class EnsembleDrivAerStarWrapper(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, + ) -> "EnsembleDrivAerStarWrapper": + """Build one GeoTransolver backbone per explicitly-listed member checkpoint.""" + if not geotransolver_available(): + raise RuntimeError( + "EnsembleDrivAerStarWrapper requires physicsnemo (GeoTransolver, " + "TransolverDataPipe, load_checkpoint)." + ) + 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( + "EnsembleDrivAerStarWrapper 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( + "EnsembleDrivAerStarWrapper: 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("EnsembleDrivAerStarWrapper: 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) -> RawOutput: + """Run one member's forward + target unscaling on the shared per-case batch.""" + 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, + ) + 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[list[RawOutput]]: + """Run every member on the shared batch; return one raw output per member. + + ``n`` (``run.uq.num_samples``) is intentionally ignored: the ensemble has a fixed number of + members, and the engine iterates over exactly the returned list. Only one member's forward is + resident on the device at a time; the returned raw outputs are the (host-bound) unscaled + predictions the engine immediately folds into its running mean/variance. + """ + if not self._models or self._datapipe is None: + raise RuntimeError("EnsembleDrivAerStarWrapper: call load() first") + return [self._forward_member(model, model_input) for model in self._models] + + 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("EnsembleDrivAerStarWrapper: 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/workflows/benchmarking/conf/config_uq_multiclass_surface.yaml b/workflows/benchmarking/conf/config_uq_multiclass_surface.yaml new file mode 100644 index 0000000..6b2baf3 --- /dev/null +++ b/workflows/benchmarking/conf/config_uq_multiclass_surface.yaml @@ -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. + +# Probabilistic (UQ) surface benchmark ACROSS DrivAerStar body styles — four models +# (deterministic / analytic-GP / MC weight-perturbation proxy / 5-member snapshot ensemble) scored on +# a 100-geometry validation sample from EACH class: estateback (class_E), fastback (class_F), +# notchback (class_N). See docs/design/uq_metrics_integration.md and conf/config_uq_surface.yaml. +# +# Rerun-in-place: this config writes to the SAME run.output_dir as before. With +# run.metrics_cache.enabled: true, the already-scored rows (deterministic / GP / MC across the 3 +# classes) are served from the per-case cache, so only the NEW ensemble_surface x 3 classes actually +# runs inference. Re-run tools/combine_class_results.py afterward to refresh the all_classes rows. +# +# Per-class vs. combined: +# * Per-class is native: each dataset entry below is a separate matrix row, so the report shows +# one row per (model x class) with its own pooled UQ metrics. +# * "All classes combined" is produced post-hoc (no extra inference) by pooling the per-case rows: +# python tools/combine_class_results.py uq_multiclass_surface/benchmark_results.json +# which appends an ``all_classes`` dataset row per model (correct pooled NLPD / coverage / zRMS / +# AUSE, and case-mean L2 / drag / lift). +# +# Prereq — materialize the sampled per-class symlink sets (100 val geometries each, seed 42): +# python tools/make_drivaerstar_class_sample.py \ +# --zarr-root /data/datasets/drivaerstar/surface_files_zarr \ +# --vtk-root /data/datasets/drivaerstar/surface_files --n 100 --seed 42 +# +# Run (on a GPU node): +# python main.py --config-name=config_uq_multiclass_surface +# torchrun --standalone --nnodes=1 --nproc_per_node=8 main.py --config-name=config_uq_multiclass_surface + +defaults: + - _self_ + +hydra: + run: + dir: ${run.output_dir} + output_subdir: hydra + job: + chdir: false + +case_id: null + +run: + device: "cuda:0" + output_dir: "uq_multiclass_surface" + seed: 42 + batch_size: 1 + save_inference_mesh: true + metrics_cache: + enabled: true + path: ${run.output_dir}/metrics_cache + uq: + enabled: true + num_samples: 32 + retain_samples: false + device_metrics: false + +# In-container mounts from interactive_job.sh: +# /workspace -> /lustre/fsw/portfolios/coreai/users/ktangsali +# /data -> /lustre/fsw/portfolios/coreai/projects/coreai_modulus_cae +_paths: + transformer_src: &transformer_src /workspace/gp_fields/physicsnemo/examples/cfd/external_aerodynamics/transformer_models + surface_stats: &surface_stats /workspace/gp_fields/physicsnemo/examples/cfd/external_aerodynamics/transformer_models/src/surface_fields_normalization.npz + baseline_ckpt: &baseline_ckpt /workspace/gp_fields/physicsnemo/examples/cfd/external_aerodynamics/transformer_models/runs/geotransolver/surface/baseline_fastback_bf/checkpoints/checkpoint.0.100.pt + +benchmark: + mode: matrix + # All three rows are transformer_models checkpoints (physical, standardize-only targets). They use + # the DrivAerStar GeoTransolver wrappers (REDIMENSIONALIZE_OUTPUTS=False). NOTE: these checkpoints + # were trained on FASTBACK (class_F) only, so estateback / notchback are out-of-distribution — the + # per-class rows quantify exactly that generalization gap. + models: + - name: geotransolver_drivaerstar_surface + inference_domain: surface + checkpoint: *baseline_ckpt + stats_path: *surface_stats + kwargs: { batch_resolution: 60000, geometry_sampling: 300000 } + + - name: geotransolver_gp_surface + inference_domain: surface + checkpoint: /workspace/gp_fields/physicsnemo/examples/cfd/external_aerodynamics/transformer_models/runs/geotransolver/surface/field_gp_v4_long_fastback_bf/checkpoints_field_gp/GeoTransolver.0.30.mdlus + stats_path: *surface_stats + kwargs: + gp_feature_norm: "l2_radial" + gp_n_inducing: 256 + gp_mlp_hidden: [128, 16] + gp_lengthscale_range: [0.01, 0.5] + gp_lengthscale_prior: [2.0, 6.0] + gp_outputscale_prior: [2.0, 0.5] + num_tasks: 4 + gp_inference_chunk_size: 51200 + geometry_sampling: 300000 + + - name: mc_perturbation_surface + inference_domain: surface + checkpoint: *baseline_ckpt + stats_path: *surface_stats + kwargs: + perturb_std: 0.02 + perturb_scope: "last_n_layers=4" + batch_resolution: 60000 + geometry_sampling: 300000 + + # 5-member SNAPSHOT ensemble: the 5 checkpoints are listed EXPLICITLY (epochs 96-100 of the + # deterministic baseline run) — one GeoTransolver backbone per member, each run as one pass. + # `checkpoint` just anchors the asset identity / cache fingerprint (point it at any one member). + # NOTE: these are consecutive epochs of ONE run (correlated) -> expect mild under-dispersion vs. + # a true deep ensemble; list members from SEPARATE runs for a genuine deep ensemble (same field). + - name: ensemble_surface + inference_domain: surface + checkpoint: *baseline_ckpt + stats_path: *surface_stats + kwargs: + member_checkpoints: + - /workspace/gp_fields/physicsnemo/examples/cfd/external_aerodynamics/transformer_models/runs/geotransolver/surface/baseline_fastback_bf/checkpoints/checkpoint.0.96.pt + - /workspace/gp_fields/physicsnemo/examples/cfd/external_aerodynamics/transformer_models/runs/geotransolver/surface/baseline_fastback_bf/checkpoints/checkpoint.0.97.pt + - /workspace/gp_fields/physicsnemo/examples/cfd/external_aerodynamics/transformer_models/runs/geotransolver/surface/baseline_fastback_bf/checkpoints/checkpoint.0.98.pt + - /workspace/gp_fields/physicsnemo/examples/cfd/external_aerodynamics/transformer_models/runs/geotransolver/surface/baseline_fastback_bf/checkpoints/checkpoint.0.99.pt + - /workspace/gp_fields/physicsnemo/examples/cfd/external_aerodynamics/transformer_models/runs/geotransolver/surface/baseline_fastback_bf/checkpoints/checkpoint.0.100.pt + batch_resolution: 60000 + geometry_sampling: 300000 + + # Three entries share the SAME adapter (`name: drivaerstar`) but different roots; `label` gives + # each a distinct row in the report (and a distinct group for the post-hoc combine). Point each + # root at the per-class symlink set from make_drivaerstar_class_sample.py. + datasets: + - name: drivaerstar + label: estateback + root: /data/datasets/drivaerstar/surface_files/eval_mix_100/class_E + kwargs: { inference_domain: surface, flip_wss_sign: false } + - name: drivaerstar + label: fastback + root: /data/datasets/drivaerstar/surface_files/eval_mix_100/class_F + kwargs: { inference_domain: surface, flip_wss_sign: false } + - name: drivaerstar + label: notchback + root: /data/datasets/drivaerstar/surface_files/eval_mix_100/class_N + 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" + 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: + - l2_pressure + - l2_shear_stress + - l2_pressure_area_weighted + - { name: drag, drag_direction: [1, 0, 0] } + - lift + - nlpd + - nlpd_epistemic + - calibration_zrms + - coverage_95 + - sharpness_std + - sharpness_epistemic_std + - uncertainty_error_spearman + - uncertainty_error_spearman_epistemic + - sparsification_ause + - sparsification_ause_epistemic + - { name: drag_uq, drag_direction: [1, 0, 0] } + +reports: + enabled: true + save_comparison_meshes: false + comparison_mesh_subdir: comparison_meshes + # Case ids are per-class (they repeat across classes) and the sets are randomly sampled, so a given + # id may not exist in every class — present ids get a VTP, others are simply skipped. Leave empty + # to skip per-case mesh export entirely on this multi-class sweep. + visual_case_ids: [] + visuals: + - sparsification_plot From a9418318ae1f622b30bf0a04574d3dec0cf39291 Mon Sep 17 00:00:00 2001 From: Kaustubh Tangsali Date: Wed, 22 Jul 2026 20:37:18 +0000 Subject: [PATCH 05/10] update after testing MC dropout style inference --- .../cfd/evaluation/benchmarks/engine.py | 40 +- .../cfd/evaluation/metrics/builtin/uq.py | 22 +- .../geotransolver_runtime.py | 14 +- .../evaluation/models/wrappers/__init__.py | 11 +- .../wrappers/ensemble_drivaerstar/wrapper.py | 8 +- .../wrappers/geotransolver_gp/wrapper.py | 15 + .../__init__.py | 6 +- .../models/wrappers/mc_dropout/wrapper.py | 285 +++++++++++++++ .../wrappers/mc_perturbation/wrapper.py | 341 ------------------ test/ci_tests/test_mc_perturbation_scope.py | 74 ---- test/ci_tests/test_uq_metrics.py | 13 +- .../conf/config_uq_multiclass_surface.yaml | 102 ++++-- .../benchmarking/conf/config_uq_surface.yaml | 44 ++- 13 files changed, 489 insertions(+), 486 deletions(-) rename physicsnemo/cfd/evaluation/models/wrappers/{mc_perturbation => mc_dropout}/__init__.py (81%) create mode 100644 physicsnemo/cfd/evaluation/models/wrappers/mc_dropout/wrapper.py delete mode 100644 physicsnemo/cfd/evaluation/models/wrappers/mc_perturbation/wrapper.py delete mode 100644 test/ci_tests/test_mc_perturbation_scope.py diff --git a/physicsnemo/cfd/evaluation/benchmarks/engine.py b/physicsnemo/cfd/evaluation/benchmarks/engine.py index 1b6e477..2f4a2d0 100644 --- a/physicsnemo/cfd/evaluation/benchmarks/engine.py +++ b/physicsnemo/cfd/evaluation/benchmarks/engine.py @@ -33,10 +33,14 @@ 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. +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`` @@ -231,6 +235,11 @@ 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: @@ -310,6 +319,7 @@ def _save_inference_mesh_if_requested( predictions: dict[str, Any], output_dir: str, dataset_name: str, + dataset_label: str | None = None, ) -> str | None: """ Write ``inference__.vtp`` or ``.vtu`` when requested. @@ -351,10 +361,16 @@ def _save_inference_mesh_if_requested( 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}…", @@ -765,6 +781,7 @@ def _run_single( predictions=predictions, output_dir=output_dir, dataset_name=dataset_config.name, + dataset_label=dataset_config.display_name, ) comparison_mesh = None @@ -862,11 +879,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()) diff --git a/physicsnemo/cfd/evaluation/metrics/builtin/uq.py b/physicsnemo/cfd/evaluation/metrics/builtin/uq.py index b5677bb..6b5c573 100644 --- a/physicsnemo/cfd/evaluation/metrics/builtin/uq.py +++ b/physicsnemo/cfd/evaluation/metrics/builtin/uq.py @@ -373,7 +373,10 @@ class _SampleAUSE: 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``. + 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: @@ -411,17 +414,27 @@ def finalize_samples( return 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]]: @@ -497,7 +510,9 @@ class _SampleDragUQ: ``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|``. + 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 them when ``output.std_mesh_field_names`` / ``epistemic_std_mesh_field_names`` are set), so it is a no-op for deterministic wrappers or when std fields are unavailable. @@ -623,8 +638,11 @@ def finalize_samples( 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]]: diff --git a/physicsnemo/cfd/evaluation/models/common_wrapper_utils/geotransolver_runtime.py b/physicsnemo/cfd/evaluation/models/common_wrapper_utils/geotransolver_runtime.py index 8f7f521..043dab8 100644 --- a/physicsnemo/cfd/evaluation/models/common_wrapper_utils/geotransolver_runtime.py +++ b/physicsnemo/cfd/evaluation/models/common_wrapper_utils/geotransolver_runtime.py @@ -242,18 +242,30 @@ def resolve_checkpoint_file(checkpoint_path: str) -> tuple[Path, int]: def build_geotransolver_backbone( - *, checkpoint_path: str, device: str, inference_mode: str + *, + 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 + ``load_checkpoint`` 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. """ model_kw = dict( DEFAULT_GEOTRANSOLVER_VOLUME_KW if inference_mode == "volume" else DEFAULT_GEOTRANSOLVER_KW ) + if concrete_dropout: + model_kw["concrete_dropout"] = True checkpoint_dir, epoch = resolve_checkpoint_file(checkpoint_path) ensure_distributed_initialized() diff --git a/physicsnemo/cfd/evaluation/models/wrappers/__init__.py b/physicsnemo/cfd/evaluation/models/wrappers/__init__.py index 8011b92..42b74f4 100644 --- a/physicsnemo/cfd/evaluation/models/wrappers/__init__.py +++ b/physicsnemo/cfd/evaluation/models/wrappers/__init__.py @@ -31,8 +31,8 @@ from physicsnemo.cfd.evaluation.models.wrappers.ensemble_drivaerstar import ( EnsembleDrivAerStarWrapper, ) -from physicsnemo.cfd.evaluation.models.wrappers.mc_perturbation import ( - MCPerturbationDrivAerStarWrapper, +from physicsnemo.cfd.evaluation.models.wrappers.mc_dropout import ( + MCDropoutDrivAerStarWrapper, ) from physicsnemo.cfd.evaluation.models.wrappers.surface_baseline import ( SurfaceBaselineWrapper, @@ -49,7 +49,10 @@ register_model("geotransolver_volume", GeoTransolverWrapper) register_model("geotransolver_drivaerstar_surface", GeoTransolverDrivAerStarWrapper) register_model("geotransolver_gp_surface", GeoTransolverGPDrivAerStarWrapper) -register_model("mc_perturbation_surface", MCPerturbationDrivAerStarWrapper) +# DUE-style bi-Lipschitz field GP: same wrapper class, distinct model name so it scores as its own +# matrix row (config supplies gp_spectral_norm_coeff / gp_dkl_residual + the DUE checkpoint). +register_model("geotransolver_gp_due_surface", GeoTransolverGPDrivAerStarWrapper) +register_model("mc_dropout_surface", MCDropoutDrivAerStarWrapper) register_model("ensemble_surface", EnsembleDrivAerStarWrapper) register_model("transolver_surface", TransolverWrapper) register_model("transolver_volume", TransolverWrapper) @@ -65,7 +68,7 @@ "GeoTransolverDrivAerStarWrapper", "GeoTransolverGPDrivAerStarWrapper", "EnsembleDrivAerStarWrapper", - "MCPerturbationDrivAerStarWrapper", + "MCDropoutDrivAerStarWrapper", "TransolverWrapper", "DominoWrapper", "SurfaceBaselineWrapper", diff --git a/physicsnemo/cfd/evaluation/models/wrappers/ensemble_drivaerstar/wrapper.py b/physicsnemo/cfd/evaluation/models/wrappers/ensemble_drivaerstar/wrapper.py index 8db4f8c..9f88875 100644 --- a/physicsnemo/cfd/evaluation/models/wrappers/ensemble_drivaerstar/wrapper.py +++ b/physicsnemo/cfd/evaluation/models/wrappers/ensemble_drivaerstar/wrapper.py @@ -16,10 +16,10 @@ """Deep-/snapshot-ensemble sampling UQ wrapper for GeoTransolver (surface). -An ensemble is the canonical "real" counterpart to the throwaway MC weight-perturbation proxy -(:mod:`..mc_perturbation`): instead of jittering one model's weights, it aggregates the predictions -of several independently-saved model checkpoints. The engine drives it through the SAME -``UQ_METHOD="sampling"`` path — but here each "pass" is a genuine model, so the across-member spread +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:: diff --git a/physicsnemo/cfd/evaluation/models/wrappers/geotransolver_gp/wrapper.py b/physicsnemo/cfd/evaluation/models/wrappers/geotransolver_gp/wrapper.py index b59c89e..1ac742c 100644 --- a/physicsnemo/cfd/evaluation/models/wrappers/geotransolver_gp/wrapper.py +++ b/physicsnemo/cfd/evaluation/models/wrappers/geotransolver_gp/wrapper.py @@ -48,9 +48,17 @@ - ``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. + +The same wrapper class is registered under two keys: ``geotransolver_gp_surface`` (plain DKL) and +``geotransolver_gp_due_surface`` (DUE variant). They differ only by ``model.kwargs`` + checkpoint, +so DUE scores as its own matrix row without a bespoke subclass. """ from contextlib import nullcontext @@ -198,6 +206,13 @@ def load( 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) diff --git a/physicsnemo/cfd/evaluation/models/wrappers/mc_perturbation/__init__.py b/physicsnemo/cfd/evaluation/models/wrappers/mc_dropout/__init__.py similarity index 81% rename from physicsnemo/cfd/evaluation/models/wrappers/mc_perturbation/__init__.py rename to physicsnemo/cfd/evaluation/models/wrappers/mc_dropout/__init__.py index 6f4d436..26668b7 100644 --- a/physicsnemo/cfd/evaluation/models/wrappers/mc_perturbation/__init__.py +++ b/physicsnemo/cfd/evaluation/models/wrappers/mc_dropout/__init__.py @@ -14,8 +14,8 @@ # See the License for the specific language governing permissions and # limitations under the License. -from physicsnemo.cfd.evaluation.models.wrappers.mc_perturbation.wrapper import ( - MCPerturbationDrivAerStarWrapper, +from physicsnemo.cfd.evaluation.models.wrappers.mc_dropout.wrapper import ( + MCDropoutDrivAerStarWrapper, ) -__all__ = ["MCPerturbationDrivAerStarWrapper"] +__all__ = ["MCDropoutDrivAerStarWrapper"] 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..b388908 --- /dev/null +++ b/physicsnemo/cfd/evaluation/models/wrappers/mc_dropout/wrapper.py @@ -0,0 +1,285 @@ +# 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. + +"""Real **MC-Dropout** sampling UQ over a Concrete-Dropout GeoTransolver (surface/volume). + +This is the MC-Dropout sampling baseline (``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, + 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 MCDropoutDrivAerStarWrapper(CFDModel): + """Real 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, + ) -> "MCDropoutDrivAerStarWrapper": + """Load the Concrete-Dropout backbone and keep its dropout layers stochastic for MC passes.""" + if not geotransolver_available(): + raise RuntimeError( + "MCDropoutDrivAerStarWrapper requires physicsnemo (GeoTransolver, " + "TransolverDataPipe, load_checkpoint)." + ) + if not _CONCRETE_DROPOUT_AVAILABLE: + raise RuntimeError( + "MCDropoutDrivAerStarWrapper 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("MCDropoutDrivAerStarWrapper: 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: + """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). + """ + if self._model is None or self._datapipe is None: + raise RuntimeError("MCDropoutDrivAerStarWrapper: 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, + ) + 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/mc_perturbation/wrapper.py b/physicsnemo/cfd/evaluation/models/wrappers/mc_perturbation/wrapper.py deleted file mode 100644 index 5078d8f..0000000 --- a/physicsnemo/cfd/evaluation/models/wrappers/mc_perturbation/wrapper.py +++ /dev/null @@ -1,341 +0,0 @@ -# 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. - -"""Monte-Carlo **weight-perturbation** proxy for a sampling UQ method (surface). - -.. warning:: - This is a **throwaway plumbing proxy**, *not* a valid UQ method. It exists only to exercise - the ``UQ_METHOD="sampling"`` engine path (N ``predict`` calls, per-pass seeding, streaming - Welford aggregation, pooled metrics, overlaid reports) *before* a real MC-Dropout wrapper is - available — **no training and no new checkpoint required** (UQ design doc §8.2). The - uncertainties it produces are **not calibrated or physically meaningful**. Swap it for the - real ``mc_dropout_surface`` wrapper (same config) once that lands. - -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` for its -deterministic forward, and makes it stochastic in :meth:`predict`: before each pass it multiplies a -subset of weights by ``(1 + eps)`` with ``eps ~ N(0, perturb_std^2)`` (a zero-cost SWAG-flavored -pseudo-ensemble), runs the forward, then restores the weights. The engine drives the N passes and -aggregates them into a -:class:`~physicsnemo.cfd.evaluation.datasets.schema.FieldDistribution` (sample mean + -across-pass std = epistemic). - -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``): - -- ``perturb_std`` (float, default ``0.02``): stddev of the multiplicative Gaussian weight noise. - ``perturb_std → 0`` collapses to the deterministic model (epistemic std → 0), a useful sanity test. -- ``perturb_scope`` (str, default ``"last_n_layers=4"``): which weights to perturb — - ``"all"``, ``"last_n_layers=K"``, or ``"fraction=f"`` (last fraction ``f`` of parameter tensors). -- plus the shared GeoTransolver runtime kwargs (``cuda_bf16_autocast``; datapipe overrides; etc.); - ``checkpoint`` / ``stats_path`` are the trained baseline. -""" - -import logging -import re -from typing import Any, ClassVar, Literal, Optional - -import torch - -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, - parse_runtime_kwargs, - unscale_targets, -) -from physicsnemo.cfd.evaluation.models.model_registry import ( - CFDModel, - ModelInput, - OutputLocation, - Predictions, - RawOutput, -) - -_LOG = logging.getLogger(__name__) - - -def _parse_perturb_scope(scope: str, named_params: list[str]) -> set[str]: - """Resolve which parameter names to perturb from a ``perturb_scope`` string. - - Supports ``"all"``, ``"last_n_layers=K"`` (params whose max embedded layer index is within - the top ``K`` distinct indices), and ``"fraction=f"`` (last fraction of parameter tensors). - Falls back to the last 25% of tensors when a layer-based scope finds no numeric indices. - """ - scope = (scope or "").strip() - if scope in ("", "all"): - return set(named_params) - - if scope.startswith("fraction="): - try: - frac = float(scope.split("=", 1)[1]) - except ValueError: - frac = 0.25 - frac = min(max(frac, 0.0), 1.0) - k = max(1, int(round(frac * len(named_params)))) - return set(named_params[-k:]) - - if scope.startswith("last_n_layers="): - try: - k = int(scope.split("=", 1)[1]) - except ValueError: - k = 4 - - def _layer_index(name: str) -> int | None: - # Largest integer among dotted segments that are pure digits (the block index). - ids = [int(t) for t in re.split(r"[.\[\]]", name) if t.isdigit()] - return max(ids) if ids else None - - indexed = {n: _layer_index(n) for n in named_params} - distinct = sorted({i for i in indexed.values() if i is not None}) - if distinct: - keep = set(distinct[-k:]) - return {n for n, i in indexed.items() if i in keep} - # No numeric layer ids -> fall back to last 25% of tensors. - - k = max(1, int(round(0.25 * len(named_params)))) - return set(named_params[-k:]) - - -class MCPerturbationDrivAerStarWrapper(CFDModel): - """Weight-perturbation sampling proxy over the baseline GeoTransolver (surface). - - 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 - # These two flags are the contract with the engine: "I produce uncertainty, and the way to - # get it is to call predict() many times". The engine then loops predict() run.uq.num_samples - # times per case and turns the spread of the passes into a mean + std (see predict()). - 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() - self._perturb_std: float = 0.02 - self._perturb_scope: str = "last_n_layers=4" - self._perturb_names: set[str] = set() - - def load( - self, - checkpoint_path: str, - stats_path: str, - device: str, - **kwargs: Any, - ) -> "MCPerturbationDrivAerStarWrapper": - """Load the baseline GeoTransolver, then pick the weights to perturb per ``perturb_scope``.""" - if not geotransolver_available(): - raise RuntimeError( - "MCPerturbationDrivAerStarWrapper requires physicsnemo (GeoTransolver, " - "TransolverDataPipe, load_checkpoint)." - ) - kw = dict(kwargs) - self._perturb_std = float(kw.pop("perturb_std", 0.02)) - self._perturb_scope = str(kw.pop("perturb_scope", "last_n_layers=4")) - 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_perturbation", 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_perturbation", 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 - # Load the ordinary trained baseline once. We never retrain or hold an ensemble of - # checkpoints in memory -- the "ensemble" is faked at inference time by jittering this one - # model's weights (that is what makes this a zero-cost proxy). - self._model = build_geotransolver_backbone( - checkpoint_path=checkpoint_path, - device=device, - inference_mode=self._inference_mode, - ) - - _LOG.warning( - "MCPerturbationDrivAerStarWrapper is a THROWAWAY sampling PROXY (weight perturbation, " - "perturb_std=%s, perturb_scope=%r); its uncertainties are NOT calibrated or " - "physically meaningful. Replace with the real MC-Dropout wrapper for valid UQ.", - self._perturb_std, - self._perturb_scope, - ) - # Decide *which* weight tensors get jittered (e.g. only the last few layers). Perturbing - # every weight usually destroys the prediction; perturbing near the output gives spread - # without wrecking accuracy. We resolve the name set once here and reuse it every pass. - names = [n for n, _ in self._model.named_parameters()] - self._perturb_names = _parse_perturb_scope(self._perturb_scope, names) - _LOG.info( - "Perturbing %d / %d parameter tensors per pass.", - len(self._perturb_names), - len(names), - ) - 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("MCPerturbationDrivAerStarWrapper: 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 _deterministic_predict(self, model_input: ModelInput) -> RawOutput: - """The unperturbed forward + target unscaling (shared runtime helpers).""" - 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 predict(self, model_input: ModelInput) -> RawOutput: - """One stochastic pass: perturb selected weights, run the forward, restore weights. - - Uses the global torch RNG (seeded per pass by the engine's sampling loop) so passes are - distinct and reproducible. ``perturb_std == 0`` reproduces the deterministic forward. - """ - if self._model is None or self._datapipe is None: - raise RuntimeError("MCPerturbationDrivAerStarWrapper: call load() first") - - # This method IS one sample of the pseudo-ensemble. The engine calls it N times per case; - # each call does three steps: (1) nudge the weights, (2) run a normal forward, (3) put the - # weights back exactly as they were. Because the noise differs each call, the N predictions - # differ slightly, and their spread is what we report as (fake) uncertainty. - - # Step 1: back up each weight tensor we are about to touch, then multiply it by (1 + eps), - # eps ~ N(0, perturb_std^2). Multiplicative noise scales the jitter to each weight's own - # magnitude. randn_like draws from the global torch RNG, which the engine re-seeds per pass, - # so passes are different from each other but reproducible run-to-run. - saved: dict[str, torch.Tensor] = {} - if self._perturb_std > 0.0 and self._perturb_names: - with torch.no_grad(): - for name, p in self._model.named_parameters(): - if name in self._perturb_names: - saved[name] = p.detach().clone() # keep the original to restore later - eps = torch.randn_like(p) * self._perturb_std - p.mul_(1.0 + eps) - try: - # Step 2: a plain deterministic forward -- but on the now-jittered weights. - return self._deterministic_predict(model_input) - finally: - # Step 3: ALWAYS restore the original weights (even if the forward raised), so the next - # pass starts from the clean baseline and perturbations never accumulate. - if saved: - with torch.no_grad(): - for name, p in self._model.named_parameters(): - if name in saved: - p.copy_(saved[name]) - - 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/test/ci_tests/test_mc_perturbation_scope.py b/test/ci_tests/test_mc_perturbation_scope.py deleted file mode 100644 index 488af93..0000000 --- a/test/ci_tests/test_mc_perturbation_scope.py +++ /dev/null @@ -1,74 +0,0 @@ -# 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 MC weight-perturbation proxy's ``perturb_scope`` parser. - -Requires torch (imported by the wrapper module); skipped in the CPU-only lint env, runs on the -GPU node. This only exercises the pure parameter-selection logic — no model / forward pass. -""" - -from __future__ import annotations - -import pytest - -pytest.importorskip("torch") - -from physicsnemo.cfd.evaluation.models.wrappers.mc_perturbation.wrapper import ( # noqa: E402 - _parse_perturb_scope, -) - -# A representative GeoTransolver-like parameter-name layout: 20 transformer blocks + heads. -_NAMES = ( - ["preprocess.weight", "preprocess.bias"] - + [f"blocks.{i}.attn.qkv.weight" for i in range(20)] - + [f"blocks.{i}.mlp.fc1.weight" for i in range(20)] - + ["head.weight", "head.bias"] -) - - -def test_scope_all_selects_everything() -> None: - assert _parse_perturb_scope("all", _NAMES) == set(_NAMES) - assert _parse_perturb_scope("", _NAMES) == set(_NAMES) - - -def test_scope_last_n_layers_selects_top_block_indices() -> None: - sel = _parse_perturb_scope("last_n_layers=2", _NAMES) - # Blocks 18 and 19 (the two highest indices) across both attn and mlp param groups. - assert sel == { - "blocks.18.attn.qkv.weight", - "blocks.19.attn.qkv.weight", - "blocks.18.mlp.fc1.weight", - "blocks.19.mlp.fc1.weight", - } - - -def test_scope_fraction_selects_last_tensors() -> None: - sel = _parse_perturb_scope("fraction=0.5", _NAMES) - assert len(sel) == round(0.5 * len(_NAMES)) - # Fraction takes the *tail* of the tensor list (includes the head params). - assert "head.weight" in sel and "preprocess.weight" not in sel - - -def test_scope_last_n_layers_falls_back_when_no_indices() -> None: - names = ["a.weight", "b.weight", "c.weight", "d.weight"] - sel = _parse_perturb_scope("last_n_layers=2", names) - # No numeric block indices -> fall back to last 25% of tensors (>=1). - assert sel == {"d.weight"} - - -def test_scope_bad_value_defaults_gracefully() -> None: - sel = _parse_perturb_scope("fraction=notanumber", _NAMES) - assert len(sel) == round(0.25 * len(_NAMES)) diff --git a/test/ci_tests/test_uq_metrics.py b/test/ci_tests/test_uq_metrics.py index 56055be..34cf23b 100644 --- a/test/ci_tests/test_uq_metrics.py +++ b/test/ci_tests/test_uq_metrics.py @@ -252,6 +252,9 @@ def _geom(scale, rank_sig): 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: @@ -321,9 +324,17 @@ def test_drag_uq_finalize_and_curves_rank_geometries() -> None: "total_std": err[::-1], # anti-correlated } fin = drag_uq.finalize_samples(collected) - assert set(fin) == {"epistemic", "total"} + 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 diff --git a/workflows/benchmarking/conf/config_uq_multiclass_surface.yaml b/workflows/benchmarking/conf/config_uq_multiclass_surface.yaml index 6b2baf3..99c97b9 100644 --- a/workflows/benchmarking/conf/config_uq_multiclass_surface.yaml +++ b/workflows/benchmarking/conf/config_uq_multiclass_surface.yaml @@ -14,21 +14,32 @@ # See the License for the specific language governing permissions and # limitations under the License. -# Probabilistic (UQ) surface benchmark ACROSS DrivAerStar body styles — four models -# (deterministic / analytic-GP / MC weight-perturbation proxy / 5-member snapshot ensemble) scored on -# a 100-geometry validation sample from EACH class: estateback (class_E), fastback (class_F), -# notchback (class_N). See docs/design/uq_metrics_integration.md and conf/config_uq_surface.yaml. +# Probabilistic (UQ) surface benchmark ACROSS DrivAerStar body styles — five models +# (deterministic / analytic-GP / DUE bi-Lipschitz GP / real MC-Dropout / 5-member snapshot ensemble) +# scored on a 100-geometry validation sample from EACH class: estateback (class_E), fastback +# (class_F), notchback (class_N). See docs/design/uq_metrics_integration.md and config_uq_surface.yaml. # -# Rerun-in-place: this config writes to the SAME run.output_dir as before. With -# run.metrics_cache.enabled: true, the already-scored rows (deterministic / GP / MC across the 3 -# classes) are served from the per-case cache, so only the NEW ensemble_surface x 3 classes actually -# runs inference. Re-run tools/combine_class_results.py afterward to refresh the all_classes rows. +# NOTE: MC-Dropout (mc_dropout_surface) is the real Concrete-Dropout baseline, and the plain GP row +# uses the from-scratch joint epoch-100 (final) checkpoint (field_gp_jointscratch_fastback): random +# init, 1024 inducing points, relaxed lengthscale cap [0.01, 2.0] for more input-dependent variance. +# The DUE row (geotransolver_gp_due_surface) is the same recipe with a spectral-norm + residual +# (bi-Lipschitz) DKL extractor to fight feature collapse. This reuses run.output_dir +# (latest_gp_mc_jointscratch): only the NEW DUE row is a cache miss and runs (300 cases); the +# deterministic / plain-GP / MC-Dropout / ensemble results are served from the existing metrics_cache +# unchanged. Drag-UQ ranking (drag AUSE / Spearman) is included (see the metrics section). Re-run +# tools/combine_class_results.py afterward to build the all_classes rows. +# +# Per-case VTP visuals: reports.visual_case_ids (below) lists one representative geometry per class, +# so each gets an inference mesh + a full comparison mesh (pred / GT / std). Because this is a fresh +# dir every case runs inference, so the meshes are written with no cache surgery needed. Outputs: +# latest_gp_mc_jointscratch/inference___.vtp +# latest_gp_mc_jointscratch/comparison_meshes/___comparison.vtp # # Per-class vs. combined: # * Per-class is native: each dataset entry below is a separate matrix row, so the report shows # one row per (model x class) with its own pooled UQ metrics. # * "All classes combined" is produced post-hoc (no extra inference) by pooling the per-case rows: -# python tools/combine_class_results.py uq_multiclass_surface/benchmark_results.json +# python tools/combine_class_results.py latest_gp_mc_jointscratch/benchmark_results.json # which appends an ``all_classes`` dataset row per model (correct pooled NLPD / coverage / zRMS / # AUSE, and case-mean L2 / drag / lift). # @@ -55,7 +66,7 @@ case_id: null run: device: "cuda:0" - output_dir: "uq_multiclass_surface" + output_dir: "latest_gp_mc_jointscratch" seed: 42 batch_size: 1 save_inference_mesh: true @@ -64,7 +75,7 @@ run: path: ${run.output_dir}/metrics_cache uq: enabled: true - num_samples: 32 + num_samples: 20 retain_samples: false device_metrics: false @@ -75,6 +86,8 @@ _paths: transformer_src: &transformer_src /workspace/gp_fields/physicsnemo/examples/cfd/external_aerodynamics/transformer_models surface_stats: &surface_stats /workspace/gp_fields/physicsnemo/examples/cfd/external_aerodynamics/transformer_models/src/surface_fields_normalization.npz baseline_ckpt: &baseline_ckpt /workspace/gp_fields/physicsnemo/examples/cfd/external_aerodynamics/transformer_models/runs/geotransolver/surface/baseline_fastback_bf/checkpoints/checkpoint.0.100.pt + # Concrete-Dropout run (model.concrete_dropout=true), same backbone/data/normalization as baseline. + concrete_dropout_ckpt: &concrete_dropout_ckpt /workspace/gp_fields/physicsnemo/examples/cfd/external_aerodynamics/transformer_models/runs/geotransolver/surface/concrete_dropout_fastback_bf/checkpoints/GeoTransolver.0.100.mdlus benchmark: mode: matrix @@ -89,28 +102,61 @@ benchmark: stats_path: *surface_stats kwargs: { batch_resolution: 60000, geometry_sampling: 300000 } + # GP field head, epoch 100 (final) of the from-scratch joint run (field_gp_jointscratch_fastback): + # trained end-to-end from random init (no warm-started backbone) with 1024 inducing points and a + # RELAXED lengthscale cap [0.01, 2.0] so the kernel can learn longer-range, input-dependent + # (heteroscedastic) variance. Hyperparameters MUST match training so the FieldGPHead reconstructs + # before the state dict loads (n_inducing, mlp_hidden, lengthscale range/prior, feature_norm). + # Sibling FieldGPHead.0.100.pt loads from the same dir. - name: geotransolver_gp_surface inference_domain: surface - checkpoint: /workspace/gp_fields/physicsnemo/examples/cfd/external_aerodynamics/transformer_models/runs/geotransolver/surface/field_gp_v4_long_fastback_bf/checkpoints_field_gp/GeoTransolver.0.30.mdlus + checkpoint: /workspace/gp_fields/physicsnemo/examples/cfd/external_aerodynamics/transformer_models/runs/geotransolver/surface/field_gp_jointscratch_fastback/checkpoints_field_gp/GeoTransolver.0.100.mdlus stats_path: *surface_stats kwargs: gp_feature_norm: "l2_radial" - gp_n_inducing: 256 + gp_n_inducing: 1024 gp_mlp_hidden: [128, 16] - gp_lengthscale_range: [0.01, 0.5] - gp_lengthscale_prior: [2.0, 6.0] + 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 - - name: mc_perturbation_surface + # DUE field GP (van Amersfoort et al., 2021): identical to the GP row above EXCEPT the DKL + # feature extractor is bi-Lipschitz (spectral-norm coeff 0.95 + residual skips) to fight DKL + # feature collapse (flat per-geometry epistemic contrast). It re-uses the same wrapper class via + # a distinct model name (geotransolver_gp_due_surface), so it is its own matrix row and cache + # namespace — the four existing rows are served from cache unchanged. The extractor is widened to + # [128, 128, 16] (vs the plain GP's [128, 16]) so the 128->128 layer forms a residual block; all + # of gp_spectral_norm_coeff / gp_dkl_residual / gp_mlp_hidden MUST match the DUE training run + # (field_gp_jointscratch_due_fastback) or the FieldGPHead state dict will not reconstruct. + # Checkpoint epoch 100 (final) — matches the plain GP row's epoch for a like-for-like comparison. + - name: geotransolver_gp_due_surface inference_domain: surface - checkpoint: *baseline_ckpt + checkpoint: /workspace/gp_fields/physicsnemo/examples/cfd/external_aerodynamics/transformer_models/runs/geotransolver/surface/field_gp_jointscratch_due_fastback/checkpoints_field_gp/GeoTransolver.0.100.mdlus + stats_path: *surface_stats + kwargs: + gp_feature_norm: "l2_radial" + gp_n_inducing: 1024 + gp_mlp_hidden: [128, 128, 16] + gp_spectral_norm_coeff: 0.95 + 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 + + # Real MC-Dropout: N stochastic forward passes over the Concrete-Dropout backbone (the dropout + # layers stay stochastic at inference). Pass count is the benchmark-wide run.uq.num_samples (20), + # NOT a model kwarg, so it matches the ensemble's budget. + - name: mc_dropout_surface + inference_domain: surface + checkpoint: *concrete_dropout_ckpt stats_path: *surface_stats kwargs: - perturb_std: 0.02 - perturb_scope: "last_n_layers=4" batch_resolution: 60000 geometry_sampling: 300000 @@ -186,15 +232,23 @@ metrics: - uncertainty_error_spearman_epistemic - sparsification_ause - sparsification_ause_epistemic + # Drag-based UQ ranking (drag AUSE / drag Spearman + the drag sparsification panels), integrated + # along the same axis as the deterministic `drag` metric. CAVEAT: drag is a coherent surface + # integral and the cheap diagonal (per-cell independent) variance propagation discards the + # spatially-correlated epistemic covariance that dominates it, so the drag-EPISTEMIC ranking in + # particular is a lower bound — interpret it alongside the field-level epistemic metrics above. - { name: drag_uq, drag_direction: [1, 0, 0] } reports: enabled: true - save_comparison_meshes: false + # Save the rich comparison meshes (pred + ground truth + std, per model) for the inspected cases. + # Mesh filenames now embed the model + dataset label, so the per-class case-id overlap is safe. + save_comparison_meshes: true comparison_mesh_subdir: comparison_meshes - # Case ids are per-class (they repeat across classes) and the sets are randomly sampled, so a given - # id may not exist in every class — present ids get a VTP, others are simply skipped. Leave empty - # to skip per-case mesh export entirely on this multi-class sweep. - visual_case_ids: [] + # One representative geometry per class (median pressure-L2 for the GP row). Case ids are per-class + # and randomly sampled, so a given id lives in exactly one class — present ids get a VTP/comparison + # mesh, others are skipped. This is a fresh output dir, so every case runs inference and these + # meshes are written directly (no cache surgery needed). + visual_case_ids: ["03275", "07251", "02456"] visuals: - sparsification_plot diff --git a/workflows/benchmarking/conf/config_uq_surface.yaml b/workflows/benchmarking/conf/config_uq_surface.yaml index 4ca3065..84a658b 100644 --- a/workflows/benchmarking/conf/config_uq_surface.yaml +++ b/workflows/benchmarking/conf/config_uq_surface.yaml @@ -14,10 +14,9 @@ # See the License for the specific language governing permissions and # limitations under the License. -# Probabilistic (UQ) surface benchmark — compares deterministic, analytic-GP, and a -# sampling (weight-perturbation) proxy model on the *same* cases, scoring the pooled UQ -# metrics (NLPD, calibration zRMS, 95% coverage, sharpness) alongside the usual L2 / drag / lift. -# See docs/design/uq_metrics_integration.md. +# Probabilistic (UQ) surface benchmark — compares deterministic, analytic-GP, and real MC-Dropout +# on the *same* cases, scoring the pooled UQ metrics (NLPD, calibration zRMS, 95% coverage, +# sharpness) alongside the usual L2 / drag / lift. See docs/design/uq_metrics_integration.md. # # Run (on a GPU node): # python main.py --config-name=config_uq_surface @@ -26,9 +25,8 @@ # * geotransolver_drivaerstar_surface -> deterministic baseline. UQ metrics finalize to NaN # (no std); included so the report shows the det-vs-UQ contrast. # * geotransolver_gp_surface -> UQ_METHOD="analytic": one forward pass, GP posterior std. -# * mc_perturbation_surface -> UQ_METHOD="sampling": run.uq.num_samples stochastic passes, -# aggregated by the engine (Welford). THROWAWAY proxy — its -# uncertainties are NOT calibrated; swap for real MC-Dropout later. +# * mc_dropout_surface -> UQ_METHOD="sampling": run.uq.num_samples stochastic passes over +# the Concrete-Dropout backbone, aggregated by the engine (Welford). defaults: - _self_ @@ -61,25 +59,26 @@ run: retain_samples: false device_metrics: false -# Shared anchors for the surface normalization (both checkpoints were trained with the same npz) -# and the deterministic baseline checkpoint (used by both the baseline and the perturbation proxy). -# Paths are the in-container mounts from interactive_job.sh: +# Shared anchors for the surface normalization (all checkpoints were trained with the same npz) and +# the deterministic baseline checkpoint. Paths are the in-container mounts from interactive_job.sh: # /workspace -> /lustre/fsw/portfolios/coreai/users/ktangsali # /data -> /lustre/fsw/portfolios/coreai/projects/coreai_modulus_cae _paths: transformer_src: &transformer_src /workspace/gp_fields/physicsnemo/examples/cfd/external_aerodynamics/transformer_models surface_stats: &surface_stats /workspace/gp_fields/physicsnemo/examples/cfd/external_aerodynamics/transformer_models/src/surface_fields_normalization.npz baseline_ckpt: &baseline_ckpt /workspace/gp_fields/physicsnemo/examples/cfd/external_aerodynamics/transformer_models/runs/geotransolver/surface/baseline_fastback_bf/checkpoints/checkpoint.0.100.pt + # Concrete-Dropout run (model.concrete_dropout=true), same backbone/data/normalization as baseline. + concrete_dropout_ckpt: &concrete_dropout_ckpt /workspace/gp_fields/physicsnemo/examples/cfd/external_aerodynamics/transformer_models/runs/geotransolver/surface/concrete_dropout_fastback_bf/checkpoints/GeoTransolver.0.100.mdlus benchmark: mode: matrix # All three rows are transformer_models checkpoints (physical, standardize-only targets: # physical = norm*std + mean). They use the DrivAerStar GeoTransolver wrappers, which do NOT # re-dimensionalize by ρu² (that is only for the non-dimensional DrivAerML/HF - # `geotransolver_surface` wrapper). The deterministic, GP, and MC-perturbation wrappers are three + # `geotransolver_surface` wrapper). The deterministic, GP, and MC-Dropout wrappers are three # completely independent CFDModel subclasses (no shared base), each with REDIMENSIONALIZE_OUTPUTS=False. models: - # Deterministic baseline (same checkpoint the perturbation proxy decorates). Epoch 100 = final. + # Deterministic baseline. Epoch 100 = final. - name: geotransolver_drivaerstar_surface inference_domain: surface checkpoint: *baseline_ckpt @@ -87,14 +86,14 @@ benchmark: kwargs: { batch_resolution: 60000, geometry_sampling: 300000 } # Analytic GP field head. Point `checkpoint` at the specific GeoTransolver backbone file for the - # epoch you want; the epoch (30 here = the best-calibrated checkpoint from the prior eval) is - # parsed from that name and the sibling FieldGPHead.0.30.pt is loaded from the same directory. - # (No `checkpoint_epoch` knob / directory scanning -> the backbone and head can't drift.) The GP - # hyperparameters below MATCH the field_gp_v4_long_fastback_bf training run (feature_norm=l2_radial, - # lengthscale_range=[0.01,0.5], lengthscale_prior=[2.0,6.0]; the rest are the conf defaults). + # epoch you want; the epoch (38 here) is parsed from that name and the sibling FieldGPHead.0.38.pt + # is loaded from the same directory. (No `checkpoint_epoch` knob / directory scanning -> the + # backbone and head can't drift.) e38 of the pinned l2_radial run won the e30/e35/e38 sweep on + # drag AUSE and calibration. The GP hyperparameters below MATCH that training run + # (feature_norm=l2_radial, lengthscale_range=[0.01,0.5], lengthscale_prior=[2.0,6.0]; rest are conf defaults). - name: geotransolver_gp_surface inference_domain: surface - checkpoint: /workspace/gp_fields/physicsnemo/examples/cfd/external_aerodynamics/transformer_models/runs/geotransolver/surface/field_gp_v4_long_fastback_bf/checkpoints_field_gp/GeoTransolver.0.30.mdlus + checkpoint: /workspace/gp_fields/physicsnemo/examples/cfd/external_aerodynamics/transformer_models/runs/geotransolver/surface/field_gp_v4_long_fastback_pin/checkpoints_field_gp/GeoTransolver.0.38.mdlus stats_path: *surface_stats kwargs: gp_feature_norm: "l2_radial" @@ -107,14 +106,13 @@ benchmark: gp_inference_chunk_size: 51200 geometry_sampling: 300000 - # Sampling proxy (weight perturbation) over the deterministic baseline — no new training. - - name: mc_perturbation_surface + # Real MC-Dropout: run.uq.num_samples stochastic passes over the Concrete-Dropout backbone + # (the ConcreteDropout layers stay stochastic at inference). + - name: mc_dropout_surface inference_domain: surface - checkpoint: *baseline_ckpt + checkpoint: *concrete_dropout_ckpt stats_path: *surface_stats kwargs: - perturb_std: 0.02 - perturb_scope: "last_n_layers=4" batch_resolution: 60000 geometry_sampling: 300000 From 6e42bc700b29f6add3e6766342a177052635461a Mon Sep 17 00:00:00 2001 From: Kaustubh Tangsali Date: Wed, 22 Jul 2026 22:12:43 +0000 Subject: [PATCH 06/10] cleanup for PR prep --- .../cfd/evaluation/benchmarks/engine.py | 2 +- physicsnemo/cfd/evaluation/datasets/schema.py | 2 +- .../cfd/evaluation/metrics/builtin/uq.py | 11 +- .../cfd/evaluation/models/model_registry.py | 2 +- .../evaluation/models/wrappers/__init__.py | 15 +- .../wrappers/ensemble_drivaerstar/__init__.py | 4 +- .../wrappers/ensemble_drivaerstar/wrapper.py | 16 +- .../wrappers/geotransolver_gp/wrapper.py | 6 +- .../models/wrappers/mc_dropout/__init__.py | 4 +- .../models/wrappers/mc_dropout/wrapper.py | 18 +- .../SKILL.md | 113 +++++++- .../references/example_wrapper.py | 236 +++++++++++++++- .../conf/config_uq_multiclass_surface.yaml | 254 ------------------ .../benchmarking/conf/config_uq_surface.yaml | 131 +++++---- 14 files changed, 457 insertions(+), 357 deletions(-) delete mode 100644 workflows/benchmarking/conf/config_uq_multiclass_surface.yaml diff --git a/physicsnemo/cfd/evaluation/benchmarks/engine.py b/physicsnemo/cfd/evaluation/benchmarks/engine.py index 2f4a2d0..990952d 100644 --- a/physicsnemo/cfd/evaluation/benchmarks/engine.py +++ b/physicsnemo/cfd/evaluation/benchmarks/engine.py @@ -752,7 +752,7 @@ def _run_single( supports_uq = bool(getattr(wrapper, "SUPPORTS_UQ", False)) uq_method = getattr(wrapper, "UQ_METHOD", "none") if supports_uq and uq_method == "sampling" and run_config.uq.enabled: - # N stochastic passes; prepare_inputs already ran once (see UQ design doc §6.2). + # N stochastic passes; prepare_inputs already ran once (only the forward is repeated). predictions = run_sampling_inference( wrapper, case, diff --git a/physicsnemo/cfd/evaluation/datasets/schema.py b/physicsnemo/cfd/evaluation/datasets/schema.py index 1f7c14a..c8e1737 100644 --- a/physicsnemo/cfd/evaluation/datasets/schema.py +++ b/physicsnemo/cfd/evaluation/datasets/schema.py @@ -24,7 +24,7 @@ import numpy as np #: Array payloads on :class:`FieldDistribution` may be NumPy arrays or (for the on-device -#: metric path, see the UQ design doc §6.8) framework tensors such as ``torch.Tensor``. We +#: 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 diff --git a/physicsnemo/cfd/evaluation/metrics/builtin/uq.py b/physicsnemo/cfd/evaluation/metrics/builtin/uq.py index 6b5c573..471af86 100644 --- a/physicsnemo/cfd/evaluation/metrics/builtin/uq.py +++ b/physicsnemo/cfd/evaluation/metrics/builtin/uq.py @@ -456,10 +456,13 @@ def curves(self, collected: dict[str, list[float]]) -> dict[str, dict[str, Any]] # # 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-point field variances (diagonal-posterior linear error propagation; a lower bound under -# spatial correlation, but exact-enough for *ranking* geometries — which is all sparsification -# needs). 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 +# 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). diff --git a/physicsnemo/cfd/evaluation/models/model_registry.py b/physicsnemo/cfd/evaluation/models/model_registry.py index 5d753a5..c4d1633 100644 --- a/physicsnemo/cfd/evaluation/models/model_registry.py +++ b/physicsnemo/cfd/evaluation/models/model_registry.py @@ -57,7 +57,7 @@ class CFDModel(ABC): #: 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 (see the UQ design doc §4.2): + #: 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 diff --git a/physicsnemo/cfd/evaluation/models/wrappers/__init__.py b/physicsnemo/cfd/evaluation/models/wrappers/__init__.py index 42b74f4..91c8a9d 100644 --- a/physicsnemo/cfd/evaluation/models/wrappers/__init__.py +++ b/physicsnemo/cfd/evaluation/models/wrappers/__init__.py @@ -29,10 +29,10 @@ GeoTransolverGPDrivAerStarWrapper, ) from physicsnemo.cfd.evaluation.models.wrappers.ensemble_drivaerstar import ( - EnsembleDrivAerStarWrapper, + GeoTransolverEnsembleDrivAerStarWrapper, ) from physicsnemo.cfd.evaluation.models.wrappers.mc_dropout import ( - MCDropoutDrivAerStarWrapper, + GeoTransolverMCDropoutDrivAerStarWrapper, ) from physicsnemo.cfd.evaluation.models.wrappers.surface_baseline import ( SurfaceBaselineWrapper, @@ -49,11 +49,8 @@ register_model("geotransolver_volume", GeoTransolverWrapper) register_model("geotransolver_drivaerstar_surface", GeoTransolverDrivAerStarWrapper) register_model("geotransolver_gp_surface", GeoTransolverGPDrivAerStarWrapper) -# DUE-style bi-Lipschitz field GP: same wrapper class, distinct model name so it scores as its own -# matrix row (config supplies gp_spectral_norm_coeff / gp_dkl_residual + the DUE checkpoint). -register_model("geotransolver_gp_due_surface", GeoTransolverGPDrivAerStarWrapper) -register_model("mc_dropout_surface", MCDropoutDrivAerStarWrapper) -register_model("ensemble_surface", EnsembleDrivAerStarWrapper) +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) @@ -67,8 +64,8 @@ "GeoTransolverWrapper", "GeoTransolverDrivAerStarWrapper", "GeoTransolverGPDrivAerStarWrapper", - "EnsembleDrivAerStarWrapper", - "MCDropoutDrivAerStarWrapper", + "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 index c21e05e..fc9dc8a 100644 --- a/physicsnemo/cfd/evaluation/models/wrappers/ensemble_drivaerstar/__init__.py +++ b/physicsnemo/cfd/evaluation/models/wrappers/ensemble_drivaerstar/__init__.py @@ -15,7 +15,7 @@ # limitations under the License. from physicsnemo.cfd.evaluation.models.wrappers.ensemble_drivaerstar.wrapper import ( - EnsembleDrivAerStarWrapper, + GeoTransolverEnsembleDrivAerStarWrapper, ) -__all__ = ["EnsembleDrivAerStarWrapper"] +__all__ = ["GeoTransolverEnsembleDrivAerStarWrapper"] diff --git a/physicsnemo/cfd/evaluation/models/wrappers/ensemble_drivaerstar/wrapper.py b/physicsnemo/cfd/evaluation/models/wrappers/ensemble_drivaerstar/wrapper.py index 9f88875..3184d93 100644 --- a/physicsnemo/cfd/evaluation/models/wrappers/ensemble_drivaerstar/wrapper.py +++ b/physicsnemo/cfd/evaluation/models/wrappers/ensemble_drivaerstar/wrapper.py @@ -92,7 +92,7 @@ _LOG = logging.getLogger(__name__) -class EnsembleDrivAerStarWrapper(CFDModel): +class GeoTransolverEnsembleDrivAerStarWrapper(CFDModel): """K-member (snapshot/deep) ensemble over GeoTransolver checkpoints (surface). Decorates ``transformer_models`` (physical-target) checkpoints, so predictions are @@ -140,11 +140,11 @@ def load( stats_path: str, device: str, **kwargs: Any, - ) -> "EnsembleDrivAerStarWrapper": + ) -> "GeoTransolverEnsembleDrivAerStarWrapper": """Build one GeoTransolver backbone per explicitly-listed member checkpoint.""" if not geotransolver_available(): raise RuntimeError( - "EnsembleDrivAerStarWrapper requires physicsnemo (GeoTransolver, " + "GeoTransolverEnsembleDrivAerStarWrapper requires physicsnemo (GeoTransolver, " "TransolverDataPipe, load_checkpoint)." ) kw = dict(kwargs) @@ -172,7 +172,7 @@ def load( if not member_checkpoints: raise ValueError( - "EnsembleDrivAerStarWrapper requires `member_checkpoints`: an explicit list of " + "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." ) @@ -191,7 +191,7 @@ def load( for member in members ] _LOG.info( - "EnsembleDrivAerStarWrapper: loaded %d members from %s", + "GeoTransolverEnsembleDrivAerStarWrapper: loaded %d members from %s", len(self._models), [Path(m).name for m in members], ) @@ -200,7 +200,7 @@ def load( 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("EnsembleDrivAerStarWrapper: call load() first") + raise RuntimeError("GeoTransolverEnsembleDrivAerStarWrapper: call load() first") result = build_transolver_batch( case=case, inference_mode=self._inference_mode, @@ -243,13 +243,13 @@ def predict_ensemble( predictions the engine immediately folds into its running mean/variance. """ if not self._models or self._datapipe is None: - raise RuntimeError("EnsembleDrivAerStarWrapper: call load() first") + raise RuntimeError("GeoTransolverEnsembleDrivAerStarWrapper: call load() first") return [self._forward_member(model, model_input) for model in self._models] 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("EnsembleDrivAerStarWrapper: call load() first") + raise RuntimeError("GeoTransolverEnsembleDrivAerStarWrapper: call load() first") return self._forward_member(self._models[0], model_input) def decode_outputs( diff --git a/physicsnemo/cfd/evaluation/models/wrappers/geotransolver_gp/wrapper.py b/physicsnemo/cfd/evaluation/models/wrappers/geotransolver_gp/wrapper.py index 1ac742c..ebb7497 100644 --- a/physicsnemo/cfd/evaluation/models/wrappers/geotransolver_gp/wrapper.py +++ b/physicsnemo/cfd/evaluation/models/wrappers/geotransolver_gp/wrapper.py @@ -32,7 +32,7 @@ ``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, §6.3), so the pooled UQ +**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). @@ -55,10 +55,6 @@ - ``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. - -The same wrapper class is registered under two keys: ``geotransolver_gp_surface`` (plain DKL) and -``geotransolver_gp_due_surface`` (DUE variant). They differ only by ``model.kwargs`` + checkpoint, -so DUE scores as its own matrix row without a bespoke subclass. """ from contextlib import nullcontext diff --git a/physicsnemo/cfd/evaluation/models/wrappers/mc_dropout/__init__.py b/physicsnemo/cfd/evaluation/models/wrappers/mc_dropout/__init__.py index 26668b7..4da746c 100644 --- a/physicsnemo/cfd/evaluation/models/wrappers/mc_dropout/__init__.py +++ b/physicsnemo/cfd/evaluation/models/wrappers/mc_dropout/__init__.py @@ -15,7 +15,7 @@ # limitations under the License. from physicsnemo.cfd.evaluation.models.wrappers.mc_dropout.wrapper import ( - MCDropoutDrivAerStarWrapper, + GeoTransolverMCDropoutDrivAerStarWrapper, ) -__all__ = ["MCDropoutDrivAerStarWrapper"] +__all__ = ["GeoTransolverMCDropoutDrivAerStarWrapper"] diff --git a/physicsnemo/cfd/evaluation/models/wrappers/mc_dropout/wrapper.py b/physicsnemo/cfd/evaluation/models/wrappers/mc_dropout/wrapper.py index b388908..847697d 100644 --- a/physicsnemo/cfd/evaluation/models/wrappers/mc_dropout/wrapper.py +++ b/physicsnemo/cfd/evaluation/models/wrappers/mc_dropout/wrapper.py @@ -14,9 +14,9 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Real **MC-Dropout** sampling UQ over a Concrete-Dropout GeoTransolver (surface/volume). +"""**MC-Dropout** sampling UQ over a Concrete-Dropout GeoTransolver (surface/volume). -This is the MC-Dropout sampling baseline (``mc_dropout_surface``). The backbone was trained with +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`` / @@ -92,8 +92,8 @@ _LOG = logging.getLogger(__name__) -class MCDropoutDrivAerStarWrapper(CFDModel): - """Real MC-Dropout sampling over a Concrete-Dropout GeoTransolver (surface/volume). +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): @@ -141,16 +141,16 @@ def load( stats_path: str, device: str, **kwargs: Any, - ) -> "MCDropoutDrivAerStarWrapper": + ) -> "GeoTransolverMCDropoutDrivAerStarWrapper": """Load the Concrete-Dropout backbone and keep its dropout layers stochastic for MC passes.""" if not geotransolver_available(): raise RuntimeError( - "MCDropoutDrivAerStarWrapper requires physicsnemo (GeoTransolver, " + "GeoTransolverMCDropoutDrivAerStarWrapper requires physicsnemo (GeoTransolver, " "TransolverDataPipe, load_checkpoint)." ) if not _CONCRETE_DROPOUT_AVAILABLE: raise RuntimeError( - "MCDropoutDrivAerStarWrapper requires physicsnemo.nn.ConcreteDropout " + "GeoTransolverMCDropoutDrivAerStarWrapper requires physicsnemo.nn.ConcreteDropout " "(train the backbone with model.concrete_dropout=true)." ) kw = dict(kwargs) @@ -223,7 +223,7 @@ def _enable_mc_dropout(self) -> None: 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("MCDropoutDrivAerStarWrapper: call load() first") + raise RuntimeError("GeoTransolverMCDropoutDrivAerStarWrapper: call load() first") result = build_transolver_batch( case=case, inference_mode=self._inference_mode, @@ -247,7 +247,7 @@ def predict(self, model_input: ModelInput) -> RawOutput: passes differ from each other yet are reproducible run-to-run). """ if self._model is None or self._datapipe is None: - raise RuntimeError("MCDropoutDrivAerStarWrapper: call load() first") + raise RuntimeError("GeoTransolverMCDropoutDrivAerStarWrapper: call load() first") raw = geotransolver_forward( model=self._model, batch=model_input["batch"], diff --git a/skills/physicsnemo-cfd-create-model-wrapper/SKILL.md b/skills/physicsnemo-cfd-create-model-wrapper/SKILL.md index 17c7720..65a5774 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 one extra hook — either +`decode_distribution` (analytic) or `predict_ensemble` (sampling, +optional). 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,105 @@ 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). Optionally implement + **`predict_ensemble(model_input, n) -> list[RawOutput] | None`** as a + fast path (ensembles return one output per member and ignore `n`; + 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 +447,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..429b129 100644 --- a/skills/physicsnemo-cfd-create-model-wrapper/references/example_wrapper.py +++ b/skills/physicsnemo-cfd-create-model-wrapper/references/example_wrapper.py @@ -17,17 +17,23 @@ """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 @@ -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,220 @@ 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: + return self.OUTPUT_LOCATION + + def load( + self, checkpoint_path: str, stats_path: str, device: str, **kwargs: Any + ) -> "ExampleAnalyticUQWrapper": + 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: + 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); + * OPTIONAL fast path: implement ``predict_ensemble(model_input, n)`` to return all + passes/members in one call (an ensemble returns one output per member and ignores + ``n``). 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: + return self.OUTPUT_LOCATION + + def load( + self, checkpoint_path: str, stats_path: str, device: str, **kwargs: Any + ) -> "ExampleSamplingUQWrapper": + 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: + 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_ensemble( + self, model_input: torch.Tensor, n: int + ) -> Optional[list[dict[str, torch.Tensor]]]: + """Optional fast path: return every pass/member in one call. + + For a deep ensemble, return one raw output per member (``n`` is ignored). For + multi-pass single-model methods you may batch ``n`` passes here. Return ``None`` to + let the engine fall back to calling ``predict`` ``n`` times. + """ + if not self._models: + return None + return [self.predict(model_input) for _ in self._models] + + 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/workflows/benchmarking/conf/config_uq_multiclass_surface.yaml b/workflows/benchmarking/conf/config_uq_multiclass_surface.yaml deleted file mode 100644 index 99c97b9..0000000 --- a/workflows/benchmarking/conf/config_uq_multiclass_surface.yaml +++ /dev/null @@ -1,254 +0,0 @@ -# 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. - -# Probabilistic (UQ) surface benchmark ACROSS DrivAerStar body styles — five models -# (deterministic / analytic-GP / DUE bi-Lipschitz GP / real MC-Dropout / 5-member snapshot ensemble) -# scored on a 100-geometry validation sample from EACH class: estateback (class_E), fastback -# (class_F), notchback (class_N). See docs/design/uq_metrics_integration.md and config_uq_surface.yaml. -# -# NOTE: MC-Dropout (mc_dropout_surface) is the real Concrete-Dropout baseline, and the plain GP row -# uses the from-scratch joint epoch-100 (final) checkpoint (field_gp_jointscratch_fastback): random -# init, 1024 inducing points, relaxed lengthscale cap [0.01, 2.0] for more input-dependent variance. -# The DUE row (geotransolver_gp_due_surface) is the same recipe with a spectral-norm + residual -# (bi-Lipschitz) DKL extractor to fight feature collapse. This reuses run.output_dir -# (latest_gp_mc_jointscratch): only the NEW DUE row is a cache miss and runs (300 cases); the -# deterministic / plain-GP / MC-Dropout / ensemble results are served from the existing metrics_cache -# unchanged. Drag-UQ ranking (drag AUSE / Spearman) is included (see the metrics section). Re-run -# tools/combine_class_results.py afterward to build the all_classes rows. -# -# Per-case VTP visuals: reports.visual_case_ids (below) lists one representative geometry per class, -# so each gets an inference mesh + a full comparison mesh (pred / GT / std). Because this is a fresh -# dir every case runs inference, so the meshes are written with no cache surgery needed. Outputs: -# latest_gp_mc_jointscratch/inference___.vtp -# latest_gp_mc_jointscratch/comparison_meshes/___comparison.vtp -# -# Per-class vs. combined: -# * Per-class is native: each dataset entry below is a separate matrix row, so the report shows -# one row per (model x class) with its own pooled UQ metrics. -# * "All classes combined" is produced post-hoc (no extra inference) by pooling the per-case rows: -# python tools/combine_class_results.py latest_gp_mc_jointscratch/benchmark_results.json -# which appends an ``all_classes`` dataset row per model (correct pooled NLPD / coverage / zRMS / -# AUSE, and case-mean L2 / drag / lift). -# -# Prereq — materialize the sampled per-class symlink sets (100 val geometries each, seed 42): -# python tools/make_drivaerstar_class_sample.py \ -# --zarr-root /data/datasets/drivaerstar/surface_files_zarr \ -# --vtk-root /data/datasets/drivaerstar/surface_files --n 100 --seed 42 -# -# Run (on a GPU node): -# python main.py --config-name=config_uq_multiclass_surface -# torchrun --standalone --nnodes=1 --nproc_per_node=8 main.py --config-name=config_uq_multiclass_surface - -defaults: - - _self_ - -hydra: - run: - dir: ${run.output_dir} - output_subdir: hydra - job: - chdir: false - -case_id: null - -run: - device: "cuda:0" - output_dir: "latest_gp_mc_jointscratch" - seed: 42 - batch_size: 1 - save_inference_mesh: true - metrics_cache: - enabled: true - path: ${run.output_dir}/metrics_cache - uq: - enabled: true - num_samples: 20 - retain_samples: false - device_metrics: false - -# In-container mounts from interactive_job.sh: -# /workspace -> /lustre/fsw/portfolios/coreai/users/ktangsali -# /data -> /lustre/fsw/portfolios/coreai/projects/coreai_modulus_cae -_paths: - transformer_src: &transformer_src /workspace/gp_fields/physicsnemo/examples/cfd/external_aerodynamics/transformer_models - surface_stats: &surface_stats /workspace/gp_fields/physicsnemo/examples/cfd/external_aerodynamics/transformer_models/src/surface_fields_normalization.npz - baseline_ckpt: &baseline_ckpt /workspace/gp_fields/physicsnemo/examples/cfd/external_aerodynamics/transformer_models/runs/geotransolver/surface/baseline_fastback_bf/checkpoints/checkpoint.0.100.pt - # Concrete-Dropout run (model.concrete_dropout=true), same backbone/data/normalization as baseline. - concrete_dropout_ckpt: &concrete_dropout_ckpt /workspace/gp_fields/physicsnemo/examples/cfd/external_aerodynamics/transformer_models/runs/geotransolver/surface/concrete_dropout_fastback_bf/checkpoints/GeoTransolver.0.100.mdlus - -benchmark: - mode: matrix - # All three rows are transformer_models checkpoints (physical, standardize-only targets). They use - # the DrivAerStar GeoTransolver wrappers (REDIMENSIONALIZE_OUTPUTS=False). NOTE: these checkpoints - # were trained on FASTBACK (class_F) only, so estateback / notchback are out-of-distribution — the - # per-class rows quantify exactly that generalization gap. - models: - - name: geotransolver_drivaerstar_surface - inference_domain: surface - checkpoint: *baseline_ckpt - stats_path: *surface_stats - kwargs: { batch_resolution: 60000, geometry_sampling: 300000 } - - # GP field head, epoch 100 (final) of the from-scratch joint run (field_gp_jointscratch_fastback): - # trained end-to-end from random init (no warm-started backbone) with 1024 inducing points and a - # RELAXED lengthscale cap [0.01, 2.0] so the kernel can learn longer-range, input-dependent - # (heteroscedastic) variance. Hyperparameters MUST match training so the FieldGPHead reconstructs - # before the state dict loads (n_inducing, mlp_hidden, lengthscale range/prior, feature_norm). - # Sibling FieldGPHead.0.100.pt loads from the same dir. - - name: geotransolver_gp_surface - inference_domain: surface - checkpoint: /workspace/gp_fields/physicsnemo/examples/cfd/external_aerodynamics/transformer_models/runs/geotransolver/surface/field_gp_jointscratch_fastback/checkpoints_field_gp/GeoTransolver.0.100.mdlus - stats_path: *surface_stats - kwargs: - gp_feature_norm: "l2_radial" - gp_n_inducing: 1024 - gp_mlp_hidden: [128, 16] - 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 - - # DUE field GP (van Amersfoort et al., 2021): identical to the GP row above EXCEPT the DKL - # feature extractor is bi-Lipschitz (spectral-norm coeff 0.95 + residual skips) to fight DKL - # feature collapse (flat per-geometry epistemic contrast). It re-uses the same wrapper class via - # a distinct model name (geotransolver_gp_due_surface), so it is its own matrix row and cache - # namespace — the four existing rows are served from cache unchanged. The extractor is widened to - # [128, 128, 16] (vs the plain GP's [128, 16]) so the 128->128 layer forms a residual block; all - # of gp_spectral_norm_coeff / gp_dkl_residual / gp_mlp_hidden MUST match the DUE training run - # (field_gp_jointscratch_due_fastback) or the FieldGPHead state dict will not reconstruct. - # Checkpoint epoch 100 (final) — matches the plain GP row's epoch for a like-for-like comparison. - - name: geotransolver_gp_due_surface - inference_domain: surface - checkpoint: /workspace/gp_fields/physicsnemo/examples/cfd/external_aerodynamics/transformer_models/runs/geotransolver/surface/field_gp_jointscratch_due_fastback/checkpoints_field_gp/GeoTransolver.0.100.mdlus - stats_path: *surface_stats - kwargs: - gp_feature_norm: "l2_radial" - gp_n_inducing: 1024 - gp_mlp_hidden: [128, 128, 16] - gp_spectral_norm_coeff: 0.95 - 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 - - # Real MC-Dropout: N stochastic forward passes over the Concrete-Dropout backbone (the dropout - # layers stay stochastic at inference). Pass count is the benchmark-wide run.uq.num_samples (20), - # NOT a model kwarg, so it matches the ensemble's budget. - - name: mc_dropout_surface - inference_domain: surface - checkpoint: *concrete_dropout_ckpt - stats_path: *surface_stats - kwargs: - batch_resolution: 60000 - geometry_sampling: 300000 - - # 5-member SNAPSHOT ensemble: the 5 checkpoints are listed EXPLICITLY (epochs 96-100 of the - # deterministic baseline run) — one GeoTransolver backbone per member, each run as one pass. - # `checkpoint` just anchors the asset identity / cache fingerprint (point it at any one member). - # NOTE: these are consecutive epochs of ONE run (correlated) -> expect mild under-dispersion vs. - # a true deep ensemble; list members from SEPARATE runs for a genuine deep ensemble (same field). - - name: ensemble_surface - inference_domain: surface - checkpoint: *baseline_ckpt - stats_path: *surface_stats - kwargs: - member_checkpoints: - - /workspace/gp_fields/physicsnemo/examples/cfd/external_aerodynamics/transformer_models/runs/geotransolver/surface/baseline_fastback_bf/checkpoints/checkpoint.0.96.pt - - /workspace/gp_fields/physicsnemo/examples/cfd/external_aerodynamics/transformer_models/runs/geotransolver/surface/baseline_fastback_bf/checkpoints/checkpoint.0.97.pt - - /workspace/gp_fields/physicsnemo/examples/cfd/external_aerodynamics/transformer_models/runs/geotransolver/surface/baseline_fastback_bf/checkpoints/checkpoint.0.98.pt - - /workspace/gp_fields/physicsnemo/examples/cfd/external_aerodynamics/transformer_models/runs/geotransolver/surface/baseline_fastback_bf/checkpoints/checkpoint.0.99.pt - - /workspace/gp_fields/physicsnemo/examples/cfd/external_aerodynamics/transformer_models/runs/geotransolver/surface/baseline_fastback_bf/checkpoints/checkpoint.0.100.pt - batch_resolution: 60000 - geometry_sampling: 300000 - - # Three entries share the SAME adapter (`name: drivaerstar`) but different roots; `label` gives - # each a distinct row in the report (and a distinct group for the post-hoc combine). Point each - # root at the per-class symlink set from make_drivaerstar_class_sample.py. - datasets: - - name: drivaerstar - label: estateback - root: /data/datasets/drivaerstar/surface_files/eval_mix_100/class_E - kwargs: { inference_domain: surface, flip_wss_sign: false } - - name: drivaerstar - label: fastback - root: /data/datasets/drivaerstar/surface_files/eval_mix_100/class_F - kwargs: { inference_domain: surface, flip_wss_sign: false } - - name: drivaerstar - label: notchback - root: /data/datasets/drivaerstar/surface_files/eval_mix_100/class_N - 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" - 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: - - l2_pressure - - l2_shear_stress - - l2_pressure_area_weighted - - { name: drag, drag_direction: [1, 0, 0] } - - lift - - nlpd - - nlpd_epistemic - - calibration_zrms - - coverage_95 - - sharpness_std - - sharpness_epistemic_std - - uncertainty_error_spearman - - uncertainty_error_spearman_epistemic - - sparsification_ause - - sparsification_ause_epistemic - # Drag-based UQ ranking (drag AUSE / drag Spearman + the drag sparsification panels), integrated - # along the same axis as the deterministic `drag` metric. CAVEAT: drag is a coherent surface - # integral and the cheap diagonal (per-cell independent) variance propagation discards the - # spatially-correlated epistemic covariance that dominates it, so the drag-EPISTEMIC ranking in - # particular is a lower bound — interpret it alongside the field-level epistemic metrics above. - - { name: drag_uq, drag_direction: [1, 0, 0] } - -reports: - enabled: true - # Save the rich comparison meshes (pred + ground truth + std, per model) for the inspected cases. - # Mesh filenames now embed the model + dataset label, so the per-class case-id overlap is safe. - save_comparison_meshes: true - comparison_mesh_subdir: comparison_meshes - # One representative geometry per class (median pressure-L2 for the GP row). Case ids are per-class - # and randomly sampled, so a given id lives in exactly one class — present ids get a VTP/comparison - # mesh, others are skipped. This is a fresh output dir, so every case runs inference and these - # meshes are written directly (no cache surgery needed). - visual_case_ids: ["03275", "07251", "02456"] - visuals: - - sparsification_plot diff --git a/workflows/benchmarking/conf/config_uq_surface.yaml b/workflows/benchmarking/conf/config_uq_surface.yaml index 84a658b..25af2fa 100644 --- a/workflows/benchmarking/conf/config_uq_surface.yaml +++ b/workflows/benchmarking/conf/config_uq_surface.yaml @@ -14,19 +14,35 @@ # See the License for the specific language governing permissions and # limitations under the License. -# Probabilistic (UQ) surface benchmark — compares deterministic, analytic-GP, and real MC-Dropout -# on the *same* cases, scoring the pooled UQ metrics (NLPD, calibration zRMS, 95% coverage, -# sharpness) alongside the usual L2 / drag / lift. See docs/design/uq_metrics_integration.md. +# ============================================================================ +# 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 -# -# The three model rows exercise every path in the UQ engine: -# * geotransolver_drivaerstar_surface -> deterministic baseline. UQ metrics finalize to NaN -# (no std); included so the report shows the det-vs-UQ contrast. -# * geotransolver_gp_surface -> UQ_METHOD="analytic": one forward pass, GP posterior std. -# * mc_dropout_surface -> UQ_METHOD="sampling": run.uq.num_samples stochastic passes over -# the Concrete-Dropout backbone, aggregated by the engine (Welford). +# torchrun --standalone --nnodes=1 --nproc_per_node=8 main.py --config-name=config_uq_surface +# ============================================================================ defaults: - _self_ @@ -53,62 +69,57 @@ run: # ---- Uncertainty-quantification controls (additive; default off) ---- uq: enabled: true - # Stochastic passes per case for UQ_METHOD="sampling" wrappers (ignored by analytic/deterministic). + # 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 -# Shared anchors for the surface normalization (all checkpoints were trained with the same npz) and -# the deterministic baseline checkpoint. Paths are the in-container mounts from interactive_job.sh: -# /workspace -> /lustre/fsw/portfolios/coreai/users/ktangsali -# /data -> /lustre/fsw/portfolios/coreai/projects/coreai_modulus_cae +# ---- 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: - transformer_src: &transformer_src /workspace/gp_fields/physicsnemo/examples/cfd/external_aerodynamics/transformer_models - surface_stats: &surface_stats /workspace/gp_fields/physicsnemo/examples/cfd/external_aerodynamics/transformer_models/src/surface_fields_normalization.npz - baseline_ckpt: &baseline_ckpt /workspace/gp_fields/physicsnemo/examples/cfd/external_aerodynamics/transformer_models/runs/geotransolver/surface/baseline_fastback_bf/checkpoints/checkpoint.0.100.pt - # Concrete-Dropout run (model.concrete_dropout=true), same backbone/data/normalization as baseline. - concrete_dropout_ckpt: &concrete_dropout_ckpt /workspace/gp_fields/physicsnemo/examples/cfd/external_aerodynamics/transformer_models/runs/geotransolver/surface/concrete_dropout_fastback_bf/checkpoints/GeoTransolver.0.100.mdlus + surface_stats: &surface_stats /path/to/surface_fields_normalization.npz + # Deterministic + ensemble members come from a standard training run's checkpoints. + baseline_ckpt: &baseline_ckpt /path/to/run/checkpoints/checkpoint.0.100.pt + # 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 - # All three rows are transformer_models checkpoints (physical, standardize-only targets: - # physical = norm*std + mean). They use the DrivAerStar GeoTransolver wrappers, which do NOT - # re-dimensionalize by ρu² (that is only for the non-dimensional DrivAerML/HF - # `geotransolver_surface` wrapper). The deterministic, GP, and MC-Dropout wrappers are three - # completely independent CFDModel subclasses (no shared base), each with REDIMENSIONALIZE_OUTPUTS=False. models: - # Deterministic baseline. Epoch 100 = final. + # (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 } - # Analytic GP field head. Point `checkpoint` at the specific GeoTransolver backbone file for the - # epoch you want; the epoch (38 here) is parsed from that name and the sibling FieldGPHead.0.38.pt - # is loaded from the same directory. (No `checkpoint_epoch` knob / directory scanning -> the - # backbone and head can't drift.) e38 of the pinned l2_radial run won the e30/e35/e38 sweep on - # drag AUSE and calibration. The GP hyperparameters below MATCH that training run - # (feature_norm=l2_radial, lengthscale_range=[0.01,0.5], lengthscale_prior=[2.0,6.0]; rest are conf defaults). + # (2) Analytic GP field head. Point `checkpoint` at the specific GeoTransolver backbone file for + # the epoch you want; the epoch is parsed from that name and the sibling FieldGPHead.0..pt + # is loaded from the same directory (no directory scanning -> backbone and head can't drift). + # 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: /workspace/gp_fields/physicsnemo/examples/cfd/external_aerodynamics/transformer_models/runs/geotransolver/surface/field_gp_v4_long_fastback_pin/checkpoints_field_gp/GeoTransolver.0.38.mdlus + checkpoint: /path/to/run_field_gp/checkpoints_field_gp/GeoTransolver.0.100.mdlus stats_path: *surface_stats kwargs: gp_feature_norm: "l2_radial" - gp_n_inducing: 256 + gp_n_inducing: 1024 gp_mlp_hidden: [128, 16] - gp_lengthscale_range: [0.01, 0.5] - gp_lengthscale_prior: [2.0, 6.0] + 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 - # Real MC-Dropout: run.uq.num_samples stochastic passes over the Concrete-Dropout backbone + # (3) MC-Dropout: run.uq.num_samples stochastic passes over the Concrete-Dropout backbone # (the ConcreteDropout layers stay stochastic at inference). - - name: mc_dropout_surface + - name: geotransolver_mc_dropout_surface inference_domain: surface checkpoint: *concrete_dropout_ckpt stats_path: *surface_stats @@ -116,14 +127,28 @@ benchmark: 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: + member_checkpoints: + - /path/to/run/checkpoints/checkpoint.0.96.pt + - /path/to/run/checkpoints/checkpoint.0.97.pt + - /path/to/run/checkpoints/checkpoint.0.98.pt + - /path/to/run/checkpoints/checkpoint.0.99.pt + - /path/to/run/checkpoints/checkpoint.0.100.pt + batch_resolution: 60000 + geometry_sampling: 300000 + datasets: - name: drivaerstar - root: /data/datasets/drivaerstar/surface_files/class_F/val - # These checkpoints were TRAINED on DrivAerStar (class_F = Fastback) zarr, whose WSS keeps the - # raw .vtk sign — so do NOT flip to the DrivAerML convention. The adapter default - # flip_wss_sign=True targets DrivAerML-trained checkpoints; leaving it on here negates the GT - # WSS relative to the model and inflates WSS relative-L2 to ~2.0 (pred ≈ -true) while pressure - # (never flipped) stays correct. Pressure/drag match at flip_wss_sign=false. + root: /path/to/drivaerstar/surface_files/class_F/val + # flip_wss_sign controls the wall-shear-stress sign convention. The adapter default (True) + # targets DrivAerML-trained checkpoints; set it to match how your checkpoints were trained. kwargs: { inference_domain: surface, flip_wss_sign: false } reproducibility: @@ -152,8 +177,8 @@ metrics: - l2_pressure - l2_shear_stress - l2_pressure_area_weighted - # coeff (Cd prefactor 2/(A*rho*U^2)) + drag_direction are optional; shown explicitly so the - # drag_uq panel below uses the SAME direction as the deterministic drag it is diagnosing. + # 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). @@ -169,9 +194,13 @@ metrics: # 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 GP uncertainty propagated into the - # Cd integral (drag std) vs |Cd_pred - Cd_true| (AUSE_epistemic / AUSE_total, lower=better). - # Takes the same coeff / drag_direction kwargs as the deterministic `drag` metric above. + # 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: @@ -181,7 +210,7 @@ reports: # 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 - # Drag Cd panel from drag_uq), with the GP and MC-perturbation models overlaid vs the oracle. + # 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 From e01afd118a23fa56b858d09a74898302c0ee921b Mon Sep 17 00:00:00 2001 From: Kaustubh Tangsali Date: Wed, 22 Jul 2026 22:52:09 +0000 Subject: [PATCH 07/10] clean up configs, add changelog entry, make speficication of gp head checpoint explicit --- .gitignore | 5 ++- CHANGELOG.md | 43 +++++++++++++++++++ .../wrappers/geotransolver_gp/wrapper.py | 40 +++++++++++++---- .../benchmarking/conf/config_uq_surface.yaml | 18 +++++--- 4 files changed, 89 insertions(+), 17 deletions(-) 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/models/wrappers/geotransolver_gp/wrapper.py b/physicsnemo/cfd/evaluation/models/wrappers/geotransolver_gp/wrapper.py index ebb7497..574395c 100644 --- a/physicsnemo/cfd/evaluation/models/wrappers/geotransolver_gp/wrapper.py +++ b/physicsnemo/cfd/evaluation/models/wrappers/geotransolver_gp/wrapper.py @@ -37,14 +37,22 @@ Surface only (the GP head predicts 4 surface tasks: pressure + 3 wall-shear components). -The ``checkpoint`` knob must point at a **specific checkpoint file** whose name encodes the epoch -(``GeoTransolver.0..mdlus``); the epoch is parsed from that name and the sibling -``FieldGPHead.0..pt`` is loaded from the same directory. Pointing at a directory (or an -epoch-less name) is rejected — this removes the old ``checkpoint_epoch`` kwarg and the risk of the -backbone and head drifting to different (or "latest") epochs. +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``. Its directory / + epoch drive the load, so the head file named here is exactly the one read (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`` at the + same epoch. 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). @@ -58,6 +66,7 @@ """ from contextlib import nullcontext +from pathlib import Path from typing import Any, ClassVar, Optional import numpy as np @@ -153,6 +162,7 @@ def __init__(self) -> None: self._gp_kw: dict[str, Any] = {} self._checkpoint_dir: Optional[str] = None 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 @@ -182,6 +192,8 @@ def load( # 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 (optional; validated against the backbone below). + 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)), @@ -234,12 +246,22 @@ def load( inference_mode="surface", ) - # Single source of truth for the epoch: parsed from the checkpoint file name (same file - # the backbone just loaded). The sibling ``FieldGPHead.0..pt`` is loaded from this - # directory at this epoch in ``_build_and_load_head`` — no drift, no "latest" fallback. - ckpt_dir, ckpt_epoch = resolve_checkpoint_file(checkpoint_path) + # The backbone + head are loaded together by ``load_checkpoint(dir, epoch)``. Derive that + # dir + epoch from ``gp_head_checkpoint`` when given (so the head file the user names is the + # one read), else from the backbone ``checkpoint``. No cross-validation between the two — + # passing matching checkpoints is the caller's responsibility. + if head_ckpt_arg is not None: + ckpt_dir, ckpt_epoch = resolve_checkpoint_file(str(head_ckpt_arg)) + self._head_checkpoint = str(Path(head_ckpt_arg)) + else: + ckpt_dir, ckpt_epoch = resolve_checkpoint_file(checkpoint_path) + self._head_checkpoint = str(Path(ckpt_dir) / f"FieldGPHead.0.{ckpt_epoch}.pt") self._checkpoint_dir = str(ckpt_dir) self._checkpoint_epoch = ckpt_epoch + log_inference( + "geotransolver_gp", + f"GP checkpoints -> backbone: {checkpoint_path} | head: {self._head_checkpoint}", + ) return self def prepare_inputs(self, case: CanonicalCase) -> ModelInput: diff --git a/workflows/benchmarking/conf/config_uq_surface.yaml b/workflows/benchmarking/conf/config_uq_surface.yaml index 25af2fa..de2f981 100644 --- a/workflows/benchmarking/conf/config_uq_surface.yaml +++ b/workflows/benchmarking/conf/config_uq_surface.yaml @@ -97,19 +97,23 @@ benchmark: stats_path: *surface_stats kwargs: { batch_resolution: 60000, geometry_sampling: 300000 } - # (2) Analytic GP field head. Point `checkpoint` at the specific GeoTransolver backbone file for - # the epoch you want; the epoch is parsed from that name and the sibling FieldGPHead.0..pt - # is loaded from the same directory (no directory scanning -> backbone and head can't drift). - # The GP hyperparameters here MUST match the training run so the FieldGPHead reconstructs before - # the state dict loads. + # (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/checkpoints_field_gp/GeoTransolver.0.100.mdlus + 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, 16] + 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] From 5ccdd6810863944b761dedbd334ee4aefa1e94a1 Mon Sep 17 00:00:00 2001 From: Kaustubh Tangsali Date: Wed, 22 Jul 2026 23:49:03 +0000 Subject: [PATCH 08/10] do a initial agentic review and address the comments --- .../benchmarks/distributed_utils.py | 8 ++ .../cfd/evaluation/benchmarks/engine.py | 36 ++++-- .../cfd/evaluation/benchmarks/uq_inference.py | 50 +++++++- .../cfd/evaluation/metrics/builtin/uq.py | 83 +++++++++----- .../cfd/evaluation/metrics/mesh_bridge.py | 28 ++++- .../geotransolver_runtime.py | 32 +++++- .../wrappers/ensemble_drivaerstar/wrapper.py | 38 +++--- .../wrappers/geotransolver_gp/wrapper.py | 15 ++- .../models/wrappers/mc_dropout/wrapper.py | 10 +- test/ci_tests/test_uq_inference.py | 104 +++++++++++++++++ test/ci_tests/test_uq_metrics.py | 40 +++++++ test/ci_tests/test_uq_wrappers.py | 108 ++++++++++++++++++ 12 files changed, 483 insertions(+), 69 deletions(-) create mode 100644 test/ci_tests/test_uq_wrappers.py diff --git a/physicsnemo/cfd/evaluation/benchmarks/distributed_utils.py b/physicsnemo/cfd/evaluation/benchmarks/distributed_utils.py index 91baac6..db50664 100644 --- a/physicsnemo/cfd/evaluation/benchmarks/distributed_utils.py +++ b/physicsnemo/cfd/evaluation/benchmarks/distributed_utils.py @@ -109,6 +109,14 @@ def merge_benchmark_result_shards(shards: list[dict[str, Any]]) -> dict[str, Any 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 { diff --git a/physicsnemo/cfd/evaluation/benchmarks/engine.py b/physicsnemo/cfd/evaluation/benchmarks/engine.py index 990952d..6ff4714 100644 --- a/physicsnemo/cfd/evaluation/benchmarks/engine.py +++ b/physicsnemo/cfd/evaluation/benchmarks/engine.py @@ -130,6 +130,7 @@ class BenchmarkPolicyError(RuntimeError): make_reducer_partial_key, make_sample_partial_key, run_sampling_inference, + select_inference_path, strip_reducer_partials, ) @@ -749,9 +750,16 @@ def _run_single( case_cache[case_key] = case seed_inference_rng(run_config.seed, cid) model_input = wrapper.prepare_inputs(case) - supports_uq = bool(getattr(wrapper, "SUPPORTS_UQ", False)) - uq_method = getattr(wrapper, "UQ_METHOD", "none") - if supports_uq and uq_method == "sampling" and run_config.uq.enabled: + # ``run.uq.enabled`` is the master switch: when off, EVERY wrapper (sampling AND analytic) + # takes the deterministic path (single ``predict`` + ``decode_outputs``), so no + # distributions are emitted and no UQ metrics are produced — enabling apples-to-apples + # deterministic comparison runs. 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, @@ -762,7 +770,7 @@ def _run_single( case_id=cid, retain_samples=run_config.uq.retain_samples, ) - elif supports_uq and uq_method == "analytic": + elif inference_path == "analytic": raw = wrapper.predict(model_input) predictions = wrapper.decode_distribution(raw, case, model_input) else: @@ -931,9 +939,16 @@ def _run_single( # 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. - metrics_summary.update(finalize_reducer_metrics(per_case, m_dom)) - metrics_summary.update(finalize_sample_metrics(per_case, m_dom)) + # 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 ( { @@ -1248,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, @@ -1257,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/uq_inference.py b/physicsnemo/cfd/evaluation/benchmarks/uq_inference.py index 074c8f9..6860321 100644 --- a/physicsnemo/cfd/evaluation/benchmarks/uq_inference.py +++ b/physicsnemo/cfd/evaluation/benchmarks/uq_inference.py @@ -70,7 +70,9 @@ def _split_reducer_partial_key(key: str) -> tuple[str, str]: def finalize_reducer_metrics( - per_case_rows: list[dict[str, Any]], domain: str | None + 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). @@ -78,6 +80,12 @@ def finalize_reducer_metrics( ``(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). @@ -99,8 +107,14 @@ def finalize_reducer_metrics( 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 summed.items(): + for metric_name, stats in stats_by_metric.items(): try: metric = get_metric(metric_name, domain=domain) except KeyError: @@ -144,7 +158,9 @@ def _split_sample_partial_key(key: str) -> tuple[str, str]: def finalize_sample_metrics( - per_case_rows: list[dict[str, Any]], domain: str | None + 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). @@ -152,6 +168,11 @@ def finalize_sample_metrics( 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, @@ -170,8 +191,12 @@ def finalize_sample_metrics( 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 collected.items(): + for metric_name, stats in stats_by_metric.items(): try: metric = get_metric(metric_name, domain=domain) except KeyError: @@ -237,6 +262,23 @@ def is_uq_partial_key(key: str) -> bool: 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. diff --git a/physicsnemo/cfd/evaluation/metrics/builtin/uq.py b/physicsnemo/cfd/evaluation/metrics/builtin/uq.py index 471af86..fb91fde 100644 --- a/physicsnemo/cfd/evaluation/metrics/builtin/uq.py +++ b/physicsnemo/cfd/evaluation/metrics/builtin/uq.py @@ -39,6 +39,7 @@ 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 @@ -47,6 +48,10 @@ #: 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") @@ -173,7 +178,10 @@ def finalize(self, summed: dict[str, float]) -> float | dict[str, float]: chan, stat = key.split("::", 1) by_channel.setdefault(chan, {})[stat] = val if not by_channel: - return float("nan") + # 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): @@ -243,30 +251,22 @@ def _zrms_finalize(d: dict[str, float]) -> float: # channel (+ headline ``mean``) and averaged over cases by the engine. -def _rankdata(x: np.ndarray) -> np.ndarray: - """Ordinal ranks (0..n-1) of ``x`` via argsort-of-argsort. +def _spearman(a: np.ndarray, b: np.ndarray) -> float: + """Tie-aware Spearman rank correlation of ``a`` and ``b`` (via ``scipy.stats.spearmanr``). - Continuous std / error values effectively never tie, so ordinal (vs average-tie) ranks are - fine here and avoid a SciPy dependency. + ``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). """ - order = np.argsort(x, kind="stable") - ranks = np.empty(x.size, dtype=np.float64) - ranks[order] = np.arange(x.size, dtype=np.float64) - return ranks - - -def _spearman(a: np.ndarray, b: np.ndarray) -> float: - """Spearman rank correlation = Pearson correlation of the ranks of ``a`` and ``b``.""" - if a.size < 2: + 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") - ra = _rankdata(a) - rb = _rankdata(b) - ra -= ra.mean() - rb -= rb.mean() - denom = math.sqrt(float((ra**2).sum()) * float((rb**2).sum())) - if denom <= 0.0: + # 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((ra * rb).sum() / denom) + return float(spearmanr(a, b).correlation) class _UncertaintyErrorSpearman: @@ -339,7 +339,7 @@ def _suffix_rmse(order: np.ndarray) -> np.ndarray: oracle = _suffix_rmse(np.argsort(-err2)) full = float(np.sqrt(err2.mean())) ause = ( - float(np.trapz((by_unc - oracle) / full, fr)) if full > 0.0 else float("nan") + float(_trapezoid((by_unc - oracle) / full, fr)) if full > 0.0 else float("nan") ) return fr, by_unc, oracle, full, ause @@ -411,7 +411,9 @@ def finalize_samples( ) -> float | dict[str, float]: by_channel = self._regroup(collected) if not by_channel: - return float("nan") + # 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] = [] @@ -516,9 +518,12 @@ class _SampleDragUQ: 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 them when - ``output.std_mesh_field_names`` / ``epistemic_std_mesh_field_names`` are set), so it is a no-op - for deterministic wrappers or when std fields are unavailable. + 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 @@ -568,6 +573,7 @@ def partial( 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( @@ -586,10 +592,27 @@ def partial( 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") - std_p = output.std_mesh_field_names.get("pressure") - std_w = output.std_mesh_field_names.get("shear_stress") - epi_p = output.epistemic_std_mesh_field_names.get("pressure") - epi_w = output.epistemic_std_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). diff --git a/physicsnemo/cfd/evaluation/metrics/mesh_bridge.py b/physicsnemo/cfd/evaluation/metrics/mesh_bridge.py index 335ae71..137068c 100644 --- a/physicsnemo/cfd/evaluation/metrics/mesh_bridge.py +++ b/physicsnemo/cfd/evaluation/metrics/mesh_bridge.py @@ -45,6 +45,21 @@ def _unwrap_prediction_means(predictions: dict[str, Any] | None) -> dict[str, An 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, @@ -67,14 +82,15 @@ def _attach_uq_std_companions( 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: - name = std_map.get(canonical) or f"{pred_name}Std" - _assign_field(mesh, preference, name, value.std) - attached.append(name) + _assign_field(mesh, preference, std_name, value.std) + attached.append(std_name) if value.epistemic_std is not None: - name = epi_std_map.get(canonical) or f"{pred_name}EpistemicStd" - _assign_field(mesh, preference, name, value.epistemic_std) - attached.append(name) + _assign_field(mesh, preference, epi_std_name, value.epistemic_std) + attached.append(epi_std_name) return attached diff --git a/physicsnemo/cfd/evaluation/models/common_wrapper_utils/geotransolver_runtime.py b/physicsnemo/cfd/evaluation/models/common_wrapper_utils/geotransolver_runtime.py index 043dab8..b927515 100644 --- a/physicsnemo/cfd/evaluation/models/common_wrapper_utils/geotransolver_runtime.py +++ b/physicsnemo/cfd/evaluation/models/common_wrapper_utils/geotransolver_runtime.py @@ -18,8 +18,8 @@ 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, and the MC weight-perturbation proxy. They share their plumbing through the free -functions here rather than class inheritance: +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`. @@ -421,6 +421,19 @@ def build_transolver_batch( ) +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", @@ -428,12 +441,22 @@ def geotransolver_forward( 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.""" + """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 = torch.randperm(n, device=batch["embeddings"].device) + 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 @@ -560,6 +583,7 @@ def decode_volume_predictions( "build_surface_datapipe", "build_volume_datapipe", "build_transolver_batch", + "make_forward_permutation", "geotransolver_forward", "unscale_targets", "decode_surface_predictions", diff --git a/physicsnemo/cfd/evaluation/models/wrappers/ensemble_drivaerstar/wrapper.py b/physicsnemo/cfd/evaluation/models/wrappers/ensemble_drivaerstar/wrapper.py index 3184d93..a3c6a71 100644 --- a/physicsnemo/cfd/evaluation/models/wrappers/ensemble_drivaerstar/wrapper.py +++ b/physicsnemo/cfd/evaluation/models/wrappers/ensemble_drivaerstar/wrapper.py @@ -34,10 +34,12 @@ 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 every member on the shared per-case batch and -returns their raw outputs; the engine's -:func:`~physicsnemo.cfd.evaluation.benchmarks.uq_inference.run_sampling_inference` aggregates them -(Welford) into mean + across-member (epistemic) std. Because ``predict_ensemble`` returns exactly -the member count, ``run.uq.num_samples`` is ignored for this row. +**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. Because ``predict_ensemble`` yields exactly the member count, +``run.uq.num_samples`` is ignored for this row. It decorates ``transformer_models`` (physical-target) checkpoints, so predictions are re-standardized only — no dynamic-pressure re-dimensionalization @@ -58,7 +60,7 @@ import logging from pathlib import Path -from typing import Any, ClassVar, Literal, Optional +from typing import Any, ClassVar, Iterator, Literal, Optional from physicsnemo.cfd.evaluation.datasets.schema import ( CanonicalCase, @@ -78,6 +80,7 @@ decode_volume_predictions, geotransolver_available, geotransolver_forward, + make_forward_permutation, parse_runtime_kwargs, unscale_targets, ) @@ -216,14 +219,21 @@ def prepare_inputs(self, case: CanonicalCase) -> ModelInput: self._volume_length_scale = result.volume_length_scale return {"batch": result.batch, "datapipe": result.datapipe} - def _forward_member(self, model: Any, model_input: ModelInput) -> RawOutput: - """Run one member's forward + target unscaling on the shared per-case batch.""" + 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"], @@ -234,17 +244,19 @@ def _forward_member(self, model: Any, model_input: ModelInput) -> RawOutput: def predict_ensemble( self, model_input: ModelInput, n: int - ) -> Optional[list[RawOutput]]: - """Run every member on the shared batch; return one raw output per member. + ) -> Optional[Iterator[RawOutput]]: + """Yield one raw output per member (lazy generator), sharing one per-case permutation. ``n`` (``run.uq.num_samples``) is intentionally ignored: the ensemble has a fixed number of - members, and the engine iterates over exactly the returned list. Only one member's forward is - resident on the device at a time; the returned raw outputs are the (host-bound) unscaled - predictions the engine immediately folds into its running mean/variance. + members and the engine iterates over exactly what is yielded. 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") - return [self._forward_member(model, model_input) for model in self._models] + perm = make_forward_permutation(model_input["batch"]) + return (self._forward_member(model, model_input, perm) for model in self._models) def predict(self, model_input: ModelInput) -> RawOutput: """Deterministic fallback (first member). The engine uses :meth:`predict_ensemble` for UQ.""" diff --git a/physicsnemo/cfd/evaluation/models/wrappers/geotransolver_gp/wrapper.py b/physicsnemo/cfd/evaluation/models/wrappers/geotransolver_gp/wrapper.py index 574395c..0506b96 100644 --- a/physicsnemo/cfd/evaluation/models/wrappers/geotransolver_gp/wrapper.py +++ b/physicsnemo/cfd/evaluation/models/wrappers/geotransolver_gp/wrapper.py @@ -192,7 +192,9 @@ def load( # 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 (optional; validated against the backbone below). + # 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)), @@ -282,7 +284,14 @@ def prepare_inputs(self, case: CanonicalCase) -> ModelInput: 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 its weights.""" + """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 head checkpoint's directory + epoch (which resolves to the + exact ``FieldGPHead.0..pt`` file named by ``gp_head_checkpoint``). Passing only the + head to ``load_checkpoint`` avoids reloading — or silently overwriting — the backbone with a + different one that happens to sit in the head's directory. + """ dev = torch.device(self._cfg.device) feature_dim = self._probe_feature_dim(batch) self._head = FieldGPHead( @@ -293,7 +302,7 @@ def _build_and_load_head(self, batch: dict) -> None: with trusted_torch_load_context(): loaded_epoch = load_checkpoint( path=self._checkpoint_dir, - models=[self._model, self._head], + models=[self._head], device=dev, epoch=self._checkpoint_epoch, ) diff --git a/physicsnemo/cfd/evaluation/models/wrappers/mc_dropout/wrapper.py b/physicsnemo/cfd/evaluation/models/wrappers/mc_dropout/wrapper.py index 847697d..b3ccae9 100644 --- a/physicsnemo/cfd/evaluation/models/wrappers/mc_dropout/wrapper.py +++ b/physicsnemo/cfd/evaluation/models/wrappers/mc_dropout/wrapper.py @@ -69,6 +69,7 @@ decode_volume_predictions, geotransolver_available, geotransolver_forward, + make_forward_permutation, parse_runtime_kwargs, unscale_targets, ) @@ -237,14 +238,18 @@ def prepare_inputs(self, case: CanonicalCase) -> ModelInput: 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} + # 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). + 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") @@ -254,6 +259,7 @@ def predict(self, model_input: ModelInput) -> RawOutput: 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"], diff --git a/test/ci_tests/test_uq_inference.py b/test/ci_tests/test_uq_inference.py index 6f7e6fa..235bffa 100644 --- a/test/ci_tests/test_uq_inference.py +++ b/test/ci_tests/test_uq_inference.py @@ -18,15 +18,19 @@ 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, ) @@ -140,3 +144,103 @@ def _case(n, sigma): 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: + # 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 diff --git a/test/ci_tests/test_uq_metrics.py b/test/ci_tests/test_uq_metrics.py index 34cf23b..955fd58 100644 --- a/test/ci_tests/test_uq_metrics.py +++ b/test/ci_tests/test_uq_metrics.py @@ -34,12 +34,15 @@ 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, ) @@ -200,6 +203,43 @@ def test_spearman_high_when_uncertainty_tracks_error() -> None: 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: spear = _make_pointwise_uq_metrics()["uncertainty_error_spearman"] spear_epi = _make_pointwise_uq_metrics()["uncertainty_error_spearman_epistemic"] diff --git a/test/ci_tests/test_uq_wrappers.py b/test/ci_tests/test_uq_wrappers.py new file mode 100644 index 0000000..5cc190c --- /dev/null +++ b/test/ci_tests/test_uq_wrappers.py @@ -0,0 +1,108 @@ +# 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, +) + + +# -------------------------------------------------------------------------------------------- +# Checkpoint routing: a specific file (epoch-in-name) -> (dir, epoch); directory / bad name reject +# -------------------------------------------------------------------------------------------- + + +def test_resolve_checkpoint_file_parses_epoch(tmp_path) -> None: + ckpt = tmp_path / "GeoTransolver.0.30.mdlus" + ckpt.write_bytes(b"") # existence not required, but harmless + directory, epoch = resolve_checkpoint_file(str(ckpt)) + assert directory == tmp_path and epoch == 30 + # Nonexistent path still resolves purely from the file name (no silent latest-epoch fallback). + d2, e2 = resolve_checkpoint_file(str(tmp_path / "checkpoint.0.100.pt")) + assert d2 == tmp_path and e2 == 100 + + +def test_resolve_checkpoint_file_rejects_directory_and_epochless(tmp_path) -> None: + 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 + + +# -------------------------------------------------------------------------------------------- +# Block-partition permutation covers the WHOLE point cloud (no subsampling of samples) +# -------------------------------------------------------------------------------------------- + + +def test_make_forward_permutation_covers_all_points() -> None: + 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: + # 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: + # 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 From 5db085b1fd0a0952aad74b2db8075c3ce0305da8 Mon Sep 17 00:00:00 2001 From: Kaustubh Tangsali Date: Thu, 23 Jul 2026 00:02:00 +0000 Subject: [PATCH 09/10] fix pre-commit errors --- .../cfd/evaluation/benchmarks/engine.py | 4 +- .../cfd/evaluation/benchmarks/report.py | 6 +- .../cfd/evaluation/benchmarks/uq_inference.py | 10 ++-- physicsnemo/cfd/evaluation/datasets/schema.py | 12 ++-- .../cfd/evaluation/metrics/builtin/uq.py | 25 +++++--- .../geotransolver_runtime.py | 14 +++-- .../evaluation/models/wrappers/__init__.py | 8 ++- .../wrappers/ensemble_drivaerstar/wrapper.py | 20 +++++-- .../wrappers/geotransolver_gp/wrapper.py | 7 ++- .../models/wrappers/mc_dropout/wrapper.py | 16 ++++-- .../reports/builtin/sparsification_plot.py | 8 +-- .../references/example_wrapper.py | 24 ++++++-- test/ci_tests/test_l2_errors.py | 3 + test/ci_tests/test_uq_inference.py | 15 ++++- test/ci_tests/test_uq_metrics.py | 57 ++++++++++++++----- test/ci_tests/test_uq_wrappers.py | 9 ++- 16 files changed, 173 insertions(+), 65 deletions(-) diff --git a/physicsnemo/cfd/evaluation/benchmarks/engine.py b/physicsnemo/cfd/evaluation/benchmarks/engine.py index 6ff4714..6ad3912 100644 --- a/physicsnemo/cfd/evaluation/benchmarks/engine.py +++ b/physicsnemo/cfd/evaluation/benchmarks/engine.py @@ -241,9 +241,7 @@ def _sanitize_path_token(token: str) -> str: 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: +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 diff --git a/physicsnemo/cfd/evaluation/benchmarks/report.py b/physicsnemo/cfd/evaluation/benchmarks/report.py index 2ee4be3..fe1ef67 100644 --- a/physicsnemo/cfd/evaluation/benchmarks/report.py +++ b/physicsnemo/cfd/evaluation/benchmarks/report.py @@ -46,7 +46,11 @@ def _sanitize_for_strict_json(obj: Any) -> Any: 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: + 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()) diff --git a/physicsnemo/cfd/evaluation/benchmarks/uq_inference.py b/physicsnemo/cfd/evaluation/benchmarks/uq_inference.py index 6860321..413adbc 100644 --- a/physicsnemo/cfd/evaluation/benchmarks/uq_inference.py +++ b/physicsnemo/cfd/evaluation/benchmarks/uq_inference.py @@ -103,9 +103,9 @@ def finalize_reducer_metrics( if not partial_key: continue summed.setdefault(metric_name, {}) - summed[metric_name][partial_key] = ( - summed[metric_name].get(partial_key, 0.0) + float(val) - ) + 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. @@ -340,7 +340,9 @@ def result( 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 = ( + 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: diff --git a/physicsnemo/cfd/evaluation/datasets/schema.py b/physicsnemo/cfd/evaluation/datasets/schema.py index c8e1737..e7a3594 100644 --- a/physicsnemo/cfd/evaluation/datasets/schema.py +++ b/physicsnemo/cfd/evaluation/datasets/schema.py @@ -158,8 +158,12 @@ class FieldDistribution: 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 + 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( @@ -198,9 +202,7 @@ def _coerce(x: ArrayLike | None) -> ArrayLike | None: ) -def as_distribution( - predictions: dict[str, Any], key: str -) -> FieldDistribution | None: +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. diff --git a/physicsnemo/cfd/evaluation/metrics/builtin/uq.py b/physicsnemo/cfd/evaluation/metrics/builtin/uq.py index fb91fde..545b195 100644 --- a/physicsnemo/cfd/evaluation/metrics/builtin/uq.py +++ b/physicsnemo/cfd/evaluation/metrics/builtin/uq.py @@ -166,7 +166,9 @@ def partial(self, gt: Any, predictions: Any, **_: Any) -> dict[str, float]: ): 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) + 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]: @@ -396,7 +398,9 @@ def partial(self, gt: Any, predictions: Any, **_: Any) -> dict[str, float]: return stats @staticmethod - def _regroup(collected: dict[str, list[float]]) -> dict[str, dict[str, list[float]]]: + 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(): @@ -464,7 +468,7 @@ def curves(self, collected: dict[str, list[float]]) -> dict[str, dict[str, Any]] # 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 +# 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). @@ -631,18 +635,21 @@ def partial( 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) - ): + 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, + 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), diff --git a/physicsnemo/cfd/evaluation/models/common_wrapper_utils/geotransolver_runtime.py b/physicsnemo/cfd/evaluation/models/common_wrapper_utils/geotransolver_runtime.py index b927515..953e4e3 100644 --- a/physicsnemo/cfd/evaluation/models/common_wrapper_utils/geotransolver_runtime.py +++ b/physicsnemo/cfd/evaluation/models/common_wrapper_utils/geotransolver_runtime.py @@ -225,7 +225,9 @@ def resolve_checkpoint_file(checkpoint_path: str) -> tuple[Path, int]: """ ckpt = Path(checkpoint_path) if not checkpoint_path: - raise ValueError("checkpoint path is empty; point it at a specific checkpoint file.") + 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 " @@ -287,7 +289,11 @@ def _move_reference_scale_to_device(dp: "TransolverDataPipe", device: str) -> No def build_surface_datapipe( - *, geometry_sampling: int, surface_factors: Any, device: str, user_kw: dict[str, Any] + *, + 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} @@ -461,9 +467,7 @@ def geotransolver_forward( preds_list: list[torch.Tensor] = [] use_full_fx = "geometry" in batch - ac_ctx = ( - cuda_bf16_autocast(device) if cuda_bf16_autocast_enabled else nullcontext() - ) + 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: diff --git a/physicsnemo/cfd/evaluation/models/wrappers/__init__.py b/physicsnemo/cfd/evaluation/models/wrappers/__init__.py index 91c8a9d..71d013c 100644 --- a/physicsnemo/cfd/evaluation/models/wrappers/__init__.py +++ b/physicsnemo/cfd/evaluation/models/wrappers/__init__.py @@ -49,8 +49,12 @@ 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( + "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) diff --git a/physicsnemo/cfd/evaluation/models/wrappers/ensemble_drivaerstar/wrapper.py b/physicsnemo/cfd/evaluation/models/wrappers/ensemble_drivaerstar/wrapper.py index a3c6a71..10d3130 100644 --- a/physicsnemo/cfd/evaluation/models/wrappers/ensemble_drivaerstar/wrapper.py +++ b/physicsnemo/cfd/evaluation/models/wrappers/ensemble_drivaerstar/wrapper.py @@ -169,7 +169,9 @@ def load( ) self._surface_factors = None else: - log_inference("ensemble", f"Loading surface normalization from {stats_path}") + log_inference( + "ensemble", f"Loading surface normalization from {stats_path}" + ) self._surface_factors = load_transolver_surface_factors(stats_path, device) self._volume_factors = None @@ -203,7 +205,9 @@ def load( 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") + raise RuntimeError( + "GeoTransolverEnsembleDrivAerStarWrapper: call load() first" + ) result = build_transolver_batch( case=case, inference_mode=self._inference_mode, @@ -254,14 +258,20 @@ def predict_ensemble( 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") + raise RuntimeError( + "GeoTransolverEnsembleDrivAerStarWrapper: call load() first" + ) perm = make_forward_permutation(model_input["batch"]) - return (self._forward_member(model, model_input, perm) for model in self._models) + return ( + self._forward_member(model, model_input, perm) for model in self._models + ) 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") + raise RuntimeError( + "GeoTransolverEnsembleDrivAerStarWrapper: call load() first" + ) return self._forward_member(self._models[0], model_input) def decode_outputs( diff --git a/physicsnemo/cfd/evaluation/models/wrappers/geotransolver_gp/wrapper.py b/physicsnemo/cfd/evaluation/models/wrappers/geotransolver_gp/wrapper.py index 0506b96..3c6fdb0 100644 --- a/physicsnemo/cfd/evaluation/models/wrappers/geotransolver_gp/wrapper.py +++ b/physicsnemo/cfd/evaluation/models/wrappers/geotransolver_gp/wrapper.py @@ -257,7 +257,9 @@ def load( self._head_checkpoint = str(Path(head_ckpt_arg)) else: ckpt_dir, ckpt_epoch = resolve_checkpoint_file(checkpoint_path) - self._head_checkpoint = str(Path(ckpt_dir) / f"FieldGPHead.0.{ckpt_epoch}.pt") + self._head_checkpoint = str( + Path(ckpt_dir) / f"FieldGPHead.0.{ckpt_epoch}.pt" + ) self._checkpoint_dir = str(ckpt_dir) self._checkpoint_epoch = ckpt_epoch log_inference( @@ -438,7 +440,8 @@ def decode_distribution( raw_output["mean"], raw_output["total_std"], raw_output["epistemic_std"] ) log_inference( - "geotransolver_gp", "Decoding GP distribution (pressure + WSS, physical units)…" + "geotransolver_gp", + "Decoding GP distribution (pressure + WSS, physical units)…", ) out: dict[str, FieldDistribution] = {} out["pressure"] = build_predictive_distribution( diff --git a/physicsnemo/cfd/evaluation/models/wrappers/mc_dropout/wrapper.py b/physicsnemo/cfd/evaluation/models/wrappers/mc_dropout/wrapper.py index b3ccae9..1501d07 100644 --- a/physicsnemo/cfd/evaluation/models/wrappers/mc_dropout/wrapper.py +++ b/physicsnemo/cfd/evaluation/models/wrappers/mc_dropout/wrapper.py @@ -163,7 +163,9 @@ def load( self._cfg = parse_runtime_kwargs(kw, device) if self._inference_mode == "volume": - log_inference("mc_dropout", f"Loading volume normalization from {stats_path}") + 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( @@ -172,7 +174,9 @@ def load( ) self._surface_factors = None else: - log_inference("mc_dropout", f"Loading surface normalization from {stats_path}") + 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 @@ -224,7 +228,9 @@ def _enable_mc_dropout(self) -> None: 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") + raise RuntimeError( + "GeoTransolverMCDropoutDrivAerStarWrapper: call load() first" + ) result = build_transolver_batch( case=case, inference_mode=self._inference_mode, @@ -252,7 +258,9 @@ def predict(self, model_input: ModelInput) -> RawOutput: 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") + raise RuntimeError( + "GeoTransolverMCDropoutDrivAerStarWrapper: call load() first" + ) raw = geotransolver_forward( model=self._model, batch=model_input["batch"], diff --git a/physicsnemo/cfd/evaluation/reports/builtin/sparsification_plot.py b/physicsnemo/cfd/evaluation/reports/builtin/sparsification_plot.py index ff115a2..5e6d90b 100644 --- a/physicsnemo/cfd/evaluation/reports/builtin/sparsification_plot.py +++ b/physicsnemo/cfd/evaluation/reports/builtin/sparsification_plot.py @@ -69,7 +69,9 @@ def _panel_title(metric_name: str, series: str) -> str: return f"{metric_name}:{series}" -def _collect_panels(runs: list[tuple[dict[str, Any], dict[str, Any]]]) -> list[tuple[str, str]]: +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: @@ -161,9 +163,7 @@ def sparsification_plot( 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) - } + 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) 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 429b129..505faa6 100644 --- a/skills/physicsnemo-cfd-create-model-wrapper/references/example_wrapper.py +++ b/skills/physicsnemo-cfd-create-model-wrapper/references/example_wrapper.py @@ -276,11 +276,13 @@ def __init__(self) -> None: @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). @@ -289,6 +291,7 @@ def load( 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() @@ -306,13 +309,15 @@ def predict(self, model_input: torch.Tensor) -> dict[str, torch.Tensor]: 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) + "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: + 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).""" @@ -347,7 +352,9 @@ def decode_distribution( "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), + epistemic_std=self._to_physical( + epi[:, 1:4], "shear_stress", is_std=True + ), ), } @@ -392,17 +399,21 @@ class ExampleSamplingUQWrapper(CFDModel): UQ_METHOD: ClassVar[str] = "sampling" def __init__(self) -> None: - self._models: list[torch.nn.Module] = [] # one for MC-Dropout, K for an ensemble + 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"]. @@ -414,6 +425,7 @@ def load( 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() @@ -452,7 +464,9 @@ def decode_outputs( ) -> 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") + 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) ) 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 index 235bffa..44d667b 100644 --- a/test/ci_tests/test_uq_inference.py +++ b/test/ci_tests/test_uq_inference.py @@ -39,6 +39,7 @@ 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) @@ -48,6 +49,7 @@ def test_reducer_partial_key_round_trip_and_strip() -> None: 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) @@ -78,9 +80,16 @@ def decode_outputs(self, raw, case, model_input=None): 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 + _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]) @@ -90,6 +99,7 @@ def test_run_sampling_inference_epistemic_and_samples() -> None: 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, @@ -102,6 +112,7 @@ def test_run_sampling_inference_zero_spread_gives_zero_epistemic() -> None: 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])), @@ -118,6 +129,7 @@ def test_run_sampling_inference_law_of_total_variance() -> None: 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): @@ -186,6 +198,7 @@ def test_finalize_sample_metrics_emits_nan_for_configured_but_absent() -> None: 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) diff --git a/test/ci_tests/test_uq_metrics.py b/test/ci_tests/test_uq_metrics.py index 955fd58..febc8b0 100644 --- a/test/ci_tests/test_uq_metrics.py +++ b/test/ci_tests/test_uq_metrics.py @@ -47,6 +47,7 @@ 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) ) @@ -56,6 +57,7 @@ def test_build_predictive_distribution_coerces_numpy_to_float32() -> None: 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 @@ -66,6 +68,7 @@ def test_as_distribution_passthrough_wrap_and_none() -> 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) @@ -93,6 +96,7 @@ def _gaussian_case(rng, n, sigma, epi=None): 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", @@ -113,7 +117,9 @@ def test_pooled_calibration_on_synthetic_gaussian() -> None: 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 + 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 @@ -133,6 +139,7 @@ def test_pooling_is_global_not_mean_of_case_means() -> None: 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 = [ ( @@ -140,15 +147,17 @@ def test_deterministic_prediction_yields_nan() -> None: {"pressure": FieldDistribution(mean=np.zeros(16, np.float32))}, ) ] - assert math.isnan(_run_reducer(metrics["nlpd"], cases)) + # 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 + # 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)) + 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) @@ -156,6 +165,7 @@ def test_epistemic_metric_requires_epistemic_std() -> None: 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 @@ -166,7 +176,9 @@ def test_vector_field_splits_into_components() -> None: 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})]) + 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"} @@ -193,7 +205,9 @@ def test_spearman_high_when_uncertainty_tracks_error() -> None: 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)}) + 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) @@ -218,10 +232,13 @@ def test_spearman_tie_aware_and_constant_input() -> None: 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 + 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]) @@ -241,12 +258,15 @@ def test_trapezoid_available_and_ause_uses_it() -> None: 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"]) + 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"]) @@ -266,9 +286,12 @@ def _collect_samples(metric, cases): 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 + assert not is_reducer_metric( + m + ) # sample metrics use finalize_samples, not finalize def test_sample_ause_informative_below_random_over_geometries() -> None: @@ -289,7 +312,9 @@ def _geom(scale, rank_sig): # 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)]) + 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. @@ -298,10 +323,13 @@ def _geom(scale, rank_sig): 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})]) + res = _collect_samples( + ause, [({"pressure": np.ones(n, np.float32)}, {"pressure": fd})] + ) assert math.isnan(res["pressure"]) @@ -397,4 +425,5 @@ def test_drag_uq_partial_skips_deterministic() -> None: 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 index 5cc190c..c184a8b 100644 --- a/test/ci_tests/test_uq_wrappers.py +++ b/test/ci_tests/test_uq_wrappers.py @@ -35,6 +35,7 @@ def test_resolve_checkpoint_file_parses_epoch(tmp_path) -> None: + """A specific checkpoint file name resolves to ``(directory, epoch)``.""" ckpt = tmp_path / "GeoTransolver.0.30.mdlus" ckpt.write_bytes(b"") # existence not required, but harmless directory, epoch = resolve_checkpoint_file(str(ckpt)) @@ -45,12 +46,15 @@ def test_resolve_checkpoint_file_parses_epoch(tmp_path) -> None: 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 + resolve_checkpoint_file( + str(tmp_path / "GeoTransolver.mdlus") + ) # no epoch in name # -------------------------------------------------------------------------------------------- @@ -59,6 +63,7 @@ def test_resolve_checkpoint_file_rejects_directory_and_epochless(tmp_path) -> No 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) @@ -73,6 +78,7 @@ def test_make_forward_permutation_covers_all_points() -> None: 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. @@ -90,6 +96,7 @@ def test_uq_std_field_names_fallback_and_override() -> None: 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 From d844c6b747132b3e8e171b113ebed0460cd98e06 Mon Sep 17 00:00:00 2001 From: Kaustubh Tangsali Date: Thu, 23 Jul 2026 00:41:44 +0000 Subject: [PATCH 10/10] address some more review comments --- .../cfd/evaluation/benchmarks/engine.py | 10 +- .../cfd/evaluation/benchmarks/uq_inference.py | 5 +- .../datasets/adapters/drivaerstar.py | 8 +- .../geotransolver_runtime.py | 40 ++++--- .../cfd/evaluation/models/model_registry.py | 50 ++++++--- .../wrappers/ensemble_drivaerstar/wrapper.py | 47 +++++--- .../geotransolver_drivaerstar/wrapper.py | 2 +- .../wrappers/geotransolver_gp/wrapper.py | 54 ++++----- .../models/wrappers/mc_dropout/wrapper.py | 25 ++++- .../SKILL.md | 24 ++-- .../references/example_wrapper.py | 51 +++++++-- test/ci_tests/test_drivaerstar_adapter.py | 106 ++++++++++++++++++ test/ci_tests/test_uq_inference.py | 36 ++++++ test/ci_tests/test_uq_wrappers.py | 97 +++++++++++++++- .../benchmarking/conf/config_uq_surface.yaml | 23 ++-- 15 files changed, 463 insertions(+), 115 deletions(-) create mode 100644 test/ci_tests/test_drivaerstar_adapter.py diff --git a/physicsnemo/cfd/evaluation/benchmarks/engine.py b/physicsnemo/cfd/evaluation/benchmarks/engine.py index 6ad3912..901dfed 100644 --- a/physicsnemo/cfd/evaluation/benchmarks/engine.py +++ b/physicsnemo/cfd/evaluation/benchmarks/engine.py @@ -749,9 +749,11 @@ def _run_single( seed_inference_rng(run_config.seed, cid) model_input = wrapper.prepare_inputs(case) # ``run.uq.enabled`` is the master switch: when off, EVERY wrapper (sampling AND analytic) - # takes the deterministic path (single ``predict`` + ``decode_outputs``), so no - # distributions are emitted and no UQ metrics are produced — enabling apples-to-apples - # deterministic comparison runs. See :func:`select_inference_path`. + # 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"), @@ -772,7 +774,7 @@ def _run_single( raw = wrapper.predict(model_input) predictions = wrapper.decode_distribution(raw, case, model_input) else: - raw = wrapper.predict(model_input) + raw = wrapper.predict_deterministic(model_input) predictions = wrapper.decode_outputs(raw, case, model_input) gt = case.ground_truth or {} diff --git a/physicsnemo/cfd/evaluation/benchmarks/uq_inference.py b/physicsnemo/cfd/evaluation/benchmarks/uq_inference.py index 413adbc..6e094dd 100644 --- a/physicsnemo/cfd/evaluation/benchmarks/uq_inference.py +++ b/physicsnemo/cfd/evaluation/benchmarks/uq_inference.py @@ -385,8 +385,9 @@ def run_sampling_inference( ) -> dict[str, FieldDistribution]: """Drive ``n`` stochastic passes and aggregate to a distribution per field. - Uses ``wrapper.predict_ensemble(model_input, n)`` when available (single batched call), - else calls ``wrapper.predict(model_input)`` ``n`` times, reseeding per pass from + 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: diff --git a/physicsnemo/cfd/evaluation/datasets/adapters/drivaerstar.py b/physicsnemo/cfd/evaluation/datasets/adapters/drivaerstar.py index 0c7ce65..3b5ffa4 100644 --- a/physicsnemo/cfd/evaluation/datasets/adapters/drivaerstar.py +++ b/physicsnemo/cfd/evaluation/datasets/adapters/drivaerstar.py @@ -110,7 +110,13 @@ class DrivAerStarAdapter(DatasetAdapter): - ``pressure_field_name``: source pressure array name (default ``"Pressure"``). - ``wss_component_names``: three source WSS scalar names (default ``("WallShearStressi", "WallShearStressj", "WallShearStressk")``). - - ``flip_wss_sign``: negate combined WSS to match DrivAerML (default ``True``). + - ``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). diff --git a/physicsnemo/cfd/evaluation/models/common_wrapper_utils/geotransolver_runtime.py b/physicsnemo/cfd/evaluation/models/common_wrapper_utils/geotransolver_runtime.py index 953e4e3..2d79f11 100644 --- a/physicsnemo/cfd/evaluation/models/common_wrapper_utils/geotransolver_runtime.py +++ b/physicsnemo/cfd/evaluation/models/common_wrapper_utils/geotransolver_runtime.py @@ -64,7 +64,7 @@ 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_checkpoint + from physicsnemo.utils import load_model_weights _PHYSICSNEMO_AVAILABLE = True except ImportError: @@ -115,7 +115,7 @@ def geotransolver_available() -> bool: - """True when physicsnemo (GeoTransolver, TransolverDataPipe, load_checkpoint) is importable.""" + """True when physicsnemo (GeoTransolver, TransolverDataPipe, load_model_weights) is importable.""" return _PHYSICSNEMO_AVAILABLE @@ -209,19 +209,23 @@ def parse_runtime_kwargs(kw: dict[str, Any], device: str) -> GeoTransolverRuntim def ensure_distributed_initialized() -> None: - """Initialize ``DistributedManager`` once (``load_checkpoint`` relies on it).""" + """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 file name**. + """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). ``physicsnemo.load_checkpoint`` - itself takes the parent directory plus the epoch, which are both derived here. + 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: @@ -240,6 +244,11 @@ def resolve_checkpoint_file(checkpoint_path: str) -> tuple[Path, int]: 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 @@ -257,9 +266,13 @@ def build_geotransolver_backbone( 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 - ``load_checkpoint`` 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. + ``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 @@ -268,15 +281,16 @@ def build_geotransolver_backbone( ) if concrete_dropout: model_kw["concrete_dropout"] = True - checkpoint_dir, epoch = resolve_checkpoint_file(checkpoint_path) + # 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_checkpoint( - path=str(checkpoint_dir), models=model, epoch=epoch, device=dev - ) + load_model_weights(model, checkpoint_path, device=dev) model = model.to(dev) model.eval() return model diff --git a/physicsnemo/cfd/evaluation/models/model_registry.py b/physicsnemo/cfd/evaluation/models/model_registry.py index c4d1633..9bba0f7 100644 --- a/physicsnemo/cfd/evaluation/models/model_registry.py +++ b/physicsnemo/cfd/evaluation/models/model_registry.py @@ -17,7 +17,7 @@ """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, @@ -106,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, @@ -139,21 +154,24 @@ def decode_distribution( def predict_ensemble( self, model_input: ModelInput, n: int - ) -> Optional[list[RawOutput]]: - """Optional fast-path for ``UQ_METHOD="sampling"`` wrappers. - - Return ``n`` raw outputs (stochastic passes or ensemble members) in one call, letting - the engine skip its per-pass Python loop. **Only implement this when all ``n`` raw - outputs fit simultaneously in the compute device's memory (typically GPU), since the - returned list holds them all at once** — for full-mesh CFD fields that is usually ``n×`` - the single-pass footprint and will OOM for large ``n``. - - Return ``None`` (the default, used by all shipped sampling wrappers) to have the engine - instead call :meth:`predict` ``n`` times and stream the passes through Welford - aggregation (:func:`~physicsnemo.cfd.evaluation.benchmarks.uq_inference.run_sampling_inference`): - only **one** raw output is resident on the GPU at a time, and the running mean/variance - accumulators are a handful of host (CPU) arrays of one field's size — i.e. memory is - O(1) in ``n``. Prefer this default unless you have measured that the batched path fits. + ) -> 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 diff --git a/physicsnemo/cfd/evaluation/models/wrappers/ensemble_drivaerstar/wrapper.py b/physicsnemo/cfd/evaluation/models/wrappers/ensemble_drivaerstar/wrapper.py index 10d3130..8fec309 100644 --- a/physicsnemo/cfd/evaluation/models/wrappers/ensemble_drivaerstar/wrapper.py +++ b/physicsnemo/cfd/evaluation/models/wrappers/ensemble_drivaerstar/wrapper.py @@ -33,13 +33,15 @@ 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 every member 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 +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. Because ``predict_ensemble`` yields exactly the member count, -``run.uq.num_samples`` is ignored for this row. +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 @@ -48,9 +50,11 @@ 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 checkpoint file whose name encodes an - epoch (``checkpoint.0..pt`` / ``.0..mdlus``). List members from a single run - (snapshot ensemble) or from separate runs (true deep ensemble) — same code path either way. + 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. @@ -148,7 +152,7 @@ def load( if not geotransolver_available(): raise RuntimeError( "GeoTransolverEnsembleDrivAerStarWrapper requires physicsnemo (GeoTransolver, " - "TransolverDataPipe, load_checkpoint)." + "TransolverDataPipe, load_model_weights)." ) kw = dict(kwargs) member_checkpoints = kw.pop("member_checkpoints", None) @@ -249,21 +253,32 @@ def _forward_member( def predict_ensemble( self, model_input: ModelInput, n: int ) -> Optional[Iterator[RawOutput]]: - """Yield one raw output per member (lazy generator), sharing one per-case permutation. + """Yield up to ``n`` member raw outputs (lazy generator), sharing one per-case permutation. - ``n`` (``run.uq.num_samples``) is intentionally ignored: the ensemble has a fixed number of - members and the engine iterates over exactly what is yielded. 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. + ``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 + self._forward_member(model, model_input, perm) for model in self._models[:k] ) def predict(self, model_input: ModelInput) -> RawOutput: diff --git a/physicsnemo/cfd/evaluation/models/wrappers/geotransolver_drivaerstar/wrapper.py b/physicsnemo/cfd/evaluation/models/wrappers/geotransolver_drivaerstar/wrapper.py index b38e34e..89145ff 100644 --- a/physicsnemo/cfd/evaluation/models/wrappers/geotransolver_drivaerstar/wrapper.py +++ b/physicsnemo/cfd/evaluation/models/wrappers/geotransolver_drivaerstar/wrapper.py @@ -112,7 +112,7 @@ def load( if not geotransolver_available(): raise RuntimeError( "GeoTransolver wrapper requires physicsnemo (GeoTransolver, " - "TransolverDataPipe, load_checkpoint)." + "TransolverDataPipe, load_model_weights)." ) kw = dict(kwargs) self._inference_mode = coerce_inference_domain_or_default( diff --git a/physicsnemo/cfd/evaluation/models/wrappers/geotransolver_gp/wrapper.py b/physicsnemo/cfd/evaluation/models/wrappers/geotransolver_gp/wrapper.py index 3c6fdb0..38f28de 100644 --- a/physicsnemo/cfd/evaluation/models/wrappers/geotransolver_gp/wrapper.py +++ b/physicsnemo/cfd/evaluation/models/wrappers/geotransolver_gp/wrapper.py @@ -44,11 +44,12 @@ (``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``. Its directory / - epoch drive the load, so the head file named here is exactly the one read (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`` at the - same epoch. Either way the exact head file it loads is logged. +- ``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: @@ -105,7 +106,7 @@ try: from physicsnemo.experimental.uq import FieldGPHead - from physicsnemo.utils import load_checkpoint + from physicsnemo.utils import load_model_weights _GP_AVAILABLE = True except ImportError: @@ -160,7 +161,6 @@ def __init__(self) -> None: self._cfg: GeoTransolverRuntimeConfig = GeoTransolverRuntimeConfig() self._head: Optional[FieldGPHead] = None self._gp_kw: dict[str, Any] = {} - self._checkpoint_dir: Optional[str] = None self._checkpoint_epoch: Optional[int] = None self._head_checkpoint: Optional[str] = None self._gp_chunk_size: int = 51200 @@ -248,20 +248,24 @@ def load( inference_mode="surface", ) - # The backbone + head are loaded together by ``load_checkpoint(dir, epoch)``. Derive that - # dir + epoch from ``gp_head_checkpoint`` when given (so the head file the user names is the - # one read), else from the backbone ``checkpoint``. No cross-validation between the two — - # passing matching checkpoints is the caller's responsibility. + # 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: - ckpt_dir, ckpt_epoch = resolve_checkpoint_file(str(head_ckpt_arg)) + # 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, ckpt_epoch = resolve_checkpoint_file(checkpoint_path) + ckpt_dir, self._checkpoint_epoch = resolve_checkpoint_file(checkpoint_path) self._head_checkpoint = str( - Path(ckpt_dir) / f"FieldGPHead.0.{ckpt_epoch}.pt" + Path(ckpt_dir) / f"FieldGPHead.0.{self._checkpoint_epoch}.pt" ) - self._checkpoint_dir = str(ckpt_dir) - self._checkpoint_epoch = ckpt_epoch + 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}", @@ -289,10 +293,10 @@ 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 head checkpoint's directory + epoch (which resolves to the - exact ``FieldGPHead.0..pt`` file named by ``gp_head_checkpoint``). Passing only the - head to ``load_checkpoint`` avoids reloading — or silently overwriting — the backbone with a - different one that happens to sit in the head's directory. + 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) @@ -302,18 +306,14 @@ def _build_and_load_head(self, batch: dict) -> None: **self._gp_kw, ).to(dev) with trusted_torch_load_context(): - loaded_epoch = load_checkpoint( - path=self._checkpoint_dir, - models=[self._head], - device=dev, - epoch=self._checkpoint_epoch, - ) + 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 (epoch {loaded_epoch}); feature_dim={feature_dim}, " + 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', '?')}", ) diff --git a/physicsnemo/cfd/evaluation/models/wrappers/mc_dropout/wrapper.py b/physicsnemo/cfd/evaluation/models/wrappers/mc_dropout/wrapper.py index 1501d07..a955a4c 100644 --- a/physicsnemo/cfd/evaluation/models/wrappers/mc_dropout/wrapper.py +++ b/physicsnemo/cfd/evaluation/models/wrappers/mc_dropout/wrapper.py @@ -147,7 +147,7 @@ def load( if not geotransolver_available(): raise RuntimeError( "GeoTransolverMCDropoutDrivAerStarWrapper requires physicsnemo (GeoTransolver, " - "TransolverDataPipe, load_checkpoint)." + "TransolverDataPipe, load_model_weights)." ) if not _CONCRETE_DROPOUT_AVAILABLE: raise RuntimeError( @@ -276,6 +276,29 @@ def predict(self, model_input: ModelInput) -> RawOutput: 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, diff --git a/skills/physicsnemo-cfd-create-model-wrapper/SKILL.md b/skills/physicsnemo-cfd-create-model-wrapper/SKILL.md index 65a5774..694f017 100644 --- a/skills/physicsnemo-cfd-create-model-wrapper/SKILL.md +++ b/skills/physicsnemo-cfd-create-model-wrapper/SKILL.md @@ -98,11 +98,11 @@ 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 one extra hook — either -`decode_distribution` (analytic) or `predict_ensemble` (sampling, -optional). This is fully additive: deterministic wrappers leave the -defaults (`SUPPORTS_UQ=False`) and are unaffected. See the -"Uncertainty quantification" section below. +(`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 @@ -357,10 +357,16 @@ distribution is produced*: 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). Optionally implement - **`predict_ensemble(model_input, n) -> list[RawOutput] | None`** as a - fast path (ensembles return one output per member and ignore `n`; - return `None` to fall back to N× `predict`). + 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). 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 505faa6..2e26dc0 100644 --- a/skills/physicsnemo-cfd-create-model-wrapper/references/example_wrapper.py +++ b/skills/physicsnemo-cfd-create-model-wrapper/references/example_wrapper.py @@ -38,7 +38,7 @@ from __future__ import annotations -from typing import Any, ClassVar, Optional +from typing import Any, ClassVar, Iterable, Optional import numpy as np import torch @@ -384,9 +384,14 @@ class ExampleSamplingUQWrapper(CFDModel): 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); - * OPTIONAL fast path: implement ``predict_ensemble(model_input, n)`` to return all - passes/members in one call (an ensemble returns one output per member and ignores - ``n``). No ``decode_distribution`` is needed — the engine builds the distribution. + * 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. @@ -443,18 +448,42 @@ def predict(self, model_input: torch.Tensor) -> dict[str, torch.Tensor]: } 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[list[dict[str, torch.Tensor]]]: - """Optional fast path: return every pass/member in one call. - - For a deep ensemble, return one raw output per member (``n`` is ignored). For - multi-pass single-model methods you may batch ``n`` passes here. Return ``None`` to - let the engine fall back to calling ``predict`` ``n`` times. + ) -> 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 - return [self.predict(model_input) for _ in self._models] + 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, 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_uq_inference.py b/test/ci_tests/test_uq_inference.py index 44d667b..8a132b1 100644 --- a/test/ci_tests/test_uq_inference.py +++ b/test/ci_tests/test_uq_inference.py @@ -257,3 +257,39 @@ def test_run_sampling_inference_consumes_generator_streaming() -> None: 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_wrappers.py b/test/ci_tests/test_uq_wrappers.py index c184a8b..f292b70 100644 --- a/test/ci_tests/test_uq_wrappers.py +++ b/test/ci_tests/test_uq_wrappers.py @@ -27,24 +27,33 @@ make_forward_permutation, resolve_checkpoint_file, ) +from physicsnemo.cfd.evaluation.models.model_registry import CFDModel # -------------------------------------------------------------------------------------------- -# Checkpoint routing: a specific file (epoch-in-name) -> (dir, epoch); directory / bad name reject +# 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: - """A specific checkpoint file name resolves to ``(directory, epoch)``.""" + """An existing checkpoint file name resolves to ``(directory, epoch)``.""" ckpt = tmp_path / "GeoTransolver.0.30.mdlus" - ckpt.write_bytes(b"") # existence not required, but harmless + ckpt.write_bytes(b"") directory, epoch = resolve_checkpoint_file(str(ckpt)) assert directory == tmp_path and epoch == 30 - # Nonexistent path still resolves purely from the file name (no silent latest-epoch fallback). - d2, e2 = resolve_checkpoint_file(str(tmp_path / "checkpoint.0.100.pt")) + 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): @@ -54,7 +63,83 @@ def test_resolve_checkpoint_file_rejects_directory_and_epochless(tmp_path) -> No with pytest.raises(ValueError): resolve_checkpoint_file( str(tmp_path / "GeoTransolver.mdlus") - ) # no epoch in name + ) # 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 # -------------------------------------------------------------------------------------------- diff --git a/workflows/benchmarking/conf/config_uq_surface.yaml b/workflows/benchmarking/conf/config_uq_surface.yaml index de2f981..775a289 100644 --- a/workflows/benchmarking/conf/config_uq_surface.yaml +++ b/workflows/benchmarking/conf/config_uq_surface.yaml @@ -82,8 +82,11 @@ run: # 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's checkpoints. - baseline_ckpt: &baseline_ckpt /path/to/run/checkpoints/checkpoint.0.100.pt + # 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 @@ -139,12 +142,13 @@ benchmark: 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/checkpoint.0.96.pt - - /path/to/run/checkpoints/checkpoint.0.97.pt - - /path/to/run/checkpoints/checkpoint.0.98.pt - - /path/to/run/checkpoints/checkpoint.0.99.pt - - /path/to/run/checkpoints/checkpoint.0.100.pt + - /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 @@ -152,7 +156,10 @@ benchmark: - 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) - # targets DrivAerML-trained checkpoints; set it to match how your checkpoints were trained. + # 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: