"""Monte Carlo and two-level multifidelity Monte Carlo workflows."""
from dataclasses import dataclass
from numbers import Integral
import os
import warnings
from typing import Optional
import numpy as np
from romtools.hpc.dispatchers import BaseDispatcher, resolve_dispatcher
from romtools.workflows.models import QoiModel
from romtools.workflows.parameter_spaces import ParameterSpace
from romtools.workflows.uq._evaluation import (
evaluate_qoi_model,
prepare_qoi_model,
validate_evaluation_concurrency,
)
from romtools.workflows.uq.sample_allocation import allocate_samples_from_pilot
from romtools.workflows.uq.statistics import (
compute_monte_carlo_statistics,
compute_multifidelity_statistics,
compute_paired_statistics,
)
[docs]
@dataclass
class MonteCarloDryRunResult:
"""Prepared samples and directory locations; no model evaluations occurred."""
parameter_samples: np.ndarray
run_directories: tuple[str, ...]
high_fidelity_run_directories: tuple[str, ...] = ()
low_fidelity_run_directories: tuple[str, ...] = ()
pilot_only: bool = False
[docs]
@dataclass
class MonteCarloResult:
"""Outputs from a Monte Carlo QoI-mean estimate.
``variance`` and ``standard_deviation`` describe the sampled QoI
distribution. ``standard_error`` describes uncertainty in the estimated
mean and equals ``standard_deviation / sqrt(number_of_samples)``. Each
array has one entry per flattened QoI component.
"""
parameter_samples: np.ndarray
qoi_values: np.ndarray
mean: np.ndarray
variance: np.ndarray
standard_deviation: np.ndarray
standard_error: np.ndarray
run_times: np.ndarray
[docs]
@dataclass
class MultifidelityMonteCarloResult:
"""Outputs from a fixed or pilot-allocated two-level MFMC estimate.
``variance`` is the estimated variance of the multifidelity *mean
estimator*, rather than the variance of the underlying QoI distribution.
The first ``high_fidelity_sample_count`` rows of
``low_fidelity_qoi_values`` are paired with the high-fidelity values.
"""
parameter_samples: np.ndarray
high_fidelity_qoi_values: np.ndarray
low_fidelity_qoi_values: np.ndarray
mean: np.ndarray
variance: np.ndarray
standard_error: np.ndarray
control_variate_coefficients: np.ndarray
paired_correlations: np.ndarray
high_fidelity_sample_count: int
low_fidelity_sample_count: int
pilot_sample_count: Optional[int]
low_to_high_fidelity_cost_ratio: float
high_fidelity_equivalent_cost: float
used_multifidelity_estimator: bool
high_fidelity_run_times: np.ndarray
low_fidelity_run_times: np.ndarray
def _require_absolute_directory(directory: str) -> None:
if not os.path.isabs(directory):
raise ValueError("absolute_work_dir must be an absolute path")
def _require_sample_count(name: str, value: int, minimum: int = 2) -> None:
if (
isinstance(value, bool)
or not isinstance(value, Integral)
or value < minimum
):
raise ValueError(f"{name} must be an integer of at least {minimum}")
def _save_monte_carlo_result(
dispatcher: BaseDispatcher,
directory: str,
result: MonteCarloResult,
random_seed: int,
) -> None:
dispatcher.np_savez(
os.path.join(directory, "uq_stats"),
parameter_samples=result.parameter_samples,
qoi_values=result.qoi_values,
mean=result.mean,
variance=result.variance,
standard_deviation=result.standard_deviation,
standard_error=result.standard_error,
run_times=result.run_times,
random_seed=np.array([random_seed]),
number_of_samples=np.array([result.qoi_values.shape[0]], dtype=int),
)
def _save_multifidelity_result(
dispatcher: BaseDispatcher,
directory: str,
result: MultifidelityMonteCarloResult,
random_seed: int,
requested_budget: float,
pilot_correlations: np.ndarray,
) -> None:
dispatcher.np_savez(
os.path.join(directory, "uq_stats"),
parameter_samples=result.parameter_samples,
high_fidelity_qoi_values=result.high_fidelity_qoi_values,
low_fidelity_qoi_values=result.low_fidelity_qoi_values,
mean=result.mean,
variance=result.variance,
standard_error=result.standard_error,
control_variate_coefficients=result.control_variate_coefficients,
paired_correlations=result.paired_correlations,
pilot_correlations=pilot_correlations,
high_fidelity_sample_count=np.array(
[result.high_fidelity_sample_count], dtype=int
),
low_fidelity_sample_count=np.array(
[result.low_fidelity_sample_count], dtype=int
),
pilot_sample_count=np.array(
[-1 if result.pilot_sample_count is None else result.pilot_sample_count],
dtype=int,
),
low_to_high_fidelity_cost_ratio=np.array(
[result.low_to_high_fidelity_cost_ratio]
),
high_fidelity_equivalent_cost=np.array(
[result.high_fidelity_equivalent_cost]
),
requested_high_fidelity_equivalent_budget=np.array([requested_budget]),
used_multifidelity_estimator=np.array(
[result.used_multifidelity_estimator], dtype=bool
),
high_fidelity_run_times=result.high_fidelity_run_times,
low_fidelity_run_times=result.low_fidelity_run_times,
random_seed=np.array([random_seed]),
)
[docs]
def run_monte_carlo(
model: QoiModel,
parameter_space: ParameterSpace,
absolute_work_dir: Optional[str] = None,
number_of_samples: Optional[int] = None,
random_seed: int = 1,
evaluation_concurrency: int = 1,
overwrite: bool = False,
dispatcher: Optional[BaseDispatcher] = None,
*,
absolute_uq_directory: Optional[str] = None,
dry_run: bool = False,
) -> MonteCarloResult | MonteCarloDryRunResult:
"""Estimate the mean of a model QoI by Monte Carlo sampling.
For every generated parameter sample, the workflow calls the three
:class:`~romtools.workflows.models.QoiModel` methods in order:
``populate_run_directory``, ``run_model``, and ``compute_qoi``. Scalar and
array-valued QoIs are accepted; array-valued outputs are flattened and
statistics are computed componentwise.
Args:
model: Model implementing the ``QoiModel`` protocol.
parameter_space: Distribution used to generate independent parameter
samples. Its parameter names define the dictionary passed to the
model methods.
absolute_work_dir: Absolute output path. Sample ``i`` is evaluated
in ``run_i`` and aggregate results are saved to ``uq_stats.npz``.
number_of_samples: Number of model evaluations. Must be at least two.
random_seed: Seed forwarded to ``parameter_space.generate_samples``.
evaluation_concurrency: Maximum simultaneous model evaluations. A
spawn-based process pool is used above one, so the model and
dispatcher must be pickleable.
overwrite: If ``False``, reuse run directories containing
``passed.txt``. If ``True``, populate and evaluate every sample.
dispatcher: Filesystem/dispatch implementation. Defaults to
:class:`~romtools.hpc.dispatchers.LocalDispatcher`.
dry_run: Prepare run directories and inputs without executing models,
computing QoIs, or writing aggregate statistics.
Returns:
MonteCarloDryRunResult when dry_run is True, otherwise MonteCarloResult: Samples, flattened QoI values, sample moments,
standard errors, and measured model run times.
Raises:
ValueError: If inputs or returned QoI dimensions are invalid.
RuntimeError: If a model evaluation returns a nonzero status.
Note:
A successful marker is written only after both ``run_model`` and
``compute_qoi`` complete. Reused evaluations have ``NaN`` run times
because their original timings are not reconstructed.
"""
if absolute_uq_directory is not None:
warnings.warn(
"'absolute_uq_directory' is deprecated; use 'absolute_work_dir' instead.",
DeprecationWarning,
stacklevel=2,
)
if absolute_work_dir is not None:
raise TypeError("Specify only 'absolute_work_dir', not both directory arguments.")
absolute_work_dir = absolute_uq_directory
if absolute_work_dir is None:
raise TypeError("'absolute_work_dir' is required")
if number_of_samples is None:
raise TypeError("'number_of_samples' is required")
_require_absolute_directory(absolute_work_dir)
_require_sample_count("number_of_samples", number_of_samples)
validate_evaluation_concurrency(evaluation_concurrency)
dispatcher = resolve_dispatcher(dispatcher)
dispatcher.create_empty_dir(absolute_work_dir)
samples = parameter_space.generate_samples(number_of_samples, seed=random_seed)
if dry_run:
directories = prepare_qoi_model(
model, parameter_space.get_names(), samples,
os.path.join(absolute_work_dir, "run_"),
overwrite=overwrite, dispatcher=dispatcher,
)
return MonteCarloDryRunResult(np.asarray(samples), directories)
qois, run_times = evaluate_qoi_model(
model,
parameter_space.get_names(),
samples,
os.path.join(absolute_work_dir, "run_"),
evaluation_concurrency=evaluation_concurrency,
overwrite=overwrite,
dispatcher=dispatcher,
)
statistics = compute_monte_carlo_statistics(qois)
result = MonteCarloResult(
parameter_samples=np.asarray(samples),
qoi_values=qois,
mean=statistics.mean,
variance=statistics.variance,
standard_deviation=statistics.standard_deviation,
standard_error=statistics.standard_error,
run_times=run_times,
)
_save_monte_carlo_result(dispatcher, absolute_work_dir, result, random_seed)
return result
def _validate_multifidelity_modes(
number_high: Optional[int],
number_low: Optional[int],
pilot_count: Optional[int],
budget: Optional[float],
) -> str:
fixed_requested = number_high is not None or number_low is not None
pilot_requested = pilot_count is not None or budget is not None
if fixed_requested and pilot_requested:
raise ValueError("fixed and pilot allocation modes are mutually exclusive")
if fixed_requested:
if number_high is None or number_low is None:
raise ValueError("fixed allocation requires both sample counts")
_require_sample_count("number_of_high_fidelity_samples", number_high)
_require_sample_count("number_of_low_fidelity_samples", number_low)
if number_low < number_high:
raise ValueError(
"fixed allocation requires low count >= high count >= 2"
)
return "fixed"
if pilot_count is None or budget is None:
raise ValueError("pilot allocation requires pilot_sample_count and budget")
_require_sample_count("pilot_sample_count", pilot_count)
return "pilot"
def _finite_mean(values: np.ndarray, name: str) -> float:
finite_values = np.asarray(values, dtype=float)
finite_values = finite_values[np.isfinite(finite_values)]
if finite_values.size == 0:
raise ValueError(
f"cannot estimate {name} from reused samples; supply the cost ratio"
)
value = float(np.mean(finite_values))
if value <= 0.0:
raise ValueError(f"estimated {name} must be positive")
return value
def _build_multifidelity_result(
parameter_samples: np.ndarray,
high_qois: np.ndarray,
low_qois: np.ndarray,
high_run_times: np.ndarray,
low_run_times: np.ndarray,
pilot_count: Optional[int],
cost_ratio: float,
equivalent_cost: float,
use_multifidelity: bool,
coefficients: Optional[np.ndarray],
) -> MultifidelityMonteCarloResult:
if use_multifidelity:
statistics = compute_multifidelity_statistics(
high_qois, low_qois, coefficients
)
mean = statistics.mean
variance = statistics.variance
standard_error = statistics.standard_error
used_coefficients = statistics.control_variate_coefficients
correlations = statistics.paired_correlations
else:
statistics = compute_monte_carlo_statistics(high_qois)
_, _, _, _, correlations = compute_paired_statistics(
high_qois[: low_qois.shape[0]], low_qois
)
mean = statistics.mean
variance = statistics.variance / high_qois.shape[0]
standard_error = statistics.standard_error
used_coefficients = np.zeros(high_qois.shape[1])
return MultifidelityMonteCarloResult(
parameter_samples=parameter_samples,
high_fidelity_qoi_values=high_qois,
low_fidelity_qoi_values=low_qois,
mean=mean,
variance=variance,
standard_error=standard_error,
control_variate_coefficients=used_coefficients,
paired_correlations=correlations,
high_fidelity_sample_count=high_qois.shape[0],
low_fidelity_sample_count=low_qois.shape[0],
pilot_sample_count=pilot_count,
low_to_high_fidelity_cost_ratio=cost_ratio,
high_fidelity_equivalent_cost=equivalent_cost,
used_multifidelity_estimator=use_multifidelity,
high_fidelity_run_times=high_run_times,
low_fidelity_run_times=low_run_times,
)
[docs]
def run_multifidelity_monte_carlo(
high_fidelity_model: QoiModel,
low_fidelity_model: QoiModel,
parameter_space: ParameterSpace,
absolute_work_dir: Optional[str] = None,
number_of_high_fidelity_samples: Optional[int] = None,
number_of_low_fidelity_samples: Optional[int] = None,
pilot_sample_count: Optional[int] = None,
high_fidelity_equivalent_budget: Optional[float] = None,
low_to_high_fidelity_cost_ratio: Optional[float] = None,
allocation_qoi_index: int = 0,
random_seed: int = 1,
high_fidelity_evaluation_concurrency: int = 1,
low_fidelity_evaluation_concurrency: int = 1,
control_variate_coefficients: Optional[np.ndarray] = None,
overwrite: bool = False,
dispatcher: Optional[BaseDispatcher] = None,
*,
absolute_uq_directory: Optional[str] = None,
dry_run: bool = False,
) -> MultifidelityMonteCarloResult | MonteCarloDryRunResult:
r"""Estimate a QoI mean using two-level multifidelity Monte Carlo.
The first ``N_H`` low-fidelity evaluations are paired with the
high-fidelity evaluations at identical parameter samples. For each QoI
component, the estimator is
.. math::
\widehat{\mu}_{MF} = \overline{Q}_H^{N_H}
+ \alpha\left(\overline{Q}_L^{N_L}
- \overline{Q}_L^{N_H}\right).
Exactly one allocation mode must be selected. Fixed mode accepts explicit
high- and low-fidelity sample counts. Pilot mode accepts a paired pilot
count and high-fidelity-equivalent budget, then estimates correlation and
coefficients and performs an integer search under
``N_H + cost_ratio * N_L <= budget``. Pilot evaluations count toward the
final allocation and budget. The workflow falls back to high-fidelity
Monte Carlo when the pilot predicts no benefit from additional
low-fidelity samples.
Args:
high_fidelity_model: High-fidelity ``QoiModel`` implementation.
low_fidelity_model: Low-fidelity ``QoiModel`` implementation. It must
return the same flattened QoI dimension as the high-fidelity model.
parameter_space: Distribution used to generate shared samples.
absolute_work_dir: Absolute output path. Evaluations are stored in
``high_fidelity`` and ``low_fidelity`` subdirectories and aggregate
results are written to ``uq_stats.npz``.
number_of_high_fidelity_samples: Final high-fidelity count in fixed
mode. Must be at least two.
number_of_low_fidelity_samples: Final low-fidelity count in fixed mode.
Must be no smaller than the high-fidelity count.
pilot_sample_count: Initial paired count in pilot mode. Must be at
least two and fit within the budget.
high_fidelity_equivalent_budget: Pilot-mode budget in high-fidelity
evaluation units.
low_to_high_fidelity_cost_ratio: Positive ratio ``c_L / c_H``. If
omitted in pilot mode, paired pilot run times estimate it.
allocation_qoi_index: Flattened QoI component used to select sample
counts. Statistics and coefficients are computed for every
component.
random_seed: Seed used to generate the paired samples.
high_fidelity_evaluation_concurrency: Maximum simultaneous
high-fidelity evaluations.
low_fidelity_evaluation_concurrency: Maximum simultaneous
low-fidelity evaluations.
control_variate_coefficients: Optional scalar or componentwise values
for fixed mode. If omitted, paired samples estimate them. Pilot
mode always estimates and freezes its coefficients.
overwrite: If ``False``, reuse successful evaluations marked by
``passed.txt``; otherwise rerun them.
dispatcher: Filesystem/dispatch implementation. Defaults to
:class:`~romtools.hpc.dispatchers.LocalDispatcher`.
dry_run: Prepare inputs without model execution or statistics.
Pilot-allocation mode prepares only the initial paired pilot;
the final allocation requires pilot QoI evaluations.
Returns:
MonteCarloDryRunResult when dry_run is True, otherwise MultifidelityMonteCarloResult: Ordered samples and QoIs, mean estimate,
estimator variance and standard error, allocation metadata,
correlations, coefficients, and run times.
Raises:
ValueError: If allocation inputs, costs, or QoI dimensions are invalid.
RuntimeError: If either model returns a nonzero status.
Note:
Confidence intervals are intentionally not constructed. Callers can
use the returned mean and standard error with assumptions appropriate
to their application.
"""
if absolute_uq_directory is not None:
warnings.warn(
"'absolute_uq_directory' is deprecated; use 'absolute_work_dir' instead.",
DeprecationWarning,
stacklevel=2,
)
if absolute_work_dir is not None:
raise TypeError("Specify only 'absolute_work_dir', not both directory arguments.")
absolute_work_dir = absolute_uq_directory
if absolute_work_dir is None:
raise TypeError("'absolute_work_dir' is required")
_require_absolute_directory(absolute_work_dir)
mode = _validate_multifidelity_modes(
number_of_high_fidelity_samples,
number_of_low_fidelity_samples,
pilot_sample_count,
high_fidelity_equivalent_budget,
)
if low_to_high_fidelity_cost_ratio is not None and (
not np.isfinite(low_to_high_fidelity_cost_ratio)
or low_to_high_fidelity_cost_ratio <= 0.0
):
raise ValueError("low_to_high_fidelity_cost_ratio must be positive")
validate_evaluation_concurrency(high_fidelity_evaluation_concurrency)
validate_evaluation_concurrency(low_fidelity_evaluation_concurrency)
if mode == "pilot":
if control_variate_coefficients is not None:
raise ValueError(
"pilot allocation estimates and freezes control-variate coefficients"
)
if (
not np.isfinite(high_fidelity_equivalent_budget)
or high_fidelity_equivalent_budget <= 0
):
raise ValueError("high_fidelity_equivalent_budget must be positive")
if low_to_high_fidelity_cost_ratio is not None:
pilot_cost = pilot_sample_count * (1.0 + low_to_high_fidelity_cost_ratio)
tolerance = np.finfo(float).eps * max(high_fidelity_equivalent_budget, 1.0) * 100.0
if pilot_cost > high_fidelity_equivalent_budget + tolerance:
raise ValueError("the paired pilot exceeds the high-fidelity-equivalent budget")
dispatcher = resolve_dispatcher(dispatcher)
dispatcher.create_empty_dir(absolute_work_dir)
high_directory_base = os.path.join(
absolute_work_dir, "high_fidelity", "run_"
)
low_directory_base = os.path.join(
absolute_work_dir, "low_fidelity", "run_"
)
parameter_names = parameter_space.get_names()
if mode == "fixed":
samples = np.asarray(
parameter_space.generate_samples(
number_of_low_fidelity_samples, seed=random_seed
)
)
if dry_run:
high_dirs = prepare_qoi_model(
high_fidelity_model, parameter_names,
samples[:number_of_high_fidelity_samples], high_directory_base,
overwrite=overwrite, dispatcher=dispatcher,
)
low_dirs = prepare_qoi_model(
low_fidelity_model, parameter_names, samples, low_directory_base,
overwrite=overwrite, dispatcher=dispatcher,
)
return MonteCarloDryRunResult(
samples, (), high_dirs, low_dirs,
)
high_qois, high_times = evaluate_qoi_model(
high_fidelity_model,
parameter_names,
samples[:number_of_high_fidelity_samples],
high_directory_base,
high_fidelity_evaluation_concurrency,
overwrite,
dispatcher,
)
low_qois, low_times = evaluate_qoi_model(
low_fidelity_model,
parameter_names,
samples,
low_directory_base,
low_fidelity_evaluation_concurrency,
overwrite,
dispatcher,
)
if low_to_high_fidelity_cost_ratio is None:
try:
cost_ratio = _finite_mean(low_times, "low-fidelity cost") / _finite_mean(
high_times, "high-fidelity cost"
)
except ValueError:
cost_ratio = np.nan
else:
cost_ratio = float(low_to_high_fidelity_cost_ratio)
equivalent_cost = (
number_of_high_fidelity_samples + cost_ratio * number_of_low_fidelity_samples
if np.isfinite(cost_ratio)
else np.nan
)
result = _build_multifidelity_result(
samples,
high_qois,
low_qois,
high_times,
low_times,
None,
cost_ratio,
equivalent_cost,
True,
control_variate_coefficients,
)
_save_multifidelity_result(
dispatcher,
absolute_work_dir,
result,
random_seed,
np.nan,
result.paired_correlations,
)
return result
# Pilot mode.
pilot_count = int(pilot_sample_count)
pilot_samples = np.asarray(
parameter_space.generate_samples(pilot_count, seed=random_seed)
)
if dry_run:
high_dirs = prepare_qoi_model(
high_fidelity_model, parameter_names, pilot_samples,
high_directory_base, overwrite=overwrite, dispatcher=dispatcher,
)
low_dirs = prepare_qoi_model(
low_fidelity_model, parameter_names, pilot_samples,
low_directory_base, overwrite=overwrite, dispatcher=dispatcher,
)
return MonteCarloDryRunResult(
pilot_samples, (), high_dirs, low_dirs, pilot_only=True,
)
pilot_high_qois, pilot_high_times = evaluate_qoi_model(
high_fidelity_model,
parameter_names,
pilot_samples,
high_directory_base,
high_fidelity_evaluation_concurrency,
overwrite,
dispatcher,
)
pilot_low_qois, pilot_low_times = evaluate_qoi_model(
low_fidelity_model,
parameter_names,
pilot_samples,
low_directory_base,
low_fidelity_evaluation_concurrency,
overwrite,
dispatcher,
)
if low_to_high_fidelity_cost_ratio is None:
cost_ratio = _finite_mean(
pilot_low_times, "low-fidelity cost"
) / _finite_mean(pilot_high_times, "high-fidelity cost")
else:
cost_ratio = float(low_to_high_fidelity_cost_ratio)
allocation = allocate_samples_from_pilot(
pilot_high_qois,
pilot_low_qois,
float(high_fidelity_equivalent_budget),
cost_ratio,
allocation_qoi_index,
)
number_high = allocation.number_high_fidelity_samples
number_low = allocation.number_low_fidelity_samples
additional_count = max(number_high, number_low) - pilot_count
if additional_count:
additional_samples = np.asarray(
parameter_space.generate_samples(additional_count, seed=random_seed + 1)
)
else:
additional_samples = np.empty((0, pilot_samples.shape[1]))
samples = np.vstack([pilot_samples, additional_samples])
high_extra_count = number_high - pilot_count
low_extra_count = number_low - pilot_count
if high_extra_count:
high_extra_qois, high_extra_times = evaluate_qoi_model(
high_fidelity_model,
parameter_names,
additional_samples[:high_extra_count],
high_directory_base,
high_fidelity_evaluation_concurrency,
overwrite,
dispatcher,
starting_sample_index=pilot_count,
)
high_qois = np.vstack([pilot_high_qois, high_extra_qois])
high_times = np.concatenate([pilot_high_times, high_extra_times])
else:
high_qois, high_times = pilot_high_qois, pilot_high_times
if low_extra_count:
low_extra_qois, low_extra_times = evaluate_qoi_model(
low_fidelity_model,
parameter_names,
additional_samples[:low_extra_count],
low_directory_base,
low_fidelity_evaluation_concurrency,
overwrite,
dispatcher,
starting_sample_index=pilot_count,
)
low_qois = np.vstack([pilot_low_qois, low_extra_qois])
low_times = np.concatenate([pilot_low_times, low_extra_times])
else:
low_qois, low_times = pilot_low_qois, pilot_low_times
result = _build_multifidelity_result(
samples,
high_qois,
low_qois,
high_times,
low_times,
pilot_count,
cost_ratio,
allocation.high_fidelity_equivalent_cost,
allocation.use_multifidelity_estimator,
allocation.control_variate_coefficients,
)
_save_multifidelity_result(
dispatcher,
absolute_work_dir,
result,
random_seed,
float(high_fidelity_equivalent_budget),
allocation.pilot_correlations,
)
return result