From 466c22cc9fa27a8d176049e76c863706b9b12f3e Mon Sep 17 00:00:00 2001 From: Kaustubh Tangsali Date: Thu, 9 Jul 2026 22:26:58 +0000 Subject: [PATCH 01/14] 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/14] 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/14] 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/14] 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/14] 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/14] 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/14] 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/14] 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/14] 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/14] 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: From b9d912e040314e8c492cd745d4c9403f385ec190 Mon Sep 17 00:00:00 2001 From: Kaustubh Tangsali Date: Wed, 29 Jul 2026 05:15:05 +0000 Subject: [PATCH 11/14] add files after latest gp model scoring --- .../cfd/evaluation/benchmarks/engine.py | 25 +++++++++++++------ physicsnemo/cfd/evaluation/config.py | 11 ++++++++ .../wrappers/geotransolver_gp/wrapper.py | 17 +++++++++++++ 3 files changed, 45 insertions(+), 8 deletions(-) diff --git a/physicsnemo/cfd/evaluation/benchmarks/engine.py b/physicsnemo/cfd/evaluation/benchmarks/engine.py index 901dfed..81b8415 100644 --- a/physicsnemo/cfd/evaluation/benchmarks/engine.py +++ b/physicsnemo/cfd/evaluation/benchmarks/engine.py @@ -364,10 +364,11 @@ def _save_inference_mesh_if_requested( # 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 "" + model_tok = _sanitize_path_token(model_config.display_name) stem = ( - f"inference_{model_config.name}_{label_tok}_{case_id}" + f"inference_{model_tok}_{label_tok}_{case_id}" if label_tok - else f"inference_{model_config.name}_{case_id}" + else f"inference_{model_tok}_{case_id}" ) out_path = Path(output_dir) / f"{stem}{ext}" log_dataset( @@ -578,6 +579,12 @@ def _run_single( # 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 + # Wrapper/assets/kwargs are resolved from ``model_config.name``; the result rows and per-case + # metric-cache blobs are keyed by ``display_name`` (== label or name) so the same wrapper can + # score several checkpoints/heads as distinct rows (e.g. two ``geotransolver_gp_surface`` rows + # labeled ``gp_due`` / ``gp_hetnoise``). The cache *fingerprint* stays on ``name`` + checkpoint + + # kwargs, so unlabeled rows keep byte-identical fingerprints (existing caches remain valid). + m_label = model_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: @@ -592,7 +599,7 @@ def _run_single( ) return ( { - "model": model_config.name, + "model": m_label, "dataset": ds_label, "skipped": True, "skip_reason": reason, @@ -630,7 +637,7 @@ def _run_single( if not cases: return ( { - "model": model_config.name, + "model": m_label, "dataset": ds_label, "cases": [], "metrics": {}, @@ -699,7 +706,7 @@ def _run_single( if ( blob is not None and blob.get("fingerprint") == fingerprint - and blob.get("model") == model_config.name + and blob.get("model") == m_label and blob.get("dataset") == dataset_config.name and blob.get("case_id") == cid ): @@ -896,7 +903,8 @@ def _run_single( # 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}" + _mlbl = _sanitize_path_token(m_label) + cmp_p = sub / f"{_mlbl}_{_lbl}_{cid}_comparison{ext}" try: comparison_mesh.save(str(cmp_p)) row["comparison_mesh_path"] = str(cmp_p.resolve()) @@ -915,7 +923,7 @@ def _run_single( write_metrics_cache( cache_file, fingerprint=fingerprint, - model=model_config.name, + model=m_label, dataset=dataset_config.name, case_id=cid, case_metrics=case_metrics, @@ -952,7 +960,7 @@ def _run_single( return ( { - "model": model_config.name, + "model": m_label, "dataset": ds_label, "cases": cases, "metrics": metrics_summary, @@ -1256,6 +1264,7 @@ def _config_to_dict(c: Config) -> dict: }, "model": { "name": c.model.name, + "label": c.model.label, "checkpoint": c.model.checkpoint, "stats_path": c.model.stats_path, "kwargs": c.model.kwargs, diff --git a/physicsnemo/cfd/evaluation/config.py b/physicsnemo/cfd/evaluation/config.py index e137ad7..7e366e3 100644 --- a/physicsnemo/cfd/evaluation/config.py +++ b/physicsnemo/cfd/evaluation/config.py @@ -161,6 +161,17 @@ class ModelConfig: # Optional surface/volume routing; omit to derive from wrappers (merged ``model.kwargs`` and # optional :meth:`~physicsnemo.cfd.evaluation.models.model_registry.CFDModel.inference_domain_from_kwargs`). inference_domain: str | None = None + #: Optional display label for benchmark result rows. Defaults to :attr:`name`. Set this to give + #: the SAME wrapper multiple distinct rows in one matrix run (e.g. two ``geotransolver_gp_surface`` + #: rows scoring different GP checkpoints/heads labeled ``gp_due`` / ``gp_hetnoise``). The wrapper + #: is always resolved from :attr:`name`; only the reported/cached model key changes. (Mirrors + #: :attr:`DatasetConfig.label`.) + label: str | None = None + + @property + def display_name(self) -> str: + """Label used as the ``model`` key in results/cache (falls back to the wrapper :attr:`name`).""" + return self.label or self.name def merged_kwargs_for_load(self) -> dict[str, Any]: """``model.kwargs`` plus ``inference_domain`` so wrappers can branch surface vs volume.""" diff --git a/physicsnemo/cfd/evaluation/models/wrappers/geotransolver_gp/wrapper.py b/physicsnemo/cfd/evaluation/models/wrappers/geotransolver_gp/wrapper.py index 38f28de..dd8facf 100644 --- a/physicsnemo/cfd/evaluation/models/wrappers/geotransolver_gp/wrapper.py +++ b/physicsnemo/cfd/evaluation/models/wrappers/geotransolver_gp/wrapper.py @@ -61,6 +61,11 @@ 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_noise_mlp_hidden`` (list, default ``None``) / ``gp_noise_std_range`` (default ``(1e-3, 10.0)``): + heteroscedastic (input-dependent) observation-noise head. A non-empty list (e.g. ``[64, 64]``) + makes the aleatoric std vary per point instead of one learned scalar per channel; it adds + ``noise_head.*`` parameters that must match training. ``None``/``[]`` keeps the homoscedastic + likelihood, so existing GP checkpoints are unaffected. - ``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. @@ -223,6 +228,18 @@ def load( # 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)), + # Input-dependent (heteroscedastic) observation-noise MLP. A non-empty list turns on a + # per-point noise head so the aleatoric std varies across the surface (vs one learned + # scalar per channel); this adds ``noise_head.*`` parameters that MUST match training or + # the ``FieldGPHead`` state dict will not reconstruct. ``None``/``[]`` keeps the + # homoscedastic likelihood, so existing GP checkpoints are unaffected. ``noise_std_range`` + # is the hard clamp on the per-point noise std (matches the training default when omitted). + "noise_mlp_hidden": ( + list(kw.pop("gp_noise_mlp_hidden")) + if kw.get("gp_noise_mlp_hidden") + else kw.pop("gp_noise_mlp_hidden", None) + ), + "noise_std_range": tuple(kw.pop("gp_noise_std_range", (1e-3, 10.0))), } self._cfg = parse_runtime_kwargs(kw, device) From 55ca05f76cb3a7c426a47aed7cc9dddf5a5d2591 Mon Sep 17 00:00:00 2001 From: Kaustubh Tangsali Date: Wed, 29 Jul 2026 05:40:11 +0000 Subject: [PATCH 12/14] Track the FieldVariationalGPHead rename and drop the DUE knobs physicsnemo renamed FieldGPHead to FieldVariationalGPHead and removed the DUE-style spectral-normalised DKL extractor. The import alias kept this wrapper importable but not constructible: it passed spectral_norm_coeff and dkl_residual unconditionally, so every GP benchmark row died with a TypeError. Rather than silently dropping those two kwargs, reject them with a clear error. Leftover model kwargs are forwarded to the datapipe splitter, which would swallow them, so a config still asking for a spectral-normalised extractor would otherwise build a plain MLP whose parameter layout does not match its checkpoint -- failing later and less obviously. The head-checkpoint fallback now tries FieldVariationalGPHead.0..pt and then the older FieldGPHead.0..pt, since save_checkpoint derives the file stem from the class name and existing runs wrote the old one. Verified the three GP checkpoints scored by the multiclass surface benchmark still load strict=True. config_uq_surface.yaml's GP row loses its DUE settings and now mirrors the settled training recipe (plain DKL MLP + heteroscedastic noise head). Its comments also spell out that stats_path must be the DrivAerStar normalization: those fields are in raw physical units, and the nondimensional stats that used to sit at the shared path belong to the DrivAer AWS runs, so confusing the two mis-scales pressure by ~1600x without failing. --- .../wrappers/geotransolver_gp/wrapper.py | 106 ++++++++++-------- .../benchmarking/conf/config_uq_surface.yaml | 28 +++-- 2 files changed, 76 insertions(+), 58 deletions(-) diff --git a/physicsnemo/cfd/evaluation/models/wrappers/geotransolver_gp/wrapper.py b/physicsnemo/cfd/evaluation/models/wrappers/geotransolver_gp/wrapper.py index dd8facf..20d7e4a 100644 --- a/physicsnemo/cfd/evaluation/models/wrappers/geotransolver_gp/wrapper.py +++ b/physicsnemo/cfd/evaluation/models/wrappers/geotransolver_gp/wrapper.py @@ -21,10 +21,10 @@ 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``). +head (:class:`physicsnemo.experimental.uq.FieldVariationalGPHead`): 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 from the ``transformer_models`` +field-GP recipe. It targets ``transformer_models`` (physical-target) checkpoints, so predictions are re-standardized only — no dynamic-pressure re-dimensionalization (:attr:`REDIMENSIONALIZE_OUTPUTS` = ``False``). @@ -37,30 +37,28 @@ Surface only (the GP head predicts 4 surface tasks: pressure + 3 wall-shear components). -This wrapper loads **two** checkpoints — the GeoTransolver backbone and the ``FieldGPHead`` — and -both are named explicitly: +This wrapper loads **two** checkpoints — the GeoTransolver backbone and the GP head — and both are +named explicitly: - ``checkpoint`` must point at a **specific backbone file** whose name encodes the epoch (``GeoTransolver.0..mdlus``); the epoch is parsed from that name. Pointing at a directory (or an epoch-less name) is rejected — this removes the old ``checkpoint_epoch`` kwarg and the risk of the backbone drifting to a "latest" epoch. -- ``gp_head_checkpoint`` should point at the matching ``FieldGPHead.0..pt``. The **exact - file named here is loaded** (via :func:`physicsnemo.utils.load_model_weights`), so a missing or - mistyped path raises rather than silently leaving the head randomly initialized (no cross- - validation against the backbone — passing matching checkpoints is the caller's responsibility). - It is **optional**: when omitted the wrapper falls back to the backbone's sibling - ``FieldGPHead.0..pt`` (which must exist). Either way the exact head file it loads is logged. +- ``gp_head_checkpoint`` should point at the matching head file. 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 head file, accepting either + ``FieldVariationalGPHead.0..pt`` (written by current runs) or the older + ``FieldGPHead.0..pt`` — the head class was renamed and ``save_checkpoint`` derives the file + stem from the class name. 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_head_checkpoint`` (path): the head checkpoint to load (see above). +- ``gp_feature_norm`` (``"none"`` | ``"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_noise_mlp_hidden`` (list, default ``None``) / ``gp_noise_std_range`` (default ``(1e-3, 10.0)``): heteroscedastic (input-dependent) observation-noise head. A non-empty list (e.g. ``[64, 64]``) makes the aleatoric std vary per point instead of one learned scalar per channel; it adds @@ -110,7 +108,7 @@ ) try: - from physicsnemo.experimental.uq import FieldGPHead + from physicsnemo.experimental.uq import FieldVariationalGPHead from physicsnemo.utils import load_model_weights _GP_AVAILABLE = True @@ -120,9 +118,13 @@ #: Surface field channels predicted by the GP head: pressure, wall-shear (x, y, z). NUM_SURFACE_TASKS = 4 +#: Head-checkpoint stems to try when falling back to the backbone's sibling file. ``save_checkpoint`` +#: derives the stem from the class name, so runs predating the rename wrote ``FieldGPHead.0.*.pt``. +_HEAD_CKPT_STEMS = ("FieldVariationalGPHead", "FieldGPHead") + class GeoTransolverGPDrivAerStarWrapper(CFDModel): - """GeoTransolver backbone + ``FieldGPHead`` for analytic per-point surface UQ. + """GeoTransolver backbone + ``FieldVariationalGPHead`` for analytic per-point surface UQ. Trained on ``transformer_models`` (physical-target) checkpoints, so predictions are re-standardized only (no dynamic-pressure re-dimensionalization): @@ -164,7 +166,7 @@ def __init__(self) -> None: self._datapipe_geometry_effective: Optional[int] = None self._surface_factors: Any = None self._cfg: GeoTransolverRuntimeConfig = GeoTransolverRuntimeConfig() - self._head: Optional[FieldGPHead] = None + self._head: Optional[FieldVariationalGPHead] = None self._gp_kw: dict[str, Any] = {} self._checkpoint_epoch: Optional[int] = None self._head_checkpoint: Optional[str] = None @@ -181,13 +183,13 @@ def load( ) -> "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). + The backbone is constructed via the shared runtime helpers; the head 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." + "physicsnemo.experimental.uq.FieldVariationalGPHead." ) kw = dict(kwargs) # Surface-only; validate and force the routing hint before building the datapipe. @@ -197,10 +199,23 @@ 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 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. + # Explicit GP-head checkpoint file (optional). When omitted we fall back to a sibling of the + # backbone, trying each stem in ``_HEAD_CKPT_STEMS``. No cross-checking against the backbone + # — passing a matching pair is the caller's responsibility. head_ckpt_arg = kw.pop("gp_head_checkpoint", None) + # The DUE-style bi-Lipschitz DKL extractor was dropped from the head, so these keys no + # longer describe anything buildable. Reject them explicitly: leftover model kwargs are + # forwarded to the datapipe splitter, which would swallow them, and a config asking for a + # spectral-normalised extractor would then silently build a plain MLP whose parameter + # layout does not match its checkpoint. + for dead_key in ("gp_spectral_norm_coeff", "gp_dkl_residual"): + if kw.pop(dead_key, None) is not None: + raise ValueError( + f"{dead_key!r} is no longer supported: the DUE-style spectral-normalised " + "DKL extractor has been removed from FieldVariationalGPHead. Drop it from " + "model.kwargs. Checkpoints trained with gp_spectral_norm_coeff > 0 cannot be " + "reconstructed and must be retrained with the plain DKL MLP." + ) self._gp_kw = { "num_tasks": int(kw.pop("num_tasks", NUM_SURFACE_TASKS)), "n_inducing": int(kw.pop("gp_n_inducing", 256)), @@ -221,19 +236,12 @@ 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)), # Input-dependent (heteroscedastic) observation-noise MLP. A non-empty list turns on a # per-point noise head so the aleatoric std varies across the surface (vs one learned # scalar per channel); this adds ``noise_head.*`` parameters that MUST match training or - # the ``FieldGPHead`` state dict will not reconstruct. ``None``/``[]`` keeps the - # homoscedastic likelihood, so existing GP checkpoints are unaffected. ``noise_std_range`` - # is the hard clamp on the per-point noise std (matches the training default when omitted). + # the head state dict will not reconstruct. ``None``/``[]`` keeps the homoscedastic + # likelihood, so existing GP checkpoints are unaffected. ``noise_std_range`` is the hard + # clamp on the per-point noise std (matches the training default when omitted). "noise_mlp_hidden": ( list(kw.pop("gp_noise_mlp_hidden")) if kw.get("gp_noise_mlp_hidden") @@ -266,7 +274,7 @@ def load( ) # 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``. + # user names is the one read), else from the first matching backbone sibling. # 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: @@ -275,14 +283,18 @@ def load( self._head_checkpoint = str(Path(head_ckpt_arg)) else: ckpt_dir, self._checkpoint_epoch = resolve_checkpoint_file(checkpoint_path) - self._head_checkpoint = str( - Path(ckpt_dir) / f"FieldGPHead.0.{self._checkpoint_epoch}.pt" - ) - if not Path(self._head_checkpoint).is_file(): + candidates = [ + Path(ckpt_dir) / f"{stem}.0.{self._checkpoint_epoch}.pt" + for stem in _HEAD_CKPT_STEMS + ] + found = next((c for c in candidates if c.is_file()), None) + if found is None: raise FileNotFoundError( - "GP head checkpoint not found next to the backbone: " - f"{self._head_checkpoint!r}. Pass model.kwargs.gp_head_checkpoint explicitly." + "GP head checkpoint not found next to the backbone; tried " + f"{[str(c) for c in candidates]}. Pass model.kwargs.gp_head_checkpoint " + "explicitly." ) + self._head_checkpoint = str(found) log_inference( "geotransolver_gp", f"GP checkpoints -> backbone: {checkpoint_path} | head: {self._head_checkpoint}", @@ -310,14 +322,14 @@ def _build_and_load_head(self, batch: dict) -> None: """Probe the backbone feature dim on ``batch``, build the GP head, load ONLY its weights. The backbone was already loaded from its own ``checkpoint`` in :meth:`load`; here we load - just the ``FieldGPHead`` from the EXACT ``self._head_checkpoint`` file (the one named by + just the GP head from the EXACT ``self._head_checkpoint`` file (the one named by ``gp_head_checkpoint``, or the validated backbone sibling). Loading only the head from its own file avoids reloading — or silently overwriting — the backbone, and a missing file raises rather than leaving the head randomly initialized. """ dev = torch.device(self._cfg.device) feature_dim = self._probe_feature_dim(batch) - self._head = FieldGPHead( + self._head = FieldVariationalGPHead( input_dim=feature_dim, n_train=1, # unused at inference **self._gp_kw, @@ -329,7 +341,7 @@ def _build_and_load_head(self, batch: dict) -> None: self._head.likelihood.eval() log_inference( "geotransolver_gp", - f"Loaded FieldGPHead from {self._head_checkpoint} " + f"Loaded FieldVariationalGPHead 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/workflows/benchmarking/conf/config_uq_surface.yaml b/workflows/benchmarking/conf/config_uq_surface.yaml index 775a289..df95629 100644 --- a/workflows/benchmarking/conf/config_uq_surface.yaml +++ b/workflows/benchmarking/conf/config_uq_surface.yaml @@ -100,26 +100,32 @@ benchmark: stats_path: *surface_stats kwargs: { batch_resolution: 60000, geometry_sampling: 300000 } - # (2) Analytic GP field head. This wrapper loads TWO checkpoints, both named - # explicitly: `checkpoint` is the GeoTransolver backbone (epoch parsed from the file name) and - # `gp_head_checkpoint` is the matching FieldGPHead (read as given; pass matching files). The GP - # hyperparameters here MUST match the training run so the FieldGPHead reconstructs before the - # state dict loads. - + # (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 FieldVariationalGPHead (read as given; pass matching + # files). The GP hyperparameters here MUST match the training run so the head reconstructs + # before the state dict loads -- these mirror the training config's defaults + # (`transformer_models/src/conf/geotransolver_surface_field_gp.yaml`). + # + # `stats_path` must be the SAME normalization the checkpoint was trained with. DrivAerStar + # surface fields are in raw physical units, so field-GP runs use + # `transformer_models/src/normalization/drivaerstar/surface_fields_normalization.npz`, NOT the + # nondimensional stats in `transformer_models/src/` (which belong to the DrivAer AWS runs). - name: geotransolver_gp_surface inference_domain: surface - checkpoint: /path/to/run_field_gp_due/checkpoints_field_gp/GeoTransolver.0.100.mdlus + checkpoint: /path/to/run_field_gp/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_head_checkpoint: /path/to/run_field_gp/checkpoints_field_gp/FieldVariationalGPHead.0.100.pt gp_feature_norm: "l2_radial" gp_n_inducing: 1024 - gp_mlp_hidden: [128, 128, 16] - gp_spectral_norm_coeff: 0.95 # > 0 -> DUE bi-Lipschitz extractor - gp_dkl_residual: true + gp_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] + # Heteroscedastic noise head -- must match training (adds `noise_head.*` parameters). + gp_noise_mlp_hidden: [64, 64] + gp_noise_std_range: [0.01, 10.0] num_tasks: 4 gp_inference_chunk_size: 51200 geometry_sampling: 300000 From 5f732a5b175c1d5c754b11ee2395f57765411378 Mon Sep 17 00:00:00 2001 From: Kaustubh Tangsali Date: Thu, 30 Jul 2026 17:59:02 +0000 Subject: [PATCH 13/14] Look for one head-checkpoint stem in the GP wrapper The pre-rename FieldGPHead spelling is gone from physicsnemo, and every GP entry in the benchmark configs names its head checkpoint explicitly, so the exact file is read and the sibling search never runs for them. Accepting both stems only made the fallback and its error harder to follow; the not-found message now names the single path it looked for. --- .../wrappers/geotransolver_gp/wrapper.py | 34 ++++++++----------- 1 file changed, 15 insertions(+), 19 deletions(-) diff --git a/physicsnemo/cfd/evaluation/models/wrappers/geotransolver_gp/wrapper.py b/physicsnemo/cfd/evaluation/models/wrappers/geotransolver_gp/wrapper.py index 20d7e4a..4abdd7a 100644 --- a/physicsnemo/cfd/evaluation/models/wrappers/geotransolver_gp/wrapper.py +++ b/physicsnemo/cfd/evaluation/models/wrappers/geotransolver_gp/wrapper.py @@ -48,10 +48,9 @@ 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 head file, accepting either - ``FieldVariationalGPHead.0..pt`` (written by current runs) or the older - ``FieldGPHead.0..pt`` — the head class was renamed and ``save_checkpoint`` derives the file - stem from the class name. Either way the exact head file it loads is logged. + omitted the wrapper falls back to the backbone's sibling ``FieldVariationalGPHead.0..pt`` + (``save_checkpoint`` derives the file stem from the class name). Either way the exact head file it + loads is logged. Model kwargs (``model.kwargs`` in config), matching the trained checkpoint's GP settings: @@ -118,9 +117,9 @@ #: Surface field channels predicted by the GP head: pressure, wall-shear (x, y, z). NUM_SURFACE_TASKS = 4 -#: Head-checkpoint stems to try when falling back to the backbone's sibling file. ``save_checkpoint`` -#: derives the stem from the class name, so runs predating the rename wrote ``FieldGPHead.0.*.pt``. -_HEAD_CKPT_STEMS = ("FieldVariationalGPHead", "FieldGPHead") +#: Head-checkpoint stem used when falling back to the backbone's sibling file. ``save_checkpoint`` +#: derives the stem from the class name. +_HEAD_CKPT_STEM = "FieldVariationalGPHead" class GeoTransolverGPDrivAerStarWrapper(CFDModel): @@ -199,8 +198,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 file (optional). When omitted we fall back to a sibling of the - # backbone, trying each stem in ``_HEAD_CKPT_STEMS``. No cross-checking against the backbone + # Explicit GP-head checkpoint file (optional). When omitted we fall back to the + # ``_HEAD_CKPT_STEM`` sibling of the backbone. No cross-checking against the backbone # — passing a matching pair is the caller's responsibility. head_ckpt_arg = kw.pop("gp_head_checkpoint", None) # The DUE-style bi-Lipschitz DKL extractor was dropped from the head, so these keys no @@ -283,18 +282,15 @@ def load( self._head_checkpoint = str(Path(head_ckpt_arg)) else: ckpt_dir, self._checkpoint_epoch = resolve_checkpoint_file(checkpoint_path) - candidates = [ - Path(ckpt_dir) / f"{stem}.0.{self._checkpoint_epoch}.pt" - for stem in _HEAD_CKPT_STEMS - ] - found = next((c for c in candidates if c.is_file()), None) - if found is None: + sibling = ( + Path(ckpt_dir) / f"{_HEAD_CKPT_STEM}.0.{self._checkpoint_epoch}.pt" + ) + if not sibling.is_file(): raise FileNotFoundError( - "GP head checkpoint not found next to the backbone; tried " - f"{[str(c) for c in candidates]}. Pass model.kwargs.gp_head_checkpoint " - "explicitly." + f"GP head checkpoint not found next to the backbone; tried {sibling}. " + "Pass model.kwargs.gp_head_checkpoint explicitly." ) - self._head_checkpoint = str(found) + self._head_checkpoint = str(sibling) log_inference( "geotransolver_gp", f"GP checkpoints -> backbone: {checkpoint_path} | head: {self._head_checkpoint}", From dfa44df4bce701e687a44d56a959abb0930dd8d1 Mon Sep 17 00:00:00 2001 From: Kaustubh Tangsali Date: Thu, 30 Jul 2026 22:56:56 +0000 Subject: [PATCH 14/14] update nomenclature from analytic to closed_form, added gpytorch install and other misc things --- CHANGELOG.md | 10 +++++----- .../cfd/evaluation/benchmarks/engine.py | 4 ++-- .../cfd/evaluation/benchmarks/uq_inference.py | 8 ++++---- physicsnemo/cfd/evaluation/config.py | 2 +- .../cfd/evaluation/metrics/builtin/uq.py | 6 +++--- .../cfd/evaluation/models/model_registry.py | 8 ++++---- .../wrappers/geotransolver_gp/wrapper.py | 8 ++++---- pyproject.toml | 8 +++++++- .../SKILL.md | 18 +++++++++--------- .../references/example_wrapper.py | 18 +++++++++--------- test/ci_tests/test_uq_inference.py | 14 +++++++++----- test/ci_tests/test_uq_metrics.py | 2 +- test/ci_tests/test_uq_wrappers.py | 6 +++--- .../benchmarking/conf/config_uq_surface.yaml | 8 ++++---- 14 files changed, 65 insertions(+), 55 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index ce15a96..c7c4fce 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -23,7 +23,7 @@ Versioning](https://semver.org/spec/v2.0.0.html). - **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 + (`"none"` | `"closed_form"` | `"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 @@ -45,22 +45,22 @@ Versioning](https://semver.org/spec/v2.0.0.html). 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), + archetype): `geotransolver_gp_surface` (closed-form 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 + scores deterministic + closed-form 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. + wrapper contract (capability flags, `FieldDistribution`, closed-form vs. sampling archetypes, the engine sampling loop, and UQ metrics/config), - with `ExampleAnalyticUQWrapper` / `ExampleSamplingUQWrapper` reference + with `ExampleClosedFormUQWrapper` / `ExampleSamplingUQWrapper` reference implementations. ### Deprecated diff --git a/physicsnemo/cfd/evaluation/benchmarks/engine.py b/physicsnemo/cfd/evaluation/benchmarks/engine.py index 81b8415..3fb62d3 100644 --- a/physicsnemo/cfd/evaluation/benchmarks/engine.py +++ b/physicsnemo/cfd/evaluation/benchmarks/engine.py @@ -755,7 +755,7 @@ def _run_single( case_cache[case_key] = case 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) + # ``run.uq.enabled`` is the master switch: when off, EVERY wrapper (sampling AND closed-form) # 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 @@ -777,7 +777,7 @@ def _run_single( case_id=cid, retain_samples=run_config.uq.retain_samples, ) - elif inference_path == "analytic": + elif inference_path == "closed_form": raw = wrapper.predict(model_input) predictions = wrapper.decode_distribution(raw, case, model_input) else: diff --git a/physicsnemo/cfd/evaluation/benchmarks/uq_inference.py b/physicsnemo/cfd/evaluation/benchmarks/uq_inference.py index 6e094dd..f394180 100644 --- a/physicsnemo/cfd/evaluation/benchmarks/uq_inference.py +++ b/physicsnemo/cfd/evaluation/benchmarks/uq_inference.py @@ -265,17 +265,17 @@ def is_uq_partial_key(key: str) -> bool: def select_inference_path( *, supports_uq: bool, uq_method: str, uq_enabled: bool ) -> str: - """Engine per-case dispatch: ``"sampling"`` | ``"analytic"`` | ``"deterministic"``. + """Engine per-case dispatch: ``"sampling"`` | ``"closed_form"`` | ``"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 + the deterministic path regardless of ``SUPPORTS_UQ`` / ``UQ_METHOD`` — so a closed-form 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" + if uq_enabled and supports_uq and uq_method == "closed_form": + return "closed_form" return "deterministic" diff --git a/physicsnemo/cfd/evaluation/config.py b/physicsnemo/cfd/evaluation/config.py index 7e366e3..bc9e6fe 100644 --- a/physicsnemo/cfd/evaluation/config.py +++ b/physicsnemo/cfd/evaluation/config.py @@ -105,7 +105,7 @@ class UQConfig: 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 + Ignored by ``UQ_METHOD="closed_form"`` 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 diff --git a/physicsnemo/cfd/evaluation/metrics/builtin/uq.py b/physicsnemo/cfd/evaluation/metrics/builtin/uq.py index 545b195..923554a 100644 --- a/physicsnemo/cfd/evaluation/metrics/builtin/uq.py +++ b/physicsnemo/cfd/evaluation/metrics/builtin/uq.py @@ -463,9 +463,9 @@ 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-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 +# EVERY UQ method (a closed-form GP posterior's marginals OR the sampling across-pass spread), so +# the same caveat applies to all of them — it is NOT specific to closed-form 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 diff --git a/physicsnemo/cfd/evaluation/models/model_registry.py b/physicsnemo/cfd/evaluation/models/model_registry.py index 9bba0f7..3f705bc 100644 --- a/physicsnemo/cfd/evaluation/models/model_registry.py +++ b/physicsnemo/cfd/evaluation/models/model_registry.py @@ -58,12 +58,12 @@ class CFDModel(ABC): #: ``NaN`` for them (consistent with the engine's recoverable-metric behavior). SUPPORTS_UQ: ClassVar[bool] = False #: How the predictive distribution is produced: - #: ``"analytic"`` — one forward pass emits the distribution/params (GP, mean-variance, + #: ``"closed_form"`` — 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" + UQ_METHOD: ClassVar[Literal["none", "closed_form", "sampling"]] = "none" @classmethod def inference_domain_from_kwargs( @@ -111,7 +111,7 @@ def predict_deterministic(self, model_input: ModelInput) -> RawOutput: 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`, which is correct for deterministic and closed-form wrappers (their :meth:`predict` is already deterministic). **Sampling** wrappers whose :meth:`predict` is stochastic (e.g. MC-Dropout keeps dropout @@ -143,7 +143,7 @@ def decode_distribution( ) -> 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 + **Closed-form** UQ wrappers (``UQ_METHOD="closed_form"``, 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 diff --git a/physicsnemo/cfd/evaluation/models/wrappers/geotransolver_gp/wrapper.py b/physicsnemo/cfd/evaluation/models/wrappers/geotransolver_gp/wrapper.py index 4abdd7a..ca14dd6 100644 --- a/physicsnemo/cfd/evaluation/models/wrappers/geotransolver_gp/wrapper.py +++ b/physicsnemo/cfd/evaluation/models/wrappers/geotransolver_gp/wrapper.py @@ -14,9 +14,9 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""GeoTransolver + Gaussian-Process field-head wrapper (analytic surface UQ). +"""GeoTransolver + Gaussian-Process field-head wrapper (closed-form surface UQ). -This is the ``UQ_METHOD="analytic"`` archetype and a **completely independent** :class:`CFDModel` +This is the ``UQ_METHOD="closed_form"`` 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 @@ -123,7 +123,7 @@ class GeoTransolverGPDrivAerStarWrapper(CFDModel): - """GeoTransolver backbone + ``FieldVariationalGPHead`` for analytic per-point surface UQ. + """GeoTransolver backbone + ``FieldVariationalGPHead`` for closed-form per-point surface UQ. Trained on ``transformer_models`` (physical-target) checkpoints, so predictions are re-standardized only (no dynamic-pressure re-dimensionalization): @@ -135,7 +135,7 @@ class GeoTransolverGPDrivAerStarWrapper(CFDModel): INFERENCE_DOMAIN: ClassVar[InferenceDomain | None] = "surface" OUTPUT_LOCATION: ClassVar[OutputLocation] = "cell" SUPPORTS_UQ: ClassVar[bool] = True - UQ_METHOD: ClassVar[str] = "analytic" + UQ_METHOD: ClassVar[str] = "closed_form" @property def output_location(self) -> OutputLocation: diff --git a/pyproject.toml b/pyproject.toml index 8849a43..e153116 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -49,8 +49,14 @@ gpu = [ "cuml-cu13>=25.10.0", "warp-lang>=1.13.0", ] +# gpytorch for the closed-form UQ wrappers, whose GP head +# (physicsnemo.experimental.uq.FieldVariationalGPHead) requires it at import time. +# Every other evaluation wrapper, and the sampling UQ path, install without it. +evaluation-uq = [ + "gpytorch>=1.11", +] all = [ - "nvidia-physicsnemo-cfd[dev,gpu,evaluation-hf]", + "nvidia-physicsnemo-cfd[dev,gpu,evaluation-hf,evaluation-uq]", ] docs = [ "sphinx>=6.0.0", diff --git a/skills/physicsnemo-cfd-create-model-wrapper/SKILL.md b/skills/physicsnemo-cfd-create-model-wrapper/SKILL.md index 694f017..6c466fb 100644 --- a/skills/physicsnemo-cfd-create-model-wrapper/SKILL.md +++ b/skills/physicsnemo-cfd-create-model-wrapper/SKILL.md @@ -99,7 +99,7 @@ geometry). **Uncertainty-quantification (UQ) models** set two more class variables (`SUPPORTS_UQ`, `UQ_METHOD`) and implement extra hooks — `decode_distribution` -(analytic), or for sampling a stochastic `predict` plus +(closed-form), or for sampling a stochastic `predict` plus `predict_deterministic` (and optionally `predict_ensemble`). This is fully additive: deterministic wrappers leave the defaults (`SUPPORTS_UQ=False`) and are unaffected. See the "Uncertainty quantification" section below. @@ -330,7 +330,7 @@ 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 +opts into the UQ path additively. Copy `ExampleClosedFormUQWrapper` 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. @@ -339,13 +339,13 @@ complete four-row example config. ```python SUPPORTS_UQ: ClassVar[bool] = True -UQ_METHOD: ClassVar[str] = "analytic" # or "sampling"; default "none" +UQ_METHOD: ClassVar[str] = "closed_form" # 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 +- **`UQ_METHOD="closed_form"`** — 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) -> @@ -373,7 +373,7 @@ Deterministic wrappers keep `SUPPORTS_UQ=False`; UQ metrics simply report ### `FieldDistribution` payload -`decode_distribution` (analytic) and the engine's aggregator (sampling) +`decode_distribution` (closed-form) 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`): @@ -392,7 +392,7 @@ FieldDistribution( 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 +- **Epistemic vs aleatoric.** Closed-form: 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 @@ -415,7 +415,7 @@ Turn UQ on in the run config and add the pooled metrics you want: run: uq: enabled: true - num_samples: 32 # passes for UQ_METHOD="sampling" (ignored by analytic/ensemble) + num_samples: 32 # passes for UQ_METHOD="sampling" (ignored by closed-form/ensemble) metrics: - nlpd - calibration_zrms @@ -430,7 +430,7 @@ 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). +four-row example (deterministic + closed-form GP + MC-Dropout + ensemble). ## Gotchas @@ -455,7 +455,7 @@ four-row example (deterministic + analytic GP + MC-Dropout + ensemble). - `references/example_wrapper.py` — complete `CFDModel` templates to copy and adapt (bundled; available without the repo on disk): deterministic - surface + volume, plus `ExampleAnalyticUQWrapper` and + surface + volume, plus `ExampleClosedFormUQWrapper` 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 2e26dc0..1837750 100644 --- a/skills/physicsnemo-cfd-create-model-wrapper/references/example_wrapper.py +++ b/skills/physicsnemo-cfd-create-model-wrapper/references/example_wrapper.py @@ -19,8 +19,8 @@ These are complete, correct ``CFDModel`` implementations to **adapt** when writing a 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. + * ``ExampleClosedFormUQWrapper`` — a single-pass UQ model (GP head / mean-variance / + evidential): ``UQ_METHOD="closed_form"``, 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. @@ -252,12 +252,12 @@ 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"``. +class ExampleClosedFormUQWrapper(CFDModel): + """Single-pass UQ model (GP head / mean-variance / evidential): ``UQ_METHOD="closed_form"``. 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"``; + * declare ``SUPPORTS_UQ = True`` and ``UQ_METHOD = "closed_form"``; * 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. @@ -267,7 +267,7 @@ class ExampleAnalyticUQWrapper(CFDModel): 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" + UQ_METHOD: ClassVar[str] = "closed_form" def __init__(self) -> None: self._model: torch.nn.Module | None = None @@ -281,13 +281,13 @@ def output_location(self) -> OutputLocation: def load( self, checkpoint_path: str, stats_path: str, device: str, **kwargs: Any - ) -> "ExampleAnalyticUQWrapper": + ) -> "ExampleClosedFormUQWrapper": """Build the architecture + UQ head, load weights and normalization stats.""" self._device = device # Build the architecture + UQ head and load weights here (e.g. a GP head whose # hyperparameters come from kwargs and MUST match training so the state dict loads). self._stats = load_global_stats(stats_path, device) - log_inference("example_analytic_uq", f"Loaded from {checkpoint_path}") + log_inference("example_closed_form_uq", f"Loaded from {checkpoint_path}") return self def prepare_inputs(self, case: CanonicalCase) -> torch.Tensor: @@ -510,5 +510,5 @@ def _std(self, field: str) -> Any: # 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_closed_form_uq", ExampleClosedFormUQWrapper) register_model("example_sampling_uq", ExampleSamplingUQWrapper) diff --git a/test/ci_tests/test_uq_inference.py b/test/ci_tests/test_uq_inference.py index 8a132b1..b2b0af6 100644 --- a/test/ci_tests/test_uq_inference.py +++ b/test/ci_tests/test_uq_inference.py @@ -199,18 +199,22 @@ 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. + # UQ enabled: sampling / closed-form 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" + select_inference_path( + supports_uq=True, uq_method="closed_form", uq_enabled=True + ) + == "closed_form" ) - # Master switch OFF: EVERY wrapper (incl. analytic GP) goes deterministic -> no UQ metrics. + # Master switch OFF: EVERY wrapper (incl. closed-form GP) goes deterministic -> no UQ metrics. assert ( - select_inference_path(supports_uq=True, uq_method="analytic", uq_enabled=False) + select_inference_path( + supports_uq=True, uq_method="closed_form", uq_enabled=False + ) == "deterministic" ) assert ( diff --git a/test/ci_tests/test_uq_metrics.py b/test/ci_tests/test_uq_metrics.py index febc8b0..5f84008 100644 --- a/test/ci_tests/test_uq_metrics.py +++ b/test/ci_tests/test_uq_metrics.py @@ -111,7 +111,7 @@ def test_uq_metrics_are_registered_as_reducers() -> None: def test_pooled_calibration_on_synthetic_gaussian() -> None: - """A perfectly calibrated Gaussian gives cov95≈0.95, zrms≈1, and the analytic NLPD.""" + """A perfectly calibrated Gaussian gives cov95≈0.95, zrms≈1, and the exact NLPD.""" rng = np.random.default_rng(0) metrics = _make_uq_metrics() n, sigma = 400_000, 1.0 diff --git a/test/ci_tests/test_uq_wrappers.py b/test/ci_tests/test_uq_wrappers.py index f292b70..a574a5a 100644 --- a/test/ci_tests/test_uq_wrappers.py +++ b/test/ci_tests/test_uq_wrappers.py @@ -127,7 +127,7 @@ def predict_deterministic(self, model_input): def test_predict_deterministic_default_delegates_to_predict() -> None: - """The base hook simply calls predict() (correct for deterministic / analytic wrappers).""" + """The base hook simply calls predict() (correct for deterministic / closed-form wrappers).""" w = _EchoWrapper() assert w.predict_deterministic("x") == w.predict("x") == "x" @@ -190,8 +190,8 @@ def test_uq_wrappers_registered_with_expected_contract() -> None: 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" + # Closed-form GP head: single-pass distribution. + assert gp.SUPPORTS_UQ is True and gp.UQ_METHOD == "closed_form" # 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" diff --git a/workflows/benchmarking/conf/config_uq_surface.yaml b/workflows/benchmarking/conf/config_uq_surface.yaml index df95629..54e199d 100644 --- a/workflows/benchmarking/conf/config_uq_surface.yaml +++ b/workflows/benchmarking/conf/config_uq_surface.yaml @@ -25,7 +25,7 @@ # # * 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): +# * geotransolver_gp_surface -> UQ_METHOD="closed_form" (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 @@ -69,7 +69,7 @@ run: # ---- Uncertainty-quantification controls (additive; default off) ---- uq: enabled: true - # Stochastic passes per case for UQ_METHOD="sampling" wrappers. Ignored by analytic / + # Stochastic passes per case for UQ_METHOD="sampling" wrappers. Ignored by closed-form / # 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). @@ -100,7 +100,7 @@ benchmark: stats_path: *surface_stats kwargs: { batch_resolution: 60000, geometry_sampling: 300000 } - # (2) Analytic GP field head. This wrapper loads TWO checkpoints, both named explicitly: + # (2) Closed-form 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 FieldVariationalGPHead (read as given; pass matching # files). The GP hyperparameters here MUST match the training run so the head reconstructs @@ -212,7 +212,7 @@ metrics: - sparsification_ause - sparsification_ause_epistemic # Decision-relevant drag sparsification: rank geometries by uncertainty propagated into the Cd - # integral (drag std) vs |Cd_pred - Cd_true|. NOTE: this applies to BOTH analytic (GP) and + # integral (drag std) vs |Cd_pred - Cd_true|. NOTE: this applies to BOTH closed-form (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)