diff --git a/autoemulate/convert/__init__.py b/autoemulate/convert/__init__.py new file mode 100644 index 000000000..8d492ca82 --- /dev/null +++ b/autoemulate/convert/__init__.py @@ -0,0 +1,3 @@ +from .pyomo import pyomofy + +__all__ = ["pyomofy"] diff --git a/autoemulate/convert/pyomo.py b/autoemulate/convert/pyomo.py new file mode 100644 index 000000000..ea7bb396b --- /dev/null +++ b/autoemulate/convert/pyomo.py @@ -0,0 +1,340 @@ +import pyomo.environ as pyo +from autoemulate.core.results import Result +from autoemulate.emulators.nn.mlp import MLP +from autoemulate.emulators.polynomials import PolynomialRegression +from autoemulate.emulators.transformed.base import TransformedEmulator +from autoemulate.transforms.pca import PCATransform +from autoemulate.transforms.standardize import StandardizeTransform +from torch import nn + +_SUPPORTED_TRANSFORMS = (StandardizeTransform, PCATransform) + +# Default beta for the Softplus approximation when the model was trained with ReLU. +# As beta -> inf, Softplus_beta -> ReLU. beta=50 is sharp enough to closely track +# ReLU while remaining smooth and differentiable everywhere. +_DEFAULT_RELU_BETA = 50 + + +def pyomofy( + model: TransformedEmulator | Result, pyomo_vars, relu_beta: int = _DEFAULT_RELU_BETA +): + """Convert a fitted emulator into Pyomo expressions. + + Parameters + ---------- + model : TransformedEmulator | Result + A fitted TransformedEmulator or Result wrapping one. + pyomo_vars : list or pyomo.core.base.var.Var or pyomo.core.base.var.IndexedVar + Input decision variables in the same order used for training. + relu_beta : int + Sharpness parameter for the Softplus approximation used when the model + was trained with ReLU. Softplus_beta(x) = (1/beta) * log(1 + exp(beta*x)). + As beta increases the approximation gets closer to ReLU but less smooth. + Defaults to 50. + + Returns + ------- + list + Pyomo expressions, one per output dimension. + + Raises + ------ + TypeError + If model is not a TransformedEmulator or Result. + NotImplementedError + If model type is not PolynomialRegression or MLP. + ValueError + If the model has not been fitted. + """ + if isinstance(model, Result): + model = model.model + if not isinstance(model, TransformedEmulator): + msg = "Input model must be a TransformedEmulator or Result instance." + raise TypeError(msg) + + # The actual emulator lives inside TransformedEmulator.model + emulator = model.model + if not isinstance(emulator, PolynomialRegression | MLP): + msg = ( + "Currently, only PolynomialRegression and MLP models are supported " + "for Pyomo conversion." + ) + raise NotImplementedError(msg) + if not emulator.is_fitted_: + msg = "Model must be fitted before conversion to Pyomo." + raise ValueError(msg) + + # Flatten possible IndexedVar or VarSet to an ordered list + if hasattr(pyomo_vars, "component_data_objects"): + var_list = list(pyomo_vars.component_data_objects(sort=True)) + + else: + var_list = list(pyomo_vars) + + # --- Validate all transforms upfront before doing any work --- + _validate_transforms(model.x_transforms) + _validate_transforms(model.y_transforms) + # Also validate any transforms baked into the base emulator itself (e.g. MLP + # standardize_x / standardize_y). These are fitted on already-outer-transformed + # data, so they are nearly identity in the standard workflow, but must still be + # included for correctness when the outer pipeline has no x/y transforms. + if emulator.x_transform is not None: + _validate_transforms([emulator.x_transform]) + if emulator.y_transform is not None: + _validate_transforms([emulator.y_transform]) + + # --- Apply outer x transforms (TransformedEmulator level) --- + activations = _apply_x_transforms(model.x_transforms, var_list) + + # --- Apply inner x transform (emulator level, e.g. + # MLP-level StandardizeTransform) --- + if emulator.x_transform is not None: + activations = _apply_x_transforms([emulator.x_transform], activations) + + # --- Emulator-specific forward pass --- + if isinstance(emulator, PolynomialRegression): + activations = _polynomial_forward(emulator, activations) + elif isinstance(emulator, MLP): + activations = _mlp_forward(emulator, activations, relu_beta) + + # --- Invert inner y transform first (emulator level), then outer --- + if emulator.y_transform is not None: + activations = _apply_y_inverse_transforms([emulator.y_transform], activations) + return _apply_y_inverse_transforms(model.y_transforms, activations) + + +# --------------------------------------------------------------------------- +# Transform validation and application +# --------------------------------------------------------------------------- + + +def _validate_transforms(transforms: list): + """Check every transform in the list is supported. Fail fast before any work.""" + for t in transforms: + if not isinstance(t, _SUPPORTED_TRANSFORMS): + supported = ", ".join(cls.__name__ for cls in _SUPPORTED_TRANSFORMS) + msg = ( + f"Unsupported transform: '{type(t).__name__}'. " + f"Only {supported} are supported for Pyomo conversion." + ) + raise NotImplementedError(msg) + + +def _apply_x_transforms(transforms: list, activations: list) -> list: + """Apply x transforms in forward order as Pyomo expressions.""" + for t in transforms: + if isinstance(t, StandardizeTransform): + activations = _standardize_forward(t, activations) + elif isinstance(t, PCATransform): + activations = _pca_forward(t, activations) + return activations + + +def _apply_y_inverse_transforms(transforms: list, activations: list) -> list: + """Apply y transforms in inverse order as Pyomo expressions. + + Mirrors TransformedEmulator._inv_transform_y_gaussian: + for transform in self.y_transforms[::-1]: ... + """ + for t in reversed(transforms): + if isinstance(t, StandardizeTransform): + activations = _standardize_inverse(t, activations) + elif isinstance(t, PCATransform): + activations = _pca_inverse(t, activations) + return activations + + +# --------------------------------------------------------------------------- +# StandardizeTransform: forward and inverse +# --------------------------------------------------------------------------- + + +def _standardize_forward(transform: StandardizeTransform, activations: list) -> list: + """(x - mean) / std.""" + mean = transform.mean.detach().cpu().numpy().flatten().tolist() + std = transform.std.detach().cpu().numpy().flatten().tolist() + + if len(mean) == 1 and len(activations) > 1: + mean = mean * len(activations) + std = std * len(activations) + + return [(activations[i] - mean[i]) / std[i] for i in range(len(activations))] + + +def _standardize_inverse(transform: StandardizeTransform, activations: list) -> list: + """Y * std + mean.""" + mean = transform.mean.detach().cpu().numpy().flatten().tolist() + std = transform.std.detach().cpu().numpy().flatten().tolist() + + if len(mean) == 1 and len(activations) > 1: + mean = mean * len(activations) + std = std * len(activations) + + return [activations[i] * std[i] + mean[i] for i in range(len(activations))] + + +# --------------------------------------------------------------------------- +# PCATransform: forward and inverse +# --------------------------------------------------------------------------- + + +def _pca_forward(transform: PCATransform, activations: list) -> list: + """(x - mean) @ components -> dot product per component.""" + mean = transform.mean.detach().cpu().numpy().flatten().tolist() + # components shape: (n_features, n_components) + components = transform.components.detach().cpu().numpy() + + # Center first + centered = [activations[i] - mean[i] for i in range(len(activations))] + + # Each output component j = sum_i( centered[i] * components[i, j] ) + n_components = components.shape[1] + return [ + sum(float(components[i, j]) * centered[i] for i in range(len(centered))) + for j in range(n_components) + ] + + +def _pca_inverse(transform: PCATransform, activations: list) -> list: + """Y @ components.T + mean -> dot product per original feature.""" + mean = transform.mean.detach().cpu().numpy().flatten().tolist() + # components shape: (n_features, n_components) + components = transform.components.detach().cpu().numpy() + + n_features = components.shape[0] + # Each output feature i = sum_j( activations[j] * components[i, j] ) + mean[i] + return [ + sum(float(components[i, j]) * activations[j] for j in range(len(activations))) + + mean[i] + for i in range(n_features) + ] + + +# --------------------------------------------------------------------------- +# PolynomialRegression forward pass +# --------------------------------------------------------------------------- + + +def _polynomial_forward(model: PolynomialRegression, activations: list) -> list: + """Expand polynomial features and apply the linear layer in Pyomo expressions. + + Mirrors PolynomialRegression.forward(): + x_poly = self.poly(x) + return self.linear(x_poly) + """ + # 1. Expand polynomial features using the stored powers matrix + # model.poly._powers has shape (n_output_features, n_input_features) + # Each row defines one monomial: x0^p0 * x1^p1 * ... + powers = model.poly._powers.detach().cpu().numpy() + poly_features = [] + for row in powers: + term = 1 + for feat_idx, power in enumerate(row): + if power > 0: + term = term * activations[feat_idx] ** int(power) + poly_features.append(term) + + # 2. Single linear layer (no bias) on the expanded features + weight = ( + model.linear.weight.detach().cpu().numpy() + ) # shape: (n_outputs, n_poly_features) + out_exprs = [] + for j in range(weight.shape[0]): + lin = sum( + float(weight[j, k]) * poly_features[k] for k in range(len(poly_features)) + ) + out_exprs.append(lin) + + return out_exprs + + +# --------------------------------------------------------------------------- +# MLP forward pass +# --------------------------------------------------------------------------- + + +def _silu_expr(val): + """Forward pass SiLU / Swish: x * sigmoid(x).""" + return val / (1 + pyo.exp(-val)) + + +def _sigmoid_expr(val): + """Forward pass Standard Sigmoid: 1 / (1 + exp(-x)).""" + return 1 / (1 + pyo.exp(-val)) + + +def _softplus_expr(val): + """Forward pass Softplus: log(1 + exp(x)) — standard form, beta=1.""" + return pyo.log(1 + pyo.exp(val)) + + +def _softplus_beta_expr(val, beta: float): + """Softplus with tunable sharpness: (1/beta) * log(1 + exp(beta * x)). + + As beta -> inf this converges to ReLU. Used as a drop-in replacement + when the model was trained with ReLU. + """ + return (1.0 / beta) * pyo.log(1 + pyo.exp(beta * val)) + + +# Activation map for natively supported activations (no extra args needed). +# ReLU is handled separately in _mlp_forward via Softplus_beta approximation. +_ACTIVATION_MAP = { + nn.SiLU: _silu_expr, + nn.Tanh: pyo.tanh, + nn.Sigmoid: _sigmoid_expr, + nn.Softplus: _softplus_expr, +} + +# Full supported list: includes ReLU (auto-approximated) + everything in _ACTIVATION_MAP +_SUPPORTED_ACTIVATIONS = (*_ACTIVATION_MAP.keys(), nn.ReLU) + + +def _mlp_forward(model: MLP, activations: list, relu_beta: float) -> list: + """Unroll the MLP layers into Pyomo expressions. + + Mirrors MLP.forward(): + return self.nn(x) + + ReLU layers are automatically replaced by Softplus_beta using the provided + relu_beta sharpness parameter. + """ + for module in model.nn: + if isinstance(module, nn.Linear): + weight = module.weight.detach().cpu().numpy() + bias = ( + module.bias.detach().cpu().numpy() if module.bias is not None else None + ) + + out_exprs = [] + for j in range(weight.shape[0]): + lin = sum( + float(weight[j, k]) * activations[k] + for k in range(len(activations)) + ) + if bias is not None: + lin += float(bias[j]) + out_exprs.append(lin) + activations = out_exprs + + elif isinstance(module, nn.Dropout): + # Dropout is only active during training — skip at inference + continue + + elif isinstance(module, nn.ReLU): + # ReLU is not smooth — replace with Softplus_beta approximation + activations = [_softplus_beta_expr(a, relu_beta) for a in activations] + + elif type(module) in _ACTIVATION_MAP: + act_fn = _ACTIVATION_MAP[type(module)] + activations = [act_fn(a) for a in activations] + + else: + supported = ", ".join(act.__name__ for act in _SUPPORTED_ACTIVATIONS) + msg = ( + f"Unsupported activation function: '{type(module).__name__}'\n" + f"Supported activations: {supported}\n" + ) + raise TypeError(msg) + + return activations diff --git a/docs/_toc.yml b/docs/_toc.yml index 8cc14cfc1..6ac5f2b24 100644 --- a/docs/_toc.yml +++ b/docs/_toc.yml @@ -26,6 +26,7 @@ chapters: - file: tutorials/advanced/index sections: - file: tutorials/advanced/01_add_emulators + - file: tutorials/advanced/02_pyomo_conversion - file: community/index sections: @@ -73,3 +74,6 @@ chapters: - file: reference/simulations/flow_problem - file: reference/simulations/projectile - file: reference/simulations/reaction_diffusion + - file: reference/convert/index + sections: + - file: reference/convert/pyomo \ No newline at end of file diff --git a/docs/reference/convert/index.md b/docs/reference/convert/index.md new file mode 100644 index 000000000..34b59b90f --- /dev/null +++ b/docs/reference/convert/index.md @@ -0,0 +1,4 @@ +Convert +======= + +Reference documentation for conversion utilities (e.g., exporting emulators to other frameworks). \ No newline at end of file diff --git a/docs/reference/convert/pyomo.rst b/docs/reference/convert/pyomo.rst new file mode 100644 index 000000000..0114ffe50 --- /dev/null +++ b/docs/reference/convert/pyomo.rst @@ -0,0 +1,9 @@ +Pyomo conversion (``autoemulate.convert.pyomo``) +================================================ + +This section documents the Pyomo conversion utilities. + +.. automodule:: autoemulate.convert.pyomo + :members: + :undoc-members: + :show-inheritance: \ No newline at end of file diff --git a/docs/tutorials/advanced/02_pyomo_conversion.ipynb b/docs/tutorials/advanced/02_pyomo_conversion.ipynb new file mode 100644 index 000000000..932b8dff3 --- /dev/null +++ b/docs/tutorials/advanced/02_pyomo_conversion.ipynb @@ -0,0 +1,305 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Converting Emulators to Pyomo Expressions\n", + "\n", + "This tutorial's purpose is to walk you through the process of converting your `AutoEmulate` Emulators to Pyomo algebraic expressions for easy optimization workflows. Currently only `PolynomialRegression` and `MLP` are supported for this conversion.\n", + "\n", + "We'll demonstrate the following steps:\n", + "1. Training an MLP Emulator on a simple toy function using `AutoEmulate`\n", + "2. Exporting the trained Emulator into algebraic expressions compatible with Pyomo\n", + "3. Validating the Pyomo expressions against Emulator prediction" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# General imports for the notebook\n", + "import warnings\n", + "import numpy as np\n", + "import torch\n", + "warnings.filterwarnings(\"ignore\")\n", + "np.random.seed(42)\n", + "torch.manual_seed(42)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Toy simulation\n", + "\n", + "Before we train an emulator with AutoEmulate, we need to get a set of input/output pairs from our simulation to use as training data.\n", + "\n", + "Below is a simple toy simulation that computes the product of two inputs: `y = x1 * x2`." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Generate data: y = x1 * x2\n", + "n_samples = 1000\n", + "x1 = np.random.uniform(-100, 100, size=n_samples)\n", + "x2 = np.random.uniform(-100, 100, size=n_samples)\n", + "\n", + "def F(x1, x2):\n", + " return x1 * x2\n", + "\n", + "x = np.column_stack((x1, x2))\n", + "y = F(x1, x2)\n", + "\n", + "# Convert to tensors for AutoEmulate\n", + "x = torch.tensor(x, dtype=torch.float32)\n", + "y = torch.tensor(y, dtype=torch.float32)\n", + "\n", + "x.shape, y.shape" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Train Emulator\n", + "\n", + "For this tutorial, we'll focus on training an MLP (Multi-Layer Perceptron) Emulator." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from autoemulate import AutoEmulate\n", + "\n", + "# Initialize AutoEmulate with MLP Emulator only\n", + "ae = AutoEmulate(x, y, log_level=\"info\", models=['MLP'])\n", + "Emulator = ae.best_result()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Converting to Pyomo\n", + "\n", + "Now that we have a trained emulator, we can convert it to Pyomo algebraic expressions. This allows us to use the emulator in optimization problems with mathematical programming solvers." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Before conversion, we will setup our optimization framework through pyomo. Pyomo has a **ConcreteModel**, which is a container that holds all the components of an optimization problem - variables, objectives, constraints. Think of it as the \"problem definition sheet.\"\n", + "\n", + "- **Pyomo variables:** are symbolic placeholders for inputs that a mathematical solver can manipulate during optimization. They correspond to the same x1, x2 input features the Emulator was trained on. There are multiple types of variable domains in pyomo, some common are explained below:\n", + "\n", + " | Domain | Use case |\n", + " |------------------------|-----------------------------------------------|\n", + " | `pyo.Reals` | Continuous variables (default for regression) |\n", + " | `pyo.NonNegativeReals` | Continuous, must be ≥ 0 |\n", + " | `pyo.NonPositiveReals` | Continuous, must be ≤ 0 |\n", + " | `pyo.Integers` | Integer variables |\n", + " | `pyo.Binary` | 0/1 decisions |\n", + " | `pyo.PositiveReals` | Strictly > 0 |\n", + " | `pyo.NegativeReals` | Strictly < 0 |" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "\n", + "- **Pyomo algebraic expressions:** are the Emulator's learned input→output relationships written as pure math, which any mathematical programming solver can evaluate and search for an optimal solution." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import pyomo.environ as pyo\n", + "\n", + "# Create a Pyomo concrete model\n", + "pyo_model = pyo.ConcreteModel()\n", + "\n", + "# Define decision variables (input variables for the autoemulate emulator)\n", + "# `pyo.Reals` are continuous real number. Because our training data is continuous, we use `pyo.Reals` for the decision variables. \n", + "# If the training data were integers, we would use `pyo.Integers` instead. Bounds are set to (-100, 100) based on the range of the training data.\n", + "pyo_model.x1 = pyo.Var(domain=pyo.Reals, bounds=(-100, 100))\n", + "pyo_model.x2 = pyo.Var(domain=pyo.Reals, bounds=(-100, 100))" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Export Emulator to Pyomo\n", + "\n", + "This is the core integration step. The `pyomofy` function converts the Emulator's weights, biases, and activation functions into explicit algebraic expressions.\n", + "\n", + "It automatically handles:\n", + "\n", + "* **Input Standardization:** Considering any scaling applied to inputs of the emulator during training.\n", + "* **Forward Pass:** inputs propagates through the architecture of the emulator and a output value is computed.\n", + "* **Output Inverse Scaling:** Inverting any scaling applied to the output of the Emulator.\n", + "\n", + "\n", + "* **ReLU Approximation:** If `nn.ReLU` was used during training of MLP Emulator, it's automatically approximated with Softplus (controlled by `relu_beta`, default `50`). if results mismatch significantly, increase `relu_beta` value (max 100) or use a smooth activation function like `SiLU` during training:\n", + "\n", + " ```\n", + " from torch import nn\n", + "\n", + " # Train with SiLU activation for better Pyomo conversion\n", + " ae = AutoEmulate(\n", + " x, y, \n", + " models=['MLP'],\n", + " model_params={'activation_cls': nn.SiLU} # Use SiLU instead of ReLU\n", + " )\n", + " ```" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from autoemulate.convert.pyomo import pyomofy\n", + "\n", + "# Convert the Emulator to Pyomo expressions\n", + "# Pass the Result object (or TransformedEmulator) and the list of Pyomo variables\n", + "# Returns a list of expressions (one per output of the Emulator.)\n", + "emulator_expressions = pyomofy(Emulator, [pyo_model.x1, pyo_model.x2], relu_beta=100)\n", + "\n", + "# Since we only have 1 output for our Emulator, we take the first element of expression list emulator_expressions[0].\n", + "# For multi-output emulators, to access nth output, use emulator_expressions[n-1]\n", + "emulator_expr = emulator_expressions[0]" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Verify Numerical Equivalence\n", + "\n", + "Before using the Pyomo expressions in optimization, we **must** verify that they yield the same values as the Emulator. This validation ensures the conversion was successful and the algebraic expressions accurately represent the Emulator." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Pick an input point to compare pyomo prediction with the Emulator's prediction. We picked first sample x[0] of the simulation data.\n", + "\n", + "x_init = x[0]\n", + "\n", + "# Set pyomo variables to the selected input values. item() convert tensors to python float for inputting to pyomo.\n", + "pyo_model.x1.set_value(x_init[0].item())\n", + "pyo_model.x2.set_value(x_init[1].item())\n", + "\n", + "# 1. Predict Pyomo at picked input point\n", + "pyomo_val = pyo.value(emulator_expr)\n", + "\n", + "# 2. Predict Emulator at the same input point\n", + "\n", + "Emulator_input = x[0].reshape(1, -1)\n", + "Emulator_val = Emulator.model.predict(Emulator_input).item()\n", + "\n", + "# Print comparison\n", + "print(f\"Input point: x1={x_init[0]:.4f}, x2={x_init[1]:.4f}\")\n", + "print(f\"Emulator prediction: {Emulator_val:.12f}\")\n", + "print(f\"Pyomo prediction: {pyomo_val:.12f}\")\n", + "print(f\"\\nPyomo vs Emulator: {abs(pyomo_val - Emulator_val):.12f}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Using the Pyomo Expression\n", + "\n", + "Now that we've validated the conversion, the Pyomo expression can be used in optimization problems. Here's a simple example of how you might set up an objective function:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Maximize the emulator output value.\n", + "# `sense=pyo.maximize` tells the solver to find inputs that make the expression as large as possible.\n", + "# `sense=pyo.minimize` tells the solver to find inputs that make the expression as small as possible.\n", + "\n", + "pyo_model.obj = pyo.Objective(expr=emulator_expr, sense=pyo.maximize)\n", + "\n", + "# Optionally, we can add constraints on the emulator output:\n", + "# pyo_model.constraint1 = pyo.Constraint(expr=emulator_expr >= 1000)\n", + "\n", + "# OR restrict inputs to the some other data range:\n", + "\n", + "# pyo_model.constraint2 = pyo.Constraint(expr=pyo.inequality(-50, pyo_model.x1, 100))\n", + "# pyo_model.constraint3 = pyo.Constraint(expr=pyo.inequality(-100, pyo_model.x2, 200))\n", + "\n", + "# Or bounded on one side:\n", + "\n", + "# pyo_model.constraint4 = pyo.Constraint(expr=pyo_model.x1 >= -100)\n", + "# pyo_model.constraint5 = pyo.Constraint(expr=pyo_model.x1 <= 100)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Solve the optimization problem using mathematical solvers.\n", + "# In this example, we will use IPOPT (a gradient-based NLP solver).\n", + "# Other options: 'glpk' (linear/integer), 'gurobi', 'bonmin' (mixed-integer nonlinear).\n", + "\n", + "solver = pyo.SolverFactory('ipopt')\n", + "results = solver.solve(pyo_model, tee=True) # tee=True prints solver log\n", + "\n", + "# Extract and print results\n", + "print(f\"Solver status: {results.solver.status}\")\n", + "print(f\"Objective value: {pyo.value(pyo_model.obj):.6f}\")\n", + "print(f\"Optimal x1: {pyo.value(pyo_model.x1):.6f}\")\n", + "print(f\"Optimal x2: {pyo.value(pyo_model.x2):.6f}\")" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.9.0" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} diff --git a/docs/tutorials/advanced/index.md b/docs/tutorials/advanced/index.md index d936c2776..4e92d10e4 100644 --- a/docs/tutorials/advanced/index.md +++ b/docs/tutorials/advanced/index.md @@ -1,3 +1,3 @@ # Advanced usage -This section covers some of the more advanced features of AutoEmulate, including adding custom emulators. \ No newline at end of file +This section covers some of the more advanced features of AutoEmulate, including adding custom emulators, converting emulator to pyomo algebraic equations. \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index b811d0ab7..55c8c1069 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -38,6 +38,8 @@ dependencies = [ "getdist>=1.7.2", "einops>=0.8.1", "harmonic>=1.3.0", + "pyomo>=6", + ] [project.urls] diff --git a/tests/convert/test_pyomo.py b/tests/convert/test_pyomo.py new file mode 100644 index 000000000..a379fe1fd --- /dev/null +++ b/tests/convert/test_pyomo.py @@ -0,0 +1,300 @@ +import math + +import numpy as np +import pyomo.environ as pyo +import pytest # type: ignore[import] +import torch +from autoemulate import AutoEmulate +from autoemulate.convert.pyomo import pyomofy +from autoemulate.transforms.pca import PCATransform +from autoemulate.transforms.standardize import StandardizeTransform +from torch import nn + + +def _make_product_data(n_samples=200, seed=0): + rng = np.random.default_rng(seed) + x1 = rng.uniform(0, 1, size=n_samples) + x2 = rng.uniform(0, 1, size=n_samples) + + x = np.column_stack((x1, x2)).astype(np.float32) + y = (x1 * x2).astype(np.float32) + + x_tensor = torch.tensor(x, dtype=torch.float32) + y_tensor = torch.tensor(y, dtype=torch.float32) + return x, x_tensor, y_tensor + + +def _pyomo_model_with_vars(x_init): + m = pyo.ConcreteModel() + m.x1 = pyo.Var(initialize=float(x_init[0]), bounds=(0, 1)) # type: ignore[assignment] + m.x2 = pyo.Var(initialize=float(x_init[1]), bounds=(0, 1)) # type: ignore[assignment] + return m + + +def _predict_like_pyomofy(result, x_tensor_2d: torch.Tensor) -> float: + """ + Predict with the same object that pyomofy exports. + + In this codebase: + - pyomofy(Result) unwraps to Result.model (a TransformedEmulator) + - TransformedEmulator is not callable, so we use .predict(...) + """ + out = result.model.predict(x_tensor_2d) + if isinstance(out, torch.Tensor): + out = out.detach().cpu().numpy() + return float(np.asarray(out).reshape(-1)[0]) + + +def test_pyomofy_matches_autoemulate_predict_mlp_relu(): + x, x_tensor, y_tensor = _make_product_data(n_samples=300, seed=42) + + ae = AutoEmulate( + x_tensor, + y_tensor, + log_level="warning", + models=["MLP"], + model_params={"activation_cls": nn.ReLU, "random_seed": 42}, + ) + result = ae.best_result() + assert result is not None + + m = _pyomo_model_with_vars(x_init=x[0]) + x1_var = m.x1 # type: ignore[assignment] + x2_var = m.x2 # type: ignore[assignment] + exprs = pyomofy(result, [x1_var, x2_var]) + assert isinstance(exprs, list) + assert len(exprs) == 1 + expr = exprs[0] + + test_pts = np.array( + [ + [0.1, 0.9], + [0.25, 0.4], + [0.7, 0.2], + ], + dtype=np.float32, + ) + for pt in test_pts: + x1_var.set_value(float(pt[0])) # type: ignore[union-attr] + x2_var.set_value(float(pt[1])) # type: ignore[union-attr] + pyomo_val = pyo.value(expr) + + torch_val = _predict_like_pyomofy(result, torch.tensor(pt.reshape(1, -1))) + + # ReLU is exported as smooth Softplus_beta, so expect approximation error. + assert pyomo_val == pytest.approx(torch_val, rel=2e-2, abs=1e-3) + + +def test_pyomofy_matches_autoemulate_predict_polynomial_regression(): + x, x_tensor, y_tensor = _make_product_data(n_samples=300, seed=123) + + ae = AutoEmulate( + x_tensor, + y_tensor, + log_level="warning", + models=["PolynomialRegression"], + model_params={"degree": 2}, + ) + result = ae.best_result() + assert result is not None + + m = _pyomo_model_with_vars(x_init=x[0]) + x1_var = m.x1 # type: ignore[assignment] + x2_var = m.x2 # type: ignore[assignment] + exprs = pyomofy(result, [x1_var, x2_var]) + assert len(exprs) == 1 + expr = exprs[0] + + test_pts = np.array( + [ + [0.15, 0.85], + [0.33, 0.44], + [0.9, 0.1], + ], + dtype=np.float32, + ) + for pt in test_pts: + x1_var.set_value(float(pt[0])) # type: ignore[union-attr] + x2_var.set_value(float(pt[1])) # type: ignore[union-attr] + pyomo_val = pyo.value(expr) + + torch_val = _predict_like_pyomofy(result, torch.tensor(pt.reshape(1, -1))) + assert pyomo_val == pytest.approx(torch_val, rel=1e-7, abs=1e-7) + + +def test_pyomofy_matches_autoemulate_predict_with_x_standardize_then_pca(): + x, x_tensor, y_tensor = _make_product_data(n_samples=400, seed=7) + + ae = AutoEmulate( + x_tensor, + y_tensor, + log_level="warning", + models=["MLP"], + model_params={"activation_cls": nn.Sigmoid, "random_seed": 7}, + x_transforms_list=[[StandardizeTransform(), PCATransform(n_components=2)]], + y_transforms_list=[[StandardizeTransform()]], + ) + result = ae.best_result() + assert result is not None + + m = _pyomo_model_with_vars(x_init=x[0]) + x1_var = m.x1 # type: ignore[assignment] + x2_var = m.x2 # type: ignore[assignment] + expr = pyomofy(result, [x1_var, x2_var])[0] + + test_pts = np.array( + [ + [0.05, 0.95], + [0.5, 0.5], + [0.8, 0.3], + ], + dtype=np.float32, + ) + for pt in test_pts: + x1_var.set_value(float(pt[0])) # type: ignore[union-attr] + x2_var.set_value(float(pt[1])) # type: ignore[union-attr] + pyomo_val = pyo.value(expr) + + torch_val = _predict_like_pyomofy(result, torch.tensor(pt.reshape(1, -1))) + + # PCA-based pipelines can accumulate more numerical mismatch than the + # simpler cases (matrix products + Pyomo exp/log in the NN). + assert pyomo_val == pytest.approx(torch_val, rel=2e-1, abs=1e-2) + + +def test_pyomofy_raises_on_unsupported_transform(): + x, x_tensor, y_tensor = _make_product_data(n_samples=200, seed=99) + + ae = AutoEmulate( + x_tensor, + y_tensor, + log_level="warning", + models=["MLP"], + model_params={"activation_cls": nn.Sigmoid, "random_seed": 99}, + ) + result = ae.best_result() + assert result is not None + + class _UnsupportedTransform: + pass + + # Use setattr to avoid type checker issues with assignment + model = result.model + model.x_transforms = [_UnsupportedTransform()] # type: ignore[assignment] + + m = _pyomo_model_with_vars(x_init=x[0]) + x1_var = m.x1 # type: ignore[assignment] + x2_var = m.x2 # type: ignore[assignment] + + with pytest.raises(NotImplementedError): + pyomofy(result, [x1_var, x2_var]) + + +def test_pyomofy_raises_on_unsupported_inner_emulator_transform(): + """Unsupported transforms on the inner emulator (e.g. MLP.x_transform) must also + be caught before any Pyomo work is done.""" + x, x_tensor, y_tensor = _make_product_data(n_samples=200, seed=99) + + ae = AutoEmulate( + x_tensor, + y_tensor, + log_level="warning", + models=["MLP"], + model_params={"activation_cls": nn.Sigmoid, "random_seed": 99}, + ) + result = ae.best_result() + assert result is not None + + class _UnsupportedTransform: + pass + + # Inject an unsupported transform at the inner emulator level + emulator = result.model.model + emulator.x_transform = _UnsupportedTransform() # type: ignore[assignment] + + m = _pyomo_model_with_vars(x_init=x[0]) + x1_var = m.x1 # type: ignore[assignment] + x2_var = m.x2 # type: ignore[assignment] + + with pytest.raises(NotImplementedError): + pyomofy(result, [x1_var, x2_var]) + + +def test_pyomofy_matches_emulator_with_inner_transforms_only(): + """Correctness test for the path where inner emulator transforms are non-trivial. + + We deliberately bypass the outer TransformedEmulator standardization by passing + x_transforms_list=[[]] and y_transforms_list=[[]], so that the MLP's own + x_transform/y_transform (standardize_x=True, standardize_y=True by default) are + the only active transforms. In this case pyomofy must apply them explicitly, + otherwise the Pyomo expression would be evaluated on the raw un-scaled inputs and + produce wrong outputs. + """ + + x, x_tensor, y_tensor = _make_product_data(n_samples=300, seed=55) + + # No outer transforms — inner MLP transforms will now have non-trivial mean/std + ae = AutoEmulate( + x_tensor, + y_tensor, + log_level="warning", + models=["MLP"], + model_params={"activation_cls": nn.Sigmoid, "random_seed": 55}, + x_transforms_list=[[]], + y_transforms_list=[[]], + ) + result = ae.best_result() + assert result is not None + + emulator = result.model.model + # Confirm the inner transforms are non-trivial (mean not ≈ 0 or std not ≈ 1) + if emulator.x_transform is not None: + inner_mean = emulator.x_transform.mean.abs().max().item() + inner_std = (emulator.x_transform.std - 1).abs().max().item() + assert inner_mean > 1e-3 or inner_std > 1e-3, ( + "Expected non-trivial inner x_transform when no outer transforms are used" + ) + + m = _pyomo_model_with_vars(x_init=x[0]) + x1_var = m.x1 # type: ignore[assignment] + x2_var = m.x2 # type: ignore[assignment] + expr = pyomofy(result, [x1_var, x2_var])[0] + + test_pts = np.array( + [[0.1, 0.9], [0.4, 0.4], [0.8, 0.2]], + dtype=np.float32, + ) + for pt in test_pts: + x1_var.set_value(float(pt[0])) # type: ignore[union-attr] + x2_var.set_value(float(pt[1])) # type: ignore[union-attr] + pyomo_val = pyo.value(expr) + torch_val = _predict_like_pyomofy(result, torch.tensor(pt.reshape(1, -1))) + assert pyomo_val == pytest.approx(torch_val, rel=1e-5, abs=1e-5) + + +def test_pyomofy_expr_is_smooth_for_relu_via_softplus_beta(): + x, x_tensor, y_tensor = _make_product_data(n_samples=250, seed=101) + + ae = AutoEmulate( + x_tensor, + y_tensor, + log_level="warning", + models=["MLP"], + model_params={"activation_cls": nn.ReLU, "random_seed": 101}, + ) + result = ae.best_result() + m = _pyomo_model_with_vars(x_init=x[0]) + x1_var = m.x1 # type: ignore[assignment] + x2_var = m.x2 # type: ignore[assignment] + expr = pyomofy(result, [x1_var, x2_var], relu_beta=50)[0] + + x2_var.set_value(0.6) # type: ignore[union-attr] + xs = [0.0, 0.1, 0.2, 0.4, 0.8, 1.0] + vals = [] + for x1 in xs: + x1_var.set_value(x1) # type: ignore[union-attr] + v = pyo.value(expr) + assert math.isfinite(v) # type: ignore[union-attr] + vals.append(v) + + assert vals[-1] >= vals[0]