implementing PR feedback

Author: ParticularlyPythonicBS

docs/source/computational_implementation.rst

@@ -389,7 +389,7 @@ The Temoa model code is organized into clear, purpose-driven packages:
   * ``monte_carlo`` - :doc:`monte_carlo` (Uncertainty quantification)
   * ``myopic`` - Sequential decision making with limited foresight
   * ``single_vector_mga`` - Focused MGA on specific variables ([!] untested in v4.0)
-  * ``stochastics`` - Stochastic programming capabilities
+  * ``stochastics`` - :doc:`stochastics` (Stochastic programming capabilities)
 
 * ``temoa._internal`` - Internal utilities (not part of public API)
 

docs/source/stochastics.rst

@@ -60,30 +60,33 @@ The stochastic configuration file defines the scenarios, their probabilities, an
    # Define perturbations for a specific scenario
    [[perturbations]]
    scenario = "low_cost"
-   variable = "cost_variable"
-   tech = "IMPHCO1"
-   # The value is a multiplier applied to the base value in the database
+   table = "cost_variable"
+   # Filter specifies which rows in the table to perturb
+   filter = { tech = "IMPHCO1" }
+   # Action can be "multiply", "add", or "set" (defaults to "set")
+   action = "multiply"
    value = 0.5
 
    [[perturbations]]
    scenario = "high_cost"
-   variable = "cost_variable"
-   tech = "IMPHCO1"
+   table = "cost_variable"
+   filter = { tech = "IMPHCO1" }
+   action = "multiply"
    value = 1.5
 
 Perturbation Options
 ^^^^^^^^^^^^^^^^^^^^
 
-Currently, the following perturbations are supported:
+Currently, the following fields are required for each perturbation:
 
-* **variable**: The Temoa parameter to perturb. Currently supports:
-    * ``cost_variable``
-    * ``cost_invest``
-    * ``cost_fixed``
-    * ``demand``
-* **tech**: The technology to which the perturbation applies (required for cost perturbations).
-* **commodity**: The commodity to which the perturbation applies (required for demand perturbations).
-* **value**: A multiplier applied to the base value found in the input database. For example, a value of ``1.5`` increases the base value by 50%.
+* **scenario**: The name of the scenario to which this perturbation applies.
+* **table**: The Temoa parameter (database table) to perturb (e.g., ``cost_variable``, ``demand``, ``capacity_factor_process``).
+* **filter**: A dictionary of column-value pairs used to identify specific rows. Since the extension uses the dynamic manifest from ``HybridLoader``, any column belonging to the table's index can be used for filtering.
+* **action**: The operation to perform. Supported values:
+    * ``multiply``: Multiply the base value by ``value``.
+    * ``add``: Add ``value`` to the base value.
+    * ``set``: Replace the base value with ``value``.
+* **value**: The numeric value used in the perturbation action.
 
 How it Works
 ------------
@@ -101,6 +104,4 @@ Limitations
 -----------
 
 * **Two-Stage Only**: While ``mpi-sppy`` supports multi-stage stochastic programming, the current Temoa integration is tailored for two-stage problems where the first time period constitutes the first stage.
-* **Cost/Demand Multipliers**: Perturbations are currently implemented as multipliers on base values. Absolute value overrides are not yet supported.
-* **Solver Support**: The extension has been primarily tested with the ``appsi_highs`` and ``cbc`` solvers via Pyomo.
 * **Result Persistence**: Currently, only the expected objective value and summary logs are produced. Detailed per-scenario result persistence to the database is under development.

temoa/extensions/stochastics/scenario_creator.py

@@ -1,31 +1,50 @@
-import sqlite3
 import logging
+import sqlite3
+from typing import TYPE_CHECKING, Any
+
 import pyomo.environ as pyo
 from mpisppy.utils.sputils import attach_root_node
 
-from temoa.core.config import TemoaConfig
-from temoa.extensions.stochastics.stochastic_config import StochasticConfig
-from temoa.data_io.hybrid_loader import HybridLoader
 from temoa._internal.run_actions import build_instance
 from temoa.components.costs import period_cost_rule
+from temoa.data_io.hybrid_loader import HybridLoader
+
+if TYPE_CHECKING:
+    from temoa.core.config import TemoaConfig
+    from temoa.extensions.stochastics.stochastic_config import StochasticConfig
 
 logger = logging.getLogger(__name__)
 
-def scenario_creator(scenario_name, **kwargs):
+
+def scenario_creator(scenario_name: str, **kwargs: Any) -> pyo.ConcreteModel:
     """
     Creator for mpi-sppy scenarios.
 
     Args:
         scenario_name (str): Name of the scenario to create.
         **kwargs: Must contain 'temoa_config' and 'stoch_config'.
     """
+    if 'temoa_config' not in kwargs or 'stoch_config' not in kwargs:
+        raise ValueError("scenario_creator requires 'temoa_config' and 'stoch_config' in kwargs")
+
     temoa_config: TemoaConfig = kwargs['temoa_config']
     stoch_config: StochasticConfig = kwargs['stoch_config']
 
     # 1. Load base data
-    with sqlite3.connect(temoa_config.input_database) as con:
-        hybrid_loader = HybridLoader(db_connection=con, config=temoa_config)
-        data_dict = hybrid_loader.create_data_dict(myopic_index=None)
+    try:
+        with sqlite3.connect(temoa_config.input_database) as con:
+            hybrid_loader = HybridLoader(db_connection=con, config=temoa_config)
+            data_dict = hybrid_loader.create_data_dict(myopic_index=None)
+
+            # Build a map of table -> index columns from the manifest
+            # For each LoadItem, the index columns are all but the last one (which is the value)
+            table_index_map = {}
+            for item in hybrid_loader.manifest:
+                if item.table not in table_index_map and item.columns:
+                    table_index_map[item.table] = item.columns[:-1]
+    except Exception as e:
+        logger.exception('Failed to connect to database %s', temoa_config.input_database)
+        raise RuntimeError(f'Failed to connect to database {temoa_config.input_database}') from e
 
     # 2. Apply perturbations for this scenario
     for p in stoch_config.perturbations:
@@ -34,30 +53,25 @@ def scenario_creator(scenario_name, **kwargs):
 
         target_param = data_dict.get(p.table)
         if target_param is None:
-            logger.warning(f"Table {p.table} not found in data_dict for scenario {scenario_name}")
+            logger.warning(
+                'Table %s not found in data_dict for scenario %s', p.table, scenario_name
+            )
             continue
 
         # target_param is {(idx...): value}
         # We need to find entries matching p.filter
+        index_cols = table_index_map.get(p.table)
+        if index_cols is None:
+            logger.warning(
+                'Table %s not found in manifest; cannot map indices for scenario %s',
+                p.table,
+                scenario_name,
+            )
+            continue
+
         for idx_tuple, current_val in list(target_param.items()):
-            # Temoa v4 parameter indices:
-            # cost_variable: (region, period, tech, vintage)
-            # cost_invest: (region, tech, period)
-            # cost_fixed: (region, period, tech, vintage)
-
-            # Map index tuple to names based on table
-            # This is a bit brittle, but handles common cases
-            index_map = {}
-            if p.table == 'cost_variable':
-                index_map = {'region': idx_tuple[0], 'period': idx_tuple[1], 'tech': idx_tuple[2], 'vintage': idx_tuple[3]}
-            elif p.table == 'cost_invest':
-                index_map = {'region': idx_tuple[0], 'tech': idx_tuple[1], 'period': idx_tuple[2]}
-            elif p.table == 'cost_fixed':
-                index_map = {'region': idx_tuple[0], 'period': idx_tuple[1], 'tech': idx_tuple[2], 'vintage': idx_tuple[3]}
-            else:
-                # Generic fallback if filter names match index positions?
-                # For now, only support these three
-                pass
+            # Map index tuple to names based on table manifest
+            index_map = dict(zip(index_cols, idx_tuple, strict=True))
 
             # Check if filter matches
             match = True
@@ -79,10 +93,16 @@ def scenario_creator(scenario_name, **kwargs):
     instance = build_instance(data_portal, silent=True)
 
     # 4. Attach root node (Stage 1)
-    periods = sorted(list(instance.time_optimize))
+    periods = sorted(instance.time_optimize)
     first_period = periods[0]
 
-    prob = stoch_config.scenarios.get(scenario_name, {}).get('probability', 1.0)
+    prob = stoch_config.scenarios.get(scenario_name)
+    if prob is None:
+        logger.warning(
+            "Scenario '%s' not found in stochastic config probabilities; defaulting to 1.0",
+            scenario_name,
+        )
+        prob = 1.0
     instance._mpisppy_probability = prob
 
     # First stage variables: v_new_capacity[*, *, first_period]
@@ -91,6 +111,15 @@ def scenario_creator(scenario_name, **kwargs):
         if p == first_period:
             first_stage_vars.append(instance.v_new_capacity[r, t, p])
 
+    if not first_stage_vars:
+        logger.error(
+            'No first-stage variables (v_new_capacity for period %s) found for scenario %s. '
+            'Stochastic optimization requires at least one first-stage decision.',
+            first_period,
+            scenario_name,
+        )
+        raise ValueError(f'No first-stage variables found for scenario {scenario_name}')
+
     # First stage cost: PeriodCost[first_period]
     # We can use the period_cost_rule directly
     first_stage_cost_expr = period_cost_rule(instance, first_period)

temoa/extensions/stochastics/stochastic_config.py

@@ -1,38 +1,83 @@
+import logging
+import tomllib
 from dataclasses import dataclass, field
 from pathlib import Path
 from typing import Any
-import tomllib
+
+logger = logging.getLogger(__name__)
+
 
 @dataclass
 class Perturbation:
     scenario: str
     table: str
-    column: str
     filter: dict[str, Any]
     action: str  # 'multiply', 'add', 'set'
     value: float
 
+    def __post_init__(self) -> None:
+        allowed_actions = {'multiply', 'add', 'set'}
+        if self.action not in allowed_actions:
+            raise ValueError(
+                f"Invalid perturbation action '{self.action}'; must be one of {allowed_actions}"
+            )
+
+
 @dataclass
 class StochasticConfig:
-    scenarios: dict[str, dict[str, float]] # scenario_name -> {probability: float}
+    scenarios: dict[str, float]  # scenario_name -> probability
     perturbations: list[Perturbation] = field(default_factory=list)
+    solver_options: dict[str, Any] = field(default_factory=dict)
 
     @classmethod
     def from_toml(cls, path: Path) -> 'StochasticConfig':
         with open(path, 'rb') as f:
             data = tomllib.load(f)
 
-        scenarios = data.get('scenarios', {})
+        scenarios_raw = data.get('scenarios', {})
+        scenarios = {}
+        for name, val in scenarios_raw.items():
+            if isinstance(val, dict):
+                scenarios[name] = float(val.get('probability', 1.0))
+            else:
+                scenarios[name] = float(val)
+
+        # Validate probability distribution
+        if scenarios:
+            total_prob = sum(scenarios.values())
+            if not (0.99 <= total_prob <= 1.01):
+                logger.warning(
+                    'Stochastic scenario probabilities sum to %s; usually they should sum to ~1.0',
+                    total_prob,
+                )
+
         perturbations_data = data.get('perturbations', [])
         perturbations = []
-        for p in perturbations_data:
-            perturbations.append(Perturbation(
-                scenario=p['scenario'],
-                table=p['table'],
-                column=p['column'],
-                filter=p['filter'],
-                action=p.get('action', 'set'),
-                value=p['value']
-            ))
-
-        return cls(scenarios=scenarios, perturbations=perturbations)
+        for i, p in enumerate(perturbations_data):
+            try:
+                scenario_name = p['scenario']
+                if scenario_name not in scenarios:
+                    raise ValueError(
+                        f'Perturbation at index {i} references nonexistent scenario: '
+                        f"'{scenario_name}'. Available scenarios: {list(scenarios.keys())}"
+                    )
+
+                perturbations.append(
+                    Perturbation(
+                        scenario=scenario_name,
+                        table=p['table'],
+                        filter=p['filter'],
+                        action=p.get('action', 'set'),
+                        value=p['value'],
+                    )
+                )
+            except KeyError as e:
+                raise ValueError(f'Perturbation at index {i} is missing required field: {e}') from e
+
+        solver_options = data.get('solver_options', {})
+
+        return cls(
+            scenarios=scenarios,
+            perturbations=perturbations,
+            solver_options=solver_options,
+        )