Add pyomofy Pyomo conversion support #963 Extension

autoemulate/convert/__init__.py

@@ -0,0 +1,3 @@
+from .pyomo import pyomofy
+
+__all__ = ["pyomofy"]

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

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

docs/reference/convert/index.md

@@ -0,0 +1,4 @@
+Convert
+=======
+
+Reference documentation for conversion utilities (e.g., exporting emulators to other frameworks).

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: