Add mkdocs tabs option to math docs (#564)

CHANGELOG.md

@@ -2,6 +2,9 @@
 
 ### User-facing changes
 
+|new| `mkdocs_tabbed` option when writing math documentation to file (`calliope.Model.math_documentation.write(...)`) which will add YAML snippets to all rendered math as a separate "tab" if writing to Markdown.
+Requires the [PyMdown tabbed extension](https://facelessuser.github.io/pymdown-extensions/extensions/tabbed/) to render the tabs correctly in an [MkDocs](https://www.mkdocs.org/) project.
+
 |new| List of pre-defined parameters given in the `pre-defined` math documentation, with references back to the constraints/variables/global expressions in which they are defined (either in the `expression` string or the `where` string).
 
 |new| Units and default values for variables and global expressions added to the math documentation.

docs/hooks/generate_math_docs.py

@@ -32,6 +32,10 @@
 The [decision variables](#decision-variables) and [parameters](#parameters) are listed at the end of the page; they also refer back to the global expressions / constraints in which they are used.
 Those parameters which are defined over time (`timesteps`) in the expressions can be defined by a user as a single, time invariant value, or as a timeseries that is [loaded from file or dataframe](../creating/data_sources.md).
 
+!!! note
+
+    For every math component in the documentation, we include the YAML snippet that was used to generate the math in a separate tab.
+
 [:fontawesome-solid-download: Download the {math_type} formulation as a YAML file]({filepath})
 """
 
@@ -110,7 +114,7 @@ def write_file(
 
     nav_reference["Pre-defined math"].append(output_file.as_posix())
 
-    math_doc = model.math_documentation.write(format="md")
+    math_doc = model.math_documentation.write(format="md", mkdocs_tabbed=True)
     file_to_download = Path("..") / filename
     output_full_filepath.write_text(
         PREPEND_SNIPPET.format(

src/calliope/attrdict.py

@@ -328,6 +328,7 @@ def to_yaml(self, path=None):
         yaml_ = ruamel_yaml.YAML()
         yaml_.indent = 2
         yaml_.block_seq_indent = 0
+        yaml_.sort_base_mapping_type_on_output = False
 
         # Numpy objects should be converted to regular Python objects,
         # so that they are properly displayed in the resulting YAML output

src/calliope/backend/backend_model.py

@@ -438,10 +438,12 @@ def _add_to_dataset(
         """
 
         add_attrs = {
-            attr: unparsed_dict.get(attr)
+            attr: unparsed_dict.pop(attr)
             for attr in self._COMPONENT_ATTR_METADATA
             if attr in unparsed_dict.keys()
         }
+        if unparsed_dict:
+            add_attrs["yaml_snippet"] = AttrDict(unparsed_dict).to_yaml()
 
         da.attrs.update(
             {

src/calliope/backend/latex_backend_model.py

@@ -55,17 +55,19 @@ def inputs(self, val: xr.Dataset):
     def write(
         self,
         filename: Literal[None] = None,
+        mkdocs_tabbed: bool = False,
         format: Optional[_ALLOWED_MATH_FILE_FORMATS] = None,
     ) -> str:
         "Expecting string if not giving filename"
 
     @overload
-    def write(self, filename: Union[str, Path]) -> None:
+    def write(self, filename: str | Path, mkdocs_tabbed: bool = False) -> None:
         "Expecting None (and format arg is not needed) if giving filename"
 
     def write(
         self,
         filename: Optional[Union[str, Path]] = None,
+        mkdocs_tabbed: bool = False,
         format: Optional[_ALLOWED_MATH_FILE_FORMATS] = None,
     ) -> Optional[str]:
         """_summary_
@@ -79,6 +81,8 @@ def write(
                 Required if expecting a string return from calling this function. The LaTeX math will be embedded in a document of the given format (tex=LaTeX, rst=reStructuredText, md=Markdown).
                 Defaults to None.
 
+                mkdocs_tabbed (bool, optional): If True and Markdown docs are being generated, the equations will be on a tab and the original YAML math definition will be on another tab.
+
         Raises:
             exceptions.ModelError: Math strings need to be built first (`build`)
             ValueError: The file format (inferred automatically from `filename` or given by `format`) must be one of ["tex", "rst", "md"].
@@ -103,7 +107,7 @@ def write(
             raise ValueError(
                 f"Math documentation format must be one of {allowed_formats}, received `{format}`"
             )
-        populated_doc = self._instance.generate_math_doc(format)
+        populated_doc = self._instance.generate_math_doc(format, mkdocs_tabbed)
 
         if filename is None:
             return populated_doc
@@ -311,11 +315,26 @@ class LatexBackendModel(backend_model.BackendModelGenerator):
     **Default**: {{ equation.default }}
     {% endif %}
     {% if equation.expression != "" %}
+    {% if mkdocs_tabbed and yaml_snippet is not none%}
+
+    === "Math"
+
+        $$
+        {{ equation.expression | trim | escape_underscores | mathify_text_in_text | indent(width=4) }}
+        $$
+
+    === "YAML"
+
+        ```yaml
+        {{ equation.yaml_snippet | trim | indent(width=4) }}
+        ```
+    {% else %}
 
     $$
-    {{ equation.expression | trim | escape_underscores | mathify_text_in_text }}
+    {{ equation.expression | trim | escape_underscores | mathify_text_in_text}}
     $$
     {% endif %}
+    {% endif %}
     {% endfor %}
     {% endfor %}
     """
@@ -460,16 +479,24 @@ def delete_component(
         if key in self._dataset and self._dataset[key].obj_type == component_type:
             del self._dataset[key]
 
-    def generate_math_doc(self, format: _ALLOWED_MATH_FILE_FORMATS = "tex") -> str:
+    def generate_math_doc(
+        self, format: _ALLOWED_MATH_FILE_FORMATS = "tex", mkdocs_tabbed: bool = False
+    ) -> str:
         """Generate the math documentation by embedding LaTeX math in a template.
 
         Args:
             format (Literal["tex", "rst", "md"]):
                 The built LaTeX math will be embedded in a document of the given format (tex=LaTeX, rst=reStructuredText, md=Markdown). Defaults to "tex".
+            mkdocs_tabbed (bool, optional): If True and format is `md`, the equations will be on a tab and the original YAML math definition will be on another tab.
 
         Returns:
             str: Generated math documentation.
         """
+        if mkdocs_tabbed and format != "md":
+            raise ModelError(
+                "Cannot use MKDocs tabs when writing math to a non-Markdown file format."
+            )
+
         doc_template = self.FORMAT_STRINGS[format]
         components = {
             objtype: [
@@ -480,6 +507,7 @@ def generate_math_doc(self, format: _ALLOWED_MATH_FILE_FORMATS = "tex") -> str:
                     "references": list(da.attrs.get("references", set())),
                     "default": da.attrs.get("default", None),
                     "unit": da.attrs.get("unit", None),
+                    "yaml_snippet": da.attrs.get("yaml_snippet", None),
                 }
                 for name, da in getattr(self, objtype).data_vars.items()
                 if "math_string" in da.attrs
@@ -496,7 +524,9 @@ def generate_math_doc(self, format: _ALLOWED_MATH_FILE_FORMATS = "tex") -> str:
         }
         if not components["parameters"]:
             del components["parameters"]
-        return self._render(doc_template, components=components)
+        return self._render(
+            doc_template, mkdocs_tabbed=mkdocs_tabbed, components=components
+        )
 
     def _add_latex_strings(
         self,
@@ -539,6 +569,7 @@ def _generate_math_string(
                 sets=sets,
             )
             where_array.attrs.update({"math_string": equation_element_string})
+
         return None
 
     @staticmethod

tests/test_backend_latex_backend.py

@@ -408,6 +408,55 @@ def test_generate_math_doc_no_params(self, dummy_model_data):
                     """
         )
 
+    def test_generate_math_doc_mkdocs_tabbed(self, dummy_model_data):
+        backend_model = latex_backend_model.LatexBackendModel(dummy_model_data)
+        backend_model.add_global_expression(
+            "expr",
+            {
+                "equations": [{"expression": "1 + 2"}],
+                "description": "foobar",
+                "default": 0,
+            },
+        )
+        doc = backend_model.generate_math_doc(format="md", mkdocs_tabbed=True)
+        assert doc == textwrap.dedent(
+            r"""
+
+                    ## Where
+
+                    ### expr
+
+                    foobar
+
+                    **Default**: 0
+
+                    === "Math"
+
+                        $$
+                        \begin{array}{l}
+                            \quad 1 + 2\\
+                        \end{array}
+                        $$
+
+                    === "YAML"
+
+                        ```yaml
+                        equations:
+                        - expression: 1 + 2
+                        ```
+                    """
+        )
+
+    def test_generate_math_doc_mkdocs_tabbed_not_in_md(self, dummy_model_data):
+        backend_model = latex_backend_model.LatexBackendModel(dummy_model_data)
+        with pytest.raises(exceptions.ModelError) as excinfo:
+            backend_model.generate_math_doc(format="rst", mkdocs_tabbed=True)
+
+        assert check_error_or_warning(
+            excinfo,
+            "Cannot use MKDocs tabs when writing math to a non-Markdown file format.",
+        )
+
     @pytest.mark.parametrize(
         ["kwargs", "expected"],
         [

tests/test_backend_pyomo.py

@@ -1800,20 +1800,28 @@ def test_new_build_get_variable(self, simple_supply):
         assert (
             var.to_series().dropna().apply(lambda x: isinstance(x, pmo.variable)).all()
         )
-        assert var.attrs == {
-            "obj_type": "variables",
-            "references": {
-                "flow_in_max",
-                "flow_out_max",
-                "cost_investment",
-                "cost_investment_flow_cap",
-                "symmetric_transmission",
-            },
-            "description": "A technology's flow capacity, also known as its nominal or nameplate capacity.",
-            "unit": "power",
-            "default": 0,
-            "coords_in_name": False,
+        expected_keys = set(
+            [
+                "obj_type",
+                "references",
+                "description",
+                "unit",
+                "default",
+                "yaml_snippet",
+                "coords_in_name",
+            ]
+        )
+        assert not expected_keys.symmetric_difference(var.attrs.keys())
+        assert var.attrs["obj_type"] == "variables"
+        assert var.attrs["references"] == {
+            "flow_in_max",
+            "flow_out_max",
+            "cost_investment",
+            "cost_investment_flow_cap",
+            "symmetric_transmission",
         }
+        assert var.attrs["default"] == 0
+        assert var.attrs["coords_in_name"] is False
 
     def test_new_build_get_variable_as_vals(self, simple_supply):
         var = simple_supply.backend.get_variable("flow_cap", as_backend_objs=False)
@@ -1855,14 +1863,22 @@ def test_new_build_get_global_expression(self, simple_supply):
             .apply(lambda x: isinstance(x, pmo.expression))
             .all()
         )
-        assert expr.attrs == {
-            "obj_type": "global_expressions",
-            "references": {"cost"},
-            "description": "The installation costs of a technology, including annualised investment costs and annual maintenance costs.",
-            "unit": "cost",
-            "default": 0,
-            "coords_in_name": False,
-        }
+        expected_keys = set(
+            [
+                "obj_type",
+                "references",
+                "description",
+                "unit",
+                "default",
+                "yaml_snippet",
+                "coords_in_name",
+            ]
+        )
+        assert not expected_keys.symmetric_difference(expr.attrs.keys())
+        assert expr.attrs["obj_type"] == "global_expressions"
+        assert expr.attrs["references"] == {"cost"}
+        assert expr.attrs["default"] == 0
+        assert expr.attrs["coords_in_name"] is False
 
     def test_new_build_get_global_expression_as_str(self, simple_supply):
         expr = simple_supply.backend.get_global_expression(
@@ -1886,12 +1902,13 @@ def test_new_build_get_constraint(self, simple_supply):
             .apply(lambda x: isinstance(x, pmo.constraint))
             .all()
         )
-        assert constr.attrs == {
-            "obj_type": "constraints",
-            "references": set(),
-            "description": "Set the global carrier balance of the optimisation problem by fixing the total production of a given carrier to equal the total consumption of that carrier at every node in every timestep.",
-            "coords_in_name": False,
-        }
+        expected_keys = set(
+            ["obj_type", "references", "description", "yaml_snippet", "coords_in_name"]
+        )
+        assert not expected_keys.symmetric_difference(constr.attrs.keys())
+        assert constr.attrs["obj_type"] == "constraints"
+        assert constr.attrs["references"] == set()
+        assert constr.attrs["coords_in_name"] is False
 
     def test_new_build_get_constraint_as_str(self, simple_supply):
         constr = simple_supply.backend.get_constraint(