diff --git a/ixmp4/core/optimization/indexset.py b/ixmp4/core/optimization/indexset.py index 9e768a7a..d9772b3b 100644 --- a/ixmp4/core/optimization/indexset.py +++ b/ixmp4/core/optimization/indexset.py @@ -23,17 +23,17 @@ def name(self) -> str: return self._model.name @property - def elements(self) -> list[float | int | str]: - return self._model.elements + def data(self) -> list[float | int | str]: + return self._model.data - def add(self, elements: float | int | list[float | int | str] | str) -> None: - """Adds elements to an existing IndexSet.""" - self.backend.optimization.indexsets.add_elements( - indexset_id=self._model.id, elements=elements + def add(self, data: float | int | list[float | int | str] | str) -> None: + """Adds data to an existing IndexSet.""" + self.backend.optimization.indexsets.add_data( + indexset_id=self._model.id, data=data ) - self._model.elements = self.backend.optimization.indexsets.get( + self._model.data = self.backend.optimization.indexsets.get( run_id=self._model.run__id, name=self._model.name - ).elements + ).data @property def run_id(self) -> int: @@ -48,21 +48,21 @@ def created_by(self) -> str | None: return self._model.created_by @property - def docs(self): + def docs(self) -> str | None: try: return self.backend.optimization.indexsets.docs.get(self.id).description except DocsModel.NotFound: return None @docs.setter - def docs(self, description): + def docs(self, description: str | None) -> None: if description is None: self.backend.optimization.indexsets.docs.delete(self.id) else: self.backend.optimization.indexsets.docs.set(self.id, description) @docs.deleter - def docs(self): + def docs(self) -> None: try: self.backend.optimization.indexsets.docs.delete(self.id) # TODO: silently failing @@ -105,7 +105,9 @@ def list(self, name: str | None = None) -> list[IndexSet]: for i in indexsets ] - def tabulate(self, name: str | None = None) -> pd.DataFrame: + def tabulate( + self, name: str | None = None, include_data: bool = False + ) -> pd.DataFrame: return self.backend.optimization.indexsets.tabulate( - run_id=self._run.id, name=name + run_id=self._run.id, name=name, include_data=include_data ) diff --git a/ixmp4/data/abstract/optimization/equation.py b/ixmp4/data/abstract/optimization/equation.py index 118708af..896534d6 100644 --- a/ixmp4/data/abstract/optimization/equation.py +++ b/ixmp4/data/abstract/optimization/equation.py @@ -177,7 +177,7 @@ def add_data(self, equation_id: int, data: dict[str, Any] | pd.DataFrame) -> Non 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 + contain values that are in the linked `IndexSet.data`. 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. diff --git a/ixmp4/data/abstract/optimization/indexset.py b/ixmp4/data/abstract/optimization/indexset.py index ccf4988f..ae0c6e2d 100644 --- a/ixmp4/data/abstract/optimization/indexset.py +++ b/ixmp4/data/abstract/optimization/indexset.py @@ -17,8 +17,8 @@ class IndexSet(base.BaseModel, Protocol): """The id of the :class:`ixmp4.data.abstract.Run` for which this IndexSet is defined. """ - elements: types.JsonList - """Unique list of str or int.""" + data: types.OptimizationDataList + """Unique list of str, int, or float.""" created_at: types.DateTime "Creation date/time. TODO" @@ -102,13 +102,18 @@ def list(self, *, name: str | None = None, **kwargs) -> list[IndexSet]: """ ... - def tabulate(self, *, name: str | None = None, **kwargs) -> pd.DataFrame: + def tabulate( + self, *, name: str | None = None, include_data: bool = False, **kwargs + ) -> pd.DataFrame: r"""Tabulate IndexSets by specified criteria. Parameters ---------- - name : str + name : str, optional The name of an IndexSet. If supplied only one result will be returned. + include_data : bool, optional + Whether to load all IndexSet data, which reduces loading speed. Defaults to + `False`. # TODO: Update kwargs \*\*kwargs: any More filter parameters as specified in @@ -120,24 +125,24 @@ def tabulate(self, *, name: str | None = None, **kwargs) -> pd.DataFrame: A data frame with the columns: - id - name - - elements + - data - run__id - created_at - created_by """ ... - def add_elements( - self, indexset_id: int, elements: float | int | List[float | int | str] | str + def add_data( + self, indexset_id: int, data: float | int | List[float | int | str] | str ) -> None: - """Adds elements to an existing IndexSet. + """Adds data to an existing IndexSet. Parameters ---------- indexset_id : int The id of the target IndexSet. - elements : float | int | List[float | int | str] | str - The elements to be added to the IndexSet. + data : float | int | List[float | int | str] | str + The data to be added to the IndexSet. Returns ------- diff --git a/ixmp4/data/abstract/optimization/parameter.py b/ixmp4/data/abstract/optimization/parameter.py index a4742b0a..86fd4567 100644 --- a/ixmp4/data/abstract/optimization/parameter.py +++ b/ixmp4/data/abstract/optimization/parameter.py @@ -176,7 +176,7 @@ def add_data(self, parameter_id: int, data: dict[str, Any] | pd.DataFrame) -> No 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 Parameter's columns. Each column can only - contain values that are in the linked `IndexSet.elements`. Each row of entries + contain values that are in the linked `IndexSet.data`. Each row of entries must be unique. No values can be missing, `None`, or `NaN`. If `data.keys()` contains names already present in `Parameter.data`, existing values will be overwritten. diff --git a/ixmp4/data/abstract/optimization/table.py b/ixmp4/data/abstract/optimization/table.py index 6aff32fc..067a252c 100644 --- a/ixmp4/data/abstract/optimization/table.py +++ b/ixmp4/data/abstract/optimization/table.py @@ -176,7 +176,7 @@ def add_data(self, table_id: int, data: dict[str, Any] | pd.DataFrame) -> None: 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 Table's columns. Each column can only - contain values that are in the linked `IndexSet.elements`. Each row of entries + contain values that are in the linked `IndexSet.data`. Each row of entries must be unique. No values can be missing, `None`, or `NaN`. If `data.keys()` contains names already present in `Table.data`, existing values will be overwritten. diff --git a/ixmp4/data/abstract/optimization/variable.py b/ixmp4/data/abstract/optimization/variable.py index 6a5a5d49..cae81378 100644 --- a/ixmp4/data/abstract/optimization/variable.py +++ b/ixmp4/data/abstract/optimization/variable.py @@ -179,7 +179,7 @@ def add_data(self, variable_id: int, data: dict[str, Any] | pd.DataFrame) -> Non 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 Variable's columns. Each column can only - contain values that are in the linked `IndexSet.elements`. Each row of entries + contain values that are in the linked `IndexSet.data`. Each row of entries must be unique. No values can be missing, `None`, or `NaN`. If `data.keys()` contains names already present in `Variable.data`, existing values will be overwritten. diff --git a/ixmp4/data/api/optimization/indexset.py b/ixmp4/data/api/optimization/indexset.py index 477bd308..890237ec 100644 --- a/ixmp4/data/api/optimization/indexset.py +++ b/ixmp4/data/api/optimization/indexset.py @@ -2,7 +2,6 @@ from typing import ClassVar, List import pandas as pd -from pydantic import StrictFloat, StrictInt, StrictStr from ixmp4.data import abstract @@ -17,13 +16,7 @@ class IndexSet(base.BaseModel): id: int name: str - elements: ( - StrictFloat - | StrictInt - | StrictStr - | list[StrictFloat | StrictInt | StrictStr] - | None - ) + data: float | int | str | list[int | float | str] | None run__id: int created_at: datetime | None @@ -64,16 +57,13 @@ def enumerate(self, **kwargs) -> list[IndexSet] | pd.DataFrame: def list(self, **kwargs) -> list[IndexSet]: return super()._list(json=kwargs) - def tabulate(self, **kwargs) -> pd.DataFrame: - return super()._tabulate(json=kwargs) + def tabulate(self, include_data: bool = False, **kwargs) -> pd.DataFrame: + return super()._tabulate(json=kwargs, params={"include_data": include_data}) - def add_elements( + def add_data( self, indexset_id: int, - elements: StrictFloat - | StrictInt - | List[StrictFloat | StrictInt | StrictStr] - | StrictStr, + data: float | int | str | List[float | int | str], ) -> None: - kwargs = {"indexset_id": indexset_id, "elements": elements} + kwargs = {"indexset_id": indexset_id, "data": data} self._request("PATCH", self.prefix + str(indexset_id) + "/", json=kwargs) diff --git a/ixmp4/data/db/optimization/equation/repository.py b/ixmp4/data/db/optimization/equation/repository.py index efe12fef..f9aa91b5 100644 --- a/ixmp4/data/db/optimization/equation/repository.py +++ b/ixmp4/data/db/optimization/equation/repository.py @@ -65,7 +65,7 @@ def _add_column( self.columns.create( name=column_name, constrained_to_indexset=indexset.id, - dtype=pd.Series(indexset.elements).dtype.name, + dtype=pd.Series(indexset.data).dtype.name, equation_id=equation_id, unique=True, **kwargs, diff --git a/ixmp4/data/db/optimization/indexset/model.py b/ixmp4/data/db/optimization/indexset/model.py index 896692a4..58afe208 100644 --- a/ixmp4/data/db/optimization/indexset/model.py +++ b/ixmp4/data/db/optimization/indexset/model.py @@ -1,6 +1,6 @@ from typing import ClassVar -from sqlalchemy.orm import validates +import numpy as np from ixmp4 import db from ixmp4.core.exceptions import OptimizationDataValidationError @@ -16,20 +16,34 @@ class IndexSet(base.BaseModel): DataInvalid: ClassVar = OptimizationDataValidationError DeletionPrevented: ClassVar = abstract.IndexSet.DeletionPrevented - elements: types.JsonList = db.Column(db.JsonType, nullable=False, default=[]) + _data_type: types.OptimizationDataType - @validates("elements") - def validate_elements(self, key, value: list[float | int | str]): - unique = set() - for element in value: - if element in unique: - raise self.DataInvalid( - f"{element} already defined for IndexSet {self.name}!" - ) - else: - unique.add(element) - return value + _data: types.Mapped[list["IndexSetData"]] = db.relationship( + back_populates="indexset" + ) + + @property + def data(self) -> list[float | int | str]: + return ( + [] + if self._data_type is None + else np.array([d.value for d in self._data], dtype=self._data_type).tolist() + ) + + @data.setter + def data(self, value: list[float | int | str]) -> None: + return None run__id: types.RunId __table_args__ = (db.UniqueConstraint("name", "run__id"),) + + +class IndexSetData(base.RootBaseModel): + table_prefix = "optimization_" + + indexset: types.Mapped["IndexSet"] = db.relationship(back_populates="_data") + indexset__id: types.IndexSetId + value: types.String = db.Column(db.String, nullable=False) + + __table_args__ = (db.UniqueConstraint("indexset__id", "value"),) diff --git a/ixmp4/data/db/optimization/indexset/repository.py b/ixmp4/data/db/optimization/indexset/repository.py index 9ec1bc81..d2281a15 100644 --- a/ixmp4/data/db/optimization/indexset/repository.py +++ b/ixmp4/data/db/optimization/indexset/repository.py @@ -8,7 +8,7 @@ from .. import base from .docs import IndexSetDocsRepository -from .model import IndexSet +from .model import IndexSet, IndexSetData class IndexSetRepository( @@ -60,22 +60,45 @@ def list(self, *args, **kwargs) -> list[IndexSet]: return super().list(*args, **kwargs) @guard("view") - def tabulate(self, *args, **kwargs) -> pd.DataFrame: - return super().tabulate(*args, **kwargs) + def tabulate(self, *args, include_data: bool = False, **kwargs) -> pd.DataFrame: + if not include_data: + return ( + super() + .tabulate(*args, **kwargs) + .rename(columns={"_data_type": "data_type"}) + ) + else: + result = super().tabulate(*args, **kwargs).drop(labels="_data_type", axis=1) + result.insert( + loc=0, + column="data", + value=[indexset.data for indexset in self.list(**kwargs)], + ) + return result @guard("edit") - def add_elements( + def add_data( self, indexset_id: int, - elements: float | int | List[float | int | str] | str, + data: float | int | List[float | int | str] | str, ) -> None: indexset = self.get_by_id(id=indexset_id) - if not isinstance(elements, list): - elements = [elements] - if indexset.elements is None: - indexset.elements = elements - else: - indexset.elements = indexset.elements + elements + if not isinstance(data, list): + data = [data] + + bulk_insert_enabled_data: list[dict[str, str]] = [ + {"value": str(d)} for d in data + ] + try: + self.session.execute( + db.insert(IndexSetData).values(indexset__id=indexset_id), + bulk_insert_enabled_data, + ) + except db.IntegrityError as e: + self.session.rollback() + raise indexset.DataInvalid from e + + indexset._data_type = type(data[0]).__name__ self.session.add(indexset) self.session.commit() diff --git a/ixmp4/data/db/optimization/parameter/repository.py b/ixmp4/data/db/optimization/parameter/repository.py index 699cfcf4..126a51fa 100644 --- a/ixmp4/data/db/optimization/parameter/repository.py +++ b/ixmp4/data/db/optimization/parameter/repository.py @@ -66,7 +66,7 @@ def _add_column( self.columns.create( name=column_name, constrained_to_indexset=indexset.id, - dtype=pd.Series(indexset.elements).dtype.name, + dtype=pd.Series(indexset.data).dtype.name, parameter_id=parameter_id, unique=True, **kwargs, diff --git a/ixmp4/data/db/optimization/table/repository.py b/ixmp4/data/db/optimization/table/repository.py index eef57e9f..717b6c51 100644 --- a/ixmp4/data/db/optimization/table/repository.py +++ b/ixmp4/data/db/optimization/table/repository.py @@ -65,7 +65,7 @@ def _add_column( self.columns.create( name=column_name, constrained_to_indexset=indexset.id, - dtype=pd.Series(indexset.elements).dtype.name, + dtype=pd.Series(indexset.data).dtype.name, table_id=table_id, unique=True, **kwargs, diff --git a/ixmp4/data/db/optimization/utils.py b/ixmp4/data/db/optimization/utils.py index 08c5c721..08b2dce2 100644 --- a/ixmp4/data/db/optimization/utils.py +++ b/ixmp4/data/db/optimization/utils.py @@ -12,10 +12,10 @@ def collect_indexsets_to_check( columns: list["Column"], ) -> dict[str, Any]: """Creates a {key:value} dict from linked Column.names and their - IndexSet.elements.""" + IndexSet.data.""" collection: dict[str, Any] = {} for column in columns: - collection[column.name] = column.indexset.elements + collection[column.name] = column.indexset.data return collection diff --git a/ixmp4/data/db/optimization/variable/repository.py b/ixmp4/data/db/optimization/variable/repository.py index 73f88aa9..b8828fd3 100644 --- a/ixmp4/data/db/optimization/variable/repository.py +++ b/ixmp4/data/db/optimization/variable/repository.py @@ -65,7 +65,7 @@ def _add_column( self.columns.create( name=column_name, constrained_to_indexset=indexset.id, - dtype=pd.Series(indexset.elements).dtype.name, + dtype=pd.Series(indexset.data).dtype.name, variable_id=variable_id, unique=True, **kwargs, diff --git a/ixmp4/data/types.py b/ixmp4/data/types.py index 67947fa0..55549c4d 100644 --- a/ixmp4/data/types.py +++ b/ixmp4/data/types.py @@ -1,5 +1,5 @@ from datetime import datetime -from typing import Any +from typing import Any, Literal from sqlalchemy.orm import Mapped as Mapped @@ -8,9 +8,11 @@ Boolean = Mapped[bool] DateTime = Mapped[datetime] Float = Mapped[float] +IndexSetId = Mapped[db.IndexSetIdType] Integer = Mapped[int] -JsonList = Mapped[list[float | int | str]] +OptimizationDataList = Mapped[list[float | int | str]] JsonDict = Mapped[dict[str, Any]] +OptimizationDataType = Mapped[Literal["float", "int", "str"] | None] String = Mapped[str] Name = Mapped[db.NameType] UniqueName = Mapped[db.UniqueNameType] diff --git a/ixmp4/db/__init__.py b/ixmp4/db/__init__.py index cb7bfeae..90a3d67d 100644 --- a/ixmp4/db/__init__.py +++ b/ixmp4/db/__init__.py @@ -50,7 +50,8 @@ update, ) from sqlalchemy.dialects.postgresql import JSONB -from sqlalchemy.exc import MultipleResultsFound +from sqlalchemy.exc import IntegrityError, MultipleResultsFound +from sqlalchemy.ext.hybrid import hybrid_property from sqlalchemy.orm import ( Relationship, Session, @@ -65,17 +66,16 @@ from . import utils Column = mapped_column +IndexSetIdType = Annotated[ + int, + Column(Integer, ForeignKey("optimization_indexset.id"), nullable=False, index=True), +] JsonType = JSON() JsonType = JsonType.with_variant(JSONB(), "postgresql") NameType = Annotated[str, Column(String(255), nullable=False, unique=False)] RunIdType = Annotated[ int, - Column( - Integer, - ForeignKey("run.id"), - nullable=False, - index=True, - ), + Column(Integer, ForeignKey("run.id"), nullable=False, index=True), ] UniqueNameType = Annotated[str, Column(String(255), nullable=False, unique=True)] UsernameType = Annotated[str, Column(String(255), nullable=True)] diff --git a/ixmp4/db/migrations/versions/0d73f7467dab_temporary_create_all_missing_.py b/ixmp4/db/migrations/versions/0d73f7467dab_temporary_create_all_missing_.py new file mode 100644 index 00000000..dd2cf996 --- /dev/null +++ b/ixmp4/db/migrations/versions/0d73f7467dab_temporary_create_all_missing_.py @@ -0,0 +1,270 @@ +# type: ignore +"""TEMPORARY Create all missing optimization items for testing + +Revision ID: 0d73f7467dab +Revises: c29289ced488 +Create Date: 2024-07-08 14:09:49.174145 + +""" + +import sqlalchemy as sa +from alembic import op +from sqlalchemy.dialects import postgresql + +# Revision identifiers, used by Alembic. +revision = "0d73f7467dab" +down_revision = "c29289ced488" +branch_labels = None +depends_on = None + + +def upgrade(): + # ### commands auto generated by Alembic - please adjust! ### + op.create_table( + "optimization_equation", + sa.Column("run__id", sa.Integer(), nullable=False), + sa.Column( + "data", + sa.JSON().with_variant( + postgresql.JSONB(astext_type=sa.Text()), "postgresql" + ), + nullable=False, + ), + sa.Column("name", sa.String(length=255), nullable=False), + sa.Column( + "id", + sa.Integer(), + sa.Identity(always=False, on_null=True, start=1, increment=1), + nullable=False, + ), + sa.Column("created_at", sa.DateTime(), nullable=True), + sa.Column("created_by", sa.String(length=255), nullable=True), + sa.ForeignKeyConstraint( + ["run__id"], ["run.id"], name=op.f("fk_optimization_equation_run__id_run") + ), + sa.PrimaryKeyConstraint("id", name=op.f("pk_optimization_equation")), + sa.UniqueConstraint( + "name", "run__id", name=op.f("uq_optimization_equation_name_run__id") + ), + ) + with op.batch_alter_table("optimization_equation", schema=None) as batch_op: + batch_op.create_index( + batch_op.f("ix_optimization_equation_run__id"), ["run__id"], unique=False + ) + + op.create_table( + "optimization_optimizationvariable", + sa.Column("run__id", sa.Integer(), nullable=False), + sa.Column( + "data", + sa.JSON().with_variant( + postgresql.JSONB(astext_type=sa.Text()), "postgresql" + ), + nullable=False, + ), + sa.Column("name", sa.String(length=255), nullable=False), + sa.Column( + "id", + sa.Integer(), + sa.Identity(always=False, on_null=True, start=1, increment=1), + nullable=False, + ), + sa.Column("created_at", sa.DateTime(), nullable=True), + sa.Column("created_by", sa.String(length=255), nullable=True), + sa.ForeignKeyConstraint( + ["run__id"], + ["run.id"], + name=op.f("fk_optimization_optimizationvariable_run__id_run"), + ), + sa.PrimaryKeyConstraint( + "id", name=op.f("pk_optimization_optimizationvariable") + ), + sa.UniqueConstraint( + "name", + "run__id", + name=op.f("uq_optimization_optimizationvariable_name_run__id"), + ), + ) + with op.batch_alter_table( + "optimization_optimizationvariable", schema=None + ) as batch_op: + batch_op.create_index( + batch_op.f("ix_optimization_optimizationvariable_run__id"), + ["run__id"], + unique=False, + ) + + op.create_table( + "optimization_parameter", + sa.Column("run__id", sa.Integer(), nullable=False), + sa.Column( + "data", + sa.JSON().with_variant( + postgresql.JSONB(astext_type=sa.Text()), "postgresql" + ), + nullable=False, + ), + sa.Column("name", sa.String(length=255), nullable=False), + sa.Column( + "id", + sa.Integer(), + sa.Identity(always=False, on_null=True, start=1, increment=1), + nullable=False, + ), + sa.Column("created_at", sa.DateTime(), nullable=True), + sa.Column("created_by", sa.String(length=255), nullable=True), + sa.ForeignKeyConstraint( + ["run__id"], ["run.id"], name=op.f("fk_optimization_parameter_run__id_run") + ), + sa.PrimaryKeyConstraint("id", name=op.f("pk_optimization_parameter")), + sa.UniqueConstraint( + "name", "run__id", name=op.f("uq_optimization_parameter_name_run__id") + ), + ) + with op.batch_alter_table("optimization_parameter", schema=None) as batch_op: + batch_op.create_index( + batch_op.f("ix_optimization_parameter_run__id"), ["run__id"], unique=False + ) + + op.create_table( + "optimization_equation_docs", + sa.Column("description", sa.Text(), nullable=False), + sa.Column("dimension__id", sa.Integer(), nullable=True), + sa.Column( + "id", + sa.Integer(), + sa.Identity(always=False, on_null=True, start=1, increment=1), + nullable=False, + ), + sa.ForeignKeyConstraint( + ["dimension__id"], + ["optimization_equation.id"], + name=op.f( + "fk_optimization_equation_docs_dimension__id_optimization_equation" + ), + ), + sa.PrimaryKeyConstraint("id", name=op.f("pk_optimization_equation_docs")), + sa.UniqueConstraint( + "dimension__id", name=op.f("uq_optimization_equation_docs_dimension__id") + ), + ) + op.create_table( + "optimization_optimizationvariable_docs", + sa.Column("description", sa.Text(), nullable=False), + sa.Column("dimension__id", sa.Integer(), nullable=True), + sa.Column( + "id", + sa.Integer(), + sa.Identity(always=False, on_null=True, start=1, increment=1), + nullable=False, + ), + sa.ForeignKeyConstraint( + ["dimension__id"], + ["optimization_optimizationvariable.id"], + name=op.f( + "fk_optimization_optimizationvariable_docs_dimension__id_optimization_optimizationvariable" + ), + ), + sa.PrimaryKeyConstraint( + "id", name=op.f("pk_optimization_optimizationvariable_docs") + ), + sa.UniqueConstraint( + "dimension__id", + name=op.f("uq_optimization_optimizationvariable_docs_dimension__id"), + ), + ) + op.create_table( + "optimization_parameter_docs", + sa.Column("description", sa.Text(), nullable=False), + sa.Column("dimension__id", sa.Integer(), nullable=True), + sa.Column( + "id", + sa.Integer(), + sa.Identity(always=False, on_null=True, start=1, increment=1), + nullable=False, + ), + sa.ForeignKeyConstraint( + ["dimension__id"], + ["optimization_parameter.id"], + name=op.f( + "fk_optimization_parameter_docs_dimension__id_optimization_parameter" + ), + ), + sa.PrimaryKeyConstraint("id", name=op.f("pk_optimization_parameter_docs")), + sa.UniqueConstraint( + "dimension__id", name=op.f("uq_optimization_parameter_docs_dimension__id") + ), + ) + with op.batch_alter_table("optimization_column", schema=None) as batch_op: + batch_op.add_column(sa.Column("equation__id", sa.Integer(), nullable=True)) + batch_op.add_column(sa.Column("parameter__id", sa.Integer(), nullable=True)) + batch_op.add_column(sa.Column("variable__id", sa.Integer(), nullable=True)) + batch_op.alter_column("table__id", existing_type=sa.INTEGER(), nullable=True) + batch_op.drop_index("ix_optimization_column_table__id") + batch_op.create_foreign_key( + batch_op.f( + "fk_optimization_column_variable__id_optimization_optimizationvariable" + ), + "optimization_optimizationvariable", + ["variable__id"], + ["id"], + ) + batch_op.create_foreign_key( + batch_op.f("fk_optimization_column_parameter__id_optimization_parameter"), + "optimization_parameter", + ["parameter__id"], + ["id"], + ) + batch_op.create_foreign_key( + batch_op.f("fk_optimization_column_equation__id_optimization_equation"), + "optimization_equation", + ["equation__id"], + ["id"], + ) + + # ### end Alembic commands ### + + +def downgrade(): + # ### commands auto generated by Alembic - please adjust! ### + with op.batch_alter_table("optimization_column", schema=None) as batch_op: + batch_op.drop_constraint( + batch_op.f("fk_optimization_column_equation__id_optimization_equation"), + type_="foreignkey", + ) + batch_op.drop_constraint( + batch_op.f("fk_optimization_column_parameter__id_optimization_parameter"), + type_="foreignkey", + ) + batch_op.drop_constraint( + batch_op.f( + "fk_optimization_column_variable__id_optimization_optimizationvariable" + ), + type_="foreignkey", + ) + batch_op.create_index( + "ix_optimization_column_table__id", ["table__id"], unique=False + ) + batch_op.alter_column("table__id", existing_type=sa.INTEGER(), nullable=False) + batch_op.drop_column("variable__id") + batch_op.drop_column("parameter__id") + batch_op.drop_column("equation__id") + + op.drop_table("optimization_parameter_docs") + op.drop_table("optimization_optimizationvariable_docs") + op.drop_table("optimization_equation_docs") + with op.batch_alter_table("optimization_parameter", schema=None) as batch_op: + batch_op.drop_index(batch_op.f("ix_optimization_parameter_run__id")) + + op.drop_table("optimization_parameter") + with op.batch_alter_table( + "optimization_optimizationvariable", schema=None + ) as batch_op: + batch_op.drop_index(batch_op.f("ix_optimization_optimizationvariable_run__id")) + + op.drop_table("optimization_optimizationvariable") + with op.batch_alter_table("optimization_equation", schema=None) as batch_op: + batch_op.drop_index(batch_op.f("ix_optimization_equation_run__id")) + + op.drop_table("optimization_equation") + # ### end Alembic commands ### diff --git a/ixmp4/db/migrations/versions/914991d09f59_normalize_indexset_data_storage.py b/ixmp4/db/migrations/versions/914991d09f59_normalize_indexset_data_storage.py new file mode 100644 index 00000000..a34d15e6 --- /dev/null +++ b/ixmp4/db/migrations/versions/914991d09f59_normalize_indexset_data_storage.py @@ -0,0 +1,77 @@ +# type: ignore +"""Normalize Indexset.data storage + +Revision ID: 914991d09f59 +Revises: 0d73f7467dab +Create Date: 2024-10-29 15:37:51.485552 + +""" + +import sqlalchemy as sa +from alembic import op +from sqlalchemy.dialects import sqlite + +# Revision identifiers, used by Alembic. +revision = "914991d09f59" +down_revision = "0d73f7467dab" +branch_labels = None +depends_on = None + + +def upgrade(): + # ### commands auto generated by Alembic - please adjust! ### + op.create_table( + "optimization_indexsetdata", + sa.Column("indexset__id", sa.Integer(), nullable=False), + sa.Column("value", sa.String(), nullable=False), + sa.Column( + "id", + sa.Integer(), + sa.Identity(always=False, on_null=True, start=1, increment=1), + nullable=False, + ), + sa.ForeignKeyConstraint( + ["indexset__id"], + ["optimization_indexset.id"], + name=op.f( + "fk_optimization_indexsetdata_indexset__id_optimization_indexset" + ), + ), + sa.PrimaryKeyConstraint("id", name=op.f("pk_optimization_indexsetdata")), + sa.UniqueConstraint( + "indexset__id", + "value", + name=op.f("uq_optimization_indexsetdata_indexset__id_value"), + ), + ) + with op.batch_alter_table("optimization_indexsetdata", schema=None) as batch_op: + batch_op.create_index( + batch_op.f("ix_optimization_indexsetdata_indexset__id"), + ["indexset__id"], + unique=False, + ) + + with op.batch_alter_table("optimization_indexset", schema=None) as batch_op: + batch_op.add_column( + sa.Column( + "_data_type", + sa.Enum("float", "int", "str", native_enum=False), + nullable=True, + ) + ) + batch_op.drop_column("elements") + + # ### end Alembic commands ### + + +def downgrade(): + # ### commands auto generated by Alembic - please adjust! ### + with op.batch_alter_table("optimization_indexset", schema=None) as batch_op: + batch_op.add_column(sa.Column("elements", sqlite.JSON(), nullable=False)) + batch_op.drop_column("_data_type") + + with op.batch_alter_table("optimization_indexsetdata", schema=None) as batch_op: + batch_op.drop_index(batch_op.f("ix_optimization_indexsetdata_indexset__id")) + + op.drop_table("optimization_indexsetdata") + # ### end Alembic commands ### diff --git a/ixmp4/server/rest/base.py b/ixmp4/server/rest/base.py index 2ade6953..17ad863a 100644 --- a/ixmp4/server/rest/base.py +++ b/ixmp4/server/rest/base.py @@ -35,8 +35,9 @@ def __init__( results: pd.DataFrame | api.DataFrame | list[EnumeratedT], **kwargs, ): - if isinstance(results, pd.DataFrame): - kwargs["results"] = api.DataFrame.model_validate(results) - else: - kwargs["results"] = results + kwargs["results"] = ( + api.DataFrame.model_validate(results) + if isinstance(results, pd.DataFrame) + else results + ) super().__init__(**kwargs) diff --git a/ixmp4/server/rest/optimization/indexset.py b/ixmp4/server/rest/optimization/indexset.py index f0ec76fc..f6471586 100644 --- a/ixmp4/server/rest/optimization/indexset.py +++ b/ixmp4/server/rest/optimization/indexset.py @@ -1,5 +1,4 @@ from fastapi import APIRouter, Body, Depends, Query -from pydantic import StrictFloat, StrictInt, StrictStr from ixmp4.data import api from ixmp4.data.backend.db import SqlAlchemyBackend as Backend @@ -20,10 +19,8 @@ class IndexSetInput(BaseModel): name: str -class ElementsInput(BaseModel): - elements: ( - StrictFloat | StrictInt | StrictStr | list[StrictFloat | StrictInt | StrictStr] - ) +class DataInput(BaseModel): + data: float | int | str | list[float | int | str] @autodoc @@ -31,6 +28,7 @@ class ElementsInput(BaseModel): def query( filter: OptimizationIndexSetFilter = Body(OptimizationIndexSetFilter()), table: bool = Query(False), + include_data: bool = Query(False), pagination: Pagination = Depends(), backend: Backend = Depends(deps.get_backend), ): @@ -40,6 +38,7 @@ def query( limit=pagination.limit, offset=pagination.offset, table=bool(table), + include_data=bool(include_data), ), total=backend.optimization.indexsets.count(_filter=filter), pagination=pagination, @@ -57,11 +56,11 @@ def create( @autodoc @router.patch("/{indexset_id}/") -def add_elements( +def add_data( indexset_id: int, - elements: ElementsInput, + data: DataInput, backend: Backend = Depends(deps.get_backend), ): - backend.optimization.indexsets.add_elements( - indexset_id=indexset_id, **elements.model_dump() + backend.optimization.indexsets.add_data( + indexset_id=indexset_id, **data.model_dump() ) diff --git a/tests/core/test_optimization_equation.py b/tests/core/test_optimization_equation.py index 56945d97..67b2fc1c 100644 --- a/tests/core/test_optimization_equation.py +++ b/tests/core/test_optimization_equation.py @@ -90,12 +90,12 @@ def test_create_equation(self, platform: ixmp4.Platform): ) # Test column.dtype is registered correctly - indexset_2.add(elements=2024) + indexset_2.add(data=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 + # If indexset doesn't have data, a generic dtype is registered assert equation_3.columns[0].dtype == "object" assert equation_3.columns[1].dtype == "int64" @@ -126,8 +126,8 @@ def test_equation_add_data(self, platform: ixmp4.Platform): 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]) + indexset.add(data=["foo", "bar", ""]) + indexset_2.add(data=[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 @@ -251,7 +251,7 @@ def test_equation_add_data(self, platform: ixmp4.Platform): 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"]) + indexset.add(data=["foo", "bar"]) test_data = { "Indexset": ["bar", "foo"], "levels": [2.0, 1], @@ -326,8 +326,8 @@ def test_tabulate_equation(self, platform: ixmp4.Platform): run.optimization.equations.tabulate(name="Equation 2"), ) - indexset.add(elements=["foo", "bar"]) - indexset_2.add(elements=[1, 2, 3]) + indexset.add(data=["foo", "bar"]) + indexset_2.add(data=[1, 2, 3]) test_data_1 = { indexset.name: ["foo"], indexset_2.name: [1], diff --git a/tests/core/test_optimization_indexset.py b/tests/core/test_optimization_indexset.py index eb0654ed..a1487e01 100644 --- a/tests/core/test_optimization_indexset.py +++ b/tests/core/test_optimization_indexset.py @@ -9,13 +9,12 @@ from ..utils import create_indexsets_for_run -def df_from_list(indexsets: list[IndexSet]): - return pd.DataFrame( +def df_from_list(indexsets: list[IndexSet], include_data: bool = False) -> pd.DataFrame: + result = pd.DataFrame( # Order is important here to avoid utils.assert_unordered_equality, # which doesn't like lists [ [ - indexset.elements, indexset.run_id, indexset.name, indexset.id, @@ -25,7 +24,6 @@ def df_from_list(indexsets: list[IndexSet]): for indexset in indexsets ], columns=[ - "elements", "run__id", "name", "id", @@ -33,6 +31,20 @@ def df_from_list(indexsets: list[IndexSet]): "created_by", ], ) + if include_data: + result.insert( + loc=0, column="data", value=[indexset.data for indexset in indexsets] + ) + else: + result.insert( + loc=0, + column="data_type", + value=[ + type(indexset.data[0]).__name__ if indexset.data != [] else None + for indexset in indexsets + ], + ) + return result class TestCoreIndexset: @@ -58,15 +70,15 @@ def test_get_indexset(self, platform: ixmp4.Platform): with pytest.raises(IndexSet.NotFound): _ = run.optimization.indexsets.get("Foo") - def test_add_elements(self, platform: ixmp4.Platform): + def test_add_data(self, platform: ixmp4.Platform): run = platform.runs.create("Model", "Scenario") - test_elements = ["foo", "bar"] + test_data = ["foo", "bar"] indexset_1 = run.optimization.indexsets.create("Indexset 1") - indexset_1.add(test_elements) # type: ignore - run.optimization.indexsets.create("Indexset 2").add(test_elements) # type: ignore + indexset_1.add(test_data) # type: ignore + run.optimization.indexsets.create("Indexset 2").add(test_data) # type: ignore indexset_2 = run.optimization.indexsets.get("Indexset 2") - assert indexset_1.elements == indexset_2.elements + assert indexset_1.data == indexset_2.data with pytest.raises(OptimizationDataValidationError): indexset_1.add(["baz", "foo"]) @@ -74,17 +86,20 @@ def test_add_elements(self, platform: ixmp4.Platform): with pytest.raises(OptimizationDataValidationError): indexset_2.add(["baz", "baz"]) - indexset_1.add(1) - indexset_3 = run.optimization.indexsets.get("Indexset 1") - indexset_2.add("1") - indexset_4 = run.optimization.indexsets.get("Indexset 2") - assert indexset_3.elements != indexset_4.elements - assert len(indexset_3.elements) == len(indexset_4.elements) + # Test data types are conserved + indexset_3 = run.optimization.indexsets.create("Indexset 3") + test_data_2: list[float | int | str] = [1.2, 3.4, 5.6] + indexset_3.add(data=test_data_2) + + assert indexset_3.data == test_data_2 + assert type(indexset_3.data[0]).__name__ == "float" - test_elements_2 = ["One", 2, 3.141] - indexset_5 = run.optimization.indexsets.create("Indexset 5") - indexset_5.add(test_elements_2) # type: ignore - assert indexset_5.elements == test_elements_2 + indexset_4 = run.optimization.indexsets.create("Indexset 4") + test_data_3: list[float | int | str] = [0, 1, 2] + indexset_4.add(data=test_data_3) + + assert indexset_4.data == test_data_3 + assert type(indexset_4.data[0]).__name__ == "int" def test_list_indexsets(self, platform: ixmp4.Platform): run = platform.runs.create("Model", "Scenario") @@ -128,6 +143,15 @@ def test_tabulate_indexsets(self, platform: ixmp4.Platform): result = run.optimization.indexsets.tabulate(name="Indexset 2") pdt.assert_frame_equal(expected, result) + # Test tabulating including the data + expected = df_from_list(indexsets=[indexset_2], include_data=True) + pdt.assert_frame_equal( + expected, + run.optimization.indexsets.tabulate( + name=indexset_2.name, include_data=True + ), + ) + def test_indexset_docs(self, platform: ixmp4.Platform): run = platform.runs.create("Model", "Scenario") (indexset_1,) = tuple( diff --git a/tests/core/test_optimization_parameter.py b/tests/core/test_optimization_parameter.py index 8602ddf1..f6baed9d 100644 --- a/tests/core/test_optimization_parameter.py +++ b/tests/core/test_optimization_parameter.py @@ -90,12 +90,12 @@ def test_create_parameter(self, platform: ixmp4.Platform): ) # Test column.dtype is registered correctly - indexset_2.add(elements=2024) + indexset_2.add(data=2024) parameter_3 = run.optimization.parameters.create( "Parameter 5", constrained_to_indexsets=[indexset.name, indexset_2.name], ) - # If indexset doesn't have elements, a generic dtype is registered + # If indexset doesn't have data, a generic dtype is registered assert parameter_3.columns[0].dtype == "object" assert parameter_3.columns[1].dtype == "int64" @@ -127,8 +127,8 @@ def test_parameter_add_data(self, platform: ixmp4.Platform): 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]) + indexset.add(data=["foo", "bar", ""]) + indexset_2.add(data=[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 @@ -308,8 +308,8 @@ def test_tabulate_parameter(self, platform: ixmp4.Platform): unit = platform.units.create("Unit") unit_2 = platform.units.create("Unit 2") - indexset.add(elements=["foo", "bar"]) - indexset_2.add(elements=[1, 2, 3]) + indexset.add(data=["foo", "bar"]) + indexset_2.add(data=[1, 2, 3]) test_data_1 = { indexset.name: ["foo"], indexset_2.name: [1], diff --git a/tests/core/test_optimization_table.py b/tests/core/test_optimization_table.py index 57110950..93f0c03e 100644 --- a/tests/core/test_optimization_table.py +++ b/tests/core/test_optimization_table.py @@ -90,12 +90,12 @@ def test_create_table(self, platform: ixmp4.Platform): ) # Test column.dtype is registered correctly - indexset_2.add(elements=2024) + indexset_2.add(data=2024) table_3 = run.optimization.tables.create( "Table 5", constrained_to_indexsets=[indexset.name, indexset_2.name], ) - # If indexset doesn't have elements, a generic dtype is registered + # If indexset doesn't have data, a generic dtype is registered assert table_3.columns[0].dtype == "object" assert table_3.columns[1].dtype == "int64" @@ -124,7 +124,7 @@ def test_table_add_data(self, platform: ixmp4.Platform): IndexSet(_backend=platform.backend, _model=model) # type: ignore for model in create_indexsets_for_run(platform=platform, run_id=run.id) ) - indexset.add(elements=["foo", "bar", ""]) + indexset.add(data=["foo", "bar", ""]) indexset_2.add([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 @@ -224,9 +224,9 @@ def test_table_add_data(self, platform: ixmp4.Platform): indexset_3 = run.optimization.indexsets.create(name="Indexset 3") test_data_5 = { indexset.name: ["foo", "foo", "bar"], - indexset_3.name: [1, "2", 3.14], + indexset_3.name: [1.0, 2.2, 3.14], } - indexset_3.add(elements=[1, "2", 3.14]) + indexset_3.add(data=[1.0, 2.2, 3.14]) table_5 = run.optimization.tables.create( name="Table 5", constrained_to_indexsets=[indexset.name, indexset_3.name], diff --git a/tests/core/test_optimization_variable.py b/tests/core/test_optimization_variable.py index 0cd096b5..560b4783 100644 --- a/tests/core/test_optimization_variable.py +++ b/tests/core/test_optimization_variable.py @@ -113,12 +113,12 @@ def test_create_variable(self, platform: ixmp4.Platform): ) # Test column.dtype is registered correctly - indexset_2.add(elements=2024) + indexset_2.add(data=2024) variable_4 = run.optimization.variables.create( "Variable 4", constrained_to_indexsets=[indexset.name, indexset_2.name], ) - # If indexset doesn't have elements, a generic dtype is registered + # If indexset doesn't have data, a generic dtype is registered assert variable_4.columns is not None assert variable_4.columns[0].dtype == "object" assert variable_4.columns[1].dtype == "int64" @@ -151,8 +151,8 @@ def test_variable_add_data(self, platform: ixmp4.Platform): 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]) + indexset.add(data=["foo", "bar", ""]) + indexset_2.add(data=[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 @@ -276,7 +276,7 @@ def test_variable_add_data(self, platform: ixmp4.Platform): def test_variable_remove_data(self, platform: ixmp4.Platform): run = platform.runs.create("Model", "Scenario") indexset = run.optimization.indexsets.create("Indexset") - indexset.add(elements=["foo", "bar"]) + indexset.add(data=["foo", "bar"]) test_data = { "Indexset": ["bar", "foo"], "levels": [2.0, 1], @@ -349,8 +349,8 @@ def test_tabulate_variable(self, platform: ixmp4.Platform): run.optimization.variables.tabulate(name="Variable 2"), ) - indexset.add(elements=["foo", "bar"]) - indexset_2.add(elements=[1, 2, 3]) + indexset.add(data=["foo", "bar"]) + indexset_2.add(data=[1, 2, 3]) test_data_1 = { indexset.name: ["foo"], indexset_2.name: [1], diff --git a/tests/data/test_optimization_equation.py b/tests/data/test_optimization_equation.py index b82df865..48aa29ae 100644 --- a/tests/data/test_optimization_equation.py +++ b/tests/data/test_optimization_equation.py @@ -91,9 +91,7 @@ def test_create_equation(self, platform: ixmp4.Platform): ) # Test column.dtype is registered correctly - platform.backend.optimization.indexsets.add_elements( - indexset_2.id, elements=2024 - ) + platform.backend.optimization.indexsets.add_data(indexset_2.id, data=2024) indexset_2 = platform.backend.optimization.indexsets.get( run.id, indexset_2.name ) @@ -102,7 +100,7 @@ def test_create_equation(self, platform: ixmp4.Platform): name="Equation 5", constrained_to_indexsets=[indexset.name, indexset_2.name], ) - # If indexset doesn't have elements, a generic dtype is registered + # If indexset doesn't have data, a generic dtype is registered assert equation_3.columns[0].dtype == "object" assert equation_3.columns[1].dtype == "int64" @@ -128,11 +126,11 @@ def test_equation_add_data(self, platform: ixmp4.Platform): 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_data( + indexset_id=indexset.id, data=["foo", "bar", ""] ) - platform.backend.optimization.indexsets.add_elements( - indexset_id=indexset_2.id, elements=[1, 2, 3] + platform.backend.optimization.indexsets.add_data( + indexset_id=indexset_2.id, data=[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 @@ -280,8 +278,8 @@ def test_equation_remove_data(self, platform: ixmp4.Platform): (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"] + platform.backend.optimization.indexsets.add_data( + indexset_id=indexset.id, data=["foo", "bar"] ) test_data = { indexset.name: ["bar", "foo"], @@ -363,11 +361,11 @@ def test_tabulate_equation(self, platform: ixmp4.Platform): 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_data( + indexset_id=indexset.id, data=["foo", "bar"] ) - platform.backend.optimization.indexsets.add_elements( - indexset_id=indexset_2.id, elements=[1, 2, 3] + platform.backend.optimization.indexsets.add_data( + indexset_id=indexset_2.id, data=[1, 2, 3] ) test_data_1 = { indexset.name: ["foo"], diff --git a/tests/data/test_optimization_indexset.py b/tests/data/test_optimization_indexset.py index 9eb4e094..d4cd5bbc 100644 --- a/tests/data/test_optimization_indexset.py +++ b/tests/data/test_optimization_indexset.py @@ -9,13 +9,12 @@ from ..utils import create_indexsets_for_run -def df_from_list(indexsets: list): - return pd.DataFrame( +def df_from_list(indexsets: list[IndexSet], include_data: bool = False) -> pd.DataFrame: + result = pd.DataFrame( # Order is important here to avoid utils.assert_unordered_equality, # which doesn't like lists [ [ - indexset.elements, indexset.run__id, indexset.name, indexset.id, @@ -25,7 +24,6 @@ def df_from_list(indexsets: list): for indexset in indexsets ], columns=[ - "elements", "run__id", "name", "id", @@ -33,6 +31,20 @@ def df_from_list(indexsets: list): "created_by", ], ) + if include_data: + result.insert( + loc=0, column="data", value=[indexset.data for indexset in indexsets] + ) + else: + result.insert( + loc=0, + column="data_type", + value=[ + type(indexset.data[0]).__name__ if indexset.data != [] else None + for indexset in indexsets + ], + ) + return result class TestDataOptimizationIndexSet: @@ -92,11 +104,11 @@ def test_tabulate_indexsets(self, platform: ixmp4.Platform): indexset_1, indexset_2 = create_indexsets_for_run( platform=platform, run_id=run.id ) - platform.backend.optimization.indexsets.add_elements( - indexset_id=indexset_1.id, elements="foo" + platform.backend.optimization.indexsets.add_data( + indexset_id=indexset_1.id, data="foo" ) - platform.backend.optimization.indexsets.add_elements( - indexset_id=indexset_2.id, elements=[1, 2] + platform.backend.optimization.indexsets.add_data( + indexset_id=indexset_2.id, data=[1, 2] ) indexset_1 = platform.backend.optimization.indexsets.get( @@ -115,76 +127,85 @@ def test_tabulate_indexsets(self, platform: ixmp4.Platform): expected, platform.backend.optimization.indexsets.tabulate(name=indexset_1.name), ) + # Test only indexsets belonging to this Run are tabulated if run_id is provided run_2 = platform.backend.runs.create("Model", "Scenario") indexset_3, indexset_4 = create_indexsets_for_run( - platform=platform, run_id=run_2.id, offset=2 + platform=platform, run_id=run_2.id, offset=3 ) expected = df_from_list(indexsets=[indexset_3, indexset_4]) pdt.assert_frame_equal( expected, platform.backend.optimization.indexsets.tabulate(run_id=run_2.id) ) - def test_add_elements(self, platform: ixmp4.Platform): - test_elements = ["foo", "bar"] + # Test tabulating including the data + expected = df_from_list(indexsets=[indexset_2], include_data=True) + pdt.assert_frame_equal( + expected, + platform.backend.optimization.indexsets.tabulate( + name=indexset_2.name, include_data=True + ), + ) + + def test_add_data(self, platform: ixmp4.Platform): + test_data = ["foo", "bar"] run = platform.backend.runs.create("Model", "Scenario") indexset_1, indexset_2 = create_indexsets_for_run( platform=platform, run_id=run.id ) - platform.backend.optimization.indexsets.add_elements( + platform.backend.optimization.indexsets.add_data( indexset_id=indexset_1.id, - elements=test_elements, # type: ignore + data=test_data, # type: ignore ) indexset_1 = platform.backend.optimization.indexsets.get( run_id=run.id, name=indexset_1.name ) - platform.backend.optimization.indexsets.add_elements( + platform.backend.optimization.indexsets.add_data( indexset_id=indexset_2.id, - elements=test_elements, # type: ignore + data=test_data, # type: ignore ) assert ( - indexset_1.elements + indexset_1.data == platform.backend.optimization.indexsets.get( run_id=run.id, name=indexset_2.name - ).elements + ).data ) with pytest.raises(OptimizationDataValidationError): - platform.backend.optimization.indexsets.add_elements( - indexset_id=indexset_1.id, elements=["baz", "foo"] + platform.backend.optimization.indexsets.add_data( + indexset_id=indexset_1.id, data=["baz", "foo"] ) with pytest.raises(OptimizationDataValidationError): - platform.backend.optimization.indexsets.add_elements( - indexset_id=indexset_2.id, elements=["baz", "baz"] + platform.backend.optimization.indexsets.add_data( + indexset_id=indexset_2.id, data=["baz", "baz"] ) - platform.backend.optimization.indexsets.add_elements( - indexset_id=indexset_1.id, elements=1 + # Test data types are conserved + indexset_3, indexset_4 = create_indexsets_for_run( + platform=platform, run_id=run.id, offset=3 + ) + + test_data_2: list[float | int | str] = [1.2, 3.4, 5.6] + platform.backend.optimization.indexsets.add_data( + indexset_id=indexset_3.id, data=test_data_2 ) indexset_3 = platform.backend.optimization.indexsets.get( - run_id=run.id, name=indexset_1.name + run_id=run.id, name=indexset_3.name ) - platform.backend.optimization.indexsets.add_elements( - indexset_id=indexset_2.id, elements="1" + + assert indexset_3.data == test_data_2 + assert type(indexset_3.data[0]).__name__ == "float" + + test_data_3: list[float | int | str] = [0, 1, 2] + platform.backend.optimization.indexsets.add_data( + indexset_id=indexset_4.id, data=test_data_3 ) indexset_4 = platform.backend.optimization.indexsets.get( - run_id=run.id, name=indexset_2.name + run_id=run.id, name=indexset_4.name ) - assert indexset_3.elements != indexset_4.elements - assert len(indexset_3.elements) == len(indexset_4.elements) - test_elements_2 = [1, "2", 3.14] - indexset_5 = platform.backend.optimization.indexsets.create( - run_id=run.id, name="IndexSet 5" - ) - platform.backend.optimization.indexsets.add_elements( - indexset_id=indexset_5.id, - elements=test_elements_2, # type:ignore - ) - indexset_5 = platform.backend.optimization.indexsets.get( - run_id=run.id, name=indexset_5.name - ) - assert indexset_5.elements == test_elements_2 + assert indexset_4.data == test_data_3 + assert type(indexset_4.data[0]).__name__ == "int" diff --git a/tests/data/test_optimization_parameter.py b/tests/data/test_optimization_parameter.py index dac17e86..7993596f 100644 --- a/tests/data/test_optimization_parameter.py +++ b/tests/data/test_optimization_parameter.py @@ -93,9 +93,7 @@ def test_create_parameter(self, platform: ixmp4.Platform): ) # Test column.dtype is registered correctly - platform.backend.optimization.indexsets.add_elements( - indexset_2.id, elements=2024 - ) + platform.backend.optimization.indexsets.add_data(indexset_2.id, data=2024) indexset_2 = platform.backend.optimization.indexsets.get( run.id, indexset_2.name ) @@ -104,7 +102,7 @@ def test_create_parameter(self, platform: ixmp4.Platform): name="Parameter 5", constrained_to_indexsets=[indexset.name, indexset_2.name], ) - # If indexset doesn't have elements, a generic dtype is registered + # If indexset doesn't have data, a generic dtype is registered assert parameter_3.columns[0].dtype == "object" assert parameter_3.columns[1].dtype == "int64" @@ -129,11 +127,11 @@ def test_parameter_add_data(self, platform: ixmp4.Platform): 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_data( + indexset_id=indexset.id, data=["foo", "bar", ""] ) - platform.backend.optimization.indexsets.add_elements( - indexset_id=indexset_2.id, elements=[1, 2, 3] + platform.backend.optimization.indexsets.add_data( + indexset_id=indexset_2.id, data=[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 @@ -348,11 +346,11 @@ def test_tabulate_parameter(self, platform: ixmp4.Platform): unit = platform.backend.units.create("Unit") unit_2 = platform.backend.units.create("Unit 2") - platform.backend.optimization.indexsets.add_elements( - indexset_id=indexset.id, elements=["foo", "bar"] + platform.backend.optimization.indexsets.add_data( + indexset_id=indexset.id, data=["foo", "bar"] ) - platform.backend.optimization.indexsets.add_elements( - indexset_id=indexset_2.id, elements=[1, 2, 3] + platform.backend.optimization.indexsets.add_data( + indexset_id=indexset_2.id, data=[1, 2, 3] ) test_data_1 = { indexset.name: ["foo"], diff --git a/tests/data/test_optimization_table.py b/tests/data/test_optimization_table.py index f4643da7..5db61cc6 100644 --- a/tests/data/test_optimization_table.py +++ b/tests/data/test_optimization_table.py @@ -89,9 +89,7 @@ def test_create_table(self, platform: ixmp4.Platform): ) # Test column.dtype is registered correctly - platform.backend.optimization.indexsets.add_elements( - indexset_2.id, elements=2024 - ) + platform.backend.optimization.indexsets.add_data(indexset_2.id, data=2024) indexset_2 = platform.backend.optimization.indexsets.get( run.id, indexset_2.name ) @@ -100,7 +98,7 @@ def test_create_table(self, platform: ixmp4.Platform): name="Table 5", constrained_to_indexsets=[indexset_1.name, indexset_2.name], ) - # If indexset doesn't have elements, a generic dtype is registered + # If indexset doesn't have data, a generic dtype is registered assert table_3.columns[0].dtype == "object" assert table_3.columns[1].dtype == "int64" @@ -122,11 +120,11 @@ def test_table_add_data(self, platform: ixmp4.Platform): indexset_1, indexset_2 = create_indexsets_for_run( platform=platform, run_id=run.id ) - platform.backend.optimization.indexsets.add_elements( - indexset_id=indexset_1.id, elements=["foo", "bar", ""] + platform.backend.optimization.indexsets.add_data( + indexset_id=indexset_1.id, data=["foo", "bar", ""] ) - platform.backend.optimization.indexsets.add_elements( - indexset_id=indexset_2.id, elements=[1, 2, 3] + platform.backend.optimization.indexsets.add_data( + indexset_id=indexset_2.id, data=[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 @@ -268,11 +266,11 @@ def test_table_add_data(self, platform: ixmp4.Platform): ) test_data_5 = { indexset_1.name: ["foo", "foo", "bar"], - indexset_3.name: [1, "2", 3.14], + indexset_3.name: [1.0, 2.2, 3.14], } - platform.backend.optimization.indexsets.add_elements( - indexset_id=indexset_3.id, elements=[1, "2", 3.14] + platform.backend.optimization.indexsets.add_data( + indexset_id=indexset_3.id, data=[1.0, 2.2, 3.14] ) table_5 = platform.backend.optimization.tables.create( run_id=run.id, @@ -341,11 +339,11 @@ def test_tabulate_table(self, platform: ixmp4.Platform): platform.backend.optimization.tables.tabulate(name="Table 2"), ) - platform.backend.optimization.indexsets.add_elements( - indexset_id=indexset_1.id, elements=["foo", "bar"] + platform.backend.optimization.indexsets.add_data( + indexset_id=indexset_1.id, data=["foo", "bar"] ) - platform.backend.optimization.indexsets.add_elements( - indexset_id=indexset_2.id, elements=[1, 2, 3] + platform.backend.optimization.indexsets.add_data( + indexset_id=indexset_2.id, data=[1, 2, 3] ) test_data_1 = {indexset_1.name: ["foo"], indexset_2.name: [1]} platform.backend.optimization.tables.add_data( diff --git a/tests/data/test_optimization_variable.py b/tests/data/test_optimization_variable.py index ff8f3012..0c61ea72 100644 --- a/tests/data/test_optimization_variable.py +++ b/tests/data/test_optimization_variable.py @@ -114,9 +114,7 @@ def test_create_variable(self, platform: ixmp4.Platform): ) # Test column.dtype is registered correctly - platform.backend.optimization.indexsets.add_elements( - indexset_2.id, elements=2024 - ) + platform.backend.optimization.indexsets.add_data(indexset_2.id, data=2024) indexset_2 = platform.backend.optimization.indexsets.get( run.id, indexset_2.name ) @@ -125,7 +123,7 @@ def test_create_variable(self, platform: ixmp4.Platform): name="Variable 4", constrained_to_indexsets=[indexset.name, indexset_2.name], ) - # If indexset doesn't have elements, a generic dtype is registered + # If indexset doesn't have data, a generic dtype is registered assert variable_4.columns is not None assert variable_4.columns[0].dtype == "object" assert variable_4.columns[1].dtype == "int64" @@ -152,11 +150,11 @@ def test_variable_add_data(self, platform: ixmp4.Platform): 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_data( + indexset_id=indexset.id, data=["foo", "bar", ""] ) - platform.backend.optimization.indexsets.add_elements( - indexset_id=indexset_2.id, elements=[1, 2, 3] + platform.backend.optimization.indexsets.add_data( + indexset_id=indexset_2.id, data=[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 @@ -302,8 +300,8 @@ def test_variable_add_data(self, platform: ixmp4.Platform): def test_variable_remove_data(self, platform: ixmp4.Platform): run = platform.backend.runs.create("Model", "Scenario") indexset = platform.backend.optimization.indexsets.create(run.id, "Indexset") - platform.backend.optimization.indexsets.add_elements( - indexset_id=indexset.id, elements=["foo", "bar"] + platform.backend.optimization.indexsets.add_data( + indexset_id=indexset.id, data=["foo", "bar"] ) test_data = { "Indexset": ["bar", "foo"], @@ -383,11 +381,11 @@ def test_tabulate_variable(self, platform: ixmp4.Platform): platform.backend.optimization.variables.tabulate(name="Variable 2"), ) - platform.backend.optimization.indexsets.add_elements( - indexset_id=indexset.id, elements=["foo", "bar"] + platform.backend.optimization.indexsets.add_data( + indexset_id=indexset.id, data=["foo", "bar"] ) - platform.backend.optimization.indexsets.add_elements( - indexset_id=indexset_2.id, elements=[1, 2, 3] + platform.backend.optimization.indexsets.add_data( + indexset_id=indexset_2.id, data=[1, 2, 3] ) test_data_1 = { indexset.name: ["foo"],