From 45bb2245a5992360b7f0faf7dddeb767917bc4d0 Mon Sep 17 00:00:00 2001 From: Fridolin Glatter <83776373+glatterf42@users.noreply.github.com> Date: Thu, 3 Oct 2024 11:16:12 +0200 Subject: [PATCH] Introduce optimization equation (#98) * Introduce optimization.Equation * Fix and test equation list and tabulate for specific runs * Introduce equation.remove_data() * Remove superfluous session.add() for equation * Use own errors for Equation --- ixmp4/core/__init__.py | 1 + ixmp4/core/optimization/data.py | 3 + ixmp4/core/optimization/equation.py | 140 ++++++ ixmp4/data/abstract/__init__.py | 2 + ixmp4/data/abstract/optimization/__init__.py | 1 + ixmp4/data/abstract/optimization/column.py | 6 +- ixmp4/data/abstract/optimization/equation.py | 217 +++++++++ ixmp4/data/api/__init__.py | 2 + ixmp4/data/api/optimization/__init__.py | 1 + ixmp4/data/api/optimization/column.py | 3 +- ixmp4/data/api/optimization/equation.py | 86 ++++ ixmp4/data/backend/api.py | 2 + ixmp4/data/backend/base.py | 2 + ixmp4/data/backend/db.py | 3 + ixmp4/data/db/__init__.py | 2 + ixmp4/data/db/filters/__init__.py | 1 + ixmp4/data/db/filters/optimizationequation.py | 17 + ixmp4/data/db/optimization/__init__.py | 1 + ixmp4/data/db/optimization/column/model.py | 8 +- .../data/db/optimization/column/repository.py | 7 + .../data/db/optimization/equation/__init__.py | 2 + ixmp4/data/db/optimization/equation/docs.py | 8 + ixmp4/data/db/optimization/equation/filter.py | 17 + ixmp4/data/db/optimization/equation/model.py | 41 ++ .../db/optimization/equation/repository.py | 186 ++++++++ ixmp4/data/db/run/model.py | 2 + ixmp4/server/rest/__init__.py | 2 + ixmp4/server/rest/docs.py | 24 + ixmp4/server/rest/optimization/__init__.py | 2 +- ixmp4/server/rest/optimization/equation.py | 86 ++++ tests/core/test_optimization_equation.py | 367 +++++++++++++++ tests/data/test_docs.py | 67 +++ tests/data/test_optimization_equation.py | 416 ++++++++++++++++++ 33 files changed, 1719 insertions(+), 6 deletions(-) create mode 100644 ixmp4/core/optimization/equation.py create mode 100644 ixmp4/data/abstract/optimization/equation.py create mode 100644 ixmp4/data/api/optimization/equation.py create mode 100644 ixmp4/data/db/filters/optimizationequation.py create mode 100644 ixmp4/data/db/optimization/equation/__init__.py create mode 100644 ixmp4/data/db/optimization/equation/docs.py create mode 100644 ixmp4/data/db/optimization/equation/filter.py create mode 100644 ixmp4/data/db/optimization/equation/model.py create mode 100644 ixmp4/data/db/optimization/equation/repository.py create mode 100644 ixmp4/server/rest/optimization/equation.py create mode 100644 tests/core/test_optimization_equation.py create mode 100644 tests/data/test_optimization_equation.py diff --git a/ixmp4/core/__init__.py b/ixmp4/core/__init__.py index a63efffc..a86c4ce2 100644 --- a/ixmp4/core/__init__.py +++ b/ixmp4/core/__init__.py @@ -1,6 +1,7 @@ # flake8: noqa from .iamc.variable import Variable as Variable from .model import Model as Model +from .optimization.equation import Equation as Equation from .optimization.indexset import IndexSet as IndexSet from .optimization.scalar import Scalar as Scalar from .optimization.table import Table as Table diff --git a/ixmp4/core/optimization/data.py b/ixmp4/core/optimization/data.py index 8d5f0680..69c18ee0 100644 --- a/ixmp4/core/optimization/data.py +++ b/ixmp4/core/optimization/data.py @@ -1,6 +1,7 @@ from ixmp4.data.abstract import Run from ..base import BaseFacade +from .equation import EquationRepository from .indexset import IndexSetRepository from .parameter import ParameterRepository from .scalar import ScalarRepository @@ -12,6 +13,7 @@ class OptimizationData(BaseFacade): """An optimization data instance, which provides access to optimization data such as IndexSet, Table, Variable, etc.""" + equations: EquationRepository indexsets: IndexSetRepository parameters: ParameterRepository scalars: ScalarRepository @@ -20,6 +22,7 @@ class OptimizationData(BaseFacade): def __init__(self, *args, run: Run, **kwargs) -> None: super().__init__(*args, **kwargs) + self.equations = EquationRepository(_backend=self.backend, _run=run) self.indexsets = IndexSetRepository(_backend=self.backend, _run=run) self.parameters = ParameterRepository(_backend=self.backend, _run=run) self.scalars = ScalarRepository(_backend=self.backend, _run=run) diff --git a/ixmp4/core/optimization/equation.py b/ixmp4/core/optimization/equation.py new file mode 100644 index 00000000..638bf7f0 --- /dev/null +++ b/ixmp4/core/optimization/equation.py @@ -0,0 +1,140 @@ +from datetime import datetime +from typing import Any, ClassVar, Iterable + +import pandas as pd + +from ixmp4.core.base import BaseFacade, BaseModelFacade +from ixmp4.data.abstract import Docs as DocsModel +from ixmp4.data.abstract import Equation as EquationModel +from ixmp4.data.abstract import Run +from ixmp4.data.abstract.optimization import Column + + +class Equation(BaseModelFacade): + _model: EquationModel + NotFound: ClassVar = EquationModel.NotFound + NotUnique: ClassVar = EquationModel.NotUnique + + @property + def id(self) -> int: + return self._model.id + + @property + def name(self) -> str: + return self._model.name + + @property + def run_id(self) -> int: + return self._model.run__id + + @property + def data(self) -> dict[str, Any]: + return self._model.data + + def add(self, data: dict[str, Any] | pd.DataFrame) -> None: + """Adds data to an existing Equation.""" + self.backend.optimization.equations.add_data( + equation_id=self._model.id, data=data + ) + self._model.data = self.backend.optimization.equations.get( + run_id=self._model.run__id, name=self._model.name + ).data + + def remove_data(self) -> None: + """Removes data from an existing Equation.""" + self.backend.optimization.equations.remove_data(equation_id=self._model.id) + self._model.data = self.backend.optimization.equations.get( + run_id=self._model.run__id, name=self._model.name + ).data + + @property + def levels(self) -> list: + return self._model.data.get("levels", []) + + @property + def marginals(self) -> list: + return self._model.data.get("marginals", []) + + @property + def constrained_to_indexsets(self) -> list[str]: + return [column.indexset.name for column in self._model.columns] + + @property + def columns(self) -> list[Column]: + return self._model.columns + + @property + def created_at(self) -> datetime | None: + return self._model.created_at + + @property + def created_by(self) -> str | None: + return self._model.created_by + + @property + def docs(self): + try: + return self.backend.optimization.equations.docs.get(self.id).description + except DocsModel.NotFound: + return None + + @docs.setter + def docs(self, description): + if description is None: + self.backend.optimization.equations.docs.delete(self.id) + else: + self.backend.optimization.equations.docs.set(self.id, description) + + @docs.deleter + def docs(self): + try: + self.backend.optimization.equations.docs.delete(self.id) + # TODO: silently failing + except DocsModel.NotFound: + return None + + def __str__(self) -> str: + return f"" + + +class EquationRepository(BaseFacade): + _run: Run + + def __init__(self, _run: Run, *args, **kwargs) -> None: + super().__init__(*args, **kwargs) + self._run = _run + + def create( + self, + name: str, + constrained_to_indexsets: list[str], + column_names: list[str] | None = None, + ) -> Equation: + model = self.backend.optimization.equations.create( + name=name, + run_id=self._run.id, + constrained_to_indexsets=constrained_to_indexsets, + column_names=column_names, + ) + return Equation(_backend=self.backend, _model=model) + + def get(self, name: str) -> Equation: + model = self.backend.optimization.equations.get(run_id=self._run.id, name=name) + return Equation(_backend=self.backend, _model=model) + + def list(self, name: str | None = None) -> Iterable[Equation]: + equations = self.backend.optimization.equations.list( + run_id=self._run.id, name=name + ) + return [ + Equation( + _backend=self.backend, + _model=i, + ) + for i in equations + ] + + def tabulate(self, name: str | None = None) -> pd.DataFrame: + return self.backend.optimization.equations.tabulate( + run_id=self._run.id, name=name + ) diff --git a/ixmp4/data/abstract/__init__.py b/ixmp4/data/abstract/__init__.py index ad3655ff..13b1994e 100644 --- a/ixmp4/data/abstract/__init__.py +++ b/ixmp4/data/abstract/__init__.py @@ -29,6 +29,8 @@ from .meta import MetaValue, RunMetaEntry, RunMetaEntryRepository, StrictMetaValue from .model import Model, ModelRepository from .optimization import ( + Equation, + EquationRepository, IndexSet, IndexSetRepository, Parameter, diff --git a/ixmp4/data/abstract/optimization/__init__.py b/ixmp4/data/abstract/optimization/__init__.py index 7a7c9660..60603b1e 100644 --- a/ixmp4/data/abstract/optimization/__init__.py +++ b/ixmp4/data/abstract/optimization/__init__.py @@ -1,4 +1,5 @@ from .column import Column +from .equation import Equation, EquationRepository from .indexset import IndexSet, IndexSetRepository from .parameter import Parameter, ParameterRepository from .scalar import Scalar, ScalarRepository diff --git a/ixmp4/data/abstract/optimization/column.py b/ixmp4/data/abstract/optimization/column.py index 260f16ba..9c42580a 100644 --- a/ixmp4/data/abstract/optimization/column.py +++ b/ixmp4/data/abstract/optimization/column.py @@ -14,10 +14,12 @@ class Column(base.BaseModel, Protocol): """Unique name of the Column.""" dtype: types.String """Type of the Column's data.""" - table__id: types.Mapped[int | None] - """Foreign unique integer id of a Table.""" + equation__id: types.Mapped[int | None] + """Foreign unique integer id of a Equation.""" parameter__id: types.Mapped[int | None] """Foreign unique integer id of a Parameter.""" + table__id: types.Mapped[int | None] + """Foreign unique integer id of a Table.""" variable__id: types.Mapped[int | None] """Foreign unique integer id of a Variable.""" indexset: types.Mapped[IndexSet] diff --git a/ixmp4/data/abstract/optimization/equation.py b/ixmp4/data/abstract/optimization/equation.py new file mode 100644 index 00000000..118708af --- /dev/null +++ b/ixmp4/data/abstract/optimization/equation.py @@ -0,0 +1,217 @@ +from typing import Any, Iterable, Protocol + +import pandas as pd + +from ixmp4.data import types + +from .. import base +from ..docs import DocsRepository +from .column import Column + + +class Equation(base.BaseModel, Protocol): + """Equation data model.""" + + name: types.String + """Unique name of the Equation.""" + data: types.JsonDict + """Data stored in the Equation.""" + columns: types.Mapped[list[Column]] + """Data specifying this Equation's Columns.""" + + run__id: types.Integer + "Foreign unique integer id of a run." + + created_at: types.DateTime + "Creation date/time. TODO" + created_by: types.String + "Creator. TODO" + + def __str__(self) -> str: + return f"" + + +class EquationRepository( + base.Creator, + base.Retriever, + base.Enumerator, + Protocol, +): + docs: DocsRepository + + def create( + self, + run_id: int, + name: str, + constrained_to_indexsets: list[str], + column_names: list[str] | None = None, + ) -> Equation: + """Creates an Equation. + + Each column of the Equation needs to be constrained to an existing + :class:ixmp4.data.abstract.optimization.IndexSet. These are specified by name + and per default, these will be the column names. They can be overwritten by + specifying `column_names`, which needs to specify a unique name for each column. + + Parameters + ---------- + run_id : int + The id of the :class:`ixmp4.data.abstract.Run` for which this Equation is + defined. + name : str + The unique name of the Equation. + constrained_to_indexsets : list[str] + List of :class:`ixmp4.data.abstract.optimization.IndexSet` names that define + the allowed contents of the Equation's columns. + column_names: list[str] | None = None + Optional list of names to use as column names. If given, overwrites the + names inferred from `constrained_to_indexsets`. + + Raises + ------ + :class:`ixmp4.data.abstract.optimization.Equation.NotUnique`: + If the Equation with `name` already exists for the Run with `run_id`. + ValueError + If `column_names` are not unique or not enough names are given. + + Returns + ------- + :class:`ixmp4.data.abstract.optimization.Equation`: + The created Equation. + """ + ... + + def get(self, run_id: int, name: str) -> Equation: + """Retrieves an Equation. + + Parameters + ---------- + run_id : int + The id of the :class:`ixmp4.data.abstract.Run` for which this Equation is + defined. + name : str + The name of the Equation. + + Raises + ------ + :class:`ixmp4.data.abstract.optimization.Equation.NotFound`: + If the Equation with `name` does not exist. + + Returns + ------- + :class:`ixmp4.data.abstract.optimization.Equation`: + The retrieved Equation. + """ + ... + + def get_by_id(self, id: int) -> Equation: + """Retrieves an Equation by its id. + + Parameters + ---------- + id : int + Unique integer id. + + Raises + ------ + :class:`ixmp4.data.abstract.optimization.Equation.NotFound`. + If the Equation with `id` does not exist. + + Returns + ------- + :class:`ixmp4.data.abstract.optimization.Equation`: + The retrieved Equation. + """ + ... + + def list(self, *, name: str | None = None, **kwargs) -> Iterable[Equation]: + r"""Lists Equations by specified criteria. + + Parameters + ---------- + name : str + The name of an Equation. If supplied only one result will be returned. + # TODO: Update kwargs + \*\*kwargs: any + More filter Equations as specified in + `ixmp4.data.db.optimization.equation.filter.OptimizationEquationFilter`. + + Returns + ------- + Iterable[:class:`ixmp4.data.abstract.optimization.Equation`]: + List of Equations. + """ + ... + + def tabulate(self, *, name: str | None = None, **kwargs) -> pd.DataFrame: + r"""Tabulate Equations by specified criteria. + + Parameters + ---------- + name : str + The name of an Equation. If supplied only one result will be returned. + # TODO: Update kwargs + \*\*kwargs: any + More filter variables as specified in + `ixmp4.data.db.optimization.equation.filter.OptimizationEquationFilter`. + + Returns + ------- + :class:`pandas.DataFrame`: + A data frame with the columns: + - id + - name + - data + - run__id + - created_at + - created_by + """ + ... + + # TODO Question for Daniel: do equations need to allow adding data manually? + # TODO Once present, state how to check which IndexSets are linked and which values + # they permit + def add_data(self, equation_id: int, data: dict[str, Any] | pd.DataFrame) -> None: + r"""Adds data to an Equation. + + The data will be validated with the linked constrained + :class:`ixmp4.data.abstract.optimization.IndexSet`s. For that, `data.keys()` + must correspond to the names of the Equation's columns. Each column can only + contain values that are in the linked `IndexSet.elements`. Each row of entries + must be unique. No values can be missing, `None`, or `NaN`. If `data.keys()` + contains names already present in `Equation.data`, existing values will be + overwritten. + + Parameters + ---------- + equation_id : int + The id of the :class:`ixmp4.data.abstract.optimization.Equation`. + data : dict[str, Any] | pandas.DataFrame + The data to be added. + + Raises + ------ + ValueError: + - If values are missing, `None`, or `NaN` + - If values are not allowed based on constraints to `Indexset`s + - If rows are not unique + + Returns + ------- + None + """ + ... + + def remove_data(self, equation_id: int) -> None: + """Removes data from an Equation. + + Parameters + ---------- + equation_id : int + The id of the :class:`ixmp4.data.abstract.optimization.Equation`. + + Returns + ------- + None + """ + ... diff --git a/ixmp4/data/api/__init__.py b/ixmp4/data/api/__init__.py index ae15c4d3..f18bdabd 100644 --- a/ixmp4/data/api/__init__.py +++ b/ixmp4/data/api/__init__.py @@ -11,6 +11,8 @@ from .meta import RunMetaEntry, RunMetaEntryRepository from .model import Model, ModelRepository from .optimization import ( + Equation, + EquationRepository, IndexSet, IndexSetRepository, Parameter, diff --git a/ixmp4/data/api/optimization/__init__.py b/ixmp4/data/api/optimization/__init__.py index 8e59f24c..e4a3747f 100644 --- a/ixmp4/data/api/optimization/__init__.py +++ b/ixmp4/data/api/optimization/__init__.py @@ -1,3 +1,4 @@ +from .equation import Equation, EquationRepository from .indexset import IndexSet, IndexSetRepository from .parameter import Parameter, ParameterRepository from .scalar import Scalar, ScalarRepository diff --git a/ixmp4/data/api/optimization/column.py b/ixmp4/data/api/optimization/column.py index 7da94fe2..d2076999 100644 --- a/ixmp4/data/api/optimization/column.py +++ b/ixmp4/data/api/optimization/column.py @@ -14,8 +14,9 @@ class Column(base.BaseModel): id: int name: str dtype: str - table__id: int | None + equation__id: int | None parameter__id: int | None + table__id: int | None variable__id: int | None indexset: IndexSet constrained_to_indexset: int diff --git a/ixmp4/data/api/optimization/equation.py b/ixmp4/data/api/optimization/equation.py new file mode 100644 index 00000000..f25aefa4 --- /dev/null +++ b/ixmp4/data/api/optimization/equation.py @@ -0,0 +1,86 @@ +from datetime import datetime +from typing import Any, ClassVar, Iterable + +import pandas as pd + +from ixmp4.data import abstract + +from .. import base +from ..docs import Docs, DocsRepository +from .column import Column + + +class Equation(base.BaseModel): + NotFound: ClassVar = abstract.Equation.NotFound + NotUnique: ClassVar = abstract.Equation.NotUnique + DeletionPrevented: ClassVar = abstract.Equation.DeletionPrevented + + id: int + name: str + data: dict[str, Any] + columns: list["Column"] + run__id: int + + created_at: datetime | None + created_by: str | None + + +class EquationDocsRepository(DocsRepository): + model_class = Docs + prefix = "docs/optimization/equations/" + + +class EquationRepository( + base.Creator[Equation], + base.Retriever[Equation], + base.Enumerator[Equation], + abstract.EquationRepository, +): + model_class = Equation + prefix = "optimization/equations/" + + def __init__(self, backend, *args, **kwargs) -> None: + super().__init__(backend, *args, **kwargs) + self.docs = EquationDocsRepository(backend) + + def create( + self, + run_id: int, + name: str, + constrained_to_indexsets: list[str], + column_names: list[str] | None = None, + ) -> Equation: + return super().create( + name=name, + run_id=run_id, + constrained_to_indexsets=constrained_to_indexsets, + column_names=column_names, + ) + + def add_data(self, equation_id: int, data: dict[str, Any] | pd.DataFrame) -> None: + if isinstance(data, pd.DataFrame): + # data will always contains str, not only Hashable + data: dict[str, Any] = data.to_dict(orient="list") # type: ignore + kwargs = {"data": data} + self._request( + method="PATCH", path=self.prefix + str(equation_id) + "/data/", json=kwargs + ) + + def remove_data(self, equation_id: int) -> None: + self._request(method="DELETE", path=self.prefix + str(equation_id) + "/data/") + + def get(self, run_id: int, name: str) -> Equation: + return super().get(run_id=run_id, name=name) + + def get_by_id(self, id: int) -> Equation: + res = self._get_by_id(id) + return Equation(**res) + + def list(self, *args, **kwargs) -> Iterable[Equation]: + return super().list(*args, **kwargs) + + def tabulate(self, *args, **kwargs) -> pd.DataFrame: + return super().tabulate(*args, **kwargs) + + def enumerate(self, *args, **kwargs) -> Iterable[Equation] | pd.DataFrame: + return super().enumerate(*args, **kwargs) diff --git a/ixmp4/data/backend/api.py b/ixmp4/data/backend/api.py index 096ece7b..beda8741 100644 --- a/ixmp4/data/backend/api.py +++ b/ixmp4/data/backend/api.py @@ -11,6 +11,7 @@ from ixmp4.core.exceptions import ImproperlyConfigured, UnknownApiError from ixmp4.data.api import ( DataPointRepository, + EquationRepository, IndexSetRepository, ModelRepository, OptimizationVariableRepository, @@ -114,6 +115,7 @@ def create_repositories(self): self.iamc.variables = VariableRepository(self) self.meta = RunMetaEntryRepository(self) self.models = ModelRepository(self) + self.optimization.equations = EquationRepository(self) self.optimization.indexsets = IndexSetRepository(self) self.optimization.parameters = ParameterRepository(self) self.optimization.scalars = ScalarRepository(self) diff --git a/ixmp4/data/backend/base.py b/ixmp4/data/backend/base.py index 2648e415..ae3fb46d 100644 --- a/ixmp4/data/backend/base.py +++ b/ixmp4/data/backend/base.py @@ -1,6 +1,7 @@ from ixmp4.conf.base import PlatformInfo from ixmp4.data.abstract import ( DataPointRepository, + EquationRepository, IndexSetRepository, ModelRepository, OptimizationVariableRepository, @@ -24,6 +25,7 @@ class IamcSubobject(object): class OptimizationSubobject(object): + equations: EquationRepository indexsets: IndexSetRepository parameters: ParameterRepository scalars: ScalarRepository diff --git a/ixmp4/data/backend/db.py b/ixmp4/data/backend/db.py index f0e629fc..67cecf40 100644 --- a/ixmp4/data/backend/db.py +++ b/ixmp4/data/backend/db.py @@ -14,6 +14,7 @@ from ixmp4.data.db import ( BaseModel, DataPointRepository, + EquationRepository, IndexSetRepository, ModelRepository, OptimizationVariableRepository, @@ -52,6 +53,7 @@ class IamcSubobject(BaseIamcSubobject): class OptimizationSubobject(BaseOptimizationSubobject): + equations: EquationRepository indexsets: IndexSetRepository parameters: ParameterRepository scalars: ScalarRepository @@ -99,6 +101,7 @@ def make_repositories(self): self.iamc.variables = VariableRepository(self) self.meta = RunMetaEntryRepository(self) self.models = ModelRepository(self) + self.optimization.equations = EquationRepository(self) self.optimization.indexsets = IndexSetRepository(self) self.optimization.parameters = ParameterRepository(self) self.optimization.scalars = ScalarRepository(self) diff --git a/ixmp4/data/db/__init__.py b/ixmp4/data/db/__init__.py index 3fee4f4e..c470395d 100644 --- a/ixmp4/data/db/__init__.py +++ b/ixmp4/data/db/__init__.py @@ -20,6 +20,8 @@ from .optimization import ( Column, ColumnRepository, + Equation, + EquationRepository, IndexSet, IndexSetRepository, Parameter, diff --git a/ixmp4/data/db/filters/__init__.py b/ixmp4/data/db/filters/__init__.py index b1727549..a337785e 100644 --- a/ixmp4/data/db/filters/__init__.py +++ b/ixmp4/data/db/filters/__init__.py @@ -1,6 +1,7 @@ from .meta import RunMetaEntryFilter from .model import ModelFilter from .optimizationcolumn import OptimizationColumnFilter +from .optimizationequation import OptimizationEquationFilter from .optimizationindexset import OptimizationIndexSetFilter from .optimizationparameter import OptimizationParameterFilter from .optimizationscalar import OptimizationScalarFilter diff --git a/ixmp4/data/db/filters/optimizationequation.py b/ixmp4/data/db/filters/optimizationequation.py new file mode 100644 index 00000000..4073247a --- /dev/null +++ b/ixmp4/data/db/filters/optimizationequation.py @@ -0,0 +1,17 @@ +from typing import ClassVar + +from ixmp4.db import filters + +from .. import Equation, Run + + +class OptimizationEquationFilter(filters.BaseFilter, metaclass=filters.FilterMeta): + id: filters.Id + name: filters.String + run__id: filters.Integer = filters.Field(None, alias="run_id") + + sqla_model: ClassVar[type] = Equation + + def join(self, exc, **kwargs): + exc = exc.join(Run, onclause=Equation.run__id == Run.id) + return exc diff --git a/ixmp4/data/db/optimization/__init__.py b/ixmp4/data/db/optimization/__init__.py index f1912256..f2988cb7 100644 --- a/ixmp4/data/db/optimization/__init__.py +++ b/ixmp4/data/db/optimization/__init__.py @@ -1,4 +1,5 @@ from .column import Column, ColumnRepository +from .equation import Equation, EquationRepository from .indexset import IndexSet, IndexSetRepository from .parameter import Parameter, ParameterRepository from .scalar import Scalar, ScalarRepository diff --git a/ixmp4/data/db/optimization/column/model.py b/ixmp4/data/db/optimization/column/model.py index 0a9bab95..74eda945 100644 --- a/ixmp4/data/db/optimization/column/model.py +++ b/ixmp4/data/db/optimization/column/model.py @@ -18,16 +18,20 @@ class Column(base.BaseModel): db.String(255), nullable=False, unique=False ) # pandas dtype - table__id: types.Mapped[int | None] = db.Column( - db.Integer, db.ForeignKey("optimization_table.id"), nullable=True + equation__id: types.Mapped[int | None] = db.Column( + db.Integer, db.ForeignKey("optimization_equation.id"), nullable=True ) parameter__id: types.Mapped[int | None] = db.Column( db.Integer, db.ForeignKey("optimization_parameter.id"), nullable=True ) + table__id: types.Mapped[int | None] = db.Column( + db.Integer, db.ForeignKey("optimization_table.id"), nullable=True + ) # TODO ... variable__id: types.Mapped[int | None] = db.Column( db.Integer, db.ForeignKey("optimization_optimizationvariable.id"), nullable=True ) + indexset: types.Mapped[IndexSet] = db.relationship(single_parent=True) constrained_to_indexset: types.Integer = db.Column( db.Integer, db.ForeignKey("optimization_indexset.id"), index=True diff --git a/ixmp4/data/db/optimization/column/repository.py b/ixmp4/data/db/optimization/column/repository.py index 94ce552b..eb1d2f72 100644 --- a/ixmp4/data/db/optimization/column/repository.py +++ b/ixmp4/data/db/optimization/column/repository.py @@ -24,6 +24,7 @@ def add( name: str, constrained_to_indexset: str, dtype: str, + equation_id: int, parameter_id: int, table_id: int, variable_id: int, @@ -33,6 +34,7 @@ def add( name=name, constrained_to_indexset=constrained_to_indexset, dtype=dtype, + equation__id=equation_id, parameter__id=parameter_id, table__id=table_id, variable__id=variable_id, @@ -47,6 +49,7 @@ def create( name: str, constrained_to_indexset: int, dtype: str, + equation_id: int | None = None, parameter_id: int | None = None, table_id: int | None = None, variable_id: int | None = None, @@ -64,6 +67,9 @@ def create( contain all values used as entries in this Column. dtype : str The pandas-inferred type of the Column's data. + equation_id : int + The unique integer id of the + :class:`ixmp4.data.abstract.optimization.Equation` this Column belongs to. parameter_id : int | None, default None The unique integer id of the :class:`ixmp4.data.abstract.optimization.Parameter` this Column belongs to, @@ -94,6 +100,7 @@ def create( name=name, constrained_to_indexset=constrained_to_indexset, dtype=dtype, + equation_id=equation_id, parameter_id=parameter_id, table_id=table_id, variable_id=variable_id, diff --git a/ixmp4/data/db/optimization/equation/__init__.py b/ixmp4/data/db/optimization/equation/__init__.py new file mode 100644 index 00000000..5af80352 --- /dev/null +++ b/ixmp4/data/db/optimization/equation/__init__.py @@ -0,0 +1,2 @@ +from .model import Equation +from .repository import EquationRepository diff --git a/ixmp4/data/db/optimization/equation/docs.py b/ixmp4/data/db/optimization/equation/docs.py new file mode 100644 index 00000000..b3403cb6 --- /dev/null +++ b/ixmp4/data/db/optimization/equation/docs.py @@ -0,0 +1,8 @@ +from ixmp4.data.db.docs import BaseDocsRepository, docs_model + +from .model import Equation + + +class EquationDocsRepository(BaseDocsRepository): + model_class = docs_model(Equation) # EquationDocs + dimension_model_class = Equation diff --git a/ixmp4/data/db/optimization/equation/filter.py b/ixmp4/data/db/optimization/equation/filter.py new file mode 100644 index 00000000..bd0a0686 --- /dev/null +++ b/ixmp4/data/db/optimization/equation/filter.py @@ -0,0 +1,17 @@ +from ixmp4.data.db import filters as base +from ixmp4.data.db.run import Run +from ixmp4.db import filters, utils + +from .model import Equation + + +class RunFilter(base.RunFilter, metaclass=filters.FilterMeta): + def join(self, exc, **kwargs): + if not utils.is_joined(exc, Run): + exc = exc.join(Run, onclause=Equation.run__id == Run.id) + return exc + + +class EquationFilter(base.OptimizationEquationFilter, metaclass=filters.FilterMeta): + def join(self, exc, session=None): + return exc diff --git a/ixmp4/data/db/optimization/equation/model.py b/ixmp4/data/db/optimization/equation/model.py new file mode 100644 index 00000000..2201d2f6 --- /dev/null +++ b/ixmp4/data/db/optimization/equation/model.py @@ -0,0 +1,41 @@ +import copy +from typing import Any, ClassVar + +from sqlalchemy.orm import validates + +from ixmp4 import db +from ixmp4.core.exceptions import OptimizationDataValidationError +from ixmp4.data import types +from ixmp4.data.abstract import optimization as abstract + +from .. import Column, base, utils + + +class Equation(base.BaseModel): + # NOTE: These might be mixin-able, but would require some abstraction + NotFound: ClassVar = abstract.Equation.NotFound + NotUnique: ClassVar = abstract.Equation.NotUnique + DataInvalid: ClassVar = OptimizationDataValidationError + DeletionPrevented: ClassVar = abstract.Equation.DeletionPrevented + + # constrained_to_indexsets: ClassVar[list[str] | None] = None + + run__id: types.RunId + columns: types.Mapped[list["Column"]] = db.relationship() + data: types.JsonDict = db.Column(db.JsonType, nullable=False, default={}) + + @validates("data") + def validate_data(self, key, data: dict[str, Any]): + if data == {}: + return data + data_to_validate = copy.deepcopy(data) + del data_to_validate["levels"] + del data_to_validate["marginals"] + _ = utils.validate_data( + host=self, + data=data_to_validate, + columns=self.columns, + ) + return data + + __table_args__ = (db.UniqueConstraint("name", "run__id"),) diff --git a/ixmp4/data/db/optimization/equation/repository.py b/ixmp4/data/db/optimization/equation/repository.py new file mode 100644 index 00000000..efe12fef --- /dev/null +++ b/ixmp4/data/db/optimization/equation/repository.py @@ -0,0 +1,186 @@ +from typing import Any, Iterable + +import pandas as pd + +from ixmp4 import db +from ixmp4.core.exceptions import OptimizationItemUsageError +from ixmp4.data.abstract import optimization as abstract +from ixmp4.data.auth.decorators import guard + +from .. import ColumnRepository, base +from .docs import EquationDocsRepository +from .model import Equation + + +class EquationRepository( + base.Creator[Equation], + base.Retriever[Equation], + base.Enumerator[Equation], + abstract.EquationRepository, +): + model_class = Equation + + UsageError = OptimizationItemUsageError + + def __init__(self, *args, **kwargs) -> None: + super().__init__(*args, **kwargs) + self.docs = EquationDocsRepository(*args, **kwargs) + self.columns = ColumnRepository(*args, **kwargs) + + from .filter import EquationFilter + + self.filter_class = EquationFilter + + def _add_column( + self, + run_id: int, + equation_id: int, + column_name: str, + indexset_name: str, + **kwargs, + ) -> None: + r"""Adds a Column to an Equation. + + Parameters + ---------- + run_id : int + The id of the :class:`ixmp4.data.abstract.Run` for which the + :class:`ixmp4.data.abstract.optimization.Equation` is defined. + equation_id : int + The id of the :class:`ixmp4.data.abstract.optimization.Equation`. + column_name : str + The name of the Column, which must be unique in connection with the names of + :class:`ixmp4.data.abstract.Run` and + :class:`ixmp4.data.abstract.optimization.Equation`. + indexset_name : str + The name of the :class:`ixmp4.data.abstract.optimization.IndexSet` the + Column will be linked to. + \*\*kwargs: any + Keyword arguments to be passed to + :func:`ixmp4.data.abstract.optimization.Column.create`. + """ + indexset = self.backend.optimization.indexsets.get( + run_id=run_id, name=indexset_name + ) + self.columns.create( + name=column_name, + constrained_to_indexset=indexset.id, + dtype=pd.Series(indexset.elements).dtype.name, + equation_id=equation_id, + unique=True, + **kwargs, + ) + + def add( + self, + run_id: int, + name: str, + ) -> Equation: + equation = Equation(name=name, run__id=run_id) + equation.set_creation_info(auth_context=self.backend.auth_context) + self.session.add(equation) + + return equation + + @guard("view") + def get(self, run_id: int, name: str) -> Equation: + exc = db.select(Equation).where( + (Equation.name == name) & (Equation.run__id == run_id) + ) + try: + return self.session.execute(exc).scalar_one() + except db.NoResultFound: + raise Equation.NotFound + + @guard("view") + def get_by_id(self, id: int) -> Equation: + obj = self.session.get(self.model_class, id) + + if obj is None: + raise Equation.NotFound(id=id) + + return obj + + @guard("edit") + def create( + self, + run_id: int, + name: str, + constrained_to_indexsets: list[str], + column_names: list[str] | None = None, + **kwargs, + ) -> Equation: + # Convert to list to avoid enumerate() splitting strings to letters + if isinstance(constrained_to_indexsets, str): + constrained_to_indexsets = list(constrained_to_indexsets) + if column_names and len(column_names) != len(constrained_to_indexsets): + raise OptimizationItemUsageError( + f"While processing Equation {name}: \n" + "`constrained_to_indexsets` and `column_names` not equal in length! " + "Please provide the same number of entries for both!" + ) + # TODO: activate something like this if each column must be indexed by a unique + # indexset + # if len(constrained_to_indexsets) != len(set(constrained_to_indexsets)): + # raise ValueError("Each dimension must be constrained to a unique indexset!") # noqa + if column_names and len(column_names) != len(set(column_names)): + raise OptimizationItemUsageError( + f"While processing Equation {name}: \n" + "The given `column_names` are not unique!" + ) + + equation = super().create( + run_id=run_id, + name=name, + **kwargs, + ) + for i, name in enumerate(constrained_to_indexsets): + self._add_column( + run_id=run_id, + equation_id=equation.id, + column_name=column_names[i] if column_names else name, + indexset_name=name, + ) + + return equation + + @guard("view") + def list(self, *args, **kwargs) -> Iterable[Equation]: + return super().list(*args, **kwargs) + + @guard("view") + def tabulate(self, *args, **kwargs) -> pd.DataFrame: + return super().tabulate(*args, **kwargs) + + @guard("edit") + def add_data(self, equation_id: int, data: dict[str, Any] | pd.DataFrame) -> None: + if isinstance(data, dict): + try: + data = pd.DataFrame.from_dict(data=data) + except ValueError as e: + raise Equation.DataInvalid(str(e)) from e + equation = self.get_by_id(id=equation_id) + + missing_columns = set(["levels", "marginals"]) - set(data.columns) + if missing_columns: + raise OptimizationItemUsageError( + f"Equation.data must include the column(s): " + f"{', '.join(missing_columns)}!" + ) + + index_list = [column.name for column in equation.columns] + existing_data = pd.DataFrame(equation.data) + if not existing_data.empty: + existing_data.set_index(index_list, inplace=True) + equation.data = ( + data.set_index(index_list).combine_first(existing_data).reset_index() + ).to_dict(orient="list") + + self.session.commit() + + @guard("edit") + def remove_data(self, equation_id: int) -> None: + equation = self.get_by_id(id=equation_id) + # TODO Is there a better way to reset .data? + equation.data = {} + self.session.commit() diff --git a/ixmp4/data/db/run/model.py b/ixmp4/data/db/run/model.py index 2a524b06..79dec265 100644 --- a/ixmp4/data/db/run/model.py +++ b/ixmp4/data/db/run/model.py @@ -3,6 +3,7 @@ from ixmp4 import db from ixmp4.data import abstract, types from ixmp4.data.db.model.model import Model +from ixmp4.data.db.optimization.equation import Equation from ixmp4.data.db.optimization.indexset import IndexSet from ixmp4.data.db.optimization.parameter import Parameter from ixmp4.data.db.optimization.scalar import Scalar @@ -42,6 +43,7 @@ class Run(base.BaseModel, mixins.HasUpdateInfo): foreign_keys=[scenario__id], ) + equations: types.Mapped[list["Equation"]] = db.relationship() indexsets: types.Mapped[list["IndexSet"]] = db.relationship() parameters: types.Mapped[list["Parameter"]] = db.relationship() scalars: types.Mapped[list["Scalar"]] = db.relationship() diff --git a/ixmp4/server/rest/__init__.py b/ixmp4/server/rest/__init__.py index bf482219..cb49e0cb 100644 --- a/ixmp4/server/rest/__init__.py +++ b/ixmp4/server/rest/__init__.py @@ -17,6 +17,7 @@ from .iamc import unit as iamc_unit from .iamc import variable as iamc_variable from .middleware import RequestSizeLoggerMiddleware, RequestTimeLoggerMiddleware +from .optimization import equation as optimization_equation from .optimization import indexset as optimization_indexset from .optimization import parameter as optimization_parameter from .optimization import scalar as optimization_scalar @@ -51,6 +52,7 @@ v1.include_router(iamc_variable.router, prefix="/iamc") v1.include_router(meta.router) v1.include_router(model.router) +v1.include_router(optimization_equation.router, prefix="/optimization") v1.include_router(optimization_indexset.router, prefix="/optimization") v1.include_router(optimization_parameter.router, prefix="/optimization") v1.include_router(optimization_scalar.router, prefix="/optimization") diff --git a/ixmp4/server/rest/docs.py b/ixmp4/server/rest/docs.py index d2a5a6b7..5423bef7 100644 --- a/ixmp4/server/rest/docs.py +++ b/ixmp4/server/rest/docs.py @@ -285,3 +285,27 @@ def delete_optimization_variables( backend: Backend = Depends(deps.get_backend), ): return backend.optimization.variables.docs.delete(dimension_id) + + +@router.get("/optimization/equations/", response_model=list[api.Docs]) +def list_equations( + dimension_id: int | None = Query(None), + backend: Backend = Depends(deps.get_backend), +): + return backend.optimization.equations.docs.list(dimension_id=dimension_id) + + +@router.post("/optimization/equations/", response_model=api.Docs) +def set_equations( + docs: DocsInput, + backend: Backend = Depends(deps.get_backend), +): + return backend.optimization.equations.docs.set(**docs.model_dump()) + + +@router.delete("/optimization/equations/{dimension_id}/") +def delete_equations( + dimension_id: int = Path(), + backend: Backend = Depends(deps.get_backend), +): + return backend.optimization.equations.docs.delete(dimension_id) diff --git a/ixmp4/server/rest/optimization/__init__.py b/ixmp4/server/rest/optimization/__init__.py index fe2cd207..10536e4c 100644 --- a/ixmp4/server/rest/optimization/__init__.py +++ b/ixmp4/server/rest/optimization/__init__.py @@ -1 +1 @@ -from . import indexset, parameter, scalar, table, variable +from . import equation, indexset, parameter, scalar, table, variable diff --git a/ixmp4/server/rest/optimization/equation.py b/ixmp4/server/rest/optimization/equation.py new file mode 100644 index 00000000..79dbdb88 --- /dev/null +++ b/ixmp4/server/rest/optimization/equation.py @@ -0,0 +1,86 @@ +from typing import Any + +from fastapi import APIRouter, Body, Depends, Query + +from ixmp4.data import api +from ixmp4.data.backend.db import SqlAlchemyBackend as Backend +from ixmp4.data.db.optimization.equation.filter import EquationFilter + +from .. import deps +from ..base import BaseModel, EnumerationOutput, Pagination +from ..decorators import autodoc + +router: APIRouter = APIRouter( + prefix="/equations", + tags=["optimization", "equations"], +) + + +class EquationCreateInput(BaseModel): + run_id: int + name: str + constrained_to_indexsets: list[str] + column_names: list[str] | None + + +class DataInput(BaseModel): + data: dict[str, Any] + + +@autodoc +@router.get("/{id}/", response_model=api.Equation) +def get_by_id( + id: int, + backend: Backend = Depends(deps.get_backend), +): + return backend.optimization.equations.get_by_id(id) + + +@autodoc +@router.patch("/", response_model=EnumerationOutput[api.Equation]) +def query( + filter: EquationFilter = Body(EquationFilter(id=None, name=None)), + table: bool = Query(False), + pagination: Pagination = Depends(), + backend: Backend = Depends(deps.get_backend), +): + return EnumerationOutput( + results=backend.optimization.equations.paginate( + _filter=filter, + limit=pagination.limit, + offset=pagination.offset, + table=bool(table), + ), + total=backend.optimization.equations.count(_filter=filter), + pagination=pagination, + ) + + +@autodoc +@router.patch("/{equation_id}/data/") +def add_data( + equation_id: int, + data: DataInput, + backend: Backend = Depends(deps.get_backend), +): + return backend.optimization.equations.add_data( + equation_id=equation_id, **data.model_dump() + ) + + +@autodoc +@router.delete("/{equation_id}/data/") +def remove_data( + equation_id: int, + backend: Backend = Depends(deps.get_backend), +): + backend.optimization.equations.remove_data(equation_id == equation_id) + + +@autodoc +@router.post("/", response_model=api.Equation) +def create( + equation: EquationCreateInput, + backend: Backend = Depends(deps.get_backend), +): + return backend.optimization.equations.create(**equation.model_dump()) diff --git a/tests/core/test_optimization_equation.py b/tests/core/test_optimization_equation.py new file mode 100644 index 00000000..56945d97 --- /dev/null +++ b/tests/core/test_optimization_equation.py @@ -0,0 +1,367 @@ +import pandas as pd +import pytest + +import ixmp4 +from ixmp4.core import Equation, IndexSet +from ixmp4.core.exceptions import ( + OptimizationDataValidationError, + OptimizationItemUsageError, +) + +from ..utils import assert_unordered_equality, create_indexsets_for_run + + +def df_from_list(equations: list): + return pd.DataFrame( + [ + [ + equation.run_id, + equation.data, + equation.name, + equation.id, + equation.created_at, + equation.created_by, + ] + for equation in equations + ], + columns=[ + "run__id", + "data", + "name", + "id", + "created_at", + "created_by", + ], + ) + + +class TestCoreEquation: + def test_create_equation(self, platform: ixmp4.Platform): + run = platform.runs.create("Model", "Scenario") + + # Test normal creation + indexset, indexset_2 = tuple( + IndexSet(_backend=platform.backend, _model=model) + for model in create_indexsets_for_run(platform=platform, run_id=run.id) + ) + equation = run.optimization.equations.create( + name="Equation", + constrained_to_indexsets=[indexset.name], + ) + + assert equation.run_id == run.id + assert equation.name == "Equation" + assert equation.data == {} # JsonDict type currently requires a dict, not None + assert equation.columns[0].name == indexset.name + assert equation.constrained_to_indexsets == [indexset.name] + assert equation.levels == [] + assert equation.marginals == [] + + # Test duplicate name raises + with pytest.raises(Equation.NotUnique): + _ = run.optimization.equations.create( + "Equation", constrained_to_indexsets=[indexset.name] + ) + + # Test mismatch in constrained_to_indexsets and column_names raises + with pytest.raises(OptimizationItemUsageError, match="not equal in length"): + _ = run.optimization.equations.create( + "Equation 2", + constrained_to_indexsets=[indexset.name], + column_names=["Dimension 1", "Dimension 2"], + ) + + # Test columns_names are used for names if given + equation_2 = run.optimization.equations.create( + "Equation 2", + constrained_to_indexsets=[indexset.name], + column_names=["Column 1"], + ) + assert equation_2.columns[0].name == "Column 1" + + # Test duplicate column_names raise + with pytest.raises( + OptimizationItemUsageError, match="`column_names` are not unique" + ): + _ = run.optimization.equations.create( + name="Equation 3", + constrained_to_indexsets=[indexset.name, indexset.name], + column_names=["Column 1", "Column 1"], + ) + + # Test column.dtype is registered correctly + indexset_2.add(elements=2024) + equation_3 = run.optimization.equations.create( + "Equation 5", + constrained_to_indexsets=[indexset.name, indexset_2.name], + ) + # If indexset doesn't have elements, a generic dtype is registered + assert equation_3.columns[0].dtype == "object" + assert equation_3.columns[1].dtype == "int64" + + def test_get_equation(self, platform: ixmp4.Platform): + run = platform.runs.create("Model", "Scenario") + (indexset,) = create_indexsets_for_run( + platform=platform, run_id=run.id, amount=1 + ) + _ = run.optimization.equations.create( + name="Equation", constrained_to_indexsets=[indexset.name] + ) + equation = run.optimization.equations.get(name="Equation") + assert equation.run_id == run.id + assert equation.id == 1 + assert equation.name == "Equation" + assert equation.data == {} + assert equation.levels == [] + assert equation.marginals == [] + assert equation.columns[0].name == indexset.name + assert equation.constrained_to_indexsets == [indexset.name] + + with pytest.raises(Equation.NotFound): + _ = run.optimization.equations.get("Equation 2") + + def test_equation_add_data(self, platform: ixmp4.Platform): + run = platform.runs.create("Model", "Scenario") + indexset, indexset_2 = tuple( + IndexSet(_backend=platform.backend, _model=model) + for model in create_indexsets_for_run(platform=platform, run_id=run.id) + ) + indexset.add(elements=["foo", "bar", ""]) + indexset_2.add(elements=[1, 2, 3]) + # pandas can only convert dicts to dataframes if the values are lists + # or if index is given. But maybe using read_json instead of from_dict + # can remedy this. Or maybe we want to catch the resulting + # "ValueError: If using all scalar values, you must pass an index" and + # reraise a custom informative error? + test_data_1 = { + indexset.name: ["foo"], + indexset_2.name: [1], + "levels": [3.14], + "marginals": [0.000314], + } + equation = run.optimization.equations.create( + "Equation", + constrained_to_indexsets=[indexset.name, indexset_2.name], + ) + equation.add(data=test_data_1) + assert equation.data == test_data_1 + assert equation.levels == test_data_1["levels"] + assert equation.marginals == test_data_1["marginals"] + + equation_2 = run.optimization.equations.create( + name="Equation 2", + constrained_to_indexsets=[indexset.name, indexset_2.name], + ) + + with pytest.raises( + OptimizationItemUsageError, + match=r"must include the column\(s\): marginals!", + ): + equation_2.add( + pd.DataFrame( + { + indexset.name: [None], + indexset_2.name: [2], + "levels": [1], + } + ), + ) + + with pytest.raises( + OptimizationItemUsageError, match=r"must include the column\(s\): levels!" + ): + equation_2.add( + data=pd.DataFrame( + { + indexset.name: [None], + indexset_2.name: [2], + "marginals": [0], + } + ), + ) + + # By converting data to pd.DataFrame, we automatically enforce equal length + # of new columns, raises All arrays must be of the same length otherwise: + with pytest.raises( + OptimizationDataValidationError, + match="All arrays must be of the same length", + ): + equation_2.add( + data={ + indexset.name: ["foo", "foo"], + indexset_2.name: [2, 2], + "levels": [1, 2], + "marginals": [3], + }, + ) + + with pytest.raises( + OptimizationDataValidationError, match="contains duplicate rows" + ): + equation_2.add( + data={ + indexset.name: ["foo", "foo"], + indexset_2.name: [2, 2], + "levels": [1, 2], + "marginals": [3.4, 5.6], + }, + ) + + # Test that order is conserved + test_data_2 = { + indexset.name: ["", "", "foo", "foo", "bar", "bar"], + indexset_2.name: [3, 1, 2, 1, 2, 3], + "levels": [6, 5, 4, 3, 2, 1], + "marginals": [1, 3, 5, 6, 4, 2], + } + equation_2.add(test_data_2) + assert equation_2.data == test_data_2 + assert equation_2.levels == test_data_2["levels"] + assert equation_2.marginals == test_data_2["marginals"] + + # Test updating of existing keys + equation_4 = run.optimization.equations.create( + name="Equation 4", + constrained_to_indexsets=[indexset.name, indexset_2.name], + ) + test_data_6 = { + indexset.name: ["foo", "foo", "bar", "bar"], + indexset_2.name: [1, 3, 1, 2], + "levels": [0.00001, "2", 2.3, 400000], + "marginals": [6, 7.8, 9, 0], + } + equation_4.add(data=test_data_6) + test_data_7 = { + indexset.name: ["foo", "foo", "bar", "bar", "bar"], + indexset_2.name: [1, 2, 3, 2, 1], + "levels": [0.00001, 2.3, 3, "400000", "5"], + "marginals": [6, 7.8, 9, "0", 3], + } + equation_4.add(data=test_data_7) + expected = ( + pd.DataFrame(test_data_7) + .set_index([indexset.name, indexset_2.name]) + .combine_first( + pd.DataFrame(test_data_6).set_index([indexset.name, indexset_2.name]) + ) + .reset_index() + ) + assert_unordered_equality(expected, pd.DataFrame(equation_4.data)) + + def test_equation_remove_data(self, platform: ixmp4.Platform): + run = platform.runs.create("Model", "Scenario") + indexset = run.optimization.indexsets.create("Indexset") + indexset.add(elements=["foo", "bar"]) + test_data = { + "Indexset": ["bar", "foo"], + "levels": [2.0, 1], + "marginals": [0, "test"], + } + equation = run.optimization.equations.create( + "Equation", + constrained_to_indexsets=[indexset.name], + ) + equation.add(test_data) + assert equation.data == test_data + + equation.remove_data() + assert equation.data == {} + + def test_list_equation(self, platform: ixmp4.Platform): + run = platform.runs.create("Model", "Scenario") + # Per default, list() lists scalars for `default` version runs: + run.set_as_default() + indexset, indexset_2 = create_indexsets_for_run( + platform=platform, run_id=run.id + ) + equation = run.optimization.equations.create( + "Equation", constrained_to_indexsets=[indexset.name] + ) + equation_2 = run.optimization.equations.create( + "Equation 2", constrained_to_indexsets=[indexset_2.name] + ) + # Create new run to test listing equations of specific run + run_2 = platform.runs.create("Model", "Scenario") + (indexset,) = create_indexsets_for_run( + platform=platform, run_id=run_2.id, amount=1 + ) + run_2.optimization.equations.create( + "Equation", constrained_to_indexsets=[indexset.name] + ) + expected_ids = [equation.id, equation_2.id] + list_ids = [equation.id for equation in run.optimization.equations.list()] + assert not (set(expected_ids) ^ set(list_ids)) + + # Test retrieving just one result by providing a name + expected_id = [equation.id] + list_id = [ + equation.id for equation in run.optimization.equations.list(name="Equation") + ] + assert not (set(expected_id) ^ set(list_id)) + + def test_tabulate_equation(self, platform: ixmp4.Platform): + run = platform.runs.create("Model", "Scenario") + indexset, indexset_2 = tuple( + IndexSet(_backend=platform.backend, _model=model) + for model in create_indexsets_for_run(platform=platform, run_id=run.id) + ) + equation = run.optimization.equations.create( + name="Equation", + constrained_to_indexsets=[indexset.name, indexset_2.name], + ) + equation_2 = run.optimization.equations.create( + name="Equation 2", + constrained_to_indexsets=[indexset.name, indexset_2.name], + ) + # Create new run to test tabulating equations of specific run + run_2 = platform.runs.create("Model", "Scenario") + (indexset_3,) = create_indexsets_for_run( + platform=platform, run_id=run_2.id, amount=1 + ) + run_2.optimization.equations.create( + "Equation", constrained_to_indexsets=[indexset_3.name] + ) + pd.testing.assert_frame_equal( + df_from_list([equation_2]), + run.optimization.equations.tabulate(name="Equation 2"), + ) + + indexset.add(elements=["foo", "bar"]) + indexset_2.add(elements=[1, 2, 3]) + test_data_1 = { + indexset.name: ["foo"], + indexset_2.name: [1], + "levels": [314], + "marginals": [2.0], + } + equation.add(data=test_data_1) + + test_data_2 = { + indexset_2.name: [2, 3], + indexset.name: ["foo", "bar"], + "levels": [1, -2.0], + "marginals": [0, 10], + } + equation_2.add(data=test_data_2) + pd.testing.assert_frame_equal( + df_from_list([equation, equation_2]), + run.optimization.equations.tabulate(), + ) + + def test_equation_docs(self, platform: ixmp4.Platform): + run = platform.runs.create("Model", "Scenario") + (indexset,) = tuple( + IndexSet(_backend=platform.backend, _model=model) + for model in create_indexsets_for_run( + platform=platform, run_id=run.id, amount=1 + ) + ) + equation_1 = run.optimization.equations.create( + "Equation 1", constrained_to_indexsets=[indexset.name] + ) + docs = "Documentation of Equation 1" + equation_1.docs = docs + assert equation_1.docs == docs + + equation_1.docs = None + assert equation_1.docs is None diff --git a/tests/data/test_docs.py b/tests/data/test_docs.py index f1b6f7f4..ae529f9b 100644 --- a/tests/data/test_docs.py +++ b/tests/data/test_docs.py @@ -501,3 +501,70 @@ def test_delete_optimizationvariabledocs(self, platform: ixmp4.Platform): with pytest.raises(Docs.NotFound): platform.backend.optimization.variables.docs.get(variable.id) + + def test_get_and_set_equationdocs(self, platform: ixmp4.Platform): + run = platform.backend.runs.create("Model", "Scenario") + _ = platform.backend.optimization.indexsets.create( + run_id=run.id, name="Indexset" + ) + equation = platform.backend.optimization.equations.create( + run_id=run.id, name="Equation", constrained_to_indexsets=["Indexset"] + ) + docs_equation = platform.backend.optimization.equations.docs.set( + equation.id, "Description of test Equation" + ) + docs_equation1 = platform.backend.optimization.equations.docs.get(equation.id) + + assert docs_equation == docs_equation1 + + def test_change_empty_equationdocs(self, platform: ixmp4.Platform): + run = platform.backend.runs.create("Model", "Scenario") + _ = platform.backend.optimization.indexsets.create( + run_id=run.id, name="Indexset" + ) + equation = platform.backend.optimization.equations.create( + run_id=run.id, name="Equation", constrained_to_indexsets=["Indexset"] + ) + + with pytest.raises(Docs.NotFound): + platform.backend.optimization.equations.docs.get(equation.id) + + docs_equation1 = platform.backend.optimization.equations.docs.set( + equation.id, "Description of test Equation" + ) + + assert ( + platform.backend.optimization.equations.docs.get(equation.id) + == docs_equation1 + ) + + docs_equation2 = platform.backend.optimization.equations.docs.set( + equation.id, "Different description of test Equation" + ) + + assert ( + platform.backend.optimization.equations.docs.get(equation.id) + == docs_equation2 + ) + + def test_delete_optimizationequationdocs(self, platform: ixmp4.Platform): + run = platform.backend.runs.create("Model", "Scenario") + _ = platform.backend.optimization.indexsets.create( + run_id=run.id, name="Indexset" + ) + equation = platform.backend.optimization.equations.create( + run_id=run.id, name="Equation", constrained_to_indexsets=["Indexset"] + ) + docs_equation = platform.backend.optimization.equations.docs.set( + equation.id, "Description of test Equation" + ) + + assert ( + platform.backend.optimization.equations.docs.get(equation.id) + == docs_equation + ) + + platform.backend.optimization.equations.docs.delete(equation.id) + + with pytest.raises(Docs.NotFound): + platform.backend.optimization.equations.docs.get(equation.id) diff --git a/tests/data/test_optimization_equation.py b/tests/data/test_optimization_equation.py new file mode 100644 index 00000000..b82df865 --- /dev/null +++ b/tests/data/test_optimization_equation.py @@ -0,0 +1,416 @@ +import pandas as pd +import pytest + +import ixmp4 +from ixmp4.core.exceptions import ( + OptimizationDataValidationError, + OptimizationItemUsageError, +) +from ixmp4.data.abstract import Equation + +from ..utils import assert_unordered_equality, create_indexsets_for_run + + +def df_from_list(equations: list): + return pd.DataFrame( + [ + [ + equation.run__id, + equation.data, + equation.name, + equation.id, + equation.created_at, + equation.created_by, + ] + for equation in equations + ], + columns=[ + "run__id", + "data", + "name", + "id", + "created_at", + "created_by", + ], + ) + + +class TestDataOptimizationEquation: + def test_create_equation(self, platform: ixmp4.Platform): + run = platform.backend.runs.create("Model", "Scenario") + + # Test normal creation + indexset, indexset_2 = create_indexsets_for_run( + platform=platform, run_id=run.id + ) + equation = platform.backend.optimization.equations.create( + run_id=run.id, + name="Equation", + constrained_to_indexsets=[indexset.name], + ) + + assert equation.run__id == run.id + assert equation.name == "Equation" + assert equation.data == {} # JsonDict type currently requires a dict, not None + assert equation.columns[0].name == indexset.name + assert equation.columns[0].constrained_to_indexset == indexset.id + + # Test duplicate name raises + with pytest.raises(Equation.NotUnique): + _ = platform.backend.optimization.equations.create( + run_id=run.id, name="Equation", constrained_to_indexsets=[indexset.name] + ) + + # Test mismatch in constrained_to_indexsets and column_names raises + with pytest.raises(OptimizationItemUsageError, match="not equal in length"): + _ = platform.backend.optimization.equations.create( + run_id=run.id, + name="Equation 2", + constrained_to_indexsets=[indexset.name], + column_names=["Dimension 1", "Dimension 2"], + ) + + # Test columns_names are used for names if given + equation_2 = platform.backend.optimization.equations.create( + run_id=run.id, + name="Equation 2", + constrained_to_indexsets=[indexset.name], + column_names=["Column 1"], + ) + assert equation_2.columns[0].name == "Column 1" + + # Test duplicate column_names raise + with pytest.raises( + OptimizationItemUsageError, match="`column_names` are not unique" + ): + _ = platform.backend.optimization.equations.create( + run_id=run.id, + name="Equation 3", + constrained_to_indexsets=[indexset.name, indexset.name], + column_names=["Column 1", "Column 1"], + ) + + # Test column.dtype is registered correctly + platform.backend.optimization.indexsets.add_elements( + indexset_2.id, elements=2024 + ) + indexset_2 = platform.backend.optimization.indexsets.get( + run.id, indexset_2.name + ) + equation_3 = platform.backend.optimization.equations.create( + run_id=run.id, + name="Equation 5", + constrained_to_indexsets=[indexset.name, indexset_2.name], + ) + # If indexset doesn't have elements, a generic dtype is registered + assert equation_3.columns[0].dtype == "object" + assert equation_3.columns[1].dtype == "int64" + + def test_get_equation(self, platform: ixmp4.Platform): + run = platform.backend.runs.create("Model", "Scenario") + (indexset,) = create_indexsets_for_run( + platform=platform, run_id=run.id, amount=1 + ) + equation = platform.backend.optimization.equations.create( + run_id=run.id, name="Equation", constrained_to_indexsets=[indexset.name] + ) + assert equation == platform.backend.optimization.equations.get( + run_id=run.id, name="Equation" + ) + + with pytest.raises(Equation.NotFound): + _ = platform.backend.optimization.equations.get( + run_id=run.id, name="Equation 2" + ) + + def test_equation_add_data(self, platform: ixmp4.Platform): + run = platform.backend.runs.create("Model", "Scenario") + indexset, indexset_2 = create_indexsets_for_run( + platform=platform, run_id=run.id + ) + platform.backend.optimization.indexsets.add_elements( + indexset_id=indexset.id, elements=["foo", "bar", ""] + ) + platform.backend.optimization.indexsets.add_elements( + indexset_id=indexset_2.id, elements=[1, 2, 3] + ) + # pandas can only convert dicts to dataframes if the values are lists + # or if index is given. But maybe using read_json instead of from_dict + # can remedy this. Or maybe we want to catch the resulting + # "ValueError: If using all scalar values, you must pass an index" and + # reraise a custom informative error? + test_data_1 = { + indexset.name: ["foo"], + indexset_2.name: [1], + "levels": [3.14], + "marginals": [-3.14], + } + equation = platform.backend.optimization.equations.create( + run_id=run.id, + name="Equation", + constrained_to_indexsets=[indexset.name, indexset_2.name], + ) + platform.backend.optimization.equations.add_data( + equation_id=equation.id, data=test_data_1 + ) + + equation = platform.backend.optimization.equations.get( + run_id=run.id, name="Equation" + ) + assert equation.data == test_data_1 + + equation_2 = platform.backend.optimization.equations.create( + run_id=run.id, + name="Equation 2", + constrained_to_indexsets=[indexset.name, indexset_2.name], + ) + + with pytest.raises( + OptimizationItemUsageError, match=r"must include the column\(s\): levels!" + ): + platform.backend.optimization.equations.add_data( + equation_id=equation_2.id, + data=pd.DataFrame( + { + indexset.name: [None], + indexset_2.name: [2], + "marginals": [1], + } + ), + ) + + with pytest.raises( + OptimizationItemUsageError, + match=r"must include the column\(s\): marginals!", + ): + platform.backend.optimization.equations.add_data( + equation_id=equation_2.id, + data=pd.DataFrame( + { + indexset.name: [None], + indexset_2.name: [2], + "levels": [1], + } + ), + ) + + # By converting data to pd.DataFrame, we automatically enforce equal length + # of new columns, raises All arrays must be of the same length otherwise: + with pytest.raises( + OptimizationDataValidationError, + match="All arrays must be of the same length", + ): + platform.backend.optimization.equations.add_data( + equation_id=equation_2.id, + data={ + indexset.name: ["foo", "foo"], + indexset_2.name: [2, 2], + "levels": [1, 2], + "marginals": [1], + }, + ) + + with pytest.raises( + OptimizationDataValidationError, match="contains duplicate rows" + ): + platform.backend.optimization.equations.add_data( + equation_id=equation_2.id, + data={ + indexset.name: ["foo", "foo"], + indexset_2.name: [2, 2], + "levels": [1, 2], + "marginals": [-1, -2], + }, + ) + + # Test that order is conserved + test_data_2 = { + indexset.name: ["", "", "foo", "foo", "bar", "bar"], + indexset_2.name: [3, 1, 2, 1, 2, 3], + "levels": [6, 5, 4, 3, 2, 1], + "marginals": [1, 3, 5, 6, 4, 2], + } + platform.backend.optimization.equations.add_data( + equation_id=equation_2.id, data=test_data_2 + ) + equation_2 = platform.backend.optimization.equations.get( + run_id=run.id, name="Equation 2" + ) + assert equation_2.data == test_data_2 + + # Test updating of existing keys + equation_4 = platform.backend.optimization.equations.create( + run_id=run.id, + name="Equation 4", + constrained_to_indexsets=[indexset.name, indexset_2.name], + ) + test_data_6 = { + indexset.name: ["foo", "foo", "bar", "bar"], + indexset_2.name: [1, 3, 1, 2], + "levels": [0.00001, "2", 2.3, 400000], + "marginals": [6, 7.8, 9, 0], + } + platform.backend.optimization.equations.add_data( + equation_id=equation_4.id, data=test_data_6 + ) + test_data_7 = { + indexset.name: ["foo", "foo", "bar", "bar", "bar"], + indexset_2.name: [1, 2, 3, 2, 1], + "levels": [0.00001, 2.3, 3, "400000", "5"], + "marginals": [6, 7.8, 9, "0", 3], + } + platform.backend.optimization.equations.add_data( + equation_id=equation_4.id, data=test_data_7 + ) + equation_4 = platform.backend.optimization.equations.get( + run_id=run.id, name="Equation 4" + ) + expected = ( + pd.DataFrame(test_data_7) + .set_index([indexset.name, indexset_2.name]) + .combine_first( + pd.DataFrame(test_data_6).set_index([indexset.name, indexset_2.name]) + ) + .reset_index() + ) + assert_unordered_equality(expected, pd.DataFrame(equation_4.data)) + + def test_equation_remove_data(self, platform: ixmp4.Platform): + run = platform.backend.runs.create("Model", "Scenario") + (indexset,) = create_indexsets_for_run( + platform=platform, run_id=run.id, amount=1 + ) + platform.backend.optimization.indexsets.add_elements( + indexset_id=indexset.id, elements=["foo", "bar"] + ) + test_data = { + indexset.name: ["bar", "foo"], + "levels": [2.0, 1], + "marginals": [0, "test"], + } + equation = platform.backend.optimization.equations.create( + run_id=run.id, + name="Equation", + constrained_to_indexsets=[indexset.name], + ) + platform.backend.optimization.equations.add_data(equation.id, test_data) + equation = platform.backend.optimization.equations.get( + run_id=run.id, name="Equation" + ) + assert equation.data == test_data + + platform.backend.optimization.equations.remove_data(equation_id=equation.id) + equation = platform.backend.optimization.equations.get( + run_id=run.id, name="Equation" + ) + assert equation.data == {} + + def test_list_equation(self, platform: ixmp4.Platform): + run = platform.backend.runs.create("Model", "Scenario") + # Per default, list() lists scalars for `default` version runs: + platform.backend.runs.set_as_default_version(run.id) + indexset, indexset_2 = create_indexsets_for_run( + platform=platform, run_id=run.id + ) + equation = platform.backend.optimization.equations.create( + run_id=run.id, name="Equation", constrained_to_indexsets=[indexset.name] + ) + equation_2 = platform.backend.optimization.equations.create( + run_id=run.id, name="Equation 2", constrained_to_indexsets=[indexset_2.name] + ) + assert [ + equation, + equation_2, + ] == platform.backend.optimization.equations.list() + + assert [equation] == platform.backend.optimization.equations.list( + name="Equation" + ) + + # Test listing Equations for specific Run + run_2 = platform.backend.runs.create("Model", "Scenario") + (indexset,) = create_indexsets_for_run( + platform=platform, run_id=run_2.id, amount=1 + ) + equation_3 = platform.backend.optimization.equations.create( + run_id=run_2.id, name="Equation", constrained_to_indexsets=[indexset.name] + ) + equation_4 = platform.backend.optimization.equations.create( + run_id=run_2.id, name="Equation 2", constrained_to_indexsets=[indexset.name] + ) + assert [ + equation_3, + equation_4, + ] == platform.backend.optimization.equations.list(run_id=run_2.id) + + def test_tabulate_equation(self, platform: ixmp4.Platform): + run = platform.backend.runs.create("Model", "Scenario") + indexset, indexset_2 = create_indexsets_for_run( + platform=platform, run_id=run.id + ) + equation = platform.backend.optimization.equations.create( + run_id=run.id, + name="Equation", + constrained_to_indexsets=[indexset.name, indexset_2.name], + ) + equation_2 = platform.backend.optimization.equations.create( + run_id=run.id, + name="Equation 2", + constrained_to_indexsets=[indexset.name, indexset_2.name], + ) + pd.testing.assert_frame_equal( + df_from_list([equation_2]), + platform.backend.optimization.equations.tabulate(name="Equation 2"), + ) + + platform.backend.optimization.indexsets.add_elements( + indexset_id=indexset.id, elements=["foo", "bar"] + ) + platform.backend.optimization.indexsets.add_elements( + indexset_id=indexset_2.id, elements=[1, 2, 3] + ) + test_data_1 = { + indexset.name: ["foo"], + indexset_2.name: [1], + "levels": [32], + "marginals": [-0], + } + platform.backend.optimization.equations.add_data( + equation_id=equation.id, data=test_data_1 + ) + equation = platform.backend.optimization.equations.get( + run_id=run.id, name="Equation" + ) + + test_data_2 = { + indexset_2.name: [2, 3], + indexset.name: ["foo", "bar"], + "levels": [1, -3.1], + "marginals": [2.0, -4], + } + platform.backend.optimization.equations.add_data( + equation_id=equation_2.id, data=test_data_2 + ) + equation_2 = platform.backend.optimization.equations.get( + run_id=run.id, name="Equation 2" + ) + pd.testing.assert_frame_equal( + df_from_list([equation, equation_2]), + platform.backend.optimization.equations.tabulate(), + ) + + # Test tabulating Equations for specific Run + run_2 = platform.backend.runs.create("Model", "Scenario") + (indexset,) = create_indexsets_for_run( + platform=platform, run_id=run_2.id, amount=1 + ) + equation_3 = platform.backend.optimization.equations.create( + run_id=run_2.id, name="Equation", constrained_to_indexsets=[indexset.name] + ) + equation_4 = platform.backend.optimization.equations.create( + run_id=run_2.id, name="Equation 2", constrained_to_indexsets=[indexset.name] + ) + pd.testing.assert_frame_equal( + df_from_list([equation_3, equation_4]), + platform.backend.optimization.equations.tabulate(run_id=run_2.id), + )