drbx 2.0.0.dev0__py3-none-any.whl

drbx/validation/autodiff_diffusion.py

@@ -0,0 +1,329 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+from ..config.boutinp import parse_bout_input
+from ..native.mesh import StructuredMesh, build_structured_mesh
+from ..native.metrics import StructuredMetrics, build_structured_metrics
+from ..native.transport import advance_anomalous_diffusion_history
+from ..runtime.run_config import RunConfiguration
+
+
+@dataclass(frozen=True)
+class DiffusionAutodiffSetup:
+    config_text: str
+    run_config: RunConfiguration
+    mesh: StructuredMesh
+    metrics: StructuredMetrics
+    timestep: float
+    steps: int
+    density_boundary: str
+    pressure_boundary: str
+    background_density: float
+
+
+@dataclass(frozen=True)
+class StrongScalingPoint:
+    backend: str
+    device_count: int
+    elapsed_seconds: float
+    speedup: float
+    efficiency: float
+
+
+def build_diffusion_autodiff_setup(
+    *,
+    nx: int = 128,
+    ny: int = 24,
+    timestep: float = 5.0,
+    steps: int = 6,
+    background_density: float = 1.0,
+) -> DiffusionAutodiffSetup:
+    config_text = f"""
+nout = {steps}
+timestep = {timestep}
+
+[mesh]
+nx = {nx}
+ny = {ny}
+nz = 1
+dx = {1.0 / nx:.12f}
+dy = {1.0 / ny:.12f}
+dz = 1.0
+J = 1
+g11 = 1
+g_22 = 1
+g33 = 1
+
+[solver]
+type = native
+
+[model]
+components = h
+
+[h]
+type = evolve_density, evolve_pressure, anomalous_diffusion
+AA = 1
+charge = 1
+anomalous_D = 1.0
+thermal_conduction = false
+
+[Nh]
+bndry_all = neumann
+
+[Ph]
+bndry_all = neumann
+"""
+    config = parse_bout_input(config_text)
+    run_config = RunConfiguration.from_config(config)
+    mesh = build_structured_mesh(config, run_config)
+    metrics = build_structured_metrics(config, run_config, mesh)
+    return DiffusionAutodiffSetup(
+        config_text=config_text,
+        run_config=run_config,
+        mesh=mesh,
+        metrics=metrics,
+        timestep=timestep,
+        steps=steps,
+        density_boundary="neumann",
+        pressure_boundary="neumann",
+        background_density=background_density,
+    )
+
+
+def theta_to_physical(theta: jnp.ndarray) -> dict[str, jnp.ndarray]:
+    theta = jnp.asarray(theta, dtype=jnp.float64)
+    diffusivity = 0.05 + jax.nn.softplus(theta[0])
+    amplitude = 0.02 + jax.nn.softplus(theta[1])
+    center = jax.nn.sigmoid(theta[2])
+    width = 0.04 + 0.26 * jax.nn.sigmoid(theta[3])
+    return {
+        "anomalous_D": diffusivity,
+        "amplitude": amplitude,
+        "center": center,
+        "width": width,
+    }
+
+
+def physical_to_theta(
+    *,
+    anomalous_D: float,
+    amplitude: float,
+    center: float,
+    width: float,
+) -> jnp.ndarray:
+    center = float(np.clip(center, 1.0e-6, 1.0 - 1.0e-6))
+    width_scaled = (float(width) - 0.04) / 0.26
+    width_scaled = float(np.clip(width_scaled, 1.0e-6, 1.0 - 1.0e-6))
+
+    def inverse_softplus(value: float) -> float:
+        shifted = max(value - 0.05, 1.0e-6)
+        return float(np.log(np.expm1(shifted)))
+
+    def inverse_softplus_amp(value: float) -> float:
+        shifted = max(value - 0.02, 1.0e-6)
+        return float(np.log(np.expm1(shifted)))
+
+    return jnp.asarray(
+        [
+            inverse_softplus(anomalous_D),
+            inverse_softplus_amp(amplitude),
+            np.log(center / (1.0 - center)),
+            np.log(width_scaled / (1.0 - width_scaled)),
+        ],
+        dtype=jnp.float64,
+    )
+
+
+def design_field(setup: DiffusionAutodiffSetup, theta: jnp.ndarray) -> jnp.ndarray:
+    physical = theta_to_physical(theta)
+    return design_field_from_physical(
+        setup,
+        anomalous_D=physical["anomalous_D"],
+        amplitude=physical["amplitude"],
+        center=physical["center"],
+        width=physical["width"],
+    )
+
+
+def design_field_from_physical(
+    setup: DiffusionAutodiffSetup,
+    *,
+    anomalous_D: jnp.ndarray,
+    amplitude: jnp.ndarray,
+    center: jnp.ndarray,
+    width: jnp.ndarray,
+) -> jnp.ndarray:
+    x = jnp.linspace(0.0, 1.0, setup.mesh.nx, dtype=jnp.float64)[:, None, None]
+    y = jnp.linspace(0.0, 1.0, setup.mesh.local_ny, dtype=jnp.float64)[None, :, None]
+    profile = jnp.exp(-0.5 * ((x - center) / width) ** 2)
+    y_modulation = 1.0 + 0.15 * jnp.cos(2.0 * jnp.pi * y)
+    return setup.background_density + amplitude * profile * y_modulation
+
+
+def simulate_density_history(setup: DiffusionAutodiffSetup, theta: jnp.ndarray) -> jnp.ndarray:
+    physical = theta_to_physical(theta)
+    return simulate_density_history_from_physical(
+        setup,
+        anomalous_D=physical["anomalous_D"],
+        amplitude=physical["amplitude"],
+        center=physical["center"],
+        width=physical["width"],
+    )
+
+
+def simulate_density_history_from_physical(
+    setup: DiffusionAutodiffSetup,
+    *,
+    anomalous_D: jnp.ndarray,
+    amplitude: jnp.ndarray,
+    center: jnp.ndarray,
+    width: jnp.ndarray,
+) -> jnp.ndarray:
+    density0 = design_field_from_physical(
+        setup,
+        anomalous_D=anomalous_D,
+        amplitude=amplitude,
+        center=center,
+        width=width,
+    )
+    history = advance_anomalous_diffusion_history(
+        density0,
+        density0,
+        mesh=setup.mesh,
+        metrics=setup.metrics,
+        anomalous_D=anomalous_D,
+        density_boundary=setup.density_boundary,
+        pressure_boundary=setup.pressure_boundary,
+        timestep=setup.timestep,
+        steps=setup.steps,
+    )
+    return history.density_history
+
+
+def active_density_slice(setup: DiffusionAutodiffSetup, field: jnp.ndarray) -> jnp.ndarray:
+    return field[setup.mesh.xstart : setup.mesh.xend + 1, setup.mesh.ystart : setup.mesh.yend + 1, 0]
+
+
+def objective_for_parameter_vector(
+    theta: jnp.ndarray,
+    setup: DiffusionAutodiffSetup,
+    *,
+    target_final_density: jnp.ndarray | None = None,
+    objective_kind: str = "target_misfit",
+) -> jnp.ndarray:
+    history = simulate_density_history(setup, theta)
+    final_density = active_density_slice(setup, history[-1])
+    if objective_kind == "variance":
+        centered = final_density - jnp.mean(final_density)
+        return jnp.mean(centered * centered)
+    if target_final_density is None:
+        raise ValueError("target_final_density is required for target_misfit objectives")
+    target = jnp.asarray(target_final_density, dtype=jnp.float64)
+    mismatch = final_density - target
+    return 0.5 * jnp.mean(mismatch * mismatch)
+
+
+def objective_for_physical_parameters(
+    parameters: jnp.ndarray,
+    setup: DiffusionAutodiffSetup,
+    *,
+    target_final_density: jnp.ndarray | None = None,
+    objective_kind: str = "target_misfit",
+) -> jnp.ndarray:
+    parameters = jnp.asarray(parameters, dtype=jnp.float64)
+    history = simulate_density_history_from_physical(
+        setup,
+        anomalous_D=parameters[0],
+        amplitude=parameters[1],
+        center=parameters[2],
+        width=parameters[3],
+    )
+    final_density = active_density_slice(setup, history[-1])
+    if objective_kind == "variance":
+        centered = final_density - jnp.mean(final_density)
+        return jnp.mean(centered * centered)
+    if target_final_density is None:
+        raise ValueError("target_final_density is required for target_misfit objectives")
+    target = jnp.asarray(target_final_density, dtype=jnp.float64)
+    mismatch = final_density - target
+    return 0.5 * jnp.mean(mismatch * mismatch)
+
+
+def finite_difference_gradient(
+    objective_fn,
+    theta: jnp.ndarray,
+    *,
+    epsilon: float = 1.0e-4,
+) -> jnp.ndarray:
+    theta = jnp.asarray(theta, dtype=jnp.float64)
+    basis = jnp.eye(theta.size, dtype=jnp.float64)
+    values = []
+    for direction in basis:
+        forward = objective_fn(theta + epsilon * direction)
+        backward = objective_fn(theta - epsilon * direction)
+        values.append((forward - backward) / (2.0 * epsilon))
+    return jnp.asarray(values, dtype=jnp.float64)
+
+
+def optimize_inverse_design(
+    objective_fn,
+    theta0: jnp.ndarray,
+    *,
+    iterations: int = 80,
+    learning_rate: float = 0.05,
+    beta1: float = 0.9,
+    beta2: float = 0.999,
+    epsilon: float = 1.0e-8,
+) -> dict[str, Any]:
+    theta = jnp.asarray(theta0, dtype=jnp.float64)
+    m = jnp.zeros_like(theta)
+    v = jnp.zeros_like(theta)
+    loss_history: list[float] = []
+    theta_history: list[np.ndarray] = [np.asarray(theta)]
+    value_and_grad_fn = jax.jit(jax.value_and_grad(objective_fn))
+    objective_only_fn = jax.jit(objective_fn)
+
+    for iteration in range(1, iterations + 1):
+        loss_value, gradient = value_and_grad_fn(theta)
+        m = beta1 * m + (1.0 - beta1) * gradient
+        v = beta2 * v + (1.0 - beta2) * (gradient * gradient)
+        m_hat = m / (1.0 - beta1**iteration)
+        v_hat = v / (1.0 - beta2**iteration)
+        theta = theta - learning_rate * m_hat / (jnp.sqrt(v_hat) + epsilon)
+        loss_history.append(float(loss_value))
+        theta_history.append(np.asarray(theta))
+
+    final_loss = float(objective_only_fn(theta))
+    return {
+        "theta": theta,
+        "loss_history": np.asarray(loss_history, dtype=np.float64),
+        "theta_history": np.asarray(theta_history, dtype=np.float64),
+        "final_loss": final_loss,
+    }
+
+
+def compute_strong_scaling_points(points: list[tuple[int, float]], *, backend: str) -> list[StrongScalingPoint]:
+    if not points:
+        return []
+    ordered = sorted(points, key=lambda item: item[0])
+    baseline_devices, baseline_time = ordered[0]
+    result: list[StrongScalingPoint] = []
+    for device_count, elapsed_seconds in ordered:
+        speedup = baseline_time / elapsed_seconds
+        efficiency = speedup / float(device_count / baseline_devices)
+        result.append(
+            StrongScalingPoint(
+                backend=backend,
+                device_count=device_count,
+                elapsed_seconds=elapsed_seconds,
+                speedup=speedup,
+                efficiency=efficiency,
+            )
+        )
+    return result

drbx/validation/autodiff_diffusion_uncertainty.py

@@ -0,0 +1,235 @@
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass
+from pathlib import Path
+
+import jax
+import jax.numpy as jnp
+from matplotlib import pyplot as plt
+import numpy as np
+
+from .autodiff_diffusion import (
+    active_density_slice,
+    build_diffusion_autodiff_setup,
+    objective_for_physical_parameters,
+    simulate_density_history_from_physical,
+)
+from .publication_plotting import save_publication_figure, style_axis
+
+
+@dataclass(frozen=True)
+class AutodiffDiffusionUncertaintyArtifacts:
+    analysis_json_path: Path
+    arrays_npz_path: Path
+    plot_png_path: Path
+
+
+def create_autodiff_diffusion_uncertainty_package(
+    *,
+    output_root: str | Path,
+    case_label: str = "autodiff_diffusion_uncertainty",
+    sample_count: int = 96,
+    random_seed: int = 7,
+) -> AutodiffDiffusionUncertaintyArtifacts:
+    root = Path(output_root)
+    data_dir = root / "data"
+    images_dir = root / "images"
+    data_dir.mkdir(parents=True, exist_ok=True)
+    images_dir.mkdir(parents=True, exist_ok=True)
+
+    report, arrays = build_autodiff_diffusion_uncertainty_report(
+        sample_count=sample_count,
+        random_seed=random_seed,
+    )
+    analysis_json_path = data_dir / f"{case_label}_analysis.json"
+    arrays_npz_path = data_dir / f"{case_label}_arrays.npz"
+    plot_png_path = images_dir / f"{case_label}.png"
+    analysis_json_path.write_text(json.dumps(report, indent=2, sort_keys=True), encoding="utf-8")
+    np.savez_compressed(arrays_npz_path, **arrays)
+    save_autodiff_diffusion_uncertainty_plot(report, arrays, plot_png_path)
+    return AutodiffDiffusionUncertaintyArtifacts(
+        analysis_json_path=analysis_json_path,
+        arrays_npz_path=arrays_npz_path,
+        plot_png_path=plot_png_path,
+    )
+
+
+def build_autodiff_diffusion_uncertainty_report(
+    *,
+    sample_count: int = 96,
+    random_seed: int = 7,
+) -> tuple[dict[str, object], dict[str, np.ndarray]]:
+    setup = build_diffusion_autodiff_setup(nx=160, ny=24, timestep=3.0, steps=8)
+
+    mean_parameters = jnp.asarray([0.34, 0.15, 0.48, 0.14], dtype=jnp.float64)
+    parameter_std = jnp.asarray([0.03, 0.025, 0.04, 0.02], dtype=jnp.float64)
+    covariance = jnp.diag(parameter_std**2)
+
+    def radial_profile(parameters: jnp.ndarray) -> jnp.ndarray:
+        history = simulate_density_history_from_physical(
+            setup,
+            anomalous_D=parameters[0],
+            amplitude=parameters[1],
+            center=parameters[2],
+            width=parameters[3],
+        )
+        final_density = active_density_slice(setup, history[-1])
+        return jnp.mean(final_density, axis=1)
+
+    def scalar_qoi(parameters: jnp.ndarray) -> jnp.ndarray:
+        return objective_for_physical_parameters(
+            parameters,
+            setup,
+            objective_kind="variance",
+        )
+
+    mean_profile = radial_profile(mean_parameters)
+    mean_qoi = scalar_qoi(mean_parameters)
+    profile_jacobian = jax.jacobian(radial_profile)(mean_parameters)
+    qoi_gradient = jax.grad(scalar_qoi)(mean_parameters)
+
+    linearized_profile_variance = jnp.einsum("ik,kl,il->i", profile_jacobian, covariance, profile_jacobian)
+    linearized_profile_sigma = jnp.sqrt(jnp.maximum(linearized_profile_variance, 0.0))
+    linearized_qoi_sigma = jnp.sqrt(jnp.maximum(qoi_gradient @ covariance @ qoi_gradient, 0.0))
+
+    rng = np.random.default_rng(random_seed)
+    raw_samples = rng.normal(
+        loc=np.asarray(mean_parameters, dtype=np.float64),
+        scale=np.asarray(parameter_std, dtype=np.float64),
+        size=(sample_count, mean_parameters.size),
+    )
+    sampled_parameters = np.asarray(raw_samples, dtype=np.float64)
+    sampled_parameters[:, 0] = np.clip(sampled_parameters[:, 0], 0.08, None)
+    sampled_parameters[:, 1] = np.clip(sampled_parameters[:, 1], 0.03, None)
+    sampled_parameters[:, 2] = np.clip(sampled_parameters[:, 2], 0.05, 0.95)
+    sampled_parameters[:, 3] = np.clip(sampled_parameters[:, 3], 0.05, 0.28)
+
+    batched_profile = jax.jit(jax.vmap(radial_profile))
+    batched_qoi = jax.jit(jax.vmap(scalar_qoi))
+    sampled_profiles = np.asarray(batched_profile(jnp.asarray(sampled_parameters, dtype=jnp.float64)), dtype=np.float64)
+    sampled_qoi = np.asarray(batched_qoi(jnp.asarray(sampled_parameters, dtype=jnp.float64)), dtype=np.float64)
+
+    monte_carlo_profile_mean = sampled_profiles.mean(axis=0)
+    monte_carlo_profile_sigma = sampled_profiles.std(axis=0, ddof=1)
+    monte_carlo_qoi_mean = float(sampled_qoi.mean())
+    monte_carlo_qoi_sigma = float(sampled_qoi.std(ddof=1))
+
+    profile_sigma_gap = float(np.max(np.abs(monte_carlo_profile_sigma - np.asarray(linearized_profile_sigma))))
+    qoi_sigma_relative_error = float(
+        abs(float(linearized_qoi_sigma) - monte_carlo_qoi_sigma) / max(monte_carlo_qoi_sigma, 1.0e-12)
+    )
+
+    report = {
+        "case": "autodiff_diffusion_uncertainty",
+        "sample_count": sample_count,
+        "random_seed": random_seed,
+        "mean_parameters": np.asarray(mean_parameters, dtype=np.float64).tolist(),
+        "parameter_std": np.asarray(parameter_std, dtype=np.float64).tolist(),
+        "scalar_qoi_name": "final_density_variance",
+        "scalar_qoi_at_mean": float(mean_qoi),
+        "linearized_qoi_sigma": float(linearized_qoi_sigma),
+        "monte_carlo_qoi_mean": monte_carlo_qoi_mean,
+        "monte_carlo_qoi_sigma": monte_carlo_qoi_sigma,
+        "qoi_sigma_relative_error": qoi_sigma_relative_error,
+        "profile_sigma_max_abs_gap": profile_sigma_gap,
+        "interpretation": [
+            "The scalar QoI uses the variance of the final active-domain density field on the compact native diffusion lane.",
+            "The profile uncertainty uses the radial mean of the final active-domain density field.",
+            "The comparison is between a first-order autodiff pushforward of the parameter covariance and a vectorized Monte Carlo estimate on the same native solve path.",
+        ],
+    }
+    arrays = {
+        "mean_parameters": np.asarray(mean_parameters, dtype=np.float64),
+        "parameter_std": np.asarray(parameter_std, dtype=np.float64),
+        "radial_coordinate": np.linspace(0.0, 1.0, mean_profile.shape[0], dtype=np.float64),
+        "mean_profile": np.asarray(mean_profile, dtype=np.float64),
+        "linearized_profile_sigma": np.asarray(linearized_profile_sigma, dtype=np.float64),
+        "monte_carlo_profile_mean": monte_carlo_profile_mean,
+        "monte_carlo_profile_sigma": monte_carlo_profile_sigma,
+        "sampled_qoi": sampled_qoi,
+        "sampled_parameters": sampled_parameters,
+        "qoi_gradient": np.asarray(qoi_gradient, dtype=np.float64),
+    }
+    return report, arrays
+
+
+def save_autodiff_diffusion_uncertainty_plot(
+    report: dict[str, object],
+    arrays: dict[str, np.ndarray],
+    path: str | Path,
+) -> Path:
+    target = Path(path)
+    target.parent.mkdir(parents=True, exist_ok=True)
+
+    figure, axes = plt.subplots(1, 3, figsize=(16.2, 4.8), constrained_layout=True)
+
+    labels = ["D", "A", "center", "width"]
+    x = np.arange(len(labels))
+    axes[0].bar(x, arrays["mean_parameters"], color="#4c78a8", alpha=0.85)
+    axes[0].errorbar(
+        x,
+        arrays["mean_parameters"],
+        yerr=arrays["parameter_std"],
+        fmt="none",
+        ecolor="#111111",
+        elinewidth=2.0,
+        capsize=5,
+    )
+    axes[0].set_xticks(x, labels)
+    style_axis(axes[0], title="Input uncertainty model", ylabel="parameter value", grid="y")
+
+    sampled_qoi = arrays["sampled_qoi"]
+    axes[1].hist(sampled_qoi, bins=18, color="#72b7b2", alpha=0.8, density=True, label="Monte Carlo")
+    qoi_mean = float(report["scalar_qoi_at_mean"])
+    qoi_sigma = float(report["linearized_qoi_sigma"])
+    qoi_x = np.linspace(sampled_qoi.min(), sampled_qoi.max(), 240, dtype=np.float64)
+    normal_pdf = np.exp(-0.5 * ((qoi_x - qoi_mean) / max(qoi_sigma, 1.0e-12)) ** 2) / max(
+        qoi_sigma * np.sqrt(2.0 * np.pi),
+        1.0e-12,
+    )
+    axes[1].plot(qoi_x, normal_pdf, color="#d1495b", linewidth=2.4, label="Linearized Gaussian")
+    axes[1].axvline(float(report["monte_carlo_qoi_mean"]), color="#111111", linestyle="--", linewidth=1.5)
+    style_axis(
+        axes[1],
+        title="Scalar uncertainty propagation",
+        xlabel="final-density variance QoI",
+        ylabel="density",
+        grid="both",
+    )
+    axes[1].legend(frameon=False)
+
+    radial = arrays["radial_coordinate"]
+    mean_profile = arrays["mean_profile"]
+    linear_sigma = arrays["linearized_profile_sigma"]
+    mc_mean = arrays["monte_carlo_profile_mean"]
+    mc_sigma = arrays["monte_carlo_profile_sigma"]
+    axes[2].plot(radial, mean_profile, color="#111111", linewidth=2.5, label="mean profile")
+    axes[2].fill_between(
+        radial,
+        mean_profile - 2.0 * linear_sigma,
+        mean_profile + 2.0 * linear_sigma,
+        color="#f58518",
+        alpha=0.28,
+        label="linearized 95% band",
+    )
+    axes[2].plot(radial, mc_mean, color="#2a9d8f", linewidth=2.0, linestyle="--", label="MC mean profile")
+    axes[2].fill_between(
+        radial,
+        mc_mean - 2.0 * mc_sigma,
+        mc_mean + 2.0 * mc_sigma,
+        color="#54a24b",
+        alpha=0.18,
+        label="Monte Carlo 95% band",
+    )
+    style_axis(
+        axes[2],
+        title="Profile uncertainty pushforward",
+        xlabel="normalized radial coordinate",
+        ylabel="radial mean density",
+        grid="both",
+    )
+    axes[2].legend(frameon=False, fontsize=9)
+
+    save_publication_figure(figure, target)
+    return target