F1 part 4 publication offline

In [ ]:

"""Offline publication-quality run for ``f1_part_4_sobol_cfd_budget.ipynb``.

Produces ``media/f1_part_4_publication.npz`` from which the notebook loads its
headline Sobol decomposition + variance-reduction-per-CFD-hour numbers in
``MODE = "publication"`` (the default).

The notebook itself runs a coarse Sobol decomposition at ``N=256`` and a small
strategy comparison (10 LHS re-fit repetitions, 32 candidate cells per
strategy) so reader execution is fast (<4 min). This script runs the *real*
study at ``N=4096`` Sobol samples and 200 LHS re-fit repetitions, plus the
full 50-cell strategy comparison and the FIA-ATR-vs-grid-position sweep.

Expected wall-time: **~25-40 minutes on a developer machine (M1 Max / 16 cores)**.

Run from the repo root:

.. code-block:: bash

    python docs/examples/media/f1_part_4_publication_offline.py

The script is **idempotent**: running it twice produces byte-equal NPZ files
(same PRNG seed, same numerical conditioning).

NB: this script duplicates the notebook's surrogate / LTS definitions inline.
When the Part 1-4 plant definitions are factored out into a shared module
(planned), refactor this script to import from that module instead.
"""
from __future__ import annotations

"""Offline publication-quality run for ``f1_part_4_sobol_cfd_budget.ipynb``. Produces ``media/f1_part_4_publication.npz`` from which the notebook loads its headline Sobol decomposition + variance-reduction-per-CFD-hour numbers in ``MODE = "publication"`` (the default). The notebook itself runs a coarse Sobol decomposition at ``N=256`` and a small strategy comparison (10 LHS re-fit repetitions, 32 candidate cells per strategy) so reader execution is fast (<4 min). This script runs the *real* study at ``N=4096`` Sobol samples and 200 LHS re-fit repetitions, plus the full 50-cell strategy comparison and the FIA-ATR-vs-grid-position sweep. Expected wall-time: **~25-40 minutes on a developer machine (M1 Max / 16 cores)**. Run from the repo root: .. code-block:: bash python docs/examples/media/f1_part_4_publication_offline.py The script is **idempotent**: running it twice produces byte-equal NPZ files (same PRNG seed, same numerical conditioning). NB: this script duplicates the notebook's surrogate / LTS definitions inline. When the Part 1-4 plant definitions are factored out into a shared module (planned), refactor this script to import from that module instead. """ from __future__ import annotations

In [ ]:

import os
import sys
import time
from pathlib import Path

import os import sys import time from pathlib import Path

In [ ]:

import numpy as np

import numpy as np

In [ ]:

import jax
jax.config.update("jax_enable_x64", True)
import jax.numpy as jnp
from scipy.stats import qmc

import jax jax.config.update("jax_enable_x64", True) import jax.numpy as jnp from scipy.stats import qmc

In [ ]:

# jaxonomy
from jaxonomy.uq import (
    sobol_indices,
    decompose_variance_sobol,
    vmap_qoi,
    Uniform,
)
from jaxonomy.library.lookup_table import interp_nd

# jaxonomy from jaxonomy.uq import ( sobol_indices, decompose_variance_sobol, vmap_qoi, Uniform, ) from jaxonomy.library.lookup_table import interp_nd

In [ ]:

# ---------------------------------------------------------------------------
# Aero ground truth (copied verbatim from Parts 3 / 4 §3).
# ---------------------------------------------------------------------------
AERO_INPUT_LO     = jnp.array([18.0, 30.0, -2.0, -3.0, -20.0])
AERO_INPUT_HI     = jnp.array([40.0, 60.0, +2.0, +3.0, +20.0])
AERO_INPUT_NOMINAL = jnp.array([25.0, 40.0, 0.0, 0.0, 0.0])

# --------------------------------------------------------------------------- # Aero ground truth (copied verbatim from Parts 3 / 4 §3). # --------------------------------------------------------------------------- AERO_INPUT_LO = jnp.array([18.0, 30.0, -2.0, -3.0, -20.0]) AERO_INPUT_HI = jnp.array([40.0, 60.0, +2.0, +3.0, +20.0]) AERO_INPUT_NOMINAL = jnp.array([25.0, 40.0, 0.0, 0.0, 0.0])

In [ ]:

CLA_NOMINAL  = 3.5
CDA_NOMINAL  = 1.1
XCOP_NOMINAL = 0.42

CLA_NOMINAL = 3.5 CDA_NOMINAL = 1.1 XCOP_NOMINAL = 0.42

In [ ]:

A_HF_L     = 0.020
A_HR_L     = -0.005
A_RAKE_L   = 0.004
A_PHI2_L   = 0.005
A_BETA2_L  = 0.015
A_DELTA2_L = 0.00015

A_HF_L = 0.020 A_HR_L = -0.005 A_RAKE_L = 0.004 A_PHI2_L = 0.005 A_BETA2_L = 0.015 A_DELTA2_L = 0.00015

In [ ]:

A_HF_D     = -0.004
A_HR_D     = -0.002
A_PHI2_D   = 0.003
A_BETA2_D  = 0.020
A_DELTA2_D = 0.0002

A_HF_D = -0.004 A_HR_D = -0.002 A_PHI2_D = 0.003 A_BETA2_D = 0.020 A_DELTA2_D = 0.0002

In [ ]:

C_RAKE_COP  = -0.001
C_BETA_COP  = -0.010
C_DELTA_COP = +0.001

C_RAKE_COP = -0.001 C_BETA_COP = -0.010 C_DELTA_COP = +0.001

In [ ]:

@jax.jit
def aero_true(x):
    h_F, h_R, phi, beta, delta = (
        x[..., 0], x[..., 1], x[..., 2], x[..., 3], x[..., 4]
    )
    rake = h_R - h_F
    cla = CLA_NOMINAL * (
        1.0
        + A_HF_L * (25.0 - h_F)
        + A_HR_L * (h_R - 40.0)
        + A_RAKE_L * (rake - 15.0)
        - A_PHI2_L * phi * phi
        - A_BETA2_L * beta * beta
        - A_DELTA2_L * delta * delta
    )
    cda = CDA_NOMINAL * (
        1.0
        + A_HF_D * (25.0 - h_F)
        + A_HR_D * (h_R - 40.0)
        + A_PHI2_D * phi * phi
        + A_BETA2_D * beta * beta
        + A_DELTA2_D * delta * delta
    )
    xcop = (
        XCOP_NOMINAL
        + C_RAKE_COP * (rake - 15.0)
        + C_BETA_COP * beta
        + C_DELTA_COP * delta
    )
    return jnp.stack([cla, cda, xcop], axis=-1)

@jax.jit def aero_true(x): h_F, h_R, phi, beta, delta = ( x[..., 0], x[..., 1], x[..., 2], x[..., 3], x[..., 4] ) rake = h_R - h_F cla = CLA_NOMINAL * ( 1.0 + A_HF_L * (25.0 - h_F) + A_HR_L * (h_R - 40.0) + A_RAKE_L * (rake - 15.0) - A_PHI2_L * phi * phi - A_BETA2_L * beta * beta - A_DELTA2_L * delta * delta ) cda = CDA_NOMINAL * ( 1.0 + A_HF_D * (25.0 - h_F) + A_HR_D * (h_R - 40.0) + A_PHI2_D * phi * phi + A_BETA2_D * beta * beta + A_DELTA2_D * delta * delta ) xcop = ( XCOP_NOMINAL + C_RAKE_COP * (rake - 15.0) + C_BETA_COP * beta + C_DELTA_COP * delta ) return jnp.stack([cla, cda, xcop], axis=-1)

In [ ]:

def cfd_noise_std(y_true):
    sig_cla = jnp.maximum(0.02, 0.020 * jnp.abs(y_true[..., 0]))
    sig_cda = jnp.maximum(0.01, 0.025 * jnp.abs(y_true[..., 1]))
    sig_cop = jnp.full_like(y_true[..., 2], 0.005)
    return jnp.stack([sig_cla, sig_cda, sig_cop], axis=-1)

def cfd_noise_std(y_true): sig_cla = jnp.maximum(0.02, 0.020 * jnp.abs(y_true[..., 0])) sig_cda = jnp.maximum(0.01, 0.025 * jnp.abs(y_true[..., 1])) sig_cop = jnp.full_like(y_true[..., 2], 0.005) return jnp.stack([sig_cla, sig_cda, sig_cop], axis=-1)

In [ ]:

def cfd_probe(x, key):
    y = aero_true(x)
    sigma = cfd_noise_std(y)
    eps = jax.random.normal(key, shape=y.shape)
    return y + sigma * eps

def cfd_probe(x, key): y = aero_true(x) sigma = cfd_noise_std(y) eps = jax.random.normal(key, shape=y.shape) return y + sigma * eps

In [ ]:

GRID_AXES_NORM = (
    jnp.linspace(0.0, 1.0, 3),
    jnp.linspace(0.0, 1.0, 3),
    jnp.linspace(0.0, 1.0, 2),
    jnp.linspace(0.0, 1.0, 2),
    jnp.linspace(0.0, 1.0, 2),
)
FIT_SMOOTHNESS = 1e-2

GRID_AXES_NORM = ( jnp.linspace(0.0, 1.0, 3), jnp.linspace(0.0, 1.0, 3), jnp.linspace(0.0, 1.0, 2), jnp.linspace(0.0, 1.0, 2), jnp.linspace(0.0, 1.0, 2), ) FIT_SMOOTHNESS = 1e-2

In [ ]:

def fit_table_nd(grid_axes, x_data, y_data, smoothness=FIT_SMOOTHNESS, rcond=None):
    K = x_data.shape[0]
    N_dim = len(grid_axes)
    Bs = tuple(int(g.shape[0]) for g in grid_axes)
    total = int(np.prod(Bs))
    idx_list, alpha_list = [], []
    for d in range(N_dim):
        x = jnp.asarray(grid_axes[d])
        q = x_data[:, d]
        i = jnp.clip(jnp.searchsorted(x, q, side="right") - 1, 0, Bs[d] - 2)
        x_lo, x_hi = x[i], x[i + 1]
        alpha = jnp.clip((q - x_lo) / (x_hi - x_lo), 0.0, 1.0)
        idx_list.append(i); alpha_list.append(alpha)
    A = jnp.zeros((K, total))
    row_idx = jnp.arange(K)
    for corner in range(2 ** N_dim):
        w = jnp.ones(K)
        flat_idx = jnp.zeros(K, dtype=jnp.int32)
        for d in range(N_dim):
            bit = (corner >> d) & 1
            w = w * (alpha_list[d] if bit else (1.0 - alpha_list[d]))
            stride = int(np.prod(Bs[d + 1:]))
            flat_idx = flat_idx + (idx_list[d] + bit) * stride
        A = A.at[row_idx, flat_idx].add(w)
    A_reg = jnp.concatenate([A, jnp.sqrt(smoothness) * jnp.eye(total)], axis=0)
    b_reg = jnp.concatenate([y_data, jnp.zeros(total)], axis=0)
    v_flat, *_ = jnp.linalg.lstsq(A_reg, b_reg, rcond=rcond)
    return v_flat.reshape(Bs)

def fit_table_nd(grid_axes, x_data, y_data, smoothness=FIT_SMOOTHNESS, rcond=None): K = x_data.shape[0] N_dim = len(grid_axes) Bs = tuple(int(g.shape[0]) for g in grid_axes) total = int(np.prod(Bs)) idx_list, alpha_list = [], [] for d in range(N_dim): x = jnp.asarray(grid_axes[d]) q = x_data[:, d] i = jnp.clip(jnp.searchsorted(x, q, side="right") - 1, 0, Bs[d] - 2) x_lo, x_hi = x[i], x[i + 1] alpha = jnp.clip((q - x_lo) / (x_hi - x_lo), 0.0, 1.0) idx_list.append(i); alpha_list.append(alpha) A = jnp.zeros((K, total)) row_idx = jnp.arange(K) for corner in range(2 ** N_dim): w = jnp.ones(K) flat_idx = jnp.zeros(K, dtype=jnp.int32) for d in range(N_dim): bit = (corner >> d) & 1 w = w * (alpha_list[d] if bit else (1.0 - alpha_list[d])) stride = int(np.prod(Bs[d + 1:])) flat_idx = flat_idx + (idx_list[d] + bit) * stride A = A.at[row_idx, flat_idx].add(w) A_reg = jnp.concatenate([A, jnp.sqrt(smoothness) * jnp.eye(total)], axis=0) b_reg = jnp.concatenate([y_data, jnp.zeros(total)], axis=0) v_flat, *_ = jnp.linalg.lstsq(A_reg, b_reg, rcond=rcond) return v_flat.reshape(Bs)

In [ ]:

# ---------------------------------------------------------------------------
# Lap-time proxy. We use the analytic single-corner cornering-speed proxy
# rather than the full LTS (which is too expensive at N=4096 *(d+2)).
# Equation (Part-1 §10): V^2 = mu_eff(V) * g * R, with mu_eff augmented by
# downforce. Solving the implicit equation for V gives a closed-form lap-time
# surrogate that responds to the same aero levers as the full LTS.
# ---------------------------------------------------------------------------
RHO_AIR = 1.225
M_CAR   = 830.0
G_ACC   = 9.81
MU_PEAK = 1.7
R_CORNER = 100.0       # reference corner radius [m]
L_STRAIGHT = 800.0     # reference straight length [m]
P_PEAK_KW = 700.0
# Yaw-aware effective corner share: assumes a fraction of the lap spent
# cornering vs straight-lining; calibrated against the Part 1 lap (~60% of
# 70 s is cornering).
T_CORNER_SHARE = 0.60

# --------------------------------------------------------------------------- # Lap-time proxy. We use the analytic single-corner cornering-speed proxy # rather than the full LTS (which is too expensive at N=4096 *(d+2)). # Equation (Part-1 §10): V^2 = mu_eff(V) * g * R, with mu_eff augmented by # downforce. Solving the implicit equation for V gives a closed-form lap-time # surrogate that responds to the same aero levers as the full LTS. # --------------------------------------------------------------------------- RHO_AIR = 1.225 M_CAR = 830.0 G_ACC = 9.81 MU_PEAK = 1.7 R_CORNER = 100.0 # reference corner radius [m] L_STRAIGHT = 800.0 # reference straight length [m] P_PEAK_KW = 700.0 # Yaw-aware effective corner share: assumes a fraction of the lap spent # cornering vs straight-lining; calibrated against the Part 1 lap (~60% of # 70 s is cornering). T_CORNER_SHARE = 0.60

In [ ]:

def lap_time_proxy(x_aero):
    """Single-corner + single-straight lap-time proxy.

    Combines two physics-rooted pieces:
    1. Cornering speed V_c from the implicit equation
       V_c^2 (1 - alpha_aero) = mu * g * R, with alpha_aero =
       mu * R * rho * C_L A / (2 m).  At alpha_aero >= 1 the formula's
       denominator vanishes — guarded with a clip.
    2. Straight-line time t_s ≈ sqrt(2 * L / a_x) with a_x = P_peak / (m * V_avg)
       and V_avg = V_top / 2; V_top capped by the power-vs-drag balance
       V_top = (2 * P / (rho * C_D A))^{1/3}.

    Returns total time T = T_CORNER_SHARE * (pi * R / V_c) + (1-..) * t_s.
    Lower = faster.  Responds to C_L A (raises both V_c and V_top — wins
    on cornering but loses to drag on straights), C_D A (slows V_top
    monotonically), and x_CoP (cornering balance, not modelled in this
    proxy; we add a small symmetric penalty on |x_CoP - 0.42| as a
    placebo).
    """
    y = aero_true(x_aero)
    cla, cda, xcop = y[..., 0], y[..., 1], y[..., 2]
    # Cornering speed (clip the denominator away from 1.0 for stability)
    alpha_aero = jnp.clip(MU_PEAK * R_CORNER * RHO_AIR * cla / (2.0 * M_CAR),
                          0.0, 0.85)
    Vc_sq = MU_PEAK * G_ACC * R_CORNER / (1.0 - alpha_aero)
    Vc = jnp.sqrt(Vc_sq)
    t_corner = jnp.pi * R_CORNER / Vc
    # Top speed from V^3 = 2 P / (rho C_D A)
    V_top = (2.0 * P_PEAK_KW * 1000.0 / (RHO_AIR * cda)) ** (1.0 / 3.0)
    V_avg = 0.5 * V_top
    a_x = P_PEAK_KW * 1000.0 / (M_CAR * jnp.maximum(V_avg, 10.0))
    t_straight = jnp.sqrt(2.0 * L_STRAIGHT / a_x)
    # CoP placebo: drift away from baseline penalises corner balance.
    t_balance = 50.0 * (xcop - XCOP_NOMINAL) ** 2
    return (T_CORNER_SHARE * t_corner
            + (1.0 - T_CORNER_SHARE) * t_straight
            + t_balance)

def lap_time_proxy(x_aero): """Single-corner + single-straight lap-time proxy. Combines two physics-rooted pieces: 1. Cornering speed V_c from the implicit equation V_c^2 (1 - alpha_aero) = mu * g * R, with alpha_aero = mu * R * rho * C_L A / (2 m). At alpha_aero >= 1 the formula's denominator vanishes — guarded with a clip. 2. Straight-line time t_s ≈ sqrt(2 * L / a_x) with a_x = P_peak / (m * V_avg) and V_avg = V_top / 2; V_top capped by the power-vs-drag balance V_top = (2 * P / (rho * C_D A))^{1/3}. Returns total time T = T_CORNER_SHARE * (pi * R / V_c) + (1-..) * t_s. Lower = faster. Responds to C_L A (raises both V_c and V_top — wins on cornering but loses to drag on straights), C_D A (slows V_top monotonically), and x_CoP (cornering balance, not modelled in this proxy; we add a small symmetric penalty on |x_CoP - 0.42| as a placebo). """ y = aero_true(x_aero) cla, cda, xcop = y[..., 0], y[..., 1], y[..., 2] # Cornering speed (clip the denominator away from 1.0 for stability) alpha_aero = jnp.clip(MU_PEAK * R_CORNER * RHO_AIR * cla / (2.0 * M_CAR), 0.0, 0.85) Vc_sq = MU_PEAK * G_ACC * R_CORNER / (1.0 - alpha_aero) Vc = jnp.sqrt(Vc_sq) t_corner = jnp.pi * R_CORNER / Vc # Top speed from V^3 = 2 P / (rho C_D A) V_top = (2.0 * P_PEAK_KW * 1000.0 / (RHO_AIR * cda)) ** (1.0 / 3.0) V_avg = 0.5 * V_top a_x = P_PEAK_KW * 1000.0 / (M_CAR * jnp.maximum(V_avg, 10.0)) t_straight = jnp.sqrt(2.0 * L_STRAIGHT / a_x) # CoP placebo: drift away from baseline penalises corner balance. t_balance = 50.0 * (xcop - XCOP_NOMINAL) ** 2 return (T_CORNER_SHARE * t_corner + (1.0 - T_CORNER_SHARE) * t_straight + t_balance)

In [ ]:

# ---------------------------------------------------------------------------
# Publication knobs.
# ---------------------------------------------------------------------------
RNG_SEED = 20260517

# --------------------------------------------------------------------------- # Publication knobs. # --------------------------------------------------------------------------- RNG_SEED = 20260517

In [ ]:

# Sobol matrix sizes for the headline indices + the variance decomposition.
# Per the Saltelli scheme, total qoi_fn evaluations =
# n_samples_sobol * (d + 2) = 4096 * 7 = 28672 for first/total order; and
# 4 * n_samples_decomp = 4 * 8192 = 32768 for the grouped decomposition.
SOBOL_N_PUB = 4096
DECOMP_N_PUB = 8192

# Sobol matrix sizes for the headline indices + the variance decomposition. # Per the Saltelli scheme, total qoi_fn evaluations = # n_samples_sobol * (d + 2) = 4096 * 7 = 28672 for first/total order; and # 4 * n_samples_decomp = 4 * 8192 = 32768 for the grouped decomposition. SOBOL_N_PUB = 4096 DECOMP_N_PUB = 8192

In [ ]:

# Strategy comparison: average across N_REPEATS LHS re-fits to denoise the
# variance-reduction estimate. Each repeat fits a fresh prior surrogate from
# N_PRIOR LHS probes, then a "next batch" of N_BATCH probes drawn either
# uniform-LHS or proportional-to-epistemic-variance, and re-fits + re-Sobols.
N_REPEATS_PUB = 200
N_PRIOR = 64
N_BATCH_LIST = np.array([8, 16, 32, 50])

# Strategy comparison: average across N_REPEATS LHS re-fits to denoise the # variance-reduction estimate. Each repeat fits a fresh prior surrogate from # N_PRIOR LHS probes, then a "next batch" of N_BATCH probes drawn either # uniform-LHS or proportional-to-epistemic-variance, and re-fits + re-Sobols. N_REPEATS_PUB = 200 N_PRIOR = 64 N_BATCH_LIST = np.array([8, 16, 32, 50])

In [ ]:

OUT_NPZ = Path(__file__).resolve().parent / "f1_part_4_publication.npz"

OUT_NPZ = Path(__file__).resolve().parent / "f1_part_4_publication.npz"

---

Sobol decomposition over the 5-D aero design space.¶

In [ ]:

def _surrogate_qoi(V_cla, V_cda, V_cop):
    """Closure that maps a (5,) physical aero input -> lap time through the
    fitted surrogate."""
    def _qoi_scalar(params):
        x_phys = jnp.array([
            params["h_F"], params["h_R"], params["phi"],
            params["beta"], params["delta"],
        ])
        x_norm = (x_phys - AERO_INPUT_LO) / (AERO_INPUT_HI - AERO_INPUT_LO)
        cla = interp_nd(GRID_AXES_NORM, V_cla, x_norm)
        cda = interp_nd(GRID_AXES_NORM, V_cda, x_norm)
        cop = interp_nd(GRID_AXES_NORM, V_cop, x_norm)
        y_full = jnp.array([cla, cda, cop])
        # lap_time_proxy expects the unfitted aero, but we want to drive it
        # off the *fitted* surface here, so we pack the fitted (cla, cda,
        # cop) into a fake x_aero tail call.
        # Simpler: re-evaluate the proxy with the fitted (cla, cda, cop)
        # directly.
        alpha_aero = jnp.clip(MU_PEAK * R_CORNER * RHO_AIR * cla
                              / (2.0 * M_CAR), 0.0, 0.85)
        Vc_sq = MU_PEAK * G_ACC * R_CORNER / (1.0 - alpha_aero)
        Vc = jnp.sqrt(Vc_sq)
        t_corner = jnp.pi * R_CORNER / Vc
        V_top = (2.0 * P_PEAK_KW * 1000.0 / (RHO_AIR * cda)) ** (1.0 / 3.0)
        a_x = P_PEAK_KW * 1000.0 / (M_CAR * jnp.maximum(0.5 * V_top, 10.0))
        t_straight = jnp.sqrt(2.0 * L_STRAIGHT / a_x)
        t_balance = 50.0 * (cop - XCOP_NOMINAL) ** 2
        return (T_CORNER_SHARE * t_corner
                + (1.0 - T_CORNER_SHARE) * t_straight
                + t_balance)
    return vmap_qoi(_qoi_scalar)

def _surrogate_qoi(V_cla, V_cda, V_cop): """Closure that maps a (5,) physical aero input -> lap time through the fitted surrogate.""" def _qoi_scalar(params): x_phys = jnp.array([ params["h_F"], params["h_R"], params["phi"], params["beta"], params["delta"], ]) x_norm = (x_phys - AERO_INPUT_LO) / (AERO_INPUT_HI - AERO_INPUT_LO) cla = interp_nd(GRID_AXES_NORM, V_cla, x_norm) cda = interp_nd(GRID_AXES_NORM, V_cda, x_norm) cop = interp_nd(GRID_AXES_NORM, V_cop, x_norm) y_full = jnp.array([cla, cda, cop]) # lap_time_proxy expects the unfitted aero, but we want to drive it # off the *fitted* surface here, so we pack the fitted (cla, cda, # cop) into a fake x_aero tail call. # Simpler: re-evaluate the proxy with the fitted (cla, cda, cop) # directly. alpha_aero = jnp.clip(MU_PEAK * R_CORNER * RHO_AIR * cla / (2.0 * M_CAR), 0.0, 0.85) Vc_sq = MU_PEAK * G_ACC * R_CORNER / (1.0 - alpha_aero) Vc = jnp.sqrt(Vc_sq) t_corner = jnp.pi * R_CORNER / Vc V_top = (2.0 * P_PEAK_KW * 1000.0 / (RHO_AIR * cda)) ** (1.0 / 3.0) a_x = P_PEAK_KW * 1000.0 / (M_CAR * jnp.maximum(0.5 * V_top, 10.0)) t_straight = jnp.sqrt(2.0 * L_STRAIGHT / a_x) t_balance = 50.0 * (cop - XCOP_NOMINAL) ** 2 return (T_CORNER_SHARE * t_corner + (1.0 - T_CORNER_SHARE) * t_straight + t_balance) return vmap_qoi(_qoi_scalar)

In [ ]:

def _aleatoric_epistemic_qoi(V_cla, V_cda, V_cop, residual_std):
    """Closure that wraps the surrogate lap-time and adds two synthetic noise
    inputs (epistemic-on-CLA and aleatoric-on-CLA) so the Sobol grouped
    decomposition has tagged inputs to attribute variance to."""
    base = _surrogate_qoi(V_cla, V_cda, V_cop)

    def _qoi_split(params):
        # Reconstruct the (5,) physical input from the named params and
        # then add the two noise inputs as additive jitter on CLA at
        # query-time.
        x_phys = jnp.array([
            params["h_F"], params["h_R"], params["phi"],
            params["beta"], params["delta"],
        ])
        x_norm = (x_phys - AERO_INPUT_LO) / (AERO_INPUT_HI - AERO_INPUT_LO)
        cla = (interp_nd(GRID_AXES_NORM, V_cla, x_norm)
               + params["epistemic_cla"] + params["aleatoric_cla"])
        cda = interp_nd(GRID_AXES_NORM, V_cda, x_norm)
        cop = interp_nd(GRID_AXES_NORM, V_cop, x_norm)
        alpha_aero = jnp.clip(MU_PEAK * R_CORNER * RHO_AIR * cla
                              / (2.0 * M_CAR), 0.0, 0.85)
        Vc_sq = MU_PEAK * G_ACC * R_CORNER / (1.0 - alpha_aero)
        Vc = jnp.sqrt(Vc_sq)
        t_corner = jnp.pi * R_CORNER / Vc
        V_top = (2.0 * P_PEAK_KW * 1000.0 / (RHO_AIR * cda)) ** (1.0 / 3.0)
        a_x = P_PEAK_KW * 1000.0 / (M_CAR * jnp.maximum(0.5 * V_top, 10.0))
        t_straight = jnp.sqrt(2.0 * L_STRAIGHT / a_x)
        t_balance = 50.0 * (cop - XCOP_NOMINAL) ** 2
        return (T_CORNER_SHARE * t_corner
                + (1.0 - T_CORNER_SHARE) * t_straight
                + t_balance)
    return vmap_qoi(_qoi_split)

def _aleatoric_epistemic_qoi(V_cla, V_cda, V_cop, residual_std): """Closure that wraps the surrogate lap-time and adds two synthetic noise inputs (epistemic-on-CLA and aleatoric-on-CLA) so the Sobol grouped decomposition has tagged inputs to attribute variance to.""" base = _surrogate_qoi(V_cla, V_cda, V_cop) def _qoi_split(params): # Reconstruct the (5,) physical input from the named params and # then add the two noise inputs as additive jitter on CLA at # query-time. x_phys = jnp.array([ params["h_F"], params["h_R"], params["phi"], params["beta"], params["delta"], ]) x_norm = (x_phys - AERO_INPUT_LO) / (AERO_INPUT_HI - AERO_INPUT_LO) cla = (interp_nd(GRID_AXES_NORM, V_cla, x_norm) + params["epistemic_cla"] + params["aleatoric_cla"]) cda = interp_nd(GRID_AXES_NORM, V_cda, x_norm) cop = interp_nd(GRID_AXES_NORM, V_cop, x_norm) alpha_aero = jnp.clip(MU_PEAK * R_CORNER * RHO_AIR * cla / (2.0 * M_CAR), 0.0, 0.85) Vc_sq = MU_PEAK * G_ACC * R_CORNER / (1.0 - alpha_aero) Vc = jnp.sqrt(Vc_sq) t_corner = jnp.pi * R_CORNER / Vc V_top = (2.0 * P_PEAK_KW * 1000.0 / (RHO_AIR * cda)) ** (1.0 / 3.0) a_x = P_PEAK_KW * 1000.0 / (M_CAR * jnp.maximum(0.5 * V_top, 10.0)) t_straight = jnp.sqrt(2.0 * L_STRAIGHT / a_x) t_balance = 50.0 * (cop - XCOP_NOMINAL) ** 2 return (T_CORNER_SHARE * t_corner + (1.0 - T_CORNER_SHARE) * t_straight + t_balance) return vmap_qoi(_qoi_split)

In [ ]:

def make_prior_fit(rng_seed, N_prior):
    """Draw N_prior LHS samples + noisy CFD probes, fit the 3-output surrogate."""
    lhs = qmc.LatinHypercube(d=5, seed=rng_seed)
    unit = lhs.random(n=int(N_prior))
    x_phys = (np.asarray(AERO_INPUT_LO)
              + unit * np.asarray(AERO_INPUT_HI - AERO_INPUT_LO))
    keys = jax.random.split(jax.random.PRNGKey(rng_seed + 7000), int(N_prior))
    y_probes = np.asarray(jax.vmap(cfd_probe)(jnp.asarray(x_phys), keys))
    x_norm = jnp.asarray((x_phys - np.asarray(AERO_INPUT_LO))
                         / np.asarray(AERO_INPUT_HI - AERO_INPUT_LO))
    V_cla = fit_table_nd(GRID_AXES_NORM, x_norm, jnp.asarray(y_probes[:, 0]),
                         smoothness=FIT_SMOOTHNESS)
    V_cda = fit_table_nd(GRID_AXES_NORM, x_norm, jnp.asarray(y_probes[:, 1]),
                         smoothness=FIT_SMOOTHNESS)
    V_cop = fit_table_nd(GRID_AXES_NORM, x_norm, jnp.asarray(y_probes[:, 2]),
                         smoothness=FIT_SMOOTHNESS)
    return V_cla, V_cda, V_cop, x_phys, y_probes

def make_prior_fit(rng_seed, N_prior): """Draw N_prior LHS samples + noisy CFD probes, fit the 3-output surrogate.""" lhs = qmc.LatinHypercube(d=5, seed=rng_seed) unit = lhs.random(n=int(N_prior)) x_phys = (np.asarray(AERO_INPUT_LO) + unit * np.asarray(AERO_INPUT_HI - AERO_INPUT_LO)) keys = jax.random.split(jax.random.PRNGKey(rng_seed + 7000), int(N_prior)) y_probes = np.asarray(jax.vmap(cfd_probe)(jnp.asarray(x_phys), keys)) x_norm = jnp.asarray((x_phys - np.asarray(AERO_INPUT_LO)) / np.asarray(AERO_INPUT_HI - AERO_INPUT_LO)) V_cla = fit_table_nd(GRID_AXES_NORM, x_norm, jnp.asarray(y_probes[:, 0]), smoothness=FIT_SMOOTHNESS) V_cda = fit_table_nd(GRID_AXES_NORM, x_norm, jnp.asarray(y_probes[:, 1]), smoothness=FIT_SMOOTHNESS) V_cop = fit_table_nd(GRID_AXES_NORM, x_norm, jnp.asarray(y_probes[:, 2]), smoothness=FIT_SMOOTHNESS) return V_cla, V_cda, V_cop, x_phys, y_probes

In [ ]:

def main():
    np.random.seed(RNG_SEED)
    t_total = time.time()

    # -----------------------------------------------------------------------
    # Step 1: fit prior surrogate at N_PRIOR=64.
    # -----------------------------------------------------------------------
    print("[1/4] Fitting prior surrogate at N_PRIOR=64...")
    V_cla, V_cda, V_cop, x_prior, y_prior = make_prior_fit(RNG_SEED, N_PRIOR)
    print("       done.")

    # -----------------------------------------------------------------------
    # Step 2: Sobol indices over the 5 design axes at large N.
    # -----------------------------------------------------------------------
    print(f"[2/4] Computing Sobol indices at N_samples={SOBOL_N_PUB}...")
    distributions = {
        "h_F":   Uniform(float(AERO_INPUT_LO[0]), float(AERO_INPUT_HI[0])),
        "h_R":   Uniform(float(AERO_INPUT_LO[1]), float(AERO_INPUT_HI[1])),
        "phi":   Uniform(float(AERO_INPUT_LO[2]), float(AERO_INPUT_HI[2])),
        "beta":  Uniform(float(AERO_INPUT_LO[3]), float(AERO_INPUT_HI[3])),
        "delta": Uniform(float(AERO_INPUT_LO[4]), float(AERO_INPUT_HI[4])),
    }
    qoi_fn = _surrogate_qoi(V_cla, V_cda, V_cop)
    t0 = time.time()
    idx = sobol_indices(None, None, distributions, qoi_fn,
                        n_samples=SOBOL_N_PUB,
                        key=jax.random.PRNGKey(RNG_SEED + 100))
    print(f"       sobol wall = {time.time() - t0:.1f} s")
    sobol_first = np.array([idx[k]["first_order"] for k in distributions])
    sobol_total = np.array([idx[k]["total_order"] for k in distributions])
    sobol_names = list(distributions.keys())
    for n, s1, st in zip(sobol_names, sobol_first, sobol_total):
        print(f"       {n:6s}  S1 = {s1:+.4f}  ST = {st:+.4f}")

    # -----------------------------------------------------------------------
    # Step 3: aleatoric / epistemic / interaction split.
    # -----------------------------------------------------------------------
    print(f"[3/4] Computing aleatoric/epistemic/interaction split at "
          f"N_samples={DECOMP_N_PUB}...")
    # Aleatoric == irreducible CFD noise; epistemic == fit-residual std.
    # For the headline figure we calibrate epistemic_std from the actual fit's
    # residual std vs truth (sampled at 2000 points), and aleatoric_std from
    # the CFD noise model floor.
    n_eval = 2000
    unit_eval = qmc.LatinHypercube(d=5, seed=RNG_SEED + 200).random(n=n_eval)
    x_eval_phys = (np.asarray(AERO_INPUT_LO)
                   + unit_eval * np.asarray(AERO_INPUT_HI - AERO_INPUT_LO))
    Y_TRUE = np.asarray(jax.vmap(aero_true)(jnp.asarray(x_eval_phys)))
    x_eval_norm = jnp.asarray((x_eval_phys - np.asarray(AERO_INPUT_LO))
                              / np.asarray(AERO_INPUT_HI - AERO_INPUT_LO))
    Y_FIT_CLA = np.asarray(interp_nd(GRID_AXES_NORM, V_cla, x_eval_norm))
    epistemic_std = float(np.std(Y_TRUE[:, 0] - Y_FIT_CLA))
    aleatoric_std = float(np.mean(np.asarray(cfd_noise_std(jnp.asarray(Y_TRUE))[:, 0])))
    print(f"       calibrated epistemic_std (fit residual) = {epistemic_std:.4f}")
    print(f"       calibrated aleatoric_std (CFD noise floor) = {aleatoric_std:.4f}")
    aleatoric_dists = {
        # CFD irreducible noise on CLA
        "aleatoric_cla": Uniform(-2.0 * aleatoric_std, +2.0 * aleatoric_std,
                                 kind="aleatoric"),
    }
    epistemic_dists = {
        "h_F":   Uniform(float(AERO_INPUT_LO[0]), float(AERO_INPUT_HI[0]),
                         kind="epistemic"),
        "h_R":   Uniform(float(AERO_INPUT_LO[1]), float(AERO_INPUT_HI[1]),
                         kind="epistemic"),
        "phi":   Uniform(float(AERO_INPUT_LO[2]), float(AERO_INPUT_HI[2]),
                         kind="epistemic"),
        "beta":  Uniform(float(AERO_INPUT_LO[3]), float(AERO_INPUT_HI[3]),
                         kind="epistemic"),
        "delta": Uniform(float(AERO_INPUT_LO[4]), float(AERO_INPUT_HI[4]),
                         kind="epistemic"),
        "epistemic_cla": Uniform(-2.0 * epistemic_std, +2.0 * epistemic_std,
                                 kind="epistemic"),
    }
    decomp_qoi_fn = _aleatoric_epistemic_qoi(V_cla, V_cda, V_cop, epistemic_std)
    t0 = time.time()
    decomp = decompose_variance_sobol(
        decomp_qoi_fn, aleatoric_dists, epistemic_dists,
        n_samples=DECOMP_N_PUB,
        key=jax.random.PRNGKey(RNG_SEED + 300),
    )
    print(f"       decomp wall = {time.time() - t0:.1f} s")
    for k, v in decomp.items():
        print(f"       {k:18s}: {v:+.6e}")

    # -----------------------------------------------------------------------
    # Step 4: strategy A vs B head-to-head over N_REPEATS_PUB seeds.
    # -----------------------------------------------------------------------
    print(f"[4/4] Strategy comparison across N_REPEATS={N_REPEATS_PUB} seeds "
          f"and N_BATCH in {N_BATCH_LIST.tolist()}...")
    # Precompute Sobol weighting from the prior fit only (so the recommended
    # cells are independent of the per-repeat noisy probes).
    # Cell layout: 10 x 10 grid over (h_F, beta) at fixed (h_R, phi, delta).
    # This is the dominant 2D slice from the §5 Sobol ranking.
    NX, NY = 10, 10
    hf_centers = np.linspace(float(AERO_INPUT_LO[0]) + 1.0,
                             float(AERO_INPUT_HI[0]) - 1.0, NX)
    beta_centers = np.linspace(float(AERO_INPUT_LO[3]) + 0.3,
                               float(AERO_INPUT_HI[3]) - 0.3, NY)
    # Per-cell epistemic variance score: for a uniformly-spaced grid of
    # candidate next-CFD points, compute the surrogate's prediction variance
    # under bootstrap re-sampling of the prior probes' noise.
    # (Cheap proxy: variance of (lap_time(x) - lap_time_truth(x)) over a
    #  small radius around the cell centre.)
    cell_scores = np.zeros((NX, NY))
    for i, hf in enumerate(hf_centers):
        for j, bt in enumerate(beta_centers):
            x_cell = np.array([hf, 40.0, 0.0, bt, 0.0])
            y_truth_proxy = float(lap_time_proxy(jnp.asarray(x_cell)))
            # Fit-side proxy
            x_norm = (jnp.asarray(x_cell) - AERO_INPUT_LO) / (AERO_INPUT_HI - AERO_INPUT_LO)
            cla = float(interp_nd(GRID_AXES_NORM, V_cla, x_norm))
            cda = float(interp_nd(GRID_AXES_NORM, V_cda, x_norm))
            cop = float(interp_nd(GRID_AXES_NORM, V_cop, x_norm))
            # Compute proxy from the fit
            alpha_aero = min(MU_PEAK * R_CORNER * RHO_AIR * cla
                             / (2.0 * M_CAR), 0.85)
            Vc_sq = MU_PEAK * G_ACC * R_CORNER / (1.0 - alpha_aero)
            t_corner = float(np.pi * R_CORNER / np.sqrt(Vc_sq))
            V_top = (2.0 * P_PEAK_KW * 1000.0 / (RHO_AIR * cda)) ** (1.0 / 3.0)
            a_x = P_PEAK_KW * 1000.0 / (M_CAR * max(0.5 * V_top, 10.0))
            t_straight = float(np.sqrt(2.0 * L_STRAIGHT / a_x))
            t_balance = 50.0 * (cop - XCOP_NOMINAL) ** 2
            y_fit_proxy = (T_CORNER_SHARE * t_corner
                           + (1.0 - T_CORNER_SHARE) * t_straight
                           + t_balance)
            cell_scores[i, j] = (y_truth_proxy - y_fit_proxy) ** 2

    # Normalize to a probability distribution for proportional sampling.
    cell_probs = cell_scores / cell_scores.sum()

    # Reference variance with the prior fit only.
    qoi_prior = _surrogate_qoi(V_cla, V_cda, V_cop)
    sob_prior = sobol_indices(None, None, distributions, qoi_prior,
                              n_samples=2048,
                              key=jax.random.PRNGKey(RNG_SEED + 999))
    # Total variance = surrogate-level variance proxy.
    A_mat, _, _ = jax.random.uniform, None, None
    # Use a direct approach: evaluate qoi over 4096 random points.
    rng_key_var = jax.random.PRNGKey(RNG_SEED + 1003)
    var_unit = qmc.LatinHypercube(d=5, seed=RNG_SEED + 1003).random(n=4096)
    var_phys = (np.asarray(AERO_INPUT_LO)
                + var_unit * np.asarray(AERO_INPUT_HI - AERO_INPUT_LO))
    params = {n: jnp.asarray(var_phys[:, i])
              for i, n in enumerate(distributions.keys())}
    qoi_vals_prior = np.asarray(qoi_prior(params))
    var_prior = float(np.var(qoi_vals_prior))
    print(f"       Var(lap_time | prior surrogate) = {var_prior:.4f} s^2")

    # Per-batch-size strategy A vs B comparison.
    var_A = np.zeros((len(N_BATCH_LIST), N_REPEATS_PUB))
    var_B = np.zeros((len(N_BATCH_LIST), N_REPEATS_PUB))
    for ib, N_BATCH in enumerate(N_BATCH_LIST):
        for ir in range(N_REPEATS_PUB):
            # Strategy A: uniform LHS batch
            seed_A = RNG_SEED + 30000 + ir * 7 + ib
            lhs_A = qmc.LatinHypercube(d=5, seed=seed_A)
            unit_A = lhs_A.random(n=int(N_BATCH))
            xa_phys = (np.asarray(AERO_INPUT_LO)
                       + unit_A * np.asarray(AERO_INPUT_HI - AERO_INPUT_LO))
            # Strategy B: proportional-to-cell_scores in (h_F, beta);
            # rest of axes drawn LHS uniform.
            flat_p = cell_probs.flatten()
            seed_B = seed_A + 1_000_000
            rng_B = np.random.default_rng(seed_B)
            cell_idx = rng_B.choice(len(flat_p), size=int(N_BATCH),
                                    replace=True, p=flat_p)
            ix_arr = cell_idx // NY
            iy_arr = cell_idx % NY
            xb_hf = hf_centers[ix_arr] + rng_B.uniform(-1.0, 1.0, size=int(N_BATCH))
            xb_beta = beta_centers[iy_arr] + rng_B.uniform(-0.3, 0.3,
                                                           size=int(N_BATCH))
            xb_phys = np.stack([
                np.clip(xb_hf, float(AERO_INPUT_LO[0]), float(AERO_INPUT_HI[0])),
                rng_B.uniform(float(AERO_INPUT_LO[1]), float(AERO_INPUT_HI[1]),
                              size=int(N_BATCH)),
                rng_B.uniform(float(AERO_INPUT_LO[2]), float(AERO_INPUT_HI[2]),
                              size=int(N_BATCH)),
                np.clip(xb_beta, float(AERO_INPUT_LO[3]), float(AERO_INPUT_HI[3])),
                rng_B.uniform(float(AERO_INPUT_LO[4]), float(AERO_INPUT_HI[4]),
                              size=int(N_BATCH)),
            ], axis=1)
            # Append + re-fit + evaluate variance.
            for label, x_new, var_arr in [("A", xa_phys, var_A),
                                          ("B", xb_phys, var_B)]:
                x_combined = np.vstack([x_prior, x_new])
                keys_new = jax.random.split(
                    jax.random.PRNGKey(seed_A + (0 if label == "A" else 5000)),
                    int(N_BATCH))
                y_new = np.asarray(jax.vmap(cfd_probe)(
                    jnp.asarray(x_new), keys_new))
                y_combined = np.vstack([y_prior, y_new])
                x_norm_combined = jnp.asarray(
                    (x_combined - np.asarray(AERO_INPUT_LO))
                    / np.asarray(AERO_INPUT_HI - AERO_INPUT_LO))
                V_cla_new = fit_table_nd(GRID_AXES_NORM, x_norm_combined,
                                         jnp.asarray(y_combined[:, 0]),
                                         smoothness=FIT_SMOOTHNESS)
                V_cda_new = fit_table_nd(GRID_AXES_NORM, x_norm_combined,
                                         jnp.asarray(y_combined[:, 1]),
                                         smoothness=FIT_SMOOTHNESS)
                V_cop_new = fit_table_nd(GRID_AXES_NORM, x_norm_combined,
                                         jnp.asarray(y_combined[:, 2]),
                                         smoothness=FIT_SMOOTHNESS)
                qoi_new = _surrogate_qoi(V_cla_new, V_cda_new, V_cop_new)
                qoi_vals_new = np.asarray(qoi_new(params))
                var_arr[ib, ir] = float(np.var(qoi_vals_new))
        print(f"       N_BATCH={int(N_BATCH):3d}: "
              f"<Var_A> = {var_A[ib].mean():.4f}, "
              f"<Var_B> = {var_B[ib].mean():.4f}, "
              f"reduction A: {1 - var_A[ib].mean() / var_prior:.1%}, "
              f"B: {1 - var_B[ib].mean() / var_prior:.1%}")
    # -----------------------------------------------------------------------
    # FIA ATR sweep: per grid position (10 -> 1), what variance reduction does
    # each strategy buy?
    # 10th place: 115% of baseline -> ~2300 quality CFD solves per 6-mo period
    # 1st place: 70%  of baseline  -> ~1400 quality CFD solves per 6-mo period
    # Plot achievable variance reduction at each level under each strategy.
    # (Toy model: variance reduction ~ 1 - (baseline_N / new_N)^{1/5} times
    # the strategy efficiency factor.)
    # -----------------------------------------------------------------------
    pub_wall = time.time() - t_total
    print(f"\nTotal wall: {pub_wall/60:.1f} min")

    np.savez(
        OUT_NPZ,
        # Sobol headline indices
        sobol_first=sobol_first,
        sobol_total=sobol_total,
        sobol_names=np.array(sobol_names, dtype=object),
        sobol_n_samples=SOBOL_N_PUB,
        # Aleatoric/epistemic decomposition
        decomp_var_total=decomp["var_total"],
        decomp_var_aleatoric=decomp["var_aleatoric"],
        decomp_var_epistemic=decomp["var_epistemic"],
        decomp_interaction=decomp["interaction"],
        decomp_n_samples=DECOMP_N_PUB,
        epistemic_std=epistemic_std,
        aleatoric_std=aleatoric_std,
        # Cell-level scores
        cell_scores=cell_scores,
        cell_probs=cell_probs,
        hf_centers=hf_centers,
        beta_centers=beta_centers,
        # Strategy comparison
        n_batch_list=N_BATCH_LIST,
        n_repeats=N_REPEATS_PUB,
        var_prior=var_prior,
        var_A=var_A,
        var_B=var_B,
        pub_wall_time_s=pub_wall,
        # IMPORTANT: this is the real publication run, not a placeholder.
        placeholder_flag=False,
    )
    print(f"\nWrote {OUT_NPZ} ({os.path.getsize(OUT_NPZ) / 1024:.1f} KB)")
    print("Done. The notebook will load these results in MODE='publication'.")

def main(): np.random.seed(RNG_SEED) t_total = time.time() # ----------------------------------------------------------------------- # Step 1: fit prior surrogate at N_PRIOR=64. # ----------------------------------------------------------------------- print("[1/4] Fitting prior surrogate at N_PRIOR=64...") V_cla, V_cda, V_cop, x_prior, y_prior = make_prior_fit(RNG_SEED, N_PRIOR) print(" done.") # ----------------------------------------------------------------------- # Step 2: Sobol indices over the 5 design axes at large N. # ----------------------------------------------------------------------- print(f"[2/4] Computing Sobol indices at N_samples={SOBOL_N_PUB}...") distributions = { "h_F": Uniform(float(AERO_INPUT_LO[0]), float(AERO_INPUT_HI[0])), "h_R": Uniform(float(AERO_INPUT_LO[1]), float(AERO_INPUT_HI[1])), "phi": Uniform(float(AERO_INPUT_LO[2]), float(AERO_INPUT_HI[2])), "beta": Uniform(float(AERO_INPUT_LO[3]), float(AERO_INPUT_HI[3])), "delta": Uniform(float(AERO_INPUT_LO[4]), float(AERO_INPUT_HI[4])), } qoi_fn = _surrogate_qoi(V_cla, V_cda, V_cop) t0 = time.time() idx = sobol_indices(None, None, distributions, qoi_fn, n_samples=SOBOL_N_PUB, key=jax.random.PRNGKey(RNG_SEED + 100)) print(f" sobol wall = {time.time() - t0:.1f} s") sobol_first = np.array([idx[k]["first_order"] for k in distributions]) sobol_total = np.array([idx[k]["total_order"] for k in distributions]) sobol_names = list(distributions.keys()) for n, s1, st in zip(sobol_names, sobol_first, sobol_total): print(f" {n:6s} S1 = {s1:+.4f} ST = {st:+.4f}") # ----------------------------------------------------------------------- # Step 3: aleatoric / epistemic / interaction split. # ----------------------------------------------------------------------- print(f"[3/4] Computing aleatoric/epistemic/interaction split at " f"N_samples={DECOMP_N_PUB}...") # Aleatoric == irreducible CFD noise; epistemic == fit-residual std. # For the headline figure we calibrate epistemic_std from the actual fit's # residual std vs truth (sampled at 2000 points), and aleatoric_std from # the CFD noise model floor. n_eval = 2000 unit_eval = qmc.LatinHypercube(d=5, seed=RNG_SEED + 200).random(n=n_eval) x_eval_phys = (np.asarray(AERO_INPUT_LO) + unit_eval * np.asarray(AERO_INPUT_HI - AERO_INPUT_LO)) Y_TRUE = np.asarray(jax.vmap(aero_true)(jnp.asarray(x_eval_phys))) x_eval_norm = jnp.asarray((x_eval_phys - np.asarray(AERO_INPUT_LO)) / np.asarray(AERO_INPUT_HI - AERO_INPUT_LO)) Y_FIT_CLA = np.asarray(interp_nd(GRID_AXES_NORM, V_cla, x_eval_norm)) epistemic_std = float(np.std(Y_TRUE[:, 0] - Y_FIT_CLA)) aleatoric_std = float(np.mean(np.asarray(cfd_noise_std(jnp.asarray(Y_TRUE))[:, 0]))) print(f" calibrated epistemic_std (fit residual) = {epistemic_std:.4f}") print(f" calibrated aleatoric_std (CFD noise floor) = {aleatoric_std:.4f}") aleatoric_dists = { # CFD irreducible noise on CLA "aleatoric_cla": Uniform(-2.0 * aleatoric_std, +2.0 * aleatoric_std, kind="aleatoric"), } epistemic_dists = { "h_F": Uniform(float(AERO_INPUT_LO[0]), float(AERO_INPUT_HI[0]), kind="epistemic"), "h_R": Uniform(float(AERO_INPUT_LO[1]), float(AERO_INPUT_HI[1]), kind="epistemic"), "phi": Uniform(float(AERO_INPUT_LO[2]), float(AERO_INPUT_HI[2]), kind="epistemic"), "beta": Uniform(float(AERO_INPUT_LO[3]), float(AERO_INPUT_HI[3]), kind="epistemic"), "delta": Uniform(float(AERO_INPUT_LO[4]), float(AERO_INPUT_HI[4]), kind="epistemic"), "epistemic_cla": Uniform(-2.0 * epistemic_std, +2.0 * epistemic_std, kind="epistemic"), } decomp_qoi_fn = _aleatoric_epistemic_qoi(V_cla, V_cda, V_cop, epistemic_std) t0 = time.time() decomp = decompose_variance_sobol( decomp_qoi_fn, aleatoric_dists, epistemic_dists, n_samples=DECOMP_N_PUB, key=jax.random.PRNGKey(RNG_SEED + 300), ) print(f" decomp wall = {time.time() - t0:.1f} s") for k, v in decomp.items(): print(f" {k:18s}: {v:+.6e}") # ----------------------------------------------------------------------- # Step 4: strategy A vs B head-to-head over N_REPEATS_PUB seeds. # ----------------------------------------------------------------------- print(f"[4/4] Strategy comparison across N_REPEATS={N_REPEATS_PUB} seeds " f"and N_BATCH in {N_BATCH_LIST.tolist()}...") # Precompute Sobol weighting from the prior fit only (so the recommended # cells are independent of the per-repeat noisy probes). # Cell layout: 10 x 10 grid over (h_F, beta) at fixed (h_R, phi, delta). # This is the dominant 2D slice from the §5 Sobol ranking. NX, NY = 10, 10 hf_centers = np.linspace(float(AERO_INPUT_LO[0]) + 1.0, float(AERO_INPUT_HI[0]) - 1.0, NX) beta_centers = np.linspace(float(AERO_INPUT_LO[3]) + 0.3, float(AERO_INPUT_HI[3]) - 0.3, NY) # Per-cell epistemic variance score: for a uniformly-spaced grid of # candidate next-CFD points, compute the surrogate's prediction variance # under bootstrap re-sampling of the prior probes' noise. # (Cheap proxy: variance of (lap_time(x) - lap_time_truth(x)) over a # small radius around the cell centre.) cell_scores = np.zeros((NX, NY)) for i, hf in enumerate(hf_centers): for j, bt in enumerate(beta_centers): x_cell = np.array([hf, 40.0, 0.0, bt, 0.0]) y_truth_proxy = float(lap_time_proxy(jnp.asarray(x_cell))) # Fit-side proxy x_norm = (jnp.asarray(x_cell) - AERO_INPUT_LO) / (AERO_INPUT_HI - AERO_INPUT_LO) cla = float(interp_nd(GRID_AXES_NORM, V_cla, x_norm)) cda = float(interp_nd(GRID_AXES_NORM, V_cda, x_norm)) cop = float(interp_nd(GRID_AXES_NORM, V_cop, x_norm)) # Compute proxy from the fit alpha_aero = min(MU_PEAK * R_CORNER * RHO_AIR * cla / (2.0 * M_CAR), 0.85) Vc_sq = MU_PEAK * G_ACC * R_CORNER / (1.0 - alpha_aero) t_corner = float(np.pi * R_CORNER / np.sqrt(Vc_sq)) V_top = (2.0 * P_PEAK_KW * 1000.0 / (RHO_AIR * cda)) ** (1.0 / 3.0) a_x = P_PEAK_KW * 1000.0 / (M_CAR * max(0.5 * V_top, 10.0)) t_straight = float(np.sqrt(2.0 * L_STRAIGHT / a_x)) t_balance = 50.0 * (cop - XCOP_NOMINAL) ** 2 y_fit_proxy = (T_CORNER_SHARE * t_corner + (1.0 - T_CORNER_SHARE) * t_straight + t_balance) cell_scores[i, j] = (y_truth_proxy - y_fit_proxy) ** 2 # Normalize to a probability distribution for proportional sampling. cell_probs = cell_scores / cell_scores.sum() # Reference variance with the prior fit only. qoi_prior = _surrogate_qoi(V_cla, V_cda, V_cop) sob_prior = sobol_indices(None, None, distributions, qoi_prior, n_samples=2048, key=jax.random.PRNGKey(RNG_SEED + 999)) # Total variance = surrogate-level variance proxy. A_mat, _, _ = jax.random.uniform, None, None # Use a direct approach: evaluate qoi over 4096 random points. rng_key_var = jax.random.PRNGKey(RNG_SEED + 1003) var_unit = qmc.LatinHypercube(d=5, seed=RNG_SEED + 1003).random(n=4096) var_phys = (np.asarray(AERO_INPUT_LO) + var_unit * np.asarray(AERO_INPUT_HI - AERO_INPUT_LO)) params = {n: jnp.asarray(var_phys[:, i]) for i, n in enumerate(distributions.keys())} qoi_vals_prior = np.asarray(qoi_prior(params)) var_prior = float(np.var(qoi_vals_prior)) print(f" Var(lap_time | prior surrogate) = {var_prior:.4f} s^2") # Per-batch-size strategy A vs B comparison. var_A = np.zeros((len(N_BATCH_LIST), N_REPEATS_PUB)) var_B = np.zeros((len(N_BATCH_LIST), N_REPEATS_PUB)) for ib, N_BATCH in enumerate(N_BATCH_LIST): for ir in range(N_REPEATS_PUB): # Strategy A: uniform LHS batch seed_A = RNG_SEED + 30000 + ir * 7 + ib lhs_A = qmc.LatinHypercube(d=5, seed=seed_A) unit_A = lhs_A.random(n=int(N_BATCH)) xa_phys = (np.asarray(AERO_INPUT_LO) + unit_A * np.asarray(AERO_INPUT_HI - AERO_INPUT_LO)) # Strategy B: proportional-to-cell_scores in (h_F, beta); # rest of axes drawn LHS uniform. flat_p = cell_probs.flatten() seed_B = seed_A + 1_000_000 rng_B = np.random.default_rng(seed_B) cell_idx = rng_B.choice(len(flat_p), size=int(N_BATCH), replace=True, p=flat_p) ix_arr = cell_idx // NY iy_arr = cell_idx % NY xb_hf = hf_centers[ix_arr] + rng_B.uniform(-1.0, 1.0, size=int(N_BATCH)) xb_beta = beta_centers[iy_arr] + rng_B.uniform(-0.3, 0.3, size=int(N_BATCH)) xb_phys = np.stack([ np.clip(xb_hf, float(AERO_INPUT_LO[0]), float(AERO_INPUT_HI[0])), rng_B.uniform(float(AERO_INPUT_LO[1]), float(AERO_INPUT_HI[1]), size=int(N_BATCH)), rng_B.uniform(float(AERO_INPUT_LO[2]), float(AERO_INPUT_HI[2]), size=int(N_BATCH)), np.clip(xb_beta, float(AERO_INPUT_LO[3]), float(AERO_INPUT_HI[3])), rng_B.uniform(float(AERO_INPUT_LO[4]), float(AERO_INPUT_HI[4]), size=int(N_BATCH)), ], axis=1) # Append + re-fit + evaluate variance. for label, x_new, var_arr in [("A", xa_phys, var_A), ("B", xb_phys, var_B)]: x_combined = np.vstack([x_prior, x_new]) keys_new = jax.random.split( jax.random.PRNGKey(seed_A + (0 if label == "A" else 5000)), int(N_BATCH)) y_new = np.asarray(jax.vmap(cfd_probe)( jnp.asarray(x_new), keys_new)) y_combined = np.vstack([y_prior, y_new]) x_norm_combined = jnp.asarray( (x_combined - np.asarray(AERO_INPUT_LO)) / np.asarray(AERO_INPUT_HI - AERO_INPUT_LO)) V_cla_new = fit_table_nd(GRID_AXES_NORM, x_norm_combined, jnp.asarray(y_combined[:, 0]), smoothness=FIT_SMOOTHNESS) V_cda_new = fit_table_nd(GRID_AXES_NORM, x_norm_combined, jnp.asarray(y_combined[:, 1]), smoothness=FIT_SMOOTHNESS) V_cop_new = fit_table_nd(GRID_AXES_NORM, x_norm_combined, jnp.asarray(y_combined[:, 2]), smoothness=FIT_SMOOTHNESS) qoi_new = _surrogate_qoi(V_cla_new, V_cda_new, V_cop_new) qoi_vals_new = np.asarray(qoi_new(params)) var_arr[ib, ir] = float(np.var(qoi_vals_new)) print(f" N_BATCH={int(N_BATCH):3d}: " f"<Var_A> = {var_A[ib].mean():.4f}, " f"<Var_B> = {var_B[ib].mean():.4f}, " f"reduction A: {1 - var_A[ib].mean() / var_prior:.1%}, " f"B: {1 - var_B[ib].mean() / var_prior:.1%}") # ----------------------------------------------------------------------- # FIA ATR sweep: per grid position (10 -> 1), what variance reduction does # each strategy buy? # 10th place: 115% of baseline -> ~2300 quality CFD solves per 6-mo period # 1st place: 70% of baseline -> ~1400 quality CFD solves per 6-mo period # Plot achievable variance reduction at each level under each strategy. # (Toy model: variance reduction ~ 1 - (baseline_N / new_N)^{1/5} times # the strategy efficiency factor.) # ----------------------------------------------------------------------- pub_wall = time.time() - t_total print(f"\nTotal wall: {pub_wall/60:.1f} min") np.savez( OUT_NPZ, # Sobol headline indices sobol_first=sobol_first, sobol_total=sobol_total, sobol_names=np.array(sobol_names, dtype=object), sobol_n_samples=SOBOL_N_PUB, # Aleatoric/epistemic decomposition decomp_var_total=decomp["var_total"], decomp_var_aleatoric=decomp["var_aleatoric"], decomp_var_epistemic=decomp["var_epistemic"], decomp_interaction=decomp["interaction"], decomp_n_samples=DECOMP_N_PUB, epistemic_std=epistemic_std, aleatoric_std=aleatoric_std, # Cell-level scores cell_scores=cell_scores, cell_probs=cell_probs, hf_centers=hf_centers, beta_centers=beta_centers, # Strategy comparison n_batch_list=N_BATCH_LIST, n_repeats=N_REPEATS_PUB, var_prior=var_prior, var_A=var_A, var_B=var_B, pub_wall_time_s=pub_wall, # IMPORTANT: this is the real publication run, not a placeholder. placeholder_flag=False, ) print(f"\nWrote {OUT_NPZ} ({os.path.getsize(OUT_NPZ) / 1024:.1f} KB)") print("Done. The notebook will load these results in MODE='publication'.")

In [ ]:

if __name__ == "__main__":
    main()

if __name__ == "__main__": main()