Add PyomoModel class

ixmp/__init__.py

@@ -9,7 +9,7 @@
 from ixmp.core.scenario import Scenario, TimeSeries
 from ixmp.model import MODELS
 from ixmp.model.base import ModelError
-from ixmp.model.dantzig import DantzigModel
+from ixmp.model.dantzig import DantzigGAMSModel, DantzigPyomoModel
 from ixmp.model.gams import GAMSModel
 from ixmp.reporting import Reporter
 from ixmp.utils import show_versions
@@ -37,14 +37,14 @@
 BACKENDS["jdbc"] = JDBCBackend
 
 # Register Models provided by ixmp
-MODELS.update(
-    {
-        "default": GAMSModel,
-        "gams": GAMSModel,
-        "dantzig": DantzigModel,
-    }
-)
-
+for name, cls in (
+    ("default", GAMSModel),
+    ("gams", GAMSModel),
+    ("dantzig", DantzigGAMSModel),
+    ("dantzig-gams", DantzigGAMSModel),
+    ("dantzig-pyomo", DantzigPyomoModel),
+):
+    MODELS[name] = cls
 
 # Configure the 'ixmp' logger: write messages to stdout, defaulting to level WARNING
 # and above

ixmp/model/dantzig.py

@@ -1,13 +1,16 @@
 from collections import ChainMap
+from functools import lru_cache
 from pathlib import Path
+from typing import Dict
 
 import pandas as pd
 
 from ixmp.utils import maybe_check_out, maybe_commit, update_par
 
 from .gams import GAMSModel
+from .pyomo import PyomoModel
 
-ITEMS = {
+ITEMS: Dict[str, dict] = {
     # Plants
     "i": dict(ix_type="set"),
     # Markets
@@ -61,22 +64,9 @@
 }
 
 
-class DantzigModel(GAMSModel):
-    """Dantzig's cannery/transport problem as a :class:`GAMSModel`.
-
-    Provided for testing :mod:`ixmp` code.
-    """
-
+class DantzigModel:
     name = "dantzig"
 
-    defaults = ChainMap(
-        {
-            # Override keys from GAMSModel
-            "model_file": Path(__file__).with_name("dantzig.gms"),
-        },
-        GAMSModel.defaults,
-    )
-
     @classmethod
     def initialize(cls, scenario, with_data=False):
         """Initialize the problem.
@@ -106,3 +96,56 @@ def initialize(cls, scenario, with_data=False):
         scenario.change_scalar("f", *DATA["f"])
 
         maybe_commit(scenario, checkout, f"{cls.__name__}.initialize")
+
+
+class DantzigGAMSModel(DantzigModel, GAMSModel):
+    """Dantzig's cannery/transport problem as a :class:`GAMSModel`.
+
+    Provided for testing :mod:`ixmp` code.
+    """
+
+    name = "dantzig-gams"
+
+    defaults = ChainMap(
+        {
+            # Override keys from GAMSModel
+            "model_file": Path(__file__).with_name("dantzig.gms"),
+        },
+        GAMSModel.defaults,
+    )
+
+
+def supply(model, i):
+    return sum(model.x[i, j] for j in model.j) <= model.a[i]
+
+
+def demand(model, j):
+    return sum(model.x[i, j] for i in model.i) >= model.b[j]
+
+
+@lru_cache()
+def c(model, i, j):
+    return model.f * model.d[i, j] / 1000
+
+
+def cost(model):
+    return sum(c(model, i, j) * model.x[i, j] for i in model.i for j in model.j)
+
+
+class DantzigPyomoModel(DantzigModel, PyomoModel):
+    """Dantzig's cannery/transport problem as a :class:`PyomoModel`.
+
+    Provided for testing :mod:`ixmp` code.
+    """
+
+    name = "dantzig-pyomo"
+    items = ITEMS
+    equation = dict(
+        supply=supply,
+        demand=demand,
+        cost=cost,
+    )
+    component_kwargs = dict(
+        x=dict(bounds=(0.0, None)),
+    )
+    objective = "cost"

ixmp/model/pyomo.py

@@ -0,0 +1,139 @@
+from inspect import signature
+from typing import Dict, Optional
+
+try:
+    from pyomo import environ as pyo
+    from pyomo import opt
+
+    has_pyomo = True
+except ImportError:
+    has_pyomo = False
+
+from ixmp.model.base import Model
+
+COMPONENT = dict(
+    par=pyo.Param,
+    set=pyo.Set,
+    var=pyo.Var,
+    equ=pyo.Constraint,
+)
+
+
+def get_sets(model, names):
+    return [model.component(idx_set) for idx_set in names]
+
+
+class PyomoModel(Model):
+    """General class for ixmp models using :mod:`pyomo`."""
+
+    name = "pyomo"
+
+    items: Dict[str, dict] = {}
+    constraints: Dict = {}
+    objective: Optional[str] = None
+    _partials: Dict = {}
+
+    def __init__(self, name=None, solver="glpk"):
+        if not has_pyomo:
+            raise ImportError("pyomo must be installed")
+
+        self.opt = opt.SolverFactory(solver)
+
+        m = pyo.AbstractModel()
+
+        for name, info in self.items.items():
+            if name == self.objective:
+                # Handle the objective separately
+                continue
+
+            Component = COMPONENT[info["ix_type"]]
+
+            kwargs = {}
+
+            if info["ix_type"] == "equ":
+                func = self.equation[name]
+                params = signature(func).parameters
+                idx_sets = list(params.keys())[1:]
+                kwargs = dict(rule=func)
+            else:
+                idx_sets = info.get("idx_sets", None) or []
+
+                # NB would like to do this, but pyomo doesn't recognize partial
+                #    objects as callable
+                # if info['ix_type'] != 'var':
+                #     kwarg = dict(
+                #         initialize=partial(self.to_pyomo, name)
+                # )
+
+            kwargs.update(self.component_kwargs.get(name, {}))
+
+            component = Component(*get_sets(m, idx_sets), **kwargs)
+            m.add_component(name, component)
+
+        obj_func = self.equation[self.objective]
+        obj = pyo.Objective(rule=obj_func, sense=pyo.minimize)
+        m.add_component(self.objective, obj)
+
+        # Store
+        self.model = m
+
+    def to_pyomo(self, name):
+        info = self.items[name]
+        ix_type = info["ix_type"]
+
+        if ix_type == "par":
+            item = self.scenario.par(name)
+
+            idx_sets = info.get("idx_sets", []) or []
+            if len(idx_sets):
+                series = item.set_index(idx_sets)["value"]
+                series.index = series.index.to_flat_index()
+                return series.to_dict()
+            else:
+                return {None: item["value"]}
+        elif ix_type == "set":
+            return {None: self.scenario.set(name).tolist()}
+
+    def all_to_pyomo(self):
+        return {
+            None: dict(
+                filter(
+                    lambda name_data: name_data[1],
+                    [(name, self.to_pyomo(name)) for name in self.items],
+                )
+            )
+        }
+
+    def all_from_pyomo(self, model):
+        for name, info in self.items.items():
+            if info["ix_type"] not in ("equ", "var"):
+                continue
+            self.from_pyomo(model, name)
+
+    def from_pyomo(self, model, name):
+        component = model.component(name)
+        component.display()
+        try:
+            data = component.get_values()
+        except Exception as exc:
+            print(exc)
+            return
+
+        # TODO add to Scenario; currently not possible because ixmp_source does
+        # not allow setting elements of 'equ' and 'var'
+        del data
+
+    def run(self, scenario):
+        self.scenario = scenario
+
+        data = self.all_to_pyomo()
+
+        m = self.model.create_instance(data=data)
+
+        assert m.is_constructed()
+
+        results = self.opt.solve(m)
+
+        self.all_from_pyomo(m)
+
+        delattr(self, "scenario")

ixmp/testing/data.py

@@ -155,7 +155,9 @@ def add_test_data(scen: Scenario):
     return t, t_foo, t_bar, x
 
 
-def make_dantzig(mp: Platform, solve: bool = False, quiet: bool = False) -> Scenario:
+def make_dantzig(
+    mp: Platform, solve: bool = False, quiet: bool = False, scheme="dantzig-gams"
+) -> Scenario:
     """Return :class:`ixmp.Scenario` of Dantzig's canning/transport problem.
 
     Parameters
@@ -191,7 +193,7 @@ def make_dantzig(mp: Platform, solve: bool = False, quiet: bool = False) -> Scen
         **models["dantzig"],  # type: ignore [arg-type]
         version="new",
         annotation=annot,
-        scheme="dantzig",
+        scheme=scheme,
         with_data=True,
     )
 

ixmp/tests/test_model.py

@@ -21,6 +21,26 @@ class M1(Model):
         M1()
 
 
+@pytest.mark.parametrize(
+    "kwargs",
+    [
+        dict(comment=None),
+        dict(equ_list=None, var_list=["x"]),
+        dict(equ_list=["demand", "supply"], var_list=[]),
+    ],
+    ids=["null-comment", "null-list", "empty-list"],
+)
+def test_GAMSModel(test_mp, test_data_path, kwargs):
+    s = make_dantzig(test_mp)
+    s.solve(model="dantzig", **kwargs)
+
+
+def test_PyomoModel(test_mp):
+    """Pyomo version of the Dantzig model builds and solves."""
+    s = make_dantzig(test_mp, scheme="dantzig-pyomo")
+    s.solve(model="dantzig-pyomo")
+
+
 def test_model_initialize(test_mp, caplog):
     # Model.initialize runs on an empty Scenario
     s = make_dantzig(test_mp)

pyproject.toml

@@ -94,6 +94,7 @@ module = [
   "memory_profiler",
   "pandas.*",
   "pyam",
+  "pyomo",
   "pretenders.*",
 ]
 ignore_missing_imports = true

setup.cfg

@@ -47,12 +47,15 @@ docs =
     sphinx >= 3.0
     sphinx_rtd_theme
     sphinxcontrib-bibtex
+pyomo =
+    pyomo
 report =
     genno[compat,graphviz]
 tutorial =
     jupyter
 tests =
     %(docs)s
+    %(pyomo)s
     %(report)s
     %(tutorial)s
     codecov