# (c) Kevin Dunn, 2010-2026. MIT License.
"""Budget allocation logic for multi-stage DOE strategies.
Implements the **25-40-55-15 framework**:
- Screening: 25-40 % of total budget
- Optimisation: 40-55 %
- Confirmation: 5-15 % (minimum 3 runs)
Sources:
- Montgomery, *Design and Analysis of Experiments*, 10th ed. (25% rule)
- Stat-Ease SCOR framework
- NIST Engineering Statistics Handbook section 5.3.3
"""
from __future__ import annotations
import math
from collections.abc import Callable
from typing import Any
from process_improve.experiments.designs_response_surface import dsd_run_count
from process_improve.experiments.designs_screening import plackett_burman_runs
# ---------------------------------------------------------------------------
# Run estimation look-up tables
# ---------------------------------------------------------------------------
# Box-Behnken run counts (excluding center points)
_BBD_RUNS: dict[int, int] = {
3: 12,
4: 24,
5: 40,
6: 48,
7: 56,
}
#: Factors from which a CCD's cube is the resolution V half fraction, ``generate_design(...,
#: cube="fractional")``, which is what ``recommend_strategy`` asks for; fewer take the full 2^k cube.
_CCD_HALF_FRACTION_FROM = 6
def _ccd_factorial_runs(n_factors: int) -> int:
"""Return the runs in a CCD's cube: 2^k up to five factors, the 2^(k-1) half fraction from six."""
return 2 ** (n_factors - 1) if n_factors >= _CCD_HALF_FRACTION_FROM else 2**n_factors
# ---------------------------------------------------------------------------
# Run estimation functions
# ---------------------------------------------------------------------------
def _fractional_factorial_runs(n_factors: int) -> int:
"""Return the runs in the smallest two-level design of resolution IV or more, as ``generate_design`` builds.
A resolution IV fraction needs at least ``2k`` runs, and a regular fraction reaches
that bound: 8 runs for four factors, 16 for five to eight, 32 for nine to sixteen, and
so on. Three or fewer factors have no resolution IV fraction, so they take the full
factorial.
"""
if n_factors <= 3:
return 2**n_factors
return 2 ** math.ceil(math.log2(2 * n_factors))
def _plackett_burman_runs(n_factors: int) -> int:
"""Smallest multiple of 4 above ``k`` with a Hadamard matrix (what ``generate_design`` builds)."""
return plackett_burman_runs(n_factors)
def _dsd_runs(n_factors: int) -> int:
"""Return the runs in the DSD ``generate_design`` builds (it needs at least 3 factors)."""
return dsd_run_count(max(n_factors, 3))
#: Run-count estimate per screening design, keyed by generate_design type (and the old DSD name).
_SCREENING_RUNS: dict[str, Callable[[int], int]] = {
"plackett_burman": _plackett_burman_runs,
"dsd": _dsd_runs,
"definitive_screening": _dsd_runs,
"fractional_factorial": _fractional_factorial_runs,
"full_factorial": lambda k: 2**k,
# Lin's half-fraction of the smallest Hadamard order N with N - 2 >= k: N / 2 runs
"supersaturated": lambda k: int(math.ceil((k + 2) / 4) * 2),
}
[docs]
def estimate_screening_runs(n_factors: int, design_type: str) -> int:
"""Estimate the number of runs for a screening design.
Parameters
----------
n_factors : int
Number of factors to screen.
design_type : str
One of ``"plackett_burman"``, ``"dsd"`` (or ``"definitive_screening"``),
``"fractional_factorial"``, ``"full_factorial"``, ``"supersaturated"``.
Any other value gets the Plackett-Burman estimate.
Returns
-------
int
Estimated run count including center points.
"""
return _SCREENING_RUNS.get(design_type, _plackett_burman_runs)(n_factors)
[docs]
def estimate_rsm_runs(n_factors: int, design_type: str, n_center_points: int = 3) -> int:
"""Estimate the number of runs for an RSM design.
Parameters
----------
n_factors : int
Number of factors (typically 2-5 after screening).
design_type : str
One of ``"ccd"``, ``"box_behnken"``, ``"ccd_face_centered"``,
``"d_optimal"``.
n_center_points : int
Number of center point replicates (default 3).
Returns
-------
int
Estimated run count. For a CCD it is the size ``generate_design`` builds with
``cube="full"`` up to five factors and ``cube="fractional"`` (a half-fraction
cube) from six, which is what ``recommend_strategy`` passes.
"""
if design_type in ("ccd", "ccd_face_centered"):
return _ccd_factorial_runs(n_factors) + 2 * n_factors + n_center_points
if design_type == "box_behnken":
base = _BBD_RUNS.get(n_factors, 0)
if base == 0:
# BBD not defined for this factor count; fall back to CCD estimate
return _ccd_factorial_runs(n_factors) + 2 * n_factors + n_center_points
return base + n_center_points
if design_type == "d_optimal":
# Rough estimate: 1.5x the number of model terms for a quadratic model
# Quadratic model terms: 1 + k + k(k-1)/2 + k = 1 + 2k + k(k-1)/2
n_terms = 1 + 2 * n_factors + n_factors * (n_factors - 1) // 2
return math.ceil(1.5 * n_terms)
# Fallback: CCD estimate
return _ccd_factorial_runs(n_factors) + 2 * n_factors + n_center_points
[docs]
def estimate_confirmation_runs(min_runs: int = 3) -> int:
"""Return the number of confirmation runs.
Parameters
----------
min_runs : int
Minimum confirmation runs (default 3).
Returns
-------
int
Confirmation run count (always at least 3).
"""
return max(3, min_runs)
# ---------------------------------------------------------------------------
# Budget allocation
# ---------------------------------------------------------------------------
[docs]
def allocate_budget( # noqa: C901, PLR0913, PLR0915
total_budget: int | None,
n_factors: int,
needs_screening: bool,
needs_rsm: bool,
screening_design: str = "plackett_burman",
rsm_design: str = "box_behnken",
domain_weights: dict[str, float] | None = None,
min_confirmation: int = 3,
n_center_points: int = 3,
) -> dict[str, Any]:
"""Allocate a total run budget across experimental stages.
Parameters
----------
total_budget : int or None
Total runs across all stages. If ``None``, computes an ideal budget.
n_factors : int
Total number of candidate factors.
needs_screening : bool
Whether a screening stage is needed.
needs_rsm : bool
Whether an RSM optimisation stage is needed.
screening_design : str
Preferred screening design type.
rsm_design : str
Preferred RSM design type.
domain_weights : dict or None
Stage-to-fraction mapping from the domain template.
min_confirmation : int
Minimum confirmation runs (domain-dependent).
n_center_points : int
Center points for RSM design.
Returns
-------
dict
Keys: ``"screening"``, ``"optimization"``, ``"confirmation"``,
``"total"``, ``"ideal_total"``, ``"is_tight"``, ``"warnings"``.
"""
weights = domain_weights or {"screening": 0.30, "optimization": 0.50, "confirmation": 0.10}
warnings: list[str] = []
# Estimate ideal runs for each stage
ideal_screening = estimate_screening_runs(n_factors, screening_design) if needs_screening else 0
# Assume screening reduces to ~3 significant factors for RSM
n_rsm_factors = min(n_factors, 3) if needs_screening else n_factors
ideal_rsm = estimate_rsm_runs(n_rsm_factors, rsm_design, n_center_points) if needs_rsm else 0
ideal_confirmation = estimate_confirmation_runs(min_confirmation)
ideal_total = ideal_screening + ideal_rsm + ideal_confirmation
if total_budget is None:
# No budget constraint: use ideal allocation
return {
"screening": ideal_screening,
"optimization": ideal_rsm,
"confirmation": ideal_confirmation,
"total": ideal_total,
"ideal_total": ideal_total,
"is_tight": False,
"warnings": [],
}
# Check if budget is feasible
minimum_feasible = 0
if needs_screening:
minimum_feasible += max(n_factors + 1, 4) # Absolute minimum: k+1 runs
if needs_rsm:
minimum_feasible += n_rsm_factors * 2 + 3 # Bare minimum RSM
minimum_feasible += 3 # Minimum confirmation
is_tight = total_budget < ideal_total
is_very_tight = total_budget < minimum_feasible
if is_very_tight:
warnings.append(
f"Budget of {total_budget} runs is very tight for {n_factors} factors. "
f"Minimum feasible is ~{minimum_feasible} runs. "
"Consider a single Definitive Screening Design to combine screening and curvature detection."
)
# Allocate proportionally, respecting minimums
if needs_screening and needs_rsm:
# Full multi-stage allocation
screening_share = weights.get("screening", 0.30)
rsm_share = weights.get("optimization", 0.50)
conf_share = weights.get("confirmation", 0.10)
alloc_confirmation = max(min_confirmation, int(total_budget * conf_share))
remaining = total_budget - alloc_confirmation
if remaining <= 0:
alloc_confirmation = min_confirmation
remaining = total_budget - alloc_confirmation
# Split remaining between screening and RSM
total_share = screening_share + rsm_share
alloc_screening = max(n_factors + 1, int(remaining * screening_share / total_share))
alloc_rsm = remaining - alloc_screening
# Ensure RSM has minimum runs
if alloc_rsm < n_rsm_factors * 2 + 3:
alloc_rsm = min(n_rsm_factors * 2 + 3, remaining)
alloc_screening = remaining - alloc_rsm
if alloc_screening < n_factors + 1:
alloc_screening = n_factors + 1
alloc_rsm = remaining - alloc_screening
elif needs_screening:
alloc_screening = total_budget - min_confirmation
alloc_rsm = 0
alloc_confirmation = min_confirmation
elif needs_rsm:
alloc_screening = 0
alloc_rsm = total_budget - min_confirmation
alloc_confirmation = min_confirmation
else:
alloc_screening = 0
alloc_rsm = 0
alloc_confirmation = total_budget
# Clamp all values to be non-negative
alloc_screening = max(0, alloc_screening)
alloc_rsm = max(0, alloc_rsm)
alloc_confirmation = max(0, alloc_confirmation)
if is_tight and not is_very_tight:
warnings.append(
f"Budget of {total_budget} runs is tight (ideal is {ideal_total}). "
"Designs may have reduced power or fewer center points."
)
return {
"screening": alloc_screening,
"optimization": alloc_rsm,
"confirmation": alloc_confirmation,
"total": alloc_screening + alloc_rsm + alloc_confirmation,
"ideal_total": ideal_total,
"is_tight": is_tight,
"warnings": warnings,
}
