import time
import numpy as np
import torch
from cann.data.market_surface import MarketSurface
from cann.engine.calibration_config import CalibrationConfig, CalibrationResult
from cann.engine.training import TrainedModel
from cann.optim.differential_evolution import minimize_global
[docs]
def calibrate_model(
trained_model: TrainedModel,
market_surface: MarketSurface,
calibration_config: CalibrationConfig,
) -> CalibrationResult:
"""
Calibrate model parameters to a market surface using a trained surrogate.
The columns of ``market_surface.market_inputs`` must match the feature ordering
specified by ``trained_model.model_config.market_inputs``.
This constructs a loss function J(θ) and passes it to the optimizer.
Args:
trained_model (TrainedModel): The trained surrogate model to be calibrated.
market_surface (MarketSurface): The market surface data for calibration.
calibration_config (CalibrationConfig): Configuration settings for the calibration process.
Returns:
CalibrationResult: The result of the calibration optimization.
Notes:
The optimization is performed with a global optimizer (Differential Evolution).
Surrogate inference is batched and will use CUDA if available, otherwise CPU.
"""
surrogate = trained_model.surrogate
model_cfg = trained_model.model_config
# 1. Identify bounds
bounds = model_cfg.bounds_array()
# 2. Prepare Market Data
market_inputs = market_surface.market_inputs # Shape (M, m)
observed_values = market_surface.observed.flatten() # Shape (M,)
if market_surface.weights is not None:
weights = market_surface.weights.flatten()
else:
weights = np.ones_like(observed_values)
M = market_inputs.shape[0]
device = "cuda" if torch.cuda.is_available() else "cpu"
# 3. Define Objective Function (Vectorized)
def objective_fn(x: np.ndarray) -> np.ndarray:
# x shape: (n_parameters, pop_size)
_, pop_size = x.shape
# Construct full parameter matrix: (pop_size, n_full)
candidates = x.T
# Expand candidates: (pop_size*M, n_full)
candidates_expanded = np.repeat(candidates, M, axis=0)
# Expand market inputs: (pop_size*M, m)
market_expanded = np.tile(market_inputs, (pop_size, 1))
# Predict
preds = surrogate.predict(
parameters=candidates_expanded,
market_inputs=market_expanded,
batch_size=calibration_config.batch_size,
device=device,
)
# Compute Loss
preds_reshaped = preds.reshape(pop_size, M)
diff = preds_reshaped - observed_values
weighted_sq_diff = (diff**2) * weights
loss = np.mean(weighted_sq_diff, axis=1)
return loss
# 4. Run Optimization
start_time = time.time()
opt_res = minimize_global(objective_fn, bounds, calibration_config)
end_time = time.time()
# 5. Package Results
optimal_parameters = opt_res["x"]
return CalibrationResult(
optimal_parameters=optimal_parameters,
final_loss=float(opt_res["fun"]),
num_iterations=opt_res["nit"],
num_function_evals=int(opt_res.get("candidate_nfev", opt_res["nfev"])),
runtime_seconds=end_time - start_time,
converged=bool(opt_res["success"]),
diagnostics={
"message": opt_res["message"],
"vectorized_objective_evaluations": int(opt_res.get("vectorized_nfev", opt_res["nfev"])),
"scipy_nfev": int(opt_res["nfev"]),
},
)
[docs]
def calibrate_model_with_fixed_parameters(
*,
trained_model: TrainedModel,
market_surface: MarketSurface,
calibration_config: CalibrationConfig,
fixed_parameters: dict[str, float],
free_parameter_names: list[str] | None = None,
) -> CalibrationResult:
"""Calibrate only a subset of parameters while holding others fixed.
This is useful when you want, e.g., a 3-parameter calibration using a surrogate
that was trained on the full 5-parameter Heston space.
Args:
trained_model: Trained surrogate model.
market_surface: Observed market surface.
calibration_config: Optimizer configuration.
fixed_parameters: Mapping of parameter name -> fixed value.
free_parameter_names: Optional explicit list of parameter names to optimize.
If not provided, all parameters not in fixed_parameters are optimized.
Returns:
CalibrationResult where optimal_parameters is the *full* parameter vector
in model-config ordering.
"""
surrogate = trained_model.surrogate
model_cfg = trained_model.model_config
name_to_index = {p.name: i for i, p in enumerate(model_cfg.parameters)}
for name in fixed_parameters:
if name not in name_to_index:
raise ValueError(f"Unknown fixed parameter '{name}'. Expected one of {list(name_to_index)}")
if free_parameter_names is None:
free_parameter_names = [p.name for p in model_cfg.parameters if p.name not in fixed_parameters]
else:
for name in free_parameter_names:
if name not in name_to_index:
raise ValueError(f"Unknown free parameter '{name}'. Expected one of {list(name_to_index)}")
if name in fixed_parameters:
raise ValueError(f"Parameter '{name}' cannot be both free and fixed")
free_indices = [name_to_index[n] for n in free_parameter_names]
# Bounds only for free params
bounds_full = model_cfg.bounds_array()
bounds_free = [bounds_full[i] for i in free_indices]
# Prepare Market Data
market_inputs = market_surface.market_inputs
observed_values = market_surface.observed.flatten()
if market_surface.weights is not None:
weights = market_surface.weights.flatten()
else:
weights = np.ones_like(observed_values)
M = market_inputs.shape[0]
device = "cuda" if torch.cuda.is_available() else "cpu"
fixed_vector = np.zeros((model_cfg.num_parameters,), dtype=float)
for p in model_cfg.parameters:
if p.name in fixed_parameters:
fixed_vector[name_to_index[p.name]] = float(fixed_parameters[p.name])
def objective_fn(x: np.ndarray) -> np.ndarray:
# x shape: (n_free, pop_size)
_, pop_size = x.shape
candidates_full = np.repeat(fixed_vector[None, :], pop_size, axis=0)
candidates_full[:, free_indices] = x.T
candidates_expanded = np.repeat(candidates_full, M, axis=0)
market_expanded = np.tile(market_inputs, (pop_size, 1))
preds = surrogate.predict(
parameters=candidates_expanded,
market_inputs=market_expanded,
batch_size=calibration_config.batch_size,
device=device,
)
preds_reshaped = preds.reshape(pop_size, M)
diff = preds_reshaped - observed_values
weighted_sq_diff = (diff**2) * weights
return np.mean(weighted_sq_diff, axis=1)
start_time = time.time()
opt_res = minimize_global(objective_fn, bounds_free, calibration_config)
end_time = time.time()
x_free = opt_res["x"]
x_full = fixed_vector.copy()
x_full[free_indices] = x_free
return CalibrationResult(
optimal_parameters=x_full,
final_loss=float(opt_res["fun"]),
num_iterations=opt_res["nit"],
num_function_evals=int(opt_res.get("candidate_nfev", opt_res["nfev"])),
runtime_seconds=end_time - start_time,
converged=bool(opt_res["success"]),
diagnostics={
"message": opt_res["message"],
"vectorized_objective_evaluations": int(opt_res.get("vectorized_nfev", opt_res["nfev"])),
"scipy_nfev": int(opt_res["nfev"]),
"free_parameter_names": free_parameter_names,
"fixed_parameters": dict(fixed_parameters),
},
)