Rework math handling

CHANGELOG.md

@@ -2,6 +2,8 @@
 
 ### User-facing changes
 
+|new| Math has been removed from `model.math`, and can now be accessed via `model.math.data`.
+
 |new| Direct interface to the Gurobi Python API using `!#yaml config.build.backend: gurobi` or `!#python model.build(backend="gurobi")`.
 Tests show that using the gurobi solver via the Python API reduces peak memory consumption and runtime by at least 30% for the combined model build and solve steps.
 This requires the `gurobipy` package which can be installed with `mamba`: `mamba install gurobi::gurobi`.
@@ -34,6 +36,10 @@ Parameter titles from the model definition schema will also propagate to the mod
 
 ### Internal changes
 
+|new| `ModelMath` is a new helper class to handle math additions, including separate methods for pre-defined math, user-defined math and validation checks.
+
+|changed| `MathDocumentation` has been extracted from `Model`/`LatexBackend`, and now is a postprocessing module which can take models as input.
+
 |new| `gurobipy` is a development dependency that will be added as an optional dependency to the conda-forge calliope feedstock recipe.
 
 |changed| Added any new math dicts defined with `calliope.Model.backend.add_[...](...)` to the backend math dict registry stored in `calliope.Model.backend.inputs.attrs["math"]`.

docs/hooks/generate_math_docs.py

@@ -9,6 +9,7 @@
 from pathlib import Path
 
 import calliope
+from calliope.postprocess.math_documentation import MathDocumentation
 from mkdocs.structure.files import File
 
 logger = logging.getLogger("mkdocs")
@@ -42,31 +43,33 @@ def on_files(files: list, config: dict, **kwargs):
     """Process documentation for pre-defined calliope math files."""
     model_config = calliope.AttrDict.from_yaml(MODEL_PATH)
 
-    base_model = generate_base_math_model()
+    base_documentation = generate_base_math_documentation()
     write_file(
-        "base.yaml",
+        "plan.yaml",
         textwrap.dedent(
             """
         Complete base mathematical formulation for a Calliope model.
         This math is _always_ applied but can be overridden with pre-defined additional math or [your own math][adding-your-own-math-to-a-model].
         """
         ),
-        base_model,
+        base_documentation,
         files,
         config,
     )
 
     for override in model_config["overrides"].keys():
-        custom_model = generate_custom_math_model(base_model, override)
+        custom_documentation = generate_custom_math_documentation(
+            base_documentation, override
+        )
         write_file(
             f"{override}.yaml",
             textwrap.dedent(
                 f"""
-            Pre-defined additional math to apply {custom_model.inputs.attrs['name']} math on top of the [base mathematical formulation][base-math].
+            Pre-defined additional math to apply {custom_documentation.name} math on top of the [base mathematical formulation][base-math].
             This math is _only_ applied if referenced in the `config.init.add_math` list as `{override}`.
             """
             ),
-            custom_model,
+            custom_documentation,
             files,
             config,
         )
@@ -77,7 +80,7 @@ def on_files(files: list, config: dict, **kwargs):
 def write_file(
     filename: str,
     description: str,
-    model: calliope.Model,
+    math_documentation: MathDocumentation,
     files: list[File],
     config: dict,
 ) -> None:
@@ -86,12 +89,10 @@ def write_file(
     Args:
         filename (str): name of produced `.md` file.
         description (str): first paragraph after title.
-        model (calliope.Model): calliope model with the given math.
+        math_documentation (MathDocumentation): calliope math documentation.
         files (list[File]): math files to parse.
         config (dict): documentation configuration.
     """
-    title = model.inputs.attrs["name"] + " math"
-
     output_file = (Path("math") / filename).with_suffix(".md")
     output_full_filepath = Path(TEMPDIR.name) / output_file
     output_full_filepath.parent.mkdir(exist_ok=True, parents=True)
@@ -122,7 +123,9 @@ def write_file(
 
     nav_reference["Pre-defined math"].append(output_file.as_posix())
 
-    math_doc = model.math_documentation.write(format="md", mkdocs_tabbed=True)
+    md_doc = math_documentation.write(format="md", mkdocs_tabbed=True)
+
+    title = math_documentation.name
     file_to_download = Path("..") / filename
     output_full_filepath.write_text(
         PREPEND_SNIPPET.format(
@@ -131,69 +134,69 @@ def write_file(
             math_type=title.lower(),
             filepath=file_to_download,
         )
-        + math_doc
+        + md_doc
     )
 
 
-def generate_base_math_model() -> calliope.Model:
-    """Generate model with documentation for the base math.
-
-    Args:
-        model_config (dict): Calliope model config.
+def generate_base_math_documentation() -> MathDocumentation:
+    """Generate model documentation for the base math.
 
     Returns:
-        calliope.Model: Base math model to use in generating math docs.
+        MathDocumentation: model math documentation with latex backend.
     """
     model = calliope.Model(model_definition=MODEL_PATH)
-    model.math_documentation.build()
-    return model
+    return MathDocumentation(model)
 
 
-def generate_custom_math_model(
-    base_model: calliope.Model, override: str
-) -> calliope.Model:
-    """Generate model with documentation for a pre-defined math file.
+def generate_custom_math_documentation(
+    base_documentation: MathDocumentation, override: str
+) -> MathDocumentation:
+    """Generate model documentation for a pre-defined math file.
 
     Only the changes made relative to the base math will be shown.
 
     Args:
-        base_model (calliope.Model): Calliope model with only the base math applied.
+        base_documentation (MathDocumentation): model documentation with only the base math applied.
         override (str): Name of override to load from the list available in the model config.
+
+    Returns:
+        MathDocumentation: model math documentation with latex backend.
     """
     model = calliope.Model(model_definition=MODEL_PATH, scenario=override)
 
     full_del = []
     expr_del = []
-    for component_group, component_group_dict in model.math.items():
+    for component_group, component_group_dict in model.math.data.items():
         for name, component_dict in component_group_dict.items():
-            if name in base_model.math[component_group]:
+            if name in base_documentation.math.data[component_group]:
                 if not component_dict.get("active", True):
                     expr_del.append(name)
                     component_dict["description"] = "|REMOVED|"
                     component_dict["active"] = True
-                elif base_model.math[component_group].get(name, {}) != component_dict:
+                elif (
+                    base_documentation.math.data[component_group].get(name, {})
+                    != component_dict
+                ):
                     _add_to_description(component_dict, "|UPDATED|")
                 else:
                     full_del.append(name)
             else:
                 _add_to_description(component_dict, "|NEW|")
 
-    model.math_documentation.build()
+    math_documentation = MathDocumentation(model)
     for key in expr_del:
-        model.math_documentation._instance._dataset[key].attrs["math_string"] = ""
+        math_documentation.backend._dataset[key].attrs["math_string"] = ""
     for key in full_del:
-        del model.math_documentation._instance._dataset[key]
-    for var in model.math_documentation._instance._dataset.values():
+        del math_documentation.backend._dataset[key]
+    for var in math_documentation.backend._dataset.values():
         var.attrs["references"] = var.attrs["references"].intersection(
-            model.math_documentation._instance._dataset.keys()
+            math_documentation.backend._dataset.keys()
         )
         var.attrs["references"] = var.attrs["references"].difference(expr_del)
 
-    logger.info(
-        model.math_documentation._instance._dataset["carrier_in"].attrs["references"]
-    )
+    logger.info(math_documentation.backend._dataset["carrier_in"].attrs["references"])
 
-    return model
+    return math_documentation
 
 
 def _add_to_description(component_dict: dict, update_string: str) -> None:

docs/user_defined_math/components.md

@@ -12,7 +12,7 @@ A decision variable in Calliope math looks like this:
 
 ```yaml
 variables:
---8<-- "src/calliope/math/base.yaml:variable"
+--8<-- "src/calliope/math/plan.yaml:variable"
 ```
 
 1. It needs a unique name (`storage_cap` in the example above).
@@ -48,7 +48,7 @@ To not clutter the objective function with all combinations of variables and par
 
 ```yaml
 global_expressions:
---8<-- "src/calliope/math/base.yaml:expression"
+--8<-- "src/calliope/math/plan.yaml:expression"
 ```
 
 Global expressions are by no means necessary to include, but can make more complex linear expressions easier to keep track of and can reduce post-processing requirements.
@@ -74,7 +74,7 @@ Here is an example:
 
 ```yaml
 constraints:
---8<-- "src/calliope/math/base.yaml:constraint"
+--8<-- "src/calliope/math/plan.yaml:constraint"
 ```
 
 1. It needs a unique name (`set_storage_initial` in the above example).
@@ -93,7 +93,7 @@ With your constrained decision variables and a global expression that binds thes
 
 ```yaml
 objectives:
---8<-- "src/calliope/math/base.yaml:objective"
+--8<-- "src/calliope/math/plan.yaml:objective"
 ```
 
 1. It needs a unique name.

src/calliope/backend/__init__.py

@@ -5,23 +5,30 @@
 import xarray as xr
 
 from calliope.backend.gurobi_backend_model import GurobiBackendModel
-from calliope.backend.latex_backend_model import MathDocumentation
+from calliope.backend.latex_backend_model import (
+    ALLOWED_MATH_FILE_FORMATS,
+    LatexBackendModel,
+)
 from calliope.backend.parsing import ParsedBackendComponent
 from calliope.backend.pyomo_backend_model import PyomoBackendModel
 from calliope.exceptions import BackendError
+from calliope.preprocess import ModelMath
 
-MODEL_BACKENDS = ("pyomo",)
+MODEL_BACKENDS = ("pyomo", "gurobi")
 
 if TYPE_CHECKING:
     from calliope.backend.backend_model import BackendModel
 
 
-def get_model_backend(name: str, data: xr.Dataset, **kwargs) -> "BackendModel":
+def get_model_backend(
+    name: str, data: xr.Dataset, math: ModelMath, **kwargs
+) -> "BackendModel":
     """Assign a backend using the given configuration.
 
     Args:
         name (str): name of the backend to use.
         data (Dataset): model data for the backend.
+        math (ModelMath): Calliope math.
         **kwargs: backend keyword arguments corresponding to model.config.build.
 
     Raises:
@@ -32,8 +39,8 @@ def get_model_backend(name: str, data: xr.Dataset, **kwargs) -> "BackendModel":
     """
     match name:
         case "pyomo":
-            return PyomoBackendModel(data, **kwargs)
+            return PyomoBackendModel(data, math, **kwargs)
         case "gurobi":
-            return GurobiBackendModel(data, **kwargs)
+            return GurobiBackendModel(data, math, **kwargs)
         case _:
             raise BackendError(f"Incorrect backend '{name}' requested.")