adasurrog 0.5.1__py3-none-any.whl

adasurrog/__about__.py

@@ -0,0 +1,3 @@
+"""Package metadata kept separate from runtime imports."""
+
+__version__ = "0.5.1"

adasurrog/__init__.py

@@ -0,0 +1,111 @@
+"""Public API for AdaSurroG."""
+
+from .__about__ import __version__
+from .active_learning import (
+    ActiveLearningConfig,
+    CandidateScores,
+    mc_dropout_coefficient_disagreement,
+    nearest_coverage_distance,
+    normalize_log_parameters,
+    score_and_select_candidates,
+    select_diverse_batch,
+    sobol_log_candidates,
+)
+from .checkpoint import CheckpointManager
+from .config import CheckpointConfig, StopConfig, TrainerConfig
+from .metrics import normalized_physical_error, normalized_rmse
+from .models import (
+    ChebyshevFeatureMap,
+    HierarchicalResidualRegressor,
+    PODCoefficientSurrogate,
+    PODFitResult,
+    PolynomialResidualPODSurrogate,
+    ResidualMLP,
+    fit_pod,
+    fit_weighted_pod,
+    load_pod_surrogate_from_checkpoint_payload,
+    metric_feature_weights,
+    total_degree_multi_indices,
+)
+from .refinement import (
+    FullBatchLBFGSRefiner,
+    LBFGSRefinementConfig,
+    LBFGSRefinementResult,
+    move_module_to_device_dtype,
+)
+from .state import FitResult, TrainerState
+from .trainer import TimeAccuracyTrainer
+from .transforms import AsinhTransform, Log1pTransform
+
+__all__ = [
+    "__version__",
+    "ActiveLearningConfig",
+    "AsinhTransform",
+    "CandidateScores",
+    "CheckpointConfig",
+    "CheckpointManager",
+    "ChebyshevFeatureMap",
+    "FitResult",
+    "FullBatchLBFGSRefiner",
+    "HierarchicalResidualRegressor",
+    "LBFGSRefinementConfig",
+    "LBFGSRefinementResult",
+    "move_module_to_device_dtype",
+    "Log1pTransform",
+    "PODCoefficientSurrogate",
+    "PODFitResult",
+    "PolynomialResidualPODSurrogate",
+    "ResidualMLP",
+    "StopConfig",
+    "TimeAccuracyTrainer",
+    "TrainerConfig",
+    "TrainerState",
+    "fit_pod",
+    "fit_weighted_pod",
+    "load_pod_surrogate_from_checkpoint_payload",
+    "mc_dropout_coefficient_disagreement",
+    "metric_feature_weights",
+    "nearest_coverage_distance",
+    "normalize_log_parameters",
+    "normalized_physical_error",
+    "normalized_rmse",
+    "score_and_select_candidates",
+    "select_diverse_batch",
+    "sobol_log_candidates",
+    "total_degree_multi_indices",
+]
+
+# Universal ODE pipeline API (0.5+)
+from .auto_transforms import MixedStateTransform, StateTransformSpec, infer_state_transform
+from .ode_data import ODEDataset, PilotDiagnostics, generate_fixed_dataset, run_pilot
+from .pipeline import (
+    AccuracyConfig,
+    AutoPipelineConfig,
+    BudgetConfig,
+    SamplingConfig,
+    SurrogateReport,
+    UniversalODEPipeline,
+)
+from .problem import ODEProblem, ODEProblemProtocol
+from .universal_model import AdaSurrogate, DomainCheck, UniversalPODSurrogate
+
+__all__ += [
+    "AccuracyConfig",
+    "AdaSurrogate",
+    "AutoPipelineConfig",
+    "BudgetConfig",
+    "DomainCheck",
+    "MixedStateTransform",
+    "ODEDataset",
+    "ODEProblem",
+    "ODEProblemProtocol",
+    "PilotDiagnostics",
+    "SamplingConfig",
+    "StateTransformSpec",
+    "SurrogateReport",
+    "UniversalODEPipeline",
+    "UniversalPODSurrogate",
+    "generate_fixed_dataset",
+    "infer_state_transform",
+    "run_pilot",
+]

adasurrog/active_learning.py

@@ -0,0 +1,287 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+from torch import Tensor, nn
+
+
+@dataclass(frozen=True, slots=True)
+class ActiveLearningConfig:
+    """Configuration for pool-based active learning in positive parameter spaces."""
+
+    candidate_pool_size: int = 4096
+    acquisition_batch_size: int = 64
+    mc_passes: int = 16
+    prediction_batch_size: int = 1024
+    uncertainty_weight: float = 1.0
+    coverage_weight: float = 0.60
+    boundary_weight: float = 0.15
+    diversity_weight: float = 0.50
+    minimum_distance: float = 0.01
+    seed: int = 0
+
+    def __post_init__(self) -> None:
+        if self.candidate_pool_size < 1:
+            raise ValueError("candidate_pool_size must be positive")
+        if not 1 <= self.acquisition_batch_size <= self.candidate_pool_size:
+            raise ValueError("acquisition_batch_size must be within the candidate pool")
+        if self.mc_passes < 2:
+            raise ValueError("mc_passes must be at least 2")
+        if self.prediction_batch_size < 1:
+            raise ValueError("prediction_batch_size must be positive")
+        if (
+            min(
+                self.uncertainty_weight,
+                self.coverage_weight,
+                self.boundary_weight,
+                self.diversity_weight,
+                self.minimum_distance,
+            )
+            < 0
+        ):
+            raise ValueError("active-learning weights and distances must be non-negative")
+
+
+@dataclass(frozen=True, slots=True)
+class CandidateScores:
+    parameters: Tensor
+    normalized_parameters: Tensor
+    uncertainty: Tensor
+    coverage: Tensor
+    boundary: Tensor
+    combined: Tensor
+    selected_indices: Tensor
+
+    @property
+    def selected_parameters(self) -> Tensor:
+        return self.parameters[self.selected_indices]
+
+
+def normalize_log_parameters(
+    parameters: Tensor,
+    parameter_lower: Tensor,
+    parameter_upper: Tensor,
+) -> Tensor:
+    """Map positive parameters into the normalized log box [-1, 1]^d."""
+    parameters = torch.as_tensor(parameters)
+    lower = torch.as_tensor(parameter_lower, device=parameters.device, dtype=parameters.dtype)
+    upper = torch.as_tensor(parameter_upper, device=parameters.device, dtype=parameters.dtype)
+    if torch.any(parameters <= 0):
+        raise ValueError("parameters must be positive")
+    if torch.any(lower <= 0) or torch.any(upper <= lower):
+        raise ValueError("parameter bounds must be positive and ordered")
+    return 2.0 * (parameters.log() - lower.log()) / (upper.log() - lower.log()) - 1.0
+
+
+def sobol_log_candidates(
+    count: int,
+    parameter_lower: Tensor,
+    parameter_upper: Tensor,
+    *,
+    seed: int,
+) -> Tensor:
+    """Generate a scrambled Sobol pool with log-uniform parameter marginals."""
+    lower = torch.as_tensor(parameter_lower, dtype=torch.float64)
+    upper = torch.as_tensor(parameter_upper, dtype=torch.float64)
+    if count < 1:
+        raise ValueError("count must be positive")
+    if torch.any(lower <= 0) or torch.any(upper <= lower):
+        raise ValueError("parameter bounds must be positive and ordered")
+    engine = torch.quasirandom.SobolEngine(
+        dimension=int(lower.numel()),
+        scramble=True,
+        seed=int(seed),
+    )
+    unit = engine.draw(count).to(torch.float64)
+    values = torch.exp(lower.log() + unit * (upper.log() - lower.log()))
+    return values.to(torch.float32)
+
+
+@torch.no_grad()
+def nearest_coverage_distance(
+    candidates_normalized: Tensor,
+    existing_normalized: Tensor,
+    *,
+    chunk_size: int = 2048,
+) -> Tensor:
+    """
+    Distance to the nearest existing point, scaled approximately to [0, 1].
+
+    Coordinates are assumed to lie in [-1, 1]. The maximum box-diagonal
+    distance is 2 * sqrt(dimension).
+    """
+    candidates = torch.as_tensor(candidates_normalized, dtype=torch.float32)
+    existing = torch.as_tensor(
+        existing_normalized,
+        device=candidates.device,
+        dtype=candidates.dtype,
+    )
+    if candidates.ndim != 2 or existing.ndim != 2:
+        raise ValueError("candidate and existing parameters must be rank-2 tensors")
+    if candidates.shape[1] != existing.shape[1]:
+        raise ValueError("candidate and existing parameter dimensions differ")
+    if existing.shape[0] == 0:
+        return torch.ones(candidates.shape[0], device=candidates.device)
+
+    scale = 2.0 * float(candidates.shape[1]) ** 0.5
+    parts: list[Tensor] = []
+    for start in range(0, candidates.shape[0], chunk_size):
+        current = candidates[start : start + chunk_size]
+        distance = torch.cdist(current, existing).amin(dim=1) / scale
+        parts.append(distance)
+    return torch.cat(parts).clamp(0.0, 1.0)
+
+
+def _robust_unit_interval(values: Tensor) -> Tensor:
+    values = torch.as_tensor(values, dtype=torch.float32)
+    if values.numel() == 0:
+        return values
+    low = torch.quantile(values, 0.05)
+    high = torch.quantile(values, 0.95)
+    if float(high - low) <= 1e-12:
+        return torch.zeros_like(values)
+    return ((values - low) / (high - low)).clamp(0.0, 1.0)
+
+
+@torch.no_grad()
+def mc_dropout_coefficient_disagreement(
+    model: nn.Module,
+    parameters: Tensor,
+    *,
+    passes: int = 16,
+    batch_size: int = 1024,
+) -> Tensor:
+    """
+    Estimate epistemic uncertainty from MC-dropout coefficient predictions.
+
+    The model must expose ``predict_standardized_coefficients``. LayerNorm and
+    residual blocks are safe in train mode; only Dropout changes predictions.
+    """
+    predictor = getattr(model, "predict_standardized_coefficients", None)
+    if predictor is None:
+        raise TypeError("model must define predict_standardized_coefficients")
+    if passes < 2:
+        raise ValueError("passes must be at least 2")
+
+    device = next(model.parameters()).device
+    was_training = model.training
+    model.train(True)
+    outputs: list[Tensor] = []
+    try:
+        for start in range(0, len(parameters), batch_size):
+            batch = parameters[start : start + batch_size].to(device)
+            samples = torch.stack([predictor(batch) for _ in range(passes)], dim=0)
+            disagreement = samples.std(dim=0, unbiased=False).square().mean(dim=-1).sqrt()
+            outputs.append(disagreement.cpu())
+    finally:
+        model.train(was_training)
+    return torch.cat(outputs)
+
+
+def select_diverse_batch(
+    normalized_candidates: Tensor,
+    base_score: Tensor,
+    *,
+    batch_size: int,
+    existing_normalized: Tensor | None = None,
+    diversity_weight: float = 0.5,
+    minimum_distance: float = 0.0,
+) -> Tensor:
+    """Greedy score-plus-distance batch selection."""
+    candidates = torch.as_tensor(normalized_candidates, dtype=torch.float32).cpu()
+    scores = torch.as_tensor(base_score, dtype=torch.float32).cpu()
+    if candidates.ndim != 2 or scores.ndim != 1 or len(candidates) != len(scores):
+        raise ValueError("candidate and score shapes are inconsistent")
+    if not 1 <= batch_size <= len(candidates):
+        raise ValueError("batch_size is outside the candidate range")
+
+    if existing_normalized is None or len(existing_normalized) == 0:
+        distance_to_set = torch.ones(len(candidates))
+    else:
+        distance_to_set = nearest_coverage_distance(
+            candidates,
+            torch.as_tensor(existing_normalized, dtype=torch.float32).cpu(),
+        )
+
+    selected: list[int] = []
+    available = torch.ones(len(candidates), dtype=torch.bool)
+    for _ in range(batch_size):
+        utility = scores + diversity_weight * _robust_unit_interval(distance_to_set)
+        utility[~available] = -torch.inf
+        if minimum_distance > 0:
+            sufficiently_far = distance_to_set >= minimum_distance
+            if torch.any(available & sufficiently_far):
+                utility[available & ~sufficiently_far] = -torch.inf
+        index = int(torch.argmax(utility))
+        if not torch.isfinite(utility[index]):
+            fallback = scores.clone()
+            fallback[~available] = -torch.inf
+            index = int(torch.argmax(fallback))
+        selected.append(index)
+        available[index] = False
+
+        scale = 2.0 * float(candidates.shape[1]) ** 0.5
+        new_distance = (
+            torch.linalg.vector_norm(
+                candidates - candidates[index],
+                dim=1,
+            )
+            / scale
+        )
+        distance_to_set = torch.minimum(distance_to_set, new_distance)
+
+    return torch.tensor(selected, dtype=torch.long)
+
+
+@torch.no_grad()
+def score_and_select_candidates(
+    model: nn.Module,
+    candidates: Tensor,
+    existing_parameters: Tensor,
+    parameter_lower: Tensor,
+    parameter_upper: Tensor,
+    config: ActiveLearningConfig,
+) -> CandidateScores:
+    """Score a candidate pool using uncertainty, coverage, boundary, and diversity."""
+    candidates_cpu = torch.as_tensor(candidates, dtype=torch.float32).cpu()
+    existing_cpu = torch.as_tensor(existing_parameters, dtype=torch.float32).cpu()
+    lower = torch.as_tensor(parameter_lower, dtype=torch.float32).cpu()
+    upper = torch.as_tensor(parameter_upper, dtype=torch.float32).cpu()
+
+    normalized = normalize_log_parameters(candidates_cpu, lower, upper)
+    existing_normalized = normalize_log_parameters(existing_cpu, lower, upper)
+    uncertainty_raw = mc_dropout_coefficient_disagreement(
+        model,
+        candidates_cpu,
+        passes=config.mc_passes,
+        batch_size=config.prediction_batch_size,
+    )
+    coverage_raw = nearest_coverage_distance(normalized, existing_normalized)
+    boundary_raw = normalized.abs().mean(dim=1)
+
+    uncertainty = _robust_unit_interval(uncertainty_raw)
+    coverage = _robust_unit_interval(coverage_raw)
+    boundary = _robust_unit_interval(boundary_raw)
+    combined = (
+        config.uncertainty_weight * uncertainty
+        + config.coverage_weight * coverage
+        + config.boundary_weight * boundary
+    )
+    selected = select_diverse_batch(
+        normalized,
+        combined,
+        batch_size=config.acquisition_batch_size,
+        existing_normalized=existing_normalized,
+        diversity_weight=config.diversity_weight,
+        minimum_distance=config.minimum_distance,
+    )
+    return CandidateScores(
+        parameters=candidates_cpu,
+        normalized_parameters=normalized,
+        uncertainty=uncertainty,
+        coverage=coverage,
+        boundary=boundary,
+        combined=combined,
+        selected_indices=selected,
+    )

adasurrog/auto_transforms.py

@@ -0,0 +1,140 @@
+"""Automatic per-state transforms for multi-scale ODE trajectories."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import numpy as np
+import torch
+from torch import Tensor, nn
+
+TRANSFORM_AFFINE = 0
+TRANSFORM_LOG1P = 1
+TRANSFORM_ASINH = 2
+
+
+@dataclass(frozen=True, slots=True)
+class StateTransformSpec:
+    codes: np.ndarray
+    scales: np.ndarray
+    names: tuple[str, ...]
+    diagnostics: tuple[dict[str, float | str], ...]
+
+    def to_dict(self) -> dict[str, object]:
+        return {
+            "codes": self.codes.tolist(),
+            "scales": self.scales.tolist(),
+            "names": list(self.names),
+            "diagnostics": [dict(item) for item in self.diagnostics],
+        }
+
+
+class MixedStateTransform(nn.Module):
+    """Apply affine, log1p, or asinh transforms independently per state."""
+
+    def __init__(self, codes: Tensor | np.ndarray, scales: Tensor | np.ndarray) -> None:
+        super().__init__()
+        codes_tensor = torch.as_tensor(codes, dtype=torch.long)
+        scales_tensor = torch.as_tensor(scales, dtype=torch.float32)
+        if codes_tensor.ndim != 1 or scales_tensor.shape != codes_tensor.shape:
+            raise ValueError("codes and scales must be matching one-dimensional arrays")
+        if torch.any(scales_tensor <= 0) or torch.any(~torch.isfinite(scales_tensor)):
+            raise ValueError("transform scales must be finite and positive")
+        if torch.any((codes_tensor < 0) | (codes_tensor > 2)):
+            raise ValueError("unsupported transform code")
+        self.register_buffer("codes", codes_tensor)
+        self.register_buffer("scales", scales_tensor)
+
+    def forward(self, values: Tensor) -> Tensor:
+        scales = self.scales.to(device=values.device, dtype=values.dtype)
+        codes = self.codes.to(values.device)
+        result = values / scales
+        if torch.any(codes == TRANSFORM_LOG1P):
+            index = codes == TRANSFORM_LOG1P
+            result[..., index] = torch.log1p(
+                torch.clamp_min(values[..., index], 0.0) / scales[index]
+            )
+        if torch.any(codes == TRANSFORM_ASINH):
+            index = codes == TRANSFORM_ASINH
+            result[..., index] = torch.asinh(values[..., index] / scales[index])
+        return result
+
+    def inverse(self, transformed: Tensor) -> Tensor:
+        scales = self.scales.to(device=transformed.device, dtype=transformed.dtype)
+        codes = self.codes.to(transformed.device)
+        result = transformed * scales
+        if torch.any(codes == TRANSFORM_LOG1P):
+            index = codes == TRANSFORM_LOG1P
+            result[..., index] = (
+                torch.expm1(torch.clamp(transformed[..., index], min=-20.0, max=30.0))
+                * scales[index]
+            )
+        if torch.any(codes == TRANSFORM_ASINH):
+            index = codes == TRANSFORM_ASINH
+            result[..., index] = (
+                torch.sinh(torch.clamp(transformed[..., index], min=-30.0, max=30.0))
+                * scales[index]
+            )
+        return result
+
+
+def infer_state_transform(
+    trajectories: np.ndarray,
+    state_names: tuple[str, ...],
+    *,
+    dynamic_range_threshold: float = 100.0,
+    nonnegative_tolerance: float = 1e-10,
+) -> StateTransformSpec:
+    """Infer robust per-state transforms from pilot trajectories."""
+    values = np.asarray(trajectories, dtype=np.float64)
+    if values.ndim != 3 or values.shape[-1] != len(state_names):
+        raise ValueError("trajectories must have shape [sample, time, state]")
+    codes: list[int] = []
+    scales: list[float] = []
+    diagnostics: list[dict[str, float | str]] = []
+    for state_index, name in enumerate(state_names):
+        state = values[..., state_index].reshape(-1)
+        finite = state[np.isfinite(state)]
+        if finite.size == 0:
+            raise ValueError(f"state {name!r} contains no finite values")
+        absolute = np.abs(finite)
+        nonzero = absolute[absolute > max(nonnegative_tolerance, 1e-14)]
+        q50 = float(np.quantile(nonzero, 0.50)) if nonzero.size else 1.0
+        q90 = float(np.quantile(absolute, 0.90))
+        q99 = float(np.quantile(absolute, 0.99))
+        lower = float(np.min(finite))
+        upper = float(np.max(finite))
+        reference = max(q50, q90 * 0.1, 1e-12)
+        smallest = float(np.quantile(nonzero, 0.10)) if nonzero.size else reference
+        dynamic_range = q99 / max(smallest, 1e-30)
+        nonnegative = lower >= -nonnegative_tolerance * max(1.0, q99)
+        if nonnegative and dynamic_range >= dynamic_range_threshold:
+            code = TRANSFORM_LOG1P
+            transform_name = "log1p"
+        elif not nonnegative and dynamic_range >= dynamic_range_threshold:
+            code = TRANSFORM_ASINH
+            transform_name = "asinh"
+        else:
+            code = TRANSFORM_AFFINE
+            transform_name = "affine"
+            reference = max(float(np.std(finite)), q90, 1e-12)
+        codes.append(code)
+        scales.append(reference)
+        diagnostics.append(
+            {
+                "state": name,
+                "minimum": lower,
+                "maximum": upper,
+                "q90_abs": q90,
+                "q99_abs": q99,
+                "dynamic_range": float(dynamic_range),
+                "transform": transform_name,
+                "scale": float(reference),
+            }
+        )
+    return StateTransformSpec(
+        codes=np.asarray(codes, dtype=np.int64),
+        scales=np.asarray(scales, dtype=np.float64),
+        names=state_names,
+        diagnostics=tuple(diagnostics),
+    )