595 lines
24 KiB
Python
595 lines
24 KiB
Python
|
"""Differential trace artifacts for OpenFOAM-vs-GPU RANS parity diagnostics.
|
||
|
|
|
||
|
|
The verifier already emits numeric split-stage artifacts. This module gives those
|
||
|
|
artifacts a stable semantic order, compact first-divergence summaries, and a
|
||
|
|
substitution contract that can materialize candidate artifacts with selected
|
||
|
|
checkpoints replaced by the OpenFOAM reference state.
|
||
|
|
"""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import json
|
||
|
|
import math
|
||
|
|
from collections import defaultdict
|
||
|
|
from collections.abc import Callable, Iterable, Mapping, Sequence
|
||
|
|
from pathlib import Path
|
||
|
|
from typing import Any
|
||
|
|
|
||
|
|
import numpy as np
|
||
|
|
|
||
|
|
TRACE_SCHEMA_VERSION = 1
|
||
|
|
DEFAULT_TOP_N = 5
|
||
|
|
ARTIFACT_ONLY_PREFIXES = ("candidate_diagnostics.",)
|
||
|
|
|
||
|
|
CHECKPOINT_ORDER: tuple[str, ...] = (
|
||
|
|
"matrix_terms.UEqn.ddt.diag",
|
||
|
|
"matrix_terms.UEqn.ddt.source",
|
||
|
|
"matrix_terms.UEqn.div.diag",
|
||
|
|
"matrix_terms.UEqn.div.upper",
|
||
|
|
"matrix_terms.UEqn.div.lower",
|
||
|
|
"matrix_terms.UEqn.div.source",
|
||
|
|
"matrix_terms.UEqn.divDevSigma.diag",
|
||
|
|
"matrix_terms.UEqn.divDevSigma.upper",
|
||
|
|
"matrix_terms.UEqn.divDevSigma.lower",
|
||
|
|
"matrix_terms.UEqn.divDevSigma.wall_dev_tau.source",
|
||
|
|
"matrix_terms.UEqn.divDevSigma.source",
|
||
|
|
"matrix_operator.UEqn.unrelaxed_diag",
|
||
|
|
"matrix_operator.UEqn.unrelaxed_source",
|
||
|
|
"matrix_operator.UEqn.diag",
|
||
|
|
"matrix_operator.UEqn.upper",
|
||
|
|
"matrix_operator.UEqn.lower",
|
||
|
|
"matrix_operator.UEqn.source",
|
||
|
|
"matrix_operator.UEqn.internal_coeffs",
|
||
|
|
"matrix_operator.UEqn.boundary_coeffs",
|
||
|
|
"solver.solve_UEqn.matrix_before.diag",
|
||
|
|
"solver.solve_UEqn.matrix_before.upper",
|
||
|
|
"solver.solve_UEqn.matrix_before.lower",
|
||
|
|
"solver.solve_UEqn.matrix_before.source",
|
||
|
|
"solver.solve_UEqn.solve_diag",
|
||
|
|
"solver.solve_UEqn.solve_source",
|
||
|
|
"solver.solve_UEqn.initial_residual",
|
||
|
|
"solver.solve_UEqn.preconditioned_residual",
|
||
|
|
"solver.solve_UEqn.level_preconditioned_residual",
|
||
|
|
"solver.solve_UEqn.operator_preconditioned_direction",
|
||
|
|
"solver.solve_UEqn.first_iteration.rho",
|
||
|
|
"solver.solve_UEqn.first_iteration.denominator",
|
||
|
|
"solver.solve_UEqn.first_iteration.alpha",
|
||
|
|
"solver.solve_UEqn.first_iteration.intermediate_residual",
|
||
|
|
"solver.solve_UEqn.first_iteration.second_preconditioned_residual",
|
||
|
|
"solver.solve_UEqn.first_iteration.operator_second_preconditioned_residual",
|
||
|
|
"solver.solve_UEqn.first_iteration.omega_numerator",
|
||
|
|
"solver.solve_UEqn.first_iteration.omega_denominator",
|
||
|
|
"solver.solve_UEqn.first_iteration.omega",
|
||
|
|
"solver.solve_UEqn.first_iteration.residual_after_omega",
|
||
|
|
"solver.solve_UEqn.first_iteration.solution_after_omega",
|
||
|
|
"solver.solve_UEqn.field_after",
|
||
|
|
"pressure_inputs.rAU",
|
||
|
|
"pressure_inputs.rAtU",
|
||
|
|
"pressure_inputs.HbyA",
|
||
|
|
"pressure_inputs.phiHbyA",
|
||
|
|
"matrix_operator.pEqn.diag",
|
||
|
|
"matrix_operator.pEqn.upper",
|
||
|
|
"matrix_operator.pEqn.lower",
|
||
|
|
"matrix_operator.pEqn.source",
|
||
|
|
"matrix_operator.pEqn.internal_coeffs",
|
||
|
|
"matrix_operator.pEqn.boundary_coeffs",
|
||
|
|
"solver.solve_pEqn.matrix_before.diag",
|
||
|
|
"solver.solve_pEqn.matrix_before.upper",
|
||
|
|
"solver.solve_pEqn.matrix_before.source",
|
||
|
|
"solver.solve_pEqn.initial_residual",
|
||
|
|
"solver.solve_pEqn.preconditioned_residual",
|
||
|
|
"solver.solve_pEqn.operator_preconditioned_direction",
|
||
|
|
"solver.solve_pEqn.p",
|
||
|
|
"solver.solve_pEqn.phi",
|
||
|
|
"final_correction.U",
|
||
|
|
"final_correction.p",
|
||
|
|
"final_correction.phi",
|
||
|
|
"turbulence.nut",
|
||
|
|
"turbulence.k",
|
||
|
|
"turbulence.omega",
|
||
|
|
"fields.U",
|
||
|
|
"fields.p",
|
||
|
|
"fields.phi",
|
||
|
|
"fields.nut",
|
||
|
|
"fields.k",
|
||
|
|
"fields.omega",
|
||
|
|
)
|
||
|
|
CHECKPOINT_ORDER_INDEX = {name: index for index, name in enumerate(CHECKPOINT_ORDER)}
|
||
|
|
REQUIRED_COMPARISON_CHECKPOINTS: frozenset[str] = frozenset(
|
||
|
|
{
|
||
|
|
"matrix_operator.UEqn.diag",
|
||
|
|
"matrix_operator.UEqn.upper",
|
||
|
|
"matrix_operator.UEqn.lower",
|
||
|
|
"matrix_operator.UEqn.source",
|
||
|
|
"matrix_operator.UEqn.psi",
|
||
|
|
"solver.solve_UEqn.matrix_before.diag",
|
||
|
|
"solver.solve_UEqn.matrix_before.upper",
|
||
|
|
"solver.solve_UEqn.matrix_before.lower",
|
||
|
|
"solver.solve_UEqn.matrix_before.source",
|
||
|
|
"solver.solve_UEqn.solve_diag",
|
||
|
|
"solver.solve_UEqn.solve_source",
|
||
|
|
"solver.solve_UEqn.initial_residual",
|
||
|
|
"solver.solve_UEqn.preconditioned_residual",
|
||
|
|
"solver.solve_UEqn.level_preconditioned_residual",
|
||
|
|
"solver.solve_UEqn.operator_preconditioned_direction",
|
||
|
|
"solver.solve_UEqn.first_iteration.rho",
|
||
|
|
"solver.solve_UEqn.first_iteration.denominator",
|
||
|
|
"solver.solve_UEqn.first_iteration.alpha",
|
||
|
|
"solver.solve_UEqn.first_iteration.intermediate_residual",
|
||
|
|
"solver.solve_UEqn.first_iteration.second_preconditioned_residual",
|
||
|
|
"solver.solve_UEqn.first_iteration.operator_second_preconditioned_residual",
|
||
|
|
"solver.solve_UEqn.first_iteration.omega_numerator",
|
||
|
|
"solver.solve_UEqn.first_iteration.omega_denominator",
|
||
|
|
"solver.solve_UEqn.first_iteration.omega",
|
||
|
|
"solver.solve_UEqn.first_iteration.residual_after_omega",
|
||
|
|
"solver.solve_UEqn.first_iteration.solution_after_omega",
|
||
|
|
"solver.solve_UEqn.field_after",
|
||
|
|
"pressure_inputs.rAU",
|
||
|
|
"pressure_inputs.rAtU",
|
||
|
|
"pressure_inputs.HbyA",
|
||
|
|
"pressure_inputs.phiHbyA",
|
||
|
|
"matrix_operator.pEqn.diag",
|
||
|
|
"matrix_operator.pEqn.upper",
|
||
|
|
"matrix_operator.pEqn.source",
|
||
|
|
"matrix_operator.pEqn.psi",
|
||
|
|
"solver.solve_pEqn.matrix_before.diag",
|
||
|
|
"solver.solve_pEqn.matrix_before.upper",
|
||
|
|
"solver.solve_pEqn.matrix_before.source",
|
||
|
|
"solver.solve_pEqn.p",
|
||
|
|
"solver.solve_pEqn.phi",
|
||
|
|
}
|
||
|
|
)
|
||
|
|
|
||
|
|
SUPPORTED_SUBSTITUTION_CHECKPOINTS: frozenset[str] = frozenset(
|
||
|
|
{
|
||
|
|
"matrix_operator.UEqn.diag",
|
||
|
|
"matrix_operator.UEqn.upper",
|
||
|
|
"matrix_operator.UEqn.lower",
|
||
|
|
"matrix_operator.UEqn.source",
|
||
|
|
"solver.solve_UEqn.matrix_before.diag",
|
||
|
|
"solver.solve_UEqn.matrix_before.upper",
|
||
|
|
"solver.solve_UEqn.matrix_before.lower",
|
||
|
|
"solver.solve_UEqn.matrix_before.source",
|
||
|
|
"solver.solve_UEqn.solve_source",
|
||
|
|
"solver.solve_UEqn.field_after",
|
||
|
|
"pressure_inputs.rAU",
|
||
|
|
"pressure_inputs.rAtU",
|
||
|
|
"pressure_inputs.HbyA",
|
||
|
|
"pressure_inputs.phiHbyA",
|
||
|
|
"matrix_operator.pEqn.diag",
|
||
|
|
"matrix_operator.pEqn.upper",
|
||
|
|
"matrix_operator.pEqn.source",
|
||
|
|
"solver.solve_pEqn.p",
|
||
|
|
"solver.solve_pEqn.phi",
|
||
|
|
}
|
||
|
|
)
|
||
|
|
|
||
|
|
LocationContext = Callable[[str, tuple[int, ...], tuple[int, ...]], Mapping[str, Any] | None]
|
||
|
|
|
||
|
|
|
||
|
|
def finite_float(value: Any) -> float | None:
|
||
|
|
try:
|
||
|
|
number = float(value)
|
||
|
|
except (TypeError, ValueError):
|
||
|
|
return None
|
||
|
|
return number if math.isfinite(number) else None
|
||
|
|
|
||
|
|
|
||
|
|
def json_ready(value: Any) -> Any:
|
||
|
|
if isinstance(value, np.generic):
|
||
|
|
return value.item()
|
||
|
|
if isinstance(value, np.ndarray):
|
||
|
|
return value.tolist()
|
||
|
|
if isinstance(value, Path):
|
||
|
|
return str(value)
|
||
|
|
if isinstance(value, Mapping):
|
||
|
|
return {str(key): json_ready(item) for key, item in value.items()}
|
||
|
|
if isinstance(value, (list, tuple)):
|
||
|
|
return [json_ready(item) for item in value]
|
||
|
|
if isinstance(value, (str, int, float, bool)) or value is None:
|
||
|
|
return value
|
||
|
|
return repr(value)
|
||
|
|
|
||
|
|
|
||
|
|
def load_npz_artifacts(path: Path | str | None) -> dict[str, np.ndarray] | None:
|
||
|
|
if path is None:
|
||
|
|
return None
|
||
|
|
artifact_path = Path(path)
|
||
|
|
if not artifact_path.exists():
|
||
|
|
return None
|
||
|
|
with np.load(artifact_path, allow_pickle=False) as data:
|
||
|
|
return {name: np.asarray(data[name]) for name in data.files}
|
||
|
|
|
||
|
|
|
||
|
|
def write_json_report(path: Path | str, report: Mapping[str, Any]) -> None:
|
||
|
|
output = Path(path)
|
||
|
|
output.parent.mkdir(parents=True, exist_ok=True)
|
||
|
|
output.write_text(json.dumps(json_ready(report), indent=2, sort_keys=True, allow_nan=False) + "\n")
|
||
|
|
|
||
|
|
|
||
|
|
def checkpoint_sort_key(name: str) -> tuple[int, str]:
|
||
|
|
return (CHECKPOINT_ORDER_INDEX.get(name, len(CHECKPOINT_ORDER)), name)
|
||
|
|
|
||
|
|
|
||
|
|
def is_trace_checkpoint(name: str) -> bool:
|
||
|
|
return not name.startswith(ARTIFACT_ONLY_PREFIXES)
|
||
|
|
|
||
|
|
|
||
|
|
def checkpoint_metadata(name: str) -> dict[str, Any]:
|
||
|
|
parts = name.split(".")
|
||
|
|
family = parts[0] if parts else "unknown"
|
||
|
|
metadata: dict[str, Any] = {
|
||
|
|
"checkpoint": name,
|
||
|
|
"family": family,
|
||
|
|
"order": CHECKPOINT_ORDER_INDEX.get(name),
|
||
|
|
"role": "candidate_or_reference_state",
|
||
|
|
"lifecycle_phase": lifecycle_phase(name),
|
||
|
|
}
|
||
|
|
if len(parts) >= 2:
|
||
|
|
metadata["object"] = parts[1]
|
||
|
|
if family == "matrix_operator" and len(parts) >= 3:
|
||
|
|
metadata.update({"field_or_matrix": parts[1], "coefficient": parts[2], "granularity": "cell_or_internal_face_array"})
|
||
|
|
elif family == "matrix_terms" and len(parts) >= 5:
|
||
|
|
metadata.update({"field_or_matrix": parts[1], "term": parts[2], "subterm": parts[3], "coefficient": parts[4]})
|
||
|
|
elif family == "matrix_terms" and len(parts) >= 4:
|
||
|
|
metadata.update({"field_or_matrix": parts[1], "term": parts[2], "coefficient": parts[3]})
|
||
|
|
elif family == "pressure_inputs" and len(parts) >= 2:
|
||
|
|
metadata.update({"field": parts[1], "granularity": "cell_or_internal_face_array"})
|
||
|
|
elif family == "solver" and len(parts) >= 3:
|
||
|
|
metadata.update({"solver_stage": parts[1], "solver_quantity": ".".join(parts[2:])})
|
||
|
|
elif family in {"final_correction", "fields", "turbulence"} and len(parts) >= 2:
|
||
|
|
metadata["field"] = parts[1]
|
||
|
|
metadata["required_comparison"] = name in REQUIRED_COMPARISON_CHECKPOINTS
|
||
|
|
metadata["substitution_supported"] = name in SUPPORTED_SUBSTITUTION_CHECKPOINTS
|
||
|
|
return metadata
|
||
|
|
|
||
|
|
|
||
|
|
def lifecycle_phase(name: str) -> str:
|
||
|
|
if name.startswith("matrix_terms."):
|
||
|
|
return "term_contribution_before_accumulation"
|
||
|
|
if name.startswith("matrix_operator."):
|
||
|
|
if ".unrelaxed_" in name:
|
||
|
|
return "persistent_matrix_before_relaxation"
|
||
|
|
if ".internal_coeffs" in name or ".boundary_coeffs" in name:
|
||
|
|
return "matrix_boundary_coefficients"
|
||
|
|
return "persistent_matrix_after_assembly"
|
||
|
|
if ".matrix_before." in name:
|
||
|
|
return "solve_time_matrix_before_solver"
|
||
|
|
if name.startswith("pressure_inputs."):
|
||
|
|
return "pressure_input_construction"
|
||
|
|
if ".initial_residual" in name or ".preconditioned_residual" in name or ".operator_preconditioned_direction" in name:
|
||
|
|
return "linear_solver_iteration_state"
|
||
|
|
if name.startswith("solver."):
|
||
|
|
return "linear_solver_result"
|
||
|
|
if name.startswith("final_correction."):
|
||
|
|
return "pressure_velocity_flux_correction"
|
||
|
|
if name.startswith("turbulence."):
|
||
|
|
return "turbulence_update"
|
||
|
|
if name.startswith("fields."):
|
||
|
|
return "final_field_state"
|
||
|
|
return "unknown"
|
||
|
|
|
||
|
|
|
||
|
|
def array_stats(array: Any) -> dict[str, Any]:
|
||
|
|
arr = np.asarray(array)
|
||
|
|
out: dict[str, Any] = {"shape": [int(dim) for dim in arr.shape], "dtype": str(arr.dtype), "size": int(arr.size)}
|
||
|
|
if arr.size == 0:
|
||
|
|
return {**out, "min": None, "max": None, "mean": None, "l2": 0.0}
|
||
|
|
finite = np.isfinite(arr)
|
||
|
|
if not np.any(finite):
|
||
|
|
return {**out, "min": None, "max": None, "mean": None, "l2": None, "nonfinite_count": int(arr.size)}
|
||
|
|
finite_values = arr[finite]
|
||
|
|
out.update(
|
||
|
|
{
|
||
|
|
"min": finite_float(np.min(finite_values)),
|
||
|
|
"max": finite_float(np.max(finite_values)),
|
||
|
|
"mean": finite_float(np.mean(finite_values)),
|
||
|
|
"l2": finite_float(np.linalg.norm(finite_values.reshape(-1))),
|
||
|
|
"nonfinite_count": int(arr.size - np.count_nonzero(finite)),
|
||
|
|
}
|
||
|
|
)
|
||
|
|
return out
|
||
|
|
|
||
|
|
|
||
|
|
def _value_at(array: np.ndarray, index: tuple[int, ...]) -> Any:
|
||
|
|
return json_ready(array[index])
|
||
|
|
|
||
|
|
|
||
|
|
def _artifact_location(name: str, index: tuple[int, ...], shape: tuple[int, ...]) -> dict[str, Any]:
|
||
|
|
entity_kind = "array"
|
||
|
|
if name.startswith(("matrix_operator.", "pressure_inputs.rAU", "pressure_inputs.rAtU", "pressure_inputs.HbyA", "solver.solve_UEqn", "solver.solve_pEqn", "final_correction.", "fields.", "turbulence.")):
|
||
|
|
entity_kind = "cell"
|
||
|
|
if ".upper" in name or ".lower" in name or name.endswith("phi") or name.endswith("phiHbyA"):
|
||
|
|
entity_kind = "internal_face"
|
||
|
|
return {
|
||
|
|
"array_index": list(index),
|
||
|
|
"entity_kind": entity_kind,
|
||
|
|
"entity_index": int(index[0]) if index else None,
|
||
|
|
"component_index": list(index[1:]) if len(index) > 1 else None,
|
||
|
|
"artifact": name,
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
def _top_differences(
|
||
|
|
name: str,
|
||
|
|
candidate: np.ndarray,
|
||
|
|
reference: np.ndarray,
|
||
|
|
diff: np.ndarray,
|
||
|
|
*,
|
||
|
|
top_n: int,
|
||
|
|
location_context: LocationContext | None,
|
||
|
|
) -> list[dict[str, Any]]:
|
||
|
|
if diff.size == 0 or top_n <= 0:
|
||
|
|
return []
|
||
|
|
flat = diff.reshape(-1)
|
||
|
|
finite = np.isfinite(flat)
|
||
|
|
if np.any(finite):
|
||
|
|
finite_indices = np.flatnonzero(finite)
|
||
|
|
finite_values = flat[finite_indices]
|
||
|
|
if finite_values.size > top_n:
|
||
|
|
local = np.argpartition(finite_values, -top_n)[-top_n:]
|
||
|
|
candidate_indices = finite_indices[local]
|
||
|
|
else:
|
||
|
|
candidate_indices = finite_indices
|
||
|
|
ordered = sorted(candidate_indices, key=lambda item: (-float(flat[int(item)]), int(item)))
|
||
|
|
else:
|
||
|
|
ordered = [int(item) for item in np.flatnonzero(~finite)[:top_n]]
|
||
|
|
|
||
|
|
out: list[dict[str, Any]] = []
|
||
|
|
for flat_index in ordered[:top_n]:
|
||
|
|
index = tuple(int(item) for item in np.unravel_index(int(flat_index), diff.shape))
|
||
|
|
location = _artifact_location(name, index, tuple(diff.shape))
|
||
|
|
context = location_context(name, index, tuple(diff.shape)) if location_context is not None else None
|
||
|
|
out.append(
|
||
|
|
{
|
||
|
|
"location": location,
|
||
|
|
"abs_difference": finite_float(diff[index]),
|
||
|
|
"candidate_value": _value_at(candidate, index),
|
||
|
|
"reference_value": _value_at(reference, index),
|
||
|
|
"local_context": json_ready(context),
|
||
|
|
}
|
||
|
|
)
|
||
|
|
return out
|
||
|
|
|
||
|
|
|
||
|
|
def compare_checkpoint(
|
||
|
|
name: str,
|
||
|
|
reference: np.ndarray | None,
|
||
|
|
candidate: np.ndarray | None,
|
||
|
|
*,
|
||
|
|
rtol: float,
|
||
|
|
atol: float,
|
||
|
|
top_n: int = DEFAULT_TOP_N,
|
||
|
|
location_context: LocationContext | None = None,
|
||
|
|
) -> dict[str, Any]:
|
||
|
|
out: dict[str, Any] = {
|
||
|
|
"name": name,
|
||
|
|
"metadata": checkpoint_metadata(name),
|
||
|
|
"reference_available": reference is not None,
|
||
|
|
"candidate_available": candidate is not None,
|
||
|
|
"rtol": float(rtol),
|
||
|
|
"atol": float(atol),
|
||
|
|
}
|
||
|
|
if reference is None or candidate is None:
|
||
|
|
return {**out, "status": "missing", "allclose": False, "reason": "missing_reference" if reference is None else "missing_candidate"}
|
||
|
|
|
||
|
|
ref = np.asarray(reference)
|
||
|
|
cand = np.asarray(candidate)
|
||
|
|
out.update(
|
||
|
|
{
|
||
|
|
"reference_stats": array_stats(ref),
|
||
|
|
"candidate_stats": array_stats(cand),
|
||
|
|
"shape_matches": ref.shape == cand.shape,
|
||
|
|
"dtype_matches": ref.dtype == cand.dtype,
|
||
|
|
}
|
||
|
|
)
|
||
|
|
if ref.shape != cand.shape:
|
||
|
|
return {**out, "status": "failed", "allclose": False, "reason": "shape_mismatch"}
|
||
|
|
if ref.size == 0:
|
||
|
|
return {**out, "status": "passed", "allclose": True, "reason": None, "difference_stats": array_stats(np.asarray([], dtype=np.float64)), "top_differences": []}
|
||
|
|
|
||
|
|
diff = np.abs(cand - ref)
|
||
|
|
finite = np.isfinite(diff)
|
||
|
|
finite_diff = diff[finite]
|
||
|
|
allclose = bool(np.allclose(cand, ref, rtol=rtol, atol=atol, equal_nan=False))
|
||
|
|
if finite_diff.size:
|
||
|
|
max_flat = int(np.argmax(np.where(finite, diff, -np.inf)))
|
||
|
|
max_index = tuple(int(item) for item in np.unravel_index(max_flat, diff.shape))
|
||
|
|
max_abs = finite_float(diff[max_index])
|
||
|
|
mean_abs = finite_float(np.mean(finite_diff))
|
||
|
|
rms_abs = finite_float(np.sqrt(np.mean(np.square(finite_diff))))
|
||
|
|
else:
|
||
|
|
max_index = tuple(int(item) for item in np.unravel_index(int(np.flatnonzero(~finite.reshape(-1))[0]), diff.shape))
|
||
|
|
max_abs = None
|
||
|
|
mean_abs = None
|
||
|
|
rms_abs = None
|
||
|
|
location = _artifact_location(name, max_index, tuple(diff.shape))
|
||
|
|
out.update(
|
||
|
|
{
|
||
|
|
"status": "passed" if allclose else "failed",
|
||
|
|
"allclose": allclose,
|
||
|
|
"reason": None if allclose else "value_mismatch",
|
||
|
|
"difference_stats": {
|
||
|
|
"max_abs": max_abs,
|
||
|
|
"mean_abs": mean_abs,
|
||
|
|
"rms_abs": rms_abs,
|
||
|
|
"nonfinite_error_count": int(diff.size - np.count_nonzero(finite)),
|
||
|
|
},
|
||
|
|
"largest_difference": {
|
||
|
|
"location": location,
|
||
|
|
"candidate_value": _value_at(cand, max_index),
|
||
|
|
"reference_value": _value_at(ref, max_index),
|
||
|
|
"local_context": json_ready(location_context(name, max_index, tuple(diff.shape)) if location_context is not None else None),
|
||
|
|
},
|
||
|
|
"top_differences": _top_differences(name, cand, ref, diff, top_n=top_n, location_context=location_context),
|
||
|
|
}
|
||
|
|
)
|
||
|
|
return out
|
||
|
|
|
||
|
|
|
||
|
|
def build_differential_trace_report(
|
||
|
|
*,
|
||
|
|
reference_artifacts: Mapping[str, np.ndarray] | None,
|
||
|
|
candidate_artifacts: Mapping[str, np.ndarray] | None,
|
||
|
|
reference_path: Path | str | None = None,
|
||
|
|
candidate_path: Path | str | None = None,
|
||
|
|
rtol: float = 0.0,
|
||
|
|
atol: float = 0.0,
|
||
|
|
top_n: int = DEFAULT_TOP_N,
|
||
|
|
location_context: LocationContext | None = None,
|
||
|
|
substitution: Mapping[str, Any] | None = None,
|
||
|
|
) -> dict[str, Any]:
|
||
|
|
reference = dict(reference_artifacts or {})
|
||
|
|
candidate = dict(candidate_artifacts or {})
|
||
|
|
keys = sorted((name for name in set(reference) | set(candidate) if is_trace_checkpoint(name)), key=checkpoint_sort_key)
|
||
|
|
comparisons = [
|
||
|
|
compare_checkpoint(
|
||
|
|
key,
|
||
|
|
reference.get(key),
|
||
|
|
candidate.get(key),
|
||
|
|
rtol=rtol,
|
||
|
|
atol=atol,
|
||
|
|
top_n=top_n,
|
||
|
|
location_context=location_context,
|
||
|
|
)
|
||
|
|
for key in keys
|
||
|
|
]
|
||
|
|
failed = [item for item in comparisons if item.get("status") == "failed"]
|
||
|
|
missing = [item for item in comparisons if item.get("status") == "missing"]
|
||
|
|
required_missing = [item for item in missing if item.get("name") in REQUIRED_COMPARISON_CHECKPOINTS]
|
||
|
|
artifacts_available = reference_artifacts is not None and candidate_artifacts is not None and bool(keys)
|
||
|
|
trace_status = "failed" if failed else "incomplete" if required_missing or not artifacts_available else "passed"
|
||
|
|
first_divergence = next((item for item in comparisons if item.get("status") == "failed"), None)
|
||
|
|
family_counts: dict[str, dict[str, int]] = defaultdict(lambda: {"passed": 0, "failed": 0, "missing": 0})
|
||
|
|
for item in comparisons:
|
||
|
|
family = str(item.get("metadata", {}).get("family", "unknown"))
|
||
|
|
status = str(item.get("status", "missing"))
|
||
|
|
if status not in family_counts[family]:
|
||
|
|
family_counts[family][status] = 0
|
||
|
|
family_counts[family][status] += 1
|
||
|
|
|
||
|
|
substitution_report = {
|
||
|
|
"schema_version": TRACE_SCHEMA_VERSION,
|
||
|
|
"supported_checkpoints": sorted(SUPPORTED_SUBSTITUTION_CHECKPOINTS, key=checkpoint_sort_key),
|
||
|
|
"replay_contract": "substitution replaces a named candidate checkpoint with the OpenFOAM reference array before downstream GPU computation; unsupported keys are artifact-only materialization targets",
|
||
|
|
**dict(substitution or {}),
|
||
|
|
}
|
||
|
|
return {
|
||
|
|
"schema_version": TRACE_SCHEMA_VERSION,
|
||
|
|
"status": trace_status,
|
||
|
|
"comparison_basis": "npz_reference_candidate_checkpoint_arrays",
|
||
|
|
"roles": {
|
||
|
|
"reference": {"path": str(reference_path) if reference_path is not None else None, "checkpoint_count": len(reference)},
|
||
|
|
"candidate": {"path": str(candidate_path) if candidate_path is not None else None, "checkpoint_count": len(candidate)},
|
||
|
|
},
|
||
|
|
"rtol": float(rtol),
|
||
|
|
"atol": float(atol),
|
||
|
|
"top_n": int(top_n),
|
||
|
|
"checkpoint_count": len(keys),
|
||
|
|
"passed_count": sum(1 for item in comparisons if item.get("status") == "passed"),
|
||
|
|
"failed_count": len(failed),
|
||
|
|
"missing_count": len(missing),
|
||
|
|
"required_missing_count": len(required_missing),
|
||
|
|
"first_divergence": compact_checkpoint(first_divergence) if first_divergence is not None else None,
|
||
|
|
"first_missing": compact_checkpoint(required_missing[0] if required_missing else missing[0]) if (required_missing or missing) else None,
|
||
|
|
"families": {family: dict(counts) for family, counts in sorted(family_counts.items())},
|
||
|
|
"checkpoints": comparisons,
|
||
|
|
"substitution": substitution_report,
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
def compact_checkpoint(item: Mapping[str, Any] | None) -> dict[str, Any] | None:
|
||
|
|
if item is None:
|
||
|
|
return None
|
||
|
|
diff = item.get("difference_stats", {}) if isinstance(item.get("difference_stats"), Mapping) else {}
|
||
|
|
return {
|
||
|
|
"name": item.get("name"),
|
||
|
|
"status": item.get("status"),
|
||
|
|
"reason": item.get("reason"),
|
||
|
|
"metadata": item.get("metadata"),
|
||
|
|
"max_abs": diff.get("max_abs"),
|
||
|
|
"mean_abs": diff.get("mean_abs"),
|
||
|
|
"rms_abs": diff.get("rms_abs"),
|
||
|
|
"largest_difference": item.get("largest_difference"),
|
||
|
|
"top_differences": item.get("top_differences", []),
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
def materialize_substituted_artifacts(
|
||
|
|
*,
|
||
|
|
reference_artifacts: Mapping[str, np.ndarray],
|
||
|
|
candidate_artifacts: Mapping[str, np.ndarray],
|
||
|
|
checkpoints: Sequence[str],
|
||
|
|
output_path: Path | str,
|
||
|
|
) -> dict[str, Any]:
|
||
|
|
requested = list(dict.fromkeys(str(item) for item in checkpoints))
|
||
|
|
substituted = {name: np.ascontiguousarray(value) for name, value in candidate_artifacts.items()}
|
||
|
|
applied: list[dict[str, Any]] = []
|
||
|
|
missing_reference: list[str] = []
|
||
|
|
missing_candidate: list[str] = []
|
||
|
|
shape_mismatches: list[dict[str, Any]] = []
|
||
|
|
|
||
|
|
for checkpoint in requested:
|
||
|
|
ref = reference_artifacts.get(checkpoint)
|
||
|
|
cand = candidate_artifacts.get(checkpoint)
|
||
|
|
if ref is None:
|
||
|
|
missing_reference.append(checkpoint)
|
||
|
|
continue
|
||
|
|
if cand is None:
|
||
|
|
missing_candidate.append(checkpoint)
|
||
|
|
continue
|
||
|
|
ref_array = np.asarray(ref)
|
||
|
|
cand_array = np.asarray(cand)
|
||
|
|
if ref_array.shape != cand_array.shape:
|
||
|
|
shape_mismatches.append({"checkpoint": checkpoint, "reference_shape": list(ref_array.shape), "candidate_shape": list(cand_array.shape)})
|
||
|
|
continue
|
||
|
|
substituted[checkpoint] = np.ascontiguousarray(ref_array)
|
||
|
|
applied.append(
|
||
|
|
{
|
||
|
|
"checkpoint": checkpoint,
|
||
|
|
"shape": [int(dim) for dim in ref_array.shape],
|
||
|
|
"dtype": str(ref_array.dtype),
|
||
|
|
"gpu_replay_supported": checkpoint in SUPPORTED_SUBSTITUTION_CHECKPOINTS,
|
||
|
|
}
|
||
|
|
)
|
||
|
|
|
||
|
|
output = Path(output_path)
|
||
|
|
output.parent.mkdir(parents=True, exist_ok=True)
|
||
|
|
np.savez_compressed(output, **substituted)
|
||
|
|
status = "ready" if applied and not missing_reference and not missing_candidate and not shape_mismatches else "partial" if applied else "failed"
|
||
|
|
return {
|
||
|
|
"schema_version": TRACE_SCHEMA_VERSION,
|
||
|
|
"status": status,
|
||
|
|
"path": str(output),
|
||
|
|
"format": "npz",
|
||
|
|
"requested": requested,
|
||
|
|
"applied": applied,
|
||
|
|
"missing_reference": missing_reference,
|
||
|
|
"missing_candidate": missing_candidate,
|
||
|
|
"shape_mismatches": shape_mismatches,
|
||
|
|
"artifact_count": len(substituted),
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
def extract_reference_substitutions(reference_artifacts: Mapping[str, np.ndarray], checkpoints: Iterable[str]) -> tuple[dict[str, np.ndarray], dict[str, Any]]:
|
||
|
|
substitutions: dict[str, np.ndarray] = {}
|
||
|
|
missing: list[str] = []
|
||
|
|
unsupported: list[str] = []
|
||
|
|
for checkpoint in dict.fromkeys(str(item) for item in checkpoints):
|
||
|
|
if checkpoint not in SUPPORTED_SUBSTITUTION_CHECKPOINTS:
|
||
|
|
unsupported.append(checkpoint)
|
||
|
|
continue
|
||
|
|
value = reference_artifacts.get(checkpoint)
|
||
|
|
if value is None:
|
||
|
|
missing.append(checkpoint)
|
||
|
|
continue
|
||
|
|
substitutions[checkpoint] = np.ascontiguousarray(np.asarray(value, dtype=np.float64))
|
||
|
|
return substitutions, {
|
||
|
|
"schema_version": TRACE_SCHEMA_VERSION,
|
||
|
|
"requested": list(dict.fromkeys(str(item) for item in checkpoints)),
|
||
|
|
"loaded": sorted(substitutions, key=checkpoint_sort_key),
|
||
|
|
"missing_reference": missing,
|
||
|
|
"unsupported": unsupported,
|
||
|
|
"status": "ready" if substitutions and not missing and not unsupported else "partial" if substitutions else "failed" if checkpoints else "not_requested",
|
||
|
|
}
|