"""Pyomo model export utilities.
This module provides functions for exporting Pyomo model components to TOML files.
"""
from __future__ import annotations
import pathlib
from logging import getLogger
from typing import TYPE_CHECKING
from pyomo import environ as pyo
from eta_ctrl.util import toml_export
from eta_ctrl.util.io_utils import get_unique_output_path
if TYPE_CHECKING:
from typing import Any
log = getLogger(__name__)
def _extract_variable_bounds(bounds: tuple) -> dict[str, Any]:
"""Extract and process variable bounds into a standardized format.
This helper function handles the common logic for extracting low and high bounds
from Pyomo variables, converting infinite bounds to None and handling edge cases.
:param bounds: Tuple of (lower_bound, upper_bound) from Pyomo variable
:return: Dictionary containing processed bound information
"""
bounds_info = {}
if bounds[0] is not None:
low_val = float(bounds[0]) if bounds[0] != float("-inf") else None
if low_val is not None:
bounds_info["low_value"] = low_val
if bounds[1] is not None:
high_val = float(bounds[1]) if bounds[1] != float("inf") else None
if high_val is not None:
bounds_info["high_value"] = high_val
return bounds_info
def _extract_variable_domain_type(variable: pyo.Var) -> str:
"""Extract the domain type from a Pyomo variable.
This helper function determines whether a variable is continuous or discrete
based on its domain attribute, with a fallback to continuous as the default.
:param variable: Pyomo variable with domain attribute
:return: String indicating "continuous" or "discrete"
"""
if hasattr(variable, "domain") and variable.domain is not None:
domain_name = str(variable.domain)
return "continuous" if "Real" in domain_name else "discrete"
return "continuous" # Default assumption
[docs]
def export_pyomo_state_config(model: pyo.ConcreteModel, model_name: str, output_path: pathlib.Path) -> None:
"""Export Pyomo model variables (observations) to a TOML file.
This method extracts the variables from the Pyomo model and exports them to a TOML file
for later use in state configuration.
ATTENTION: All variables are treated as observations, you need to separate these.
:param model: Pyomo ConcreteModel instance.
:param model_name: Name of the model for identification.
:param output_path: Full path where the TOML file should be saved (including filename).
"""
# Extract variables (observations) from the model
observations = []
for component in model.component_objects(pyo.Var):
var_name = component.name
var_info = {
"name": var_name,
"is_indexed": component.is_indexed(),
}
# Extract variable-specific information
if component.is_indexed():
var_info.update(extract_indexed_variable_info(component))
else:
var_info.update(extract_scalar_variable_info(component))
observations.append(var_info)
pyomo_data = {
"model_info": {
"name": model_name,
"type": "pyomo",
},
"observations": observations,
}
final_output_path = get_unique_output_path(output_path)
toml_export(final_output_path, pyomo_data)
log.info(f"Pyomo model variables exported to {final_output_path}")
[docs]
def export_pyomo_parameters(model: pyo.ConcreteModel, model_name: str, output_path: pathlib.Path) -> None:
"""Export Pyomo model parameters to a TOML file.
This method extracts parameter names and values from the Pyomo model and exports them to a TOML file.
For indexed parameters, all values are collected as arrays to preserve the complete parameter information.
:param model: Pyomo ConcreteModel instance.
:param model_name: Name of the model for identification.
:param output_path: Full path where the TOML file should be saved (including filename).
"""
# Extract parameters from the model - preserve all values for indexed parameters
parameters = {}
for component in model.component_objects(pyo.Param):
param_name = component.name
if component.is_indexed():
# For indexed parameters, collect all values as arrays to preserve complete information
param_values = []
param_indices = []
for index in component.index_set():
try:
value = pyo.value(component[index])
if value is not None:
param_values.append(str(value))
param_indices.append(str(index))
except (ValueError, TypeError):
# ValueError: Parameter value cannot be evaluated (e.g., symbolic expressions,
# uninitialized parameters, or mutable parameters without values)
# TypeError: Parameter index or value type incompatible with conversion
# (e.g., complex objects that can't be stringified)
# Skip invalid entries but continue processing other indices
continue
# Store as arrays if we have values
if param_values:
parameters[param_name] = {"values": param_values, "indices": param_indices, "is_indexed": True}
else:
# For scalar parameters, store the actual value
try:
value = pyo.value(component)
if value is not None:
parameters[param_name] = {"value": str(value), "is_indexed": False}
except (ValueError, TypeError):
# ValueError: Parameter value cannot be evaluated (e.g., uninitialized parameter)
# TypeError: Parameter value type incompatible with string conversion
# Skip invalid parameters but continue processing others
continue
final_output_path = get_unique_output_path(output_path)
pyomo_data = {
"parameters": parameters,
"model_info": {"name": model_name, "path": str(final_output_path), "type": "pyomo_parameters"},
}
toml_export(final_output_path, pyomo_data)
log.info(f"Pyomo model parameters exported to {final_output_path}")
log.info(f"Exported {len(parameters)} parameters with complete value arrays")
[docs]
def export_pyomo_state(model: pyo.ConcreteModel, model_name: str, output_dir: pathlib.Path | str | None = None) -> None:
"""Export Pyomo model state config and parameters files.
This is the main public interface for exporting Pyomo model data, creating both
state configuration and parameters files.
:param model: Pyomo ConcreteModel instance.
:param model_name: Name of the model for identification.
:param output_dir: Directory where files should be created. If None, uses current working directory.
"""
# Centralize output directory logic
output_directory = pathlib.Path.cwd().absolute() if output_dir is None else pathlib.Path(output_dir).absolute()
output_directory.mkdir(parents=True, exist_ok=True)
# Create specific file paths
state_config_path = output_directory / f"{model_name}_state_config.toml"
parameters_path = output_directory / f"{model_name}_parameters.toml"
# Call export functions with concrete paths
export_pyomo_state_config(model, model_name, state_config_path)
export_pyomo_parameters(model, model_name, parameters_path)
log.info(f"Created Pyomo model files for '{model_name}' in {output_directory}")
# ---------------------------------------------------------------------------
# PyomoModel-specific export (actions / observations / model_parameters split)
# ---------------------------------------------------------------------------
[docs]
def export_pyomo_model_state_config(model: pyo.ConcreteModel, model_name: str, output_path: pathlib.Path) -> None:
"""Export a PyomoModel's components to a state config TOML file.
Classification rules:
* Indexed ``pyo.Var`` components → ``[[actions]]``
* Indexed ``pyo.Param`` components → ``[[observations]]``
Scalar parameters are intentionally omitted here; use
:func:`export_pyomo_model_parameters` for those.
:param model: Pyomo ConcreteModel instance.
:param model_name: Name of the model for identification.
:param output_path: Full path (including filename) for the TOML file.
"""
actions: list[dict[str, Any]] = []
observations: list[dict[str, Any]] = []
for component in model.component_objects(pyo.Var):
if not component.is_indexed():
continue
var_info: dict[str, Any] = {"name": component.name}
# Reuse existing helper; strip index metadata that belongs to internal bookkeeping
raw = extract_indexed_variable_info(component)
raw.pop("index_length", None)
raw.pop("index_set", None)
raw.pop("type", None)
var_info.update(raw)
# Warn if bounds are absent — StateVar.model_post_init enforces that both
# low_value and high_value must be set for action variables, so the generated
# TOML must be completed manually before it can be loaded.
if "low_value" not in var_info or "high_value" not in var_info:
log.warning(
f"Action variable '{component.name}' has no explicit bounds in the Pyomo model. "
"You must set 'low_value' and 'high_value' manually in the generated state config TOML "
"before loading it as a StateConfig."
)
actions.append(var_info)
for component in model.component_objects(pyo.Param):
if not component.is_indexed():
continue
observations.append({"name": component.name})
pyomo_data: dict[str, Any] = {}
if actions:
pyomo_data["actions"] = actions
if observations:
pyomo_data["observations"] = observations
final_output_path = get_unique_output_path(output_path)
toml_export(final_output_path, pyomo_data)
log.info(f"PyomoModel state config exported to {final_output_path}")
[docs]
def export_pyomo_model_parameters(model: pyo.ConcreteModel, model_name: str, output_path: pathlib.Path) -> None:
"""Export scalar (non-indexed) Pyomo parameters to a model parameters TOML file.
The output corresponds to the ``[agent_specific.model_parameters]`` section
of a run config.
:param model: Pyomo ConcreteModel instance.
:param model_name: Name of the model for identification.
:param output_path: Full path (including filename) for the TOML file.
"""
model_parameters: dict[str, Any] = {}
for component in model.component_objects(pyo.Param):
if component.is_indexed():
continue
try:
value = pyo.value(component)
if value is not None:
model_parameters[component.name] = value
except (ValueError, TypeError):
# Skip parameters that cannot be evaluated (e.g. uninitialized mutable params)
continue
pyomo_data: dict[str, Any] = {
"model_info": {"name": model_name, "type": "pyomo_model_parameters"},
"model_parameters": model_parameters,
}
final_output_path = get_unique_output_path(output_path)
toml_export(final_output_path, pyomo_data)
log.info(
f"PyomoModel parameters exported to {final_output_path}. "
"Place these values under [agent_specific.model_parameters] in your run config."
)
[docs]
def export_pyomo_model_state(
model: pyo.ConcreteModel, model_name: str, output_dir: pathlib.Path | str | None = None
) -> None:
"""Export a PyomoModel's state config and model parameters to TOML files.
This is the main public interface for :class:`~eta_ctrl.simulators.PyomoModel`
state generation. It writes two files:
* ``{model_name}_state_config.toml`` — indexed Vars as actions, indexed
Params as observations.
* ``{model_name}_model_parameters.toml`` — scalar Params that belong in
``[agent_specific.model_parameters]`` of the run config.
:param model: Pyomo ConcreteModel instance.
:param model_name: Name of the model used for file naming.
:param output_dir: Target directory. Defaults to the current working directory.
"""
output_directory = pathlib.Path.cwd().absolute() if output_dir is None else pathlib.Path(output_dir).absolute()
output_directory.mkdir(parents=True, exist_ok=True)
export_pyomo_model_state_config(model, model_name, output_directory / f"{model_name}_state_config.toml")
export_pyomo_model_parameters(model, model_name, output_directory / f"{model_name}_model_parameters.toml")
log.info(f"Created PyomoModel files for '{model_name}' in {output_directory}")