Skip to content

Commit 0deac5b

Browse files
authored
DRY optimization core layer (#96)
* Add abstraction layer for core optimization reusability * Add annotations
1 parent 26abf0e commit 0deac5b

File tree

16 files changed

+232
-185
lines changed

16 files changed

+232
-185
lines changed

ixmp4/core/optimization/base.py

+94
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,94 @@
1+
from typing import TYPE_CHECKING, Generic, TypeVar
2+
3+
if TYPE_CHECKING:
4+
from . import InitKwargs
5+
6+
7+
import pandas as pd
8+
9+
# TODO Import this from typing when dropping Python 3.11
10+
from typing_extensions import TypedDict, Unpack
11+
12+
from ixmp4.core.base import BaseFacade, BaseModelFacade
13+
from ixmp4.data import abstract
14+
15+
FacadeOptimizationModelType = TypeVar(
16+
"FacadeOptimizationModelType", bound=BaseModelFacade
17+
)
18+
AbstractOptimizationModelType = TypeVar(
19+
"AbstractOptimizationModelType", bound=abstract.BaseModel
20+
)
21+
22+
23+
class OptimizationBaseRepository(
24+
BaseFacade, Generic[FacadeOptimizationModelType, AbstractOptimizationModelType]
25+
):
26+
_run: abstract.Run
27+
_backend_repository: abstract.BackendBaseRepository[AbstractOptimizationModelType]
28+
_model_type: type[FacadeOptimizationModelType]
29+
30+
def __init__(self, _run: abstract.Run, **kwargs: Unpack["InitKwargs"]) -> None:
31+
super().__init__(**kwargs)
32+
self._run = _run
33+
34+
35+
class Creator(
36+
OptimizationBaseRepository[
37+
FacadeOptimizationModelType, AbstractOptimizationModelType
38+
],
39+
abstract.Creator,
40+
):
41+
def create(
42+
self, name: str, **kwargs: Unpack["abstract.optimization.base.CreateKwargs"]
43+
) -> FacadeOptimizationModelType:
44+
model = self._backend_repository.create(
45+
run_id=self._run.id, name=name, **kwargs
46+
)
47+
return self._model_type(_backend=self.backend, _model=model)
48+
49+
50+
class Retriever(
51+
OptimizationBaseRepository[
52+
FacadeOptimizationModelType, AbstractOptimizationModelType
53+
],
54+
abstract.Retriever,
55+
):
56+
def get(self, name: str) -> FacadeOptimizationModelType:
57+
model = self._backend_repository.get(run_id=self._run.id, name=name)
58+
return self._model_type(_backend=self.backend, _model=model)
59+
60+
61+
class Lister(
62+
OptimizationBaseRepository[
63+
FacadeOptimizationModelType, AbstractOptimizationModelType
64+
],
65+
abstract.Lister,
66+
):
67+
def list(self, name: str | None = None) -> list[FacadeOptimizationModelType]:
68+
models = self._backend_repository.list(run_id=self._run.id, name=name)
69+
return [self._model_type(_backend=self.backend, _model=m) for m in models]
70+
71+
72+
class Tabulator(
73+
OptimizationBaseRepository[
74+
FacadeOptimizationModelType, AbstractOptimizationModelType
75+
],
76+
abstract.Tabulator,
77+
):
78+
def tabulate(self, name: str | None = None) -> pd.DataFrame:
79+
return self._backend_repository.tabulate(run_id=self._run.id, name=name)
80+
81+
82+
class EnumerateKwargs(TypedDict, total=False):
83+
name: str | None
84+
85+
86+
class Enumerator(
87+
Lister[FacadeOptimizationModelType, AbstractOptimizationModelType],
88+
Tabulator[FacadeOptimizationModelType, AbstractOptimizationModelType],
89+
abstract.Enumerator,
90+
):
91+
def enumerate(
92+
self, table: bool = False, **kwargs: Unpack[EnumerateKwargs]
93+
) -> list[FacadeOptimizationModelType] | pd.DataFrame:
94+
return self.tabulate(**kwargs) if table else self.list(**kwargs)

ixmp4/core/optimization/equation.py

+13-31
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,3 @@
1-
from collections.abc import Iterable
21
from datetime import datetime
32
from typing import TYPE_CHECKING, Any, ClassVar
43

@@ -10,12 +9,14 @@
109
# TODO Import this from typing when dropping Python 3.11
1110
from typing_extensions import Unpack
1211

13-
from ixmp4.core.base import BaseFacade, BaseModelFacade
12+
from ixmp4.core.base import BaseModelFacade
1413
from ixmp4.data.abstract import Docs as DocsModel
1514
from ixmp4.data.abstract import Equation as EquationModel
1615
from ixmp4.data.abstract import Run
1716
from ixmp4.data.abstract.optimization import Column
1817

18+
from .base import Creator, Lister, Retriever, Tabulator
19+
1920

2021
class Equation(BaseModelFacade):
2122
_model: EquationModel
@@ -106,44 +107,25 @@ def __str__(self) -> str:
106107
return f"<Equation {self.id} name={self.name}>"
107108

108109

109-
class EquationRepository(BaseFacade):
110-
_run: Run
111-
110+
class EquationRepository(
111+
Creator[Equation, EquationModel],
112+
Retriever[Equation, EquationModel],
113+
Lister[Equation, EquationModel],
114+
Tabulator[Equation, EquationModel],
115+
):
112116
def __init__(self, _run: Run, **kwargs: Unpack["InitKwargs"]) -> None:
113-
super().__init__(**kwargs)
114-
self._run = _run
117+
super().__init__(_run=_run, **kwargs)
118+
self._backend_repository = self.backend.optimization.equations
119+
self._model_type = Equation
115120

116121
def create(
117122
self,
118123
name: str,
119124
constrained_to_indexsets: list[str],
120125
column_names: list[str] | None = None,
121126
) -> Equation:
122-
model = self.backend.optimization.equations.create(
127+
return super().create(
123128
name=name,
124-
run_id=self._run.id,
125129
constrained_to_indexsets=constrained_to_indexsets,
126130
column_names=column_names,
127131
)
128-
return Equation(_backend=self.backend, _model=model)
129-
130-
def get(self, name: str) -> Equation:
131-
model = self.backend.optimization.equations.get(run_id=self._run.id, name=name)
132-
return Equation(_backend=self.backend, _model=model)
133-
134-
def list(self, name: str | None = None) -> Iterable[Equation]:
135-
equations = self.backend.optimization.equations.list(
136-
run_id=self._run.id, name=name
137-
)
138-
return [
139-
Equation(
140-
_backend=self.backend,
141-
_model=i,
142-
)
143-
for i in equations
144-
]
145-
146-
def tabulate(self, name: str | None = None) -> pd.DataFrame:
147-
return self.backend.optimization.equations.tabulate(
148-
run_id=self._run.id, name=name
149-
)

ixmp4/core/optimization/indexset.py

+13-36
Original file line numberDiff line numberDiff line change
@@ -4,16 +4,16 @@
44
if TYPE_CHECKING:
55
from . import InitKwargs
66

7-
import pandas as pd
8-
97
# TODO Import this from typing when dropping Python 3.11
108
from typing_extensions import Unpack
119

12-
from ixmp4.core.base import BaseFacade, BaseModelFacade
10+
from ixmp4.core.base import BaseModelFacade
1311
from ixmp4.data.abstract import Docs as DocsModel
1412
from ixmp4.data.abstract import IndexSet as IndexSetModel
1513
from ixmp4.data.abstract import Run
1614

15+
from .base import Creator, Lister, Retriever, Tabulator
16+
1717

1818
class IndexSet(BaseModelFacade):
1919
_model: IndexSetModel
@@ -81,39 +81,16 @@ def __str__(self) -> str:
8181
return f"<IndexSet {self.id} name={self.name}>"
8282

8383

84-
class IndexSetRepository(BaseFacade):
85-
_run: Run
86-
84+
class IndexSetRepository(
85+
Creator[IndexSet, IndexSetModel],
86+
Retriever[IndexSet, IndexSetModel],
87+
Lister[IndexSet, IndexSetModel],
88+
Tabulator[IndexSet, IndexSetModel],
89+
):
8790
def __init__(self, _run: Run, **kwargs: Unpack["InitKwargs"]) -> None:
88-
super().__init__(**kwargs)
89-
self._run = _run
91+
super().__init__(_run=_run, **kwargs)
92+
self._backend_repository = self.backend.optimization.indexsets
93+
self._model_type = IndexSet
9094

9195
def create(self, name: str) -> IndexSet:
92-
indexset = self.backend.optimization.indexsets.create(
93-
run_id=self._run.id,
94-
name=name,
95-
)
96-
return IndexSet(_backend=self.backend, _model=indexset)
97-
98-
def get(self, name: str) -> IndexSet:
99-
indexset = self.backend.optimization.indexsets.get(
100-
run_id=self._run.id, name=name
101-
)
102-
return IndexSet(_backend=self.backend, _model=indexset)
103-
104-
def list(self, name: str | None = None) -> list[IndexSet]:
105-
indexsets = self.backend.optimization.indexsets.list(
106-
run_id=self._run.id, name=name
107-
)
108-
return [
109-
IndexSet(
110-
_backend=self.backend,
111-
_model=i,
112-
)
113-
for i in indexsets
114-
]
115-
116-
def tabulate(self, name: str | None = None) -> pd.DataFrame:
117-
return self.backend.optimization.indexsets.tabulate(
118-
run_id=self._run.id, name=name
119-
)
96+
return super().create(name=name)

ixmp4/core/optimization/parameter.py

+13-31
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,3 @@
1-
from collections.abc import Iterable
21
from datetime import datetime
32
from typing import TYPE_CHECKING, Any, ClassVar
43

@@ -10,12 +9,14 @@
109
# TODO Import this from typing when dropping Python 3.11
1110
from typing_extensions import Unpack
1211

13-
from ixmp4.core.base import BaseFacade, BaseModelFacade
12+
from ixmp4.core.base import BaseModelFacade
1413
from ixmp4.data.abstract import Docs as DocsModel
1514
from ixmp4.data.abstract import Parameter as ParameterModel
1615
from ixmp4.data.abstract import Run, Unit
1716
from ixmp4.data.abstract.optimization import Column
1817

18+
from .base import Creator, Lister, Retriever, Tabulator
19+
1920

2021
class Parameter(BaseModelFacade):
2122
_model: ParameterModel
@@ -99,44 +100,25 @@ def __str__(self) -> str:
99100
return f"<Parameter {self.id} name={self.name}>"
100101

101102

102-
class ParameterRepository(BaseFacade):
103-
_run: Run
104-
103+
class ParameterRepository(
104+
Creator[Parameter, ParameterModel],
105+
Retriever[Parameter, ParameterModel],
106+
Lister[Parameter, ParameterModel],
107+
Tabulator[Parameter, ParameterModel],
108+
):
105109
def __init__(self, _run: Run, **kwargs: Unpack["InitKwargs"]) -> None:
106-
super().__init__(**kwargs)
107-
self._run = _run
110+
super().__init__(_run=_run, **kwargs)
111+
self._backend_repository = self.backend.optimization.parameters
112+
self._model_type = Parameter
108113

109114
def create(
110115
self,
111116
name: str,
112117
constrained_to_indexsets: list[str],
113118
column_names: list[str] | None = None,
114119
) -> Parameter:
115-
model = self.backend.optimization.parameters.create(
120+
return super().create(
116121
name=name,
117-
run_id=self._run.id,
118122
constrained_to_indexsets=constrained_to_indexsets,
119123
column_names=column_names,
120124
)
121-
return Parameter(_backend=self.backend, _model=model)
122-
123-
def get(self, name: str) -> Parameter:
124-
model = self.backend.optimization.parameters.get(run_id=self._run.id, name=name)
125-
return Parameter(_backend=self.backend, _model=model)
126-
127-
def list(self, name: str | None = None) -> Iterable[Parameter]:
128-
parameters = self.backend.optimization.parameters.list(
129-
run_id=self._run.id, name=name
130-
)
131-
return [
132-
Parameter(
133-
_backend=self.backend,
134-
_model=i,
135-
)
136-
for i in parameters
137-
]
138-
139-
def tabulate(self, name: str | None = None) -> pd.DataFrame:
140-
return self.backend.optimization.parameters.tabulate(
141-
run_id=self._run.id, name=name
142-
)

ixmp4/core/optimization/scalar.py

+11-28
Original file line numberDiff line numberDiff line change
@@ -1,22 +1,21 @@
1-
from collections.abc import Iterable
21
from datetime import datetime
32
from typing import TYPE_CHECKING, ClassVar
43

54
if TYPE_CHECKING:
65
from . import InitKwargs
76

8-
import pandas as pd
9-
107
# TODO Import this from typing when dropping Python 3.11
118
from typing_extensions import Unpack
129

13-
from ixmp4.core.base import BaseFacade, BaseModelFacade
10+
from ixmp4.core.base import BaseModelFacade
1411
from ixmp4.core.unit import Unit
1512
from ixmp4.data.abstract import Docs as DocsModel
1613
from ixmp4.data.abstract import Run
1714
from ixmp4.data.abstract import Scalar as ScalarModel
1815
from ixmp4.data.abstract import Unit as UnitModel
1916

17+
from .base import Lister, Retriever, Tabulator
18+
2019

2120
class Scalar(BaseModelFacade):
2221
_model: ScalarModel
@@ -99,12 +98,15 @@ def __str__(self) -> str:
9998
return f"<Scalar {self.id} name={self.name}>"
10099

101100

102-
class ScalarRepository(BaseFacade):
103-
_run: Run
104-
101+
class ScalarRepository(
102+
Retriever[Scalar, ScalarModel],
103+
Lister[Scalar, ScalarModel],
104+
Tabulator[Scalar, ScalarModel],
105+
):
105106
def __init__(self, _run: Run, **kwargs: Unpack["InitKwargs"]) -> None:
106-
super().__init__(**kwargs)
107-
self._run = _run
107+
super().__init__(_run=_run, **kwargs)
108+
self._backend_repository = self.backend.optimization.scalars
109+
self._model_type = Scalar
108110

109111
def create(self, name: str, value: float, unit: str | Unit | None = None) -> Scalar:
110112
if isinstance(unit, Unit):
@@ -127,22 +129,3 @@ def create(self, name: str, value: float, unit: str | Unit | None = None) -> Sca
127129
"run.optimization.scalars.update()?"
128130
) from e
129131
return Scalar(_backend=self.backend, _model=model)
130-
131-
def get(self, name: str) -> Scalar:
132-
model = self.backend.optimization.scalars.get(run_id=self._run.id, name=name)
133-
return Scalar(_backend=self.backend, _model=model)
134-
135-
def list(self, name: str | None = None) -> Iterable[Scalar]:
136-
scalars = self.backend.optimization.scalars.list(run_id=self._run.id, name=name)
137-
return [
138-
Scalar(
139-
_backend=self.backend,
140-
_model=i,
141-
)
142-
for i in scalars
143-
]
144-
145-
def tabulate(self, name: str | None = None) -> pd.DataFrame:
146-
return self.backend.optimization.scalars.tabulate(
147-
run_id=self._run.id, name=name
148-
)

0 commit comments

Comments
 (0)