diff --git a/linopy/__init__.py b/linopy/__init__.py index b813f71d..602e597f 100644 --- a/linopy/__init__.py +++ b/linopy/__init__.py @@ -27,6 +27,7 @@ Constraints, CSRConstraint, ) +from linopy.declarative.build import declarative_model from linopy.expressions import LinearExpression, QuadraticExpression, merge from linopy.io import read_netcdf from linopy.model import Model, Variable, Variables @@ -76,4 +77,5 @@ "read_netcdf", "segments", "tangent_lines", + "declarative_model", ) diff --git a/linopy/declarative/__init__.py b/linopy/declarative/__init__.py new file mode 100644 index 00000000..5000cf44 --- /dev/null +++ b/linopy/declarative/__init__.py @@ -0,0 +1,23 @@ +""" +Linopy declarative math interface. + +Build a linopy model from a declarative math definition and an xarray dataset of input data, via :func:`declarative_model`. + +This directory is adapted from the calliope Apache-2.0 licensed math backend module: +https://github.com/calliope-project/calliope/tree/9916116a06ec8c1feaf3c2606bdb8941b916ce85/src/calliope/backend +""" + +from linopy.declarative.build import DeclarativeModelBuilder, declarative_model +from linopy.declarative.helpers import HelperFunction +from linopy.declarative.latex import LatexModelBuilder, latex_math_doc +from linopy.declarative.schema import ConfigModel, MathModel + +__all__ = [ + "ConfigModel", + "DeclarativeModelBuilder", + "HelperFunction", + "LatexModelBuilder", + "MathModel", + "declarative_model", + "latex_math_doc", +] diff --git a/linopy/declarative/build.py b/linopy/declarative/build.py new file mode 100644 index 00000000..140992d0 --- /dev/null +++ b/linopy/declarative/build.py @@ -0,0 +1,415 @@ +""" +Linopy declarative model-build module. + +This module contains the entry point to build a linopy optimisation model from a declarative math definition and an +xarray dataset of input data. + +This module is adapted from the calliope Apache-2.0 licensed math backend model: +https://github.com/calliope-project/calliope/blob/9916116a06ec8c1feaf3c2606bdb8941b916ce85/src/calliope/backend/backend_model.py +""" + +from __future__ import annotations + +import logging +import time +from collections.abc import Iterable, Iterator, Mapping +from dataclasses import replace +from typing import Any + +import xarray as xr +from tqdm.auto import tqdm + +from linopy.declarative import parsing +from linopy.declarative.helpers import HelperFunction, build_registry +from linopy.declarative.nodes import Component, Context, find_refs +from linopy.declarative.schema import ( + BUILD_ORDER, + DTYPE_OPTIONS, + EQUATION_GROUP_T, + ConfigModel, + ConstraintDef, + ExpressionDef, + MathModel, + ObjectiveDef, + VariableDef, +) +from linopy.expressions import LinearExpression, merge +from linopy.io import TQDM_COLOR +from linopy.model import Model + +LOGGER = logging.getLogger(__name__) + +_SKIP_MESSAGE = "No valid data points after applying mask. Not added to model." + + +def declarative_model( + math_def: dict, + input_data: xr.Dataset, + config: dict, + helpers: Iterable[type[HelperFunction]] = (), +) -> Model: + """ + Build a linopy Model from a declarative math definition and input data. + + Parameters + ---------- + math_def : dict + Declarative math definition (see + :class:`linopy.declarative.schema.MathModel` for the expected structure). + input_data : xr.Dataset + Model input data (parameters, lookups, dimensions). + config : dict + Build configuration options. + helpers : Iterable[type[HelperFunction]], optional + User-defined helper functions to make available in math strings, in + addition to the built-in ones. + + Returns + ------- + Model + The built linopy model, ready to solve. + """ + return DeclarativeModelBuilder(math_def, input_data, config, helpers).build() + + +class _DeclarativeBase: + """Shared validation and context setup of the model and LaTeX builders.""" + + def __init__( + self, + math_def: dict, + input_data: xr.Dataset | None, + config: dict | None, + helpers: Iterable[type[HelperFunction]] = (), + *, + math_reprs: dict[str, str] | None = None, + ) -> None: + """ + Validate the math definition, input data, and config. + + Parameters + ---------- + math_def : dict + Declarative math definition. + input_data : xr.Dataset, optional + Model input data. + config : dict, optional + Build configuration options. + helpers : Iterable[type[HelperFunction]], optional + User-defined helper functions, in addition to the built-in ones. + math_reprs : dict[str, str], optional + Custom LaTeX representations per component name. + """ + self.model = Model() + self.math = MathModel.model_validate(math_def) + self.parsed = parsing.parse_math(self.math) + self.input_data = input_data if input_data is not None else xr.Dataset() + self.config = ConfigModel.model_validate(config or {}) + self._ctx = Context( + model=self.model, + input_data=self.input_data, + math=self.math, + config=self.config, + helpers=build_registry(helpers), + math_reprs=math_reprs or {}, + ) + + def _references(self, parsed: parsing.ParsedComponent) -> list[str]: + """Return the sorted names of all math components a component references.""" + refs = find_refs(parsed.mask, Component) + for equation in parsed.equations: + refs |= equation.references() + return sorted(refs) + + +class DeclarativeModelBuilder(_DeclarativeBase): + """Builder turning a declarative math definition into a linopy Model.""" + + def __init__( + self, + math_def: dict, + input_data: xr.Dataset, + config: dict, + helpers: Iterable[type[HelperFunction]] = (), + ) -> None: + """ + Validate the math definition, input data, and config, ready to `build()`. + + Parameters + ---------- + math_def : dict + Declarative math definition. + input_data : xr.Dataset + Model input data. + config : dict + Build configuration options. + helpers : Iterable[type[HelperFunction]], optional + User-defined helper functions, in addition to the built-in ones. + """ + super().__init__(math_def, input_data, config, helpers) + self.input_data = self._update_dtypes(self.input_data) + self._ctx = replace(self._ctx, input_data=self.input_data) + self._check_inputs() + + def _update_dtypes(self, ds: xr.Dataset, id_: str = "") -> xr.Dataset: + """ + Coerce dataset variables to the dtypes given by their math definitions. + + Variables not defined in the math are left unchanged (with an INFO log); + datetime/date variables pass through uncoerced. + """ + prefix = f"{id_} | " if id_ else "" + for var_name, var_data in ds.items(): + try: + math_def = self.math.find( + str(var_name), subset=["lookups", "parameters", "dimensions"] + ) + except KeyError: + LOGGER.info( + f"{prefix}input data `{var_name}` not defined in model math; " + "it will not be available in the optimisation problem." + ) + continue + + dtype_str: str = math_def["dtype"] + if dtype_str in ("datetime", "date"): + continue + dtype = DTYPE_OPTIONS[dtype_str] + LOGGER.debug( + f"{prefix}{math_def._group} | Updating values of `{var_name}` to {dtype_str} type" + ) + match dtype_str: + case "string": + updated_var = ( + var_data.astype(dtype) + .where(var_data.notnull()) + .where(var_data != "") + ) + case "bool": + updated_var = var_data.fillna(False).astype(dtype) + case _: + updated_var = var_data.astype(dtype) + + ds[var_name] = updated_var + return ds + + def _check_inputs(self) -> None: + """Run the math's input-data checks, warning or raising on triggered ones.""" + warn_msgs: list[str] = [] + error_msgs: list[str] = [] + active = self.input_data.get("active", xr.DataArray(True)) + for name in self.math.checks._active: + check = self.math.checks[name] + mask_node = self.parsed.checks[name] + check_ctx = replace(self._ctx, mode="mask", equation_name=name) + evaluated = mask_node.evaluate(check_ctx) + if (evaluated & active).any(): + messages = error_msgs if check.errors == "raise" else warn_msgs + messages.append(check.message) + + for param in self.math.parameters._active: + param_math = self.math.parameters[param] + param_data = self.input_data.get(param, None) + if param_data is None or param_math.dims is None: + continue + if set(param_data.dims).difference(param_math.dims): + error_msgs.append( + f"Parameter `{param}` has dimensions {param_data.dims}, " + f"but math definition expects max {param_math.dims}." + ) + + for lookup in self.math.lookups._active: + lookup_math = self.math.lookups[lookup] + lookup_data = self.input_data.get(lookup, None) + if lookup_data is None or lookup_math.dims is None: + continue + if set(lookup_data.dims).difference(lookup_math.dims): + error_msgs.append( + f"Lookup `{lookup}` has dimensions {lookup_data.dims}, " + f"but math definition expects max {lookup_math.dims}." + ) + + if warn_msgs: + bullets = "\n".join(f" * {msg}" for msg in sorted(set(warn_msgs))) + LOGGER.info( + f"Possible issues found during model input data checks:\n{bullets}" + ) + if error_msgs: + bullets = "\n".join(f" * {msg}" for msg in sorted(set(error_msgs))) + raise ValueError(f"Errors during model input data checks:\n{bullets}") + + @staticmethod + def _sorted_by_order(root: Mapping[str, Any]) -> list[tuple[str, Any]]: + """Return (name, definition) pairs from a root mapping, sorted by definition order.""" + return sorted(root.items(), key=lambda item: getattr(item[1], "order", 0)) + + def _iter_equations( + self, + equations: list[parsing.Equation], + group: EQUATION_GROUP_T, + mask: xr.DataArray, + ) -> Iterator[tuple[parsing.Equation, xr.DataArray]]: + """ + Yield each parsed equation with its evaluated, foreach-aligned sub-mask. + + Equations whose mask leaves no valid data point are skipped (with an INFO log). + """ + for equation in equations: + sub_mask = parsing.as_mask(equation, self._ctx, initial_mask=mask) + if not sub_mask.any(): + LOGGER.info(f"{group}:{equation.name} | {_SKIP_MESSAGE}") + continue + yield equation, parsing.drop_dims_not_in_foreach(sub_mask, equation.sets) + + def add_variable(self, name: str, definition: VariableDef) -> None: + """Add a decision variable to the model, masked by its math definition.""" + parsed = self.parsed["variables"][name] + mask = parsing.component_mask( + "variables", name, definition, parsed.mask, self._ctx + ) + if not mask.any(): + LOGGER.info(f"variables:{name} | {_SKIP_MESSAGE}") + return + self.model.add_variables( + coords=mask.coords, + name=name, + mask=mask, + upper=definition.bounds.upper, + lower=definition.bounds.lower, + integer=definition.domain == "integer", + ) + # Variable.attrs values are typed Hashable, but a sorted list serializes best. + self.model.variables[name].attrs["references"] = self._references(parsed) # type: ignore[assignment] + + def add_expression(self, name: str, definition: ExpressionDef) -> None: + """Add a named expression to the model, merging its equation variants.""" + parsed = self.parsed["expressions"][name] + mask = parsing.component_mask( + "expressions", name, definition, parsed.mask, self._ctx + ) + if not mask.any(): + LOGGER.info(f"expressions:{name} | {_SKIP_MESSAGE}") + return + expr: Any = LinearExpression(float("nan"), self.model).where(mask) + filled = xr.DataArray(False) + for equation, sub_mask in self._iter_equations( + parsed.equations, "expressions", mask + ): + if (filled & sub_mask).any(): + raise ValueError( + f"expressions:{name} | Overlapping 'mask' conditions between " + "equations are not allowed. Please revise the 'mask' conditions " + "to ensure they are mutually exclusive." + ) + filled = filled | sub_mask + expr_to_fill = parsing.as_expression(equation, self._ctx, mask=sub_mask) + expr = merge([expr, expr_to_fill.where(sub_mask)]) + if not filled.any(): + LOGGER.info(f"expressions:{name} | {_SKIP_MESSAGE}") + return + self.model.add_expressions(name=name, data=expr, mask=mask) + self.model.expressions[name].attrs["references"] = self._references(parsed) + + def add_constraint(self, name: str, definition: ConstraintDef) -> None: + """Add a constraint to the model, merging its equation variants.""" + parsed = self.parsed["constraints"][name] + mask = parsing.component_mask( + "constraints", name, definition, parsed.mask, self._ctx + ) + if not mask.any(): + LOGGER.info(f"constraints:{name} | {_SKIP_MESSAGE}") + return + lhs: Any = LinearExpression(float("nan"), self.model).where(mask) + rhs: Any = LinearExpression(float("nan"), self.model).where(mask) + sign = xr.DataArray().where(mask) + for equation, sub_mask in self._iter_equations( + parsed.equations, "constraints", mask + ): + if (sign.notnull() & sub_mask).any(): + raise ValueError( + f"constraints:{name} | Overlapping 'mask' conditions between " + "equations are not allowed. Please revise the 'mask' conditions " + "to ensure they are mutually exclusive." + ) + lhs_to_fill, sign_to_fill, rhs_to_fill = parsing.as_constraint( + equation, self._ctx, mask=sub_mask + ) + lhs = merge([lhs, lhs_to_fill]) + rhs = merge([rhs, rhs_to_fill]) + sign = sign.fillna(sign_to_fill) + + if sign.isnull().all(): + LOGGER.info(f"constraints:{name} | {_SKIP_MESSAGE}") + return + self.model.add_constraints( + coords=mask.coords, + name=name, + lhs=lhs, + # Default to equality to avoid errors on masked-out points. + sign=sign.fillna("=="), + rhs=rhs, + mask=mask, + ) + self.model.constraints[name].attrs["references"] = self._references(parsed) + + def add_objective(self, name: str, definition: ObjectiveDef) -> None: + """Set the model objective, merging its equation variants.""" + parsed = self.parsed["objectives"][name] + mask = parsing.component_mask( + "objectives", name, definition, parsed.mask, self._ctx + ) + if not mask.any(): + LOGGER.info(f"objectives:{name} | {_SKIP_MESSAGE}") + return + pieces: list[tuple[LinearExpression, xr.DataArray]] = [] + filled = xr.DataArray(False) + for equation, sub_mask in self._iter_equations( + parsed.equations, "objectives", mask + ): + if (filled & sub_mask).any(): + raise ValueError( + f"objectives:{name} | Overlapping 'mask' conditions between " + "equations are not allowed. Please revise the 'mask' conditions " + "to ensure they are mutually exclusive." + ) + filled = filled | sub_mask + pieces.append( + (parsing.as_expression(equation, self._ctx, mask=sub_mask), sub_mask) + ) + if not pieces: + LOGGER.info(f"objectives:{name} | {_SKIP_MESSAGE}") + return + expr: Any = pieces[0][0] + if len(pieces) > 1: + expr = merge([piece.where(sub_mask) for piece, sub_mask in pieces]) + self.model.add_objective(expr=expr, sense=definition.sense) + self.model.objective.attrs["references"] = self._references(parsed) + + def build(self) -> Model: + """ + Build all math components into the linopy model. + + Components are built in group order (variables, expressions, constraints, + objectives) and, within a group, by their `order` attribute where defined. + + Returns + ------- + Model + The built linopy model. + """ + active_objectives = list(self.math.objectives._active) + if len(active_objectives) > 1: + raise ValueError( + f"Only one active objective is supported, found: {active_objectives}" + ) + for group in BUILD_ORDER: + component = group.removesuffix("s") + ordered_items = self._sorted_by_order(self.math[group]._active) + for name, definition in tqdm( + ordered_items, desc=f"Building {group}.", colour=TQDM_COLOR + ): + start = time.time() + getattr(self, f"add_{component}")(name, definition) + LOGGER.debug(f"{group}:{name} | Built in {time.time() - start:.4f}s") + LOGGER.info(f"{group} | Generated.") + return self.model diff --git a/linopy/declarative/grammar.py b/linopy/declarative/grammar.py new file mode 100644 index 00000000..ea7a13c1 --- /dev/null +++ b/linopy/declarative/grammar.py @@ -0,0 +1,436 @@ +""" +Linopy declarative math grammar module. + +This module contains the pyparsing grammars that turn declarative math strings into ASTs of :mod:`linopy.declarative.nodes` node objects. +Parse actions are the node classes' `from_tokens` classmethods; all evaluation logic lives on the nodes themselves. + +The infix-notation grammar structure is adapted from the pyparsing MIT licensed `eval_arith.py` example: +https://github.com/pyparsing/pyparsing/blob/master/examples/eval_arith.py. + +This module is adapted from the calliope Apache-2.0 licensed math parsers: +- https://github.com/calliope-project/calliope/blob/9916116a06ec8c1feaf3c2606bdb8941b916ce85/src/calliope/backend/expression_parser.py +- https://github.com/calliope-project/calliope/blob/9916116a06ec8c1feaf3c2606bdb8941b916ce85/src/calliope/backend/where_parser.py +""" + +from __future__ import annotations + +from functools import cache + +import pyparsing as pp + +from linopy.declarative.nodes import ( + COMPONENT_CATEGORY_T, + Arith, + Call, + Compare, + Component, + ConfigRef, + Constant, + ListNode, + Node, + Sliced, + SliceRef, + SubExprRef, + Subset, + Unary, + find_refs, + iter_nodes, +) + +__all__ = [ + "COMPONENT_CATEGORY_T", + "EQUATION_OPERATORS", + "MASK_OPERATORS", + "REFERENCE_CLASSIFIER", + "Arith", + "Call", + "Compare", + "Component", + "ConfigRef", + "Constant", + "ListNode", + "Node", + "SliceRef", + "Sliced", + "SubExprRef", + "Subset", + "Unary", + "arithmetic_grammar", + "equation_grammar", + "find_refs", + "iter_nodes", + "mask_grammar", + "slice_grammar", + "sub_expression_grammar", +] + +pp.ParserElement.enable_packrat() + +REFERENCE_CLASSIFIER = "$" +"""Prefix marking a reference to a sub-expression (`$foo`) or slicer (`x[dim=$foo]`).""" + +EQUATION_OPERATORS = ("<=", ">=", "=") +"""Comparison operators allowed in constraint equations.""" + +MASK_OPERATORS = ("<", ">", "==", ">=", "<=") +"""Comparison operators allowed in mask strings.""" + + +# --------------------------------------------------------------------------- +# Grammar primitives +# --------------------------------------------------------------------------- + + +def _base_elements() -> tuple[pp.ParserElement, pp.ParserElement]: + """Return the (number, identifier) primitives shared by all grammars.""" + inf_kw = pp.Combine(pp.Opt(pp.Suppress(".")) + pp.Keyword("inf", caseless=True)) + number = (pp.pyparsing_common.number | inf_kw).set_parse_action(Constant.number) + identifier = ~inf_kw + pp.Word(pp.alphas, pp.alphanums + "_") + return number, identifier + + +def _names_parser(names: frozenset[str]) -> pp.ParserElement: + """Return a keyword parser matching any of `names` (or nothing if empty).""" + if not names: + return pp.NoMatch() + return pp.one_of(sorted(names), as_keyword=True) + + +def _component_parser( + names: frozenset[str], category: COMPONENT_CATEGORY_T = "any" +) -> pp.ParserElement: + """Return a parser matching any name in `names` as a :class:`Component`.""" + return _names_parser(names).set_parse_action(Component.from_tokens_as(category)) + + +def _string_parser( + identifier: pp.ParserElement, excluded_names: frozenset[str] +) -> pp.ParserElement: + """Return a parser for generic strings that are not in `excluded_names`.""" + return (~_names_parser(excluded_names) + identifier).set_parse_action( + Constant.string + ) + + +def _list_parser(*items: pp.ParserElement) -> pp.ParserElement: + """Return a parser for `[item, item, ...]` lists of the given item parsers.""" + element = pp.MatchFirst(items) + id_list = pp.Suppress("[") + pp.DelimitedList(element) + pp.Suppress("]") + return id_list.set_parse_action(ListNode.from_tokens) + + +def _call_parser( + *args: pp.ParserElement, + identifier: pp.ParserElement, + allow_nested_calls: bool = False, +) -> pp.ParserElement: + """ + Return a parser for helper-function calls `name(*args, **kwargs)`. + + Parameters + ---------- + *args : pp.ParserElement + Parsers for allowed argument values, matched in the given order. + identifier : pp.ParserElement + Parser for the function name (no parse action attached). + allow_nested_calls : bool, default: False + If True, calls may appear directly as arguments of other calls. + (Calls nested via arithmetic are enabled by passing the arithmetic + parser in `args` instead.) + """ + call = pp.Forward() + allowed_args = list(args) + if allow_nested_calls: + allowed_args.insert(0, call) + + func_name = pp.Combine(identifier + pp.Suppress("("))("func") + arg_value = pp.MatchFirst(allowed_args) + pp.NotAny("=") + arg_list = pp.Group(pp.DelimitedList(arg_value.copy()))("args") + key = identifier + pp.Suppress("=") + kwarg_list = pp.Group(pp.DelimitedList(pp.dict_of(key, arg_value)))("kwargs") + call_args = arg_list + pp.Suppress(",") + kwarg_list | pp.Opt( + arg_list, default=[] + ) + pp.Opt(kwarg_list, default={}) + + call <<= func_name + call_args + pp.Suppress(")") + return call.set_parse_action(Call.from_tokens) + + +def _sliced_component_parser( + slicers: list[pp.ParserElement], + identifier: pp.ParserElement, + component: pp.ParserElement, + allow_slice_references: bool = True, +) -> pp.ParserElement: + """ + Return a parser for sliced components `name[dim=slicer, ...]`. + + Parameters + ---------- + slicers : list[pp.ParserElement] + Parsers for allowed slicer values, matched in the given order. + identifier : pp.ParserElement + Parser for the slice dimension name (no parse action attached). + component : pp.ParserElement + Parser for the sliced component name. + allow_slice_references : bool, default: True + If True, allow `$name` slicer references (e.g. `$bar` in `foo[bars=$bar]`). + """ + slicer: pp.ParserElement = pp.MatchFirst(slicers) + if allow_slice_references: + slice_ref = pp.Suppress(REFERENCE_CLASSIFIER) + identifier + slice_ref.set_parse_action(SliceRef.from_tokens) + slicer = slice_ref | slicer + + one_slice = pp.Group( + identifier("set_name") + pp.Suppress("=") + pp.Group(slicer)("slicer") + ) + slices = pp.Group(pp.DelimitedList(one_slice))("slices") + sliced = pp.Combine(component("obj") + pp.Suppress("[")) + slices + pp.Suppress("]") + return sliced.set_parse_action(Sliced.from_tokens) + + +def _sub_expression_ref_parser(identifier: pp.ParserElement) -> pp.ParserElement: + """Return a parser for `$name` sub-expression references.""" + ref = pp.Combine(pp.Suppress(REFERENCE_CLASSIFIER) + identifier) + return ref.set_parse_action(SubExprRef.from_tokens) + + +def _arithmetic_rules( + *operands: pp.ParserElement, arithmetic: pp.Forward | None = None +) -> pp.Forward: + """ + Return an infix-notation parser combining `operands` with `+ - * / **`. + + Parameters + ---------- + *operands : pp.ParserElement + Parsers for allowed operands, matched in the given order. + arithmetic : pp.Forward, optional + If given, attach the rules to this existing forward-declared rule so + that operands (e.g. function calls) can recursively contain arithmetic. + """ + signop = pp.one_of(["+", "-"]) + multop = pp.one_of(["*", "/"]) + expop = pp.Literal("**") + if arithmetic is None: + arithmetic = pp.Forward() + arithmetic <<= pp.infix_notation( + # the order matters if two could capture the same string, e.g. "inf". + pp.MatchFirst(operands), + [ + (signop, 1, pp.opAssoc.RIGHT, Unary.from_tokens), + (expop, 2, pp.opAssoc.LEFT, Arith.from_tokens), + (multop, 2, pp.opAssoc.LEFT, Arith.from_tokens), + (signop, 2, pp.opAssoc.LEFT, Arith.from_tokens), + ], + ) + return arithmetic + + +# --------------------------------------------------------------------------- +# Grammar entry points +# --------------------------------------------------------------------------- + + +@cache +def _expression_grammar( + component_names: frozenset[str], + *, + arithmetic: bool = True, + slice_refs: bool = True, + sub_expr_refs: bool = False, +) -> pp.ParserElement: + """ + Return an expression-family grammar for the given component names. + + Parameters + ---------- + component_names : frozenset[str] + Valid math component names, to separate them from generic strings. + arithmetic : bool, default: True + If True, combine operands with infix `+ - * / **` rules (with calls + allowed to recursively contain arithmetic). If False, return a flat + single-operand grammar (used for slicer strings). + slice_refs : bool, default: True + If True, allow `$name` slicer references inside slice brackets. + sub_expr_refs : bool, default: False + If True, allow `$name` sub-expression references as operands. + """ + number, identifier = _base_elements() + string = _string_parser(identifier, component_names) + component = _component_parser(component_names) + call_list = _list_parser(number, string, component) + slicer_list = _list_parser(number, string) + sliced = _sliced_component_parser( + [number, string, slicer_list], + identifier, + component, + allow_slice_references=slice_refs, + ) + if not arithmetic: + call = _call_parser( + sliced, + component, + number, + call_list, + string, + identifier=identifier, + allow_nested_calls=True, + ) + return call | sliced | component | number | slicer_list | string + arith = pp.Forward() + call = _call_parser(arith, call_list, string, identifier=identifier) + operands = [call] + if sub_expr_refs: + operands.append(_sub_expression_ref_parser(identifier)) + operands += [sliced, number, component] + return _arithmetic_rules(*operands, arithmetic=arith) + + +def slice_grammar(component_names: frozenset[str]) -> pp.ParserElement: + """ + Return the grammar for named slicer expressions. + + Slicers are linked into equations by `$name` references inside slice brackets + (e.g. `$bar` in `foo[bars=$bar]`). Unlike sub-expressions and equations, + slicer strings allow neither arithmetic nor references to other slicers. + + Parameters + ---------- + component_names : frozenset[str] + Valid math component names, to separate them from generic strings. + """ + return _expression_grammar(component_names, arithmetic=False, slice_refs=False) + + +def sub_expression_grammar(component_names: frozenset[str]) -> pp.ParserElement: + """ + Return the grammar for named sub-expressions. + + Sub-expressions are linked into equations by `$name` references. They allow + arbitrarily nested arithmetic and function calls and `$name` slicer + references, but no references to other sub-expressions. + + Parameters + ---------- + component_names : frozenset[str] + Valid math component names, to separate them from generic strings. + """ + return _expression_grammar(component_names) + + +def arithmetic_grammar(component_names: frozenset[str]) -> pp.ParserElement: + """ + Return the grammar for arithmetic expressions (`+ - * / **`). + + Allows arbitrarily nested arithmetic and function calls, and references to + sub-expressions (`$name`) and slicers (`x[dim=$name]`). + + Parameters + ---------- + component_names : frozenset[str] + Valid math component names, to separate them from generic strings. + """ + return _expression_grammar(component_names, sub_expr_refs=True) + + +@cache +def equation_grammar(component_names: frozenset[str]) -> pp.ParserElement: + """ + Return the grammar for equations of the form `LHS OPERATOR RHS`. + + Each side is an arithmetic expression (see :func:`arithmetic_grammar`) and + the operator is one of `<=`, `>=`, `=`. + + Parameters + ---------- + component_names : frozenset[str] + Valid math component names, to separate them from generic strings. + """ + arithmetic = arithmetic_grammar(component_names) + equation = arithmetic + pp.one_of(list(EQUATION_OPERATORS)) + arithmetic + return equation.set_parse_action(Compare.from_tokens) + + +@cache +def mask_grammar( + dimensions: frozenset[str], inputs: frozenset[str], results: frozenset[str] +) -> pp.ParserElement: + """ + Return the grammar for boolean mask ("where") strings. + + Masks combine existence conditions on model data, comparisons, dimension + subsets, helper functions, booleans, and config options with + `not` / `and` / `or` operators. + + Parameters + ---------- + dimensions : frozenset[str] + Valid dimension names. + inputs : frozenset[str] + Valid parameter/lookup names. + results : frozenset[str] + Valid variable/expression names. + """ + all_names = dimensions | inputs | results + number, identifier = _base_elements() + dimension = _component_parser(dimensions, "dimension") + input_ = _component_parser(inputs, "input") + result = _component_parser(results, "result") + config_option = (pp.Suppress("config.") + identifier).set_parse_action( + ConfigRef.from_tokens + ) + bool_operand = ( + pp.Keyword("True", caseless=True) | pp.Keyword("False", caseless=True) + ).set_parse_action(Constant.boolean) + unique_string = _string_parser(identifier, all_names) + general_string = _string_parser(identifier, frozenset()) + id_list = _list_parser(number, unique_string, dimension) + + subset_items = pp.Group( + pp.DelimitedList(pp.MatchFirst([config_option, number, general_string])) + ) + subset = ( + pp.Suppress("[") + + subset_items + + pp.Suppress("]") + + pp.Suppress(pp.White(" ", min=1)) + + pp.Suppress("in") + + pp.Suppress(pp.White(" ", min=1)) + + pp.MatchFirst([dimension, input_]) + ).set_parse_action(Subset.from_tokens) + + arithmetic = pp.Forward() + comparison_call = _call_parser( + unique_string, number, id_list, arithmetic, identifier=identifier + ) + _arithmetic_rules( + comparison_call, number, dimension, input_, config_option, arithmetic=arithmetic + ) + comparison = ( + arithmetic + + pp.one_of(list(MASK_OPERATORS)) + + pp.MatchFirst([comparison_call, bool_operand, number, general_string]) + ).set_parse_action(Compare.from_tokens) + + call = _call_parser( + unique_string, + number, + id_list, + dimension, + input_, + result, + config_option, + identifier=identifier, + ) + + notop = pp.Keyword("not", caseless=True) + andorop = pp.Keyword("and", caseless=True) | pp.Keyword("or", caseless=True) + return pp.infix_notation( + pp.MatchFirst([bool_operand, comparison, call, subset, input_, result]), + [ + (notop, 1, pp.opAssoc.RIGHT, Unary.from_tokens), + (andorop, 2, pp.opAssoc.LEFT, Arith.from_tokens), + ], + ) diff --git a/linopy/declarative/helpers.py b/linopy/declarative/helpers.py new file mode 100644 index 00000000..1ef3140e --- /dev/null +++ b/linopy/declarative/helpers.py @@ -0,0 +1,661 @@ +""" +Linopy declarative math helper-functions module. + +This module contains: +- the helper functions that can be called in declarative math `mask` and `expression` strings using their `NAME`; +- the abstract base class from which users can define their own helpers, and; +- the registry builder that makes them available to a model build. + +This module is adapted from the calliope Apache-2.0 licensed helper function module: +- https://github.com/calliope-project/calliope/blob/9916116a06ec8c1feaf3c2606bdb8941b916ce85/src/calliope/backend/helper_functions.py +""" + +from __future__ import annotations + +import functools +import re +from abc import ABC, abstractmethod +from collections.abc import Iterable, Mapping +from typing import TYPE_CHECKING, Any, ClassVar, Literal + +import numpy as np +import xarray as xr + +from linopy.declarative.schema import DTYPE_OPTIONS, MathModel +from linopy.expressions import LinearExpression + +if TYPE_CHECKING: + from linopy.declarative.nodes import Context + +KIND_T = Literal["mask", "expression"] + + +def dim_iterator(math: MathModel, dim: str) -> str: + """ + Return the LaTeX iterator name of dimension `dim`. + + Falls back to the dimension name itself when the dimension is not declared + in the math or declares no iterator. + """ + if dim in math.dimensions.root: + iterator = math.dimensions[dim].iterator + if iterator != "NEEDS_ITERATOR": + return iterator + return dim + + +def _to_str_list(vals: Any) -> list[str]: + """Force `vals` to a list of strings, extracting names from any DataArray items.""" + if not isinstance(vals, list): + vals = [vals] + return [str(i.name) if isinstance(i, xr.DataArray) else str(i) for i in vals] + + +def _update_iterator( + instring: str, + iterator_converter: dict[str, str], + method: Literal["add", "replace"], +) -> str: + r""" + Update iterators in the iterator substring of a LaTeX component string. + + Find an iterator in the iterator substring of the component string (anything + wrapped in `_\text{}`, e.g. the standalone `foo` in + `\textit{my_param}_\text{bar,foo,foo=bar,foo+1}`) and append to it + (`method="add"`) or replace it (`method="replace"`). + + Parameters + ---------- + instring : str + String in which the iterator substring can be found. + iterator_converter : dict[str, str] + Mapping from the iterator to search for to the string to append/replace. + method : Literal["add", "replace"] + Whether to append to the iterator or replace it entirely. + """ + + def _replace_in_iterator(matched: re.Match) -> str: + new_iterators = [] + for it in matched.group(2).split(","): + if it in iterator_converter: + it = ( + it + iterator_converter[it] + if method == "add" + else iterator_converter[it] + ) + new_iterators.append(it) + return matched.group(1) + ",".join(new_iterators) + matched.group(3) + + return re.sub(r"(_\\text{)([^{}]*?)(})", _replace_in_iterator, instring) + + +class HelperFunction(ABC): + """ + Abstract base class of all declarative math helper functions. + + Subclasses must define the class attributes `NAME` (the function name used in + math strings) and `ALLOWED_IN` (the string types the function can be called + from), and implement :meth:`as_math_string` and :meth:`as_raw`. To make a + custom helper available to a model build, pass the subclass to + :func:`linopy.declarative.build.declarative_model` via its `helpers` argument. + """ + + NAME: ClassVar[str] + """Helper function name used in math `mask`/`expression` strings.""" + + ALLOWED_IN: ClassVar[list[KIND_T]] + """The parseable string types this function can be called from.""" + + ignore_mask: ClassVar[bool] = False + """If True, `mask` arrays are not applied to the function's incoming arguments.""" + + def __init__(self, context: Context) -> None: + """ + Initialise the helper. + + Parameters + ---------- + context : Context + The evaluation context (input data, math definition, config, ...). + """ + self._context = context + + @abstractmethod + def as_math_string(self, *args: Any, **kwargs: Any) -> str: + """Return a LaTeX math string that includes the action applied by this function.""" + + @abstractmethod + def as_raw(self, *args: Any, **kwargs: Any) -> LinearExpression | xr.DataArray: + """Apply the helper function to its evaluated argument arrays.""" + + def as_expr(self, *args: Any, **kwargs: Any) -> LinearExpression | xr.DataArray: + """ + Apply the helper function on the expression route. + + By default this delegates to :meth:`as_raw`; helpers that need distinct + behaviour on the expression route (e.g. to return a linopy expression) + should override it. + """ + return self.as_raw(*args, **kwargs) + + def _dim_iterator(self, dim: str) -> str: + """Return the LaTeX iterator name of dimension `dim`.""" + return dim_iterator(self._context.math, dim) + + def _instr(self, dim: str) -> str: + r"""Return the LaTeX "iterator in dimension" string (e.g. `\text{n} \in \text{node}`).""" + return rf"\text{{{self._dim_iterator(dim)}}} \in \text{{{dim}}}" + + def _get_dims_from_iterators(self, instring: str) -> list[str]: + """Return the dimensions whose iterators appear in a LaTeX component string.""" + math = self._context.math + + def _extract_dims(matched: re.Match) -> str: + return ",".join( + dim_name + for i in matched.group(2).split(",") + for dim_name in math.dimensions.root + if dim_iterator(math, dim_name) == i + ) + + dims = re.sub(r"^.*(_\\text{)([^{}]*?)(})", _extract_dims, instring) + return [dim for dim in dims.split(",") if dim] + + +class MaskAny(HelperFunction): + """Apply `any` over dimension(s) in a `mask` string.""" + + # Class name doesn't match NAME to avoid a clash with typing.Any + NAME = "any" + ALLOWED_IN: ClassVar[list[KIND_T]] = ["mask"] + + def as_math_string( # noqa: D102, override + self, array: str, *, over: str | list[str | xr.DataArray] + ) -> str: + overstring = r" \\ ".join(self._instr(i) for i in _to_str_list(over)) + # Using bigvee for "collective-or" + return rf"\bigvee\limits_{{\substack{{{overstring}}}}} ({array})" + + def as_raw( + self, input_component: xr.DataArray, *, over: xr.DataArray | list[xr.DataArray] + ) -> xr.DataArray: + """ + Reduce a boolean mask array by applying `any` over some dimension(s). + + Parameters + ---------- + input_component : xr.DataArray + Boolean array to reduce. + over : xr.DataArray | list[xr.DataArray] + Dimension(s) over which to apply `any`. + + Returns + ------- + xr.DataArray + Boolean array with dimensions reduced by applying a boolean OR + operation along the dimensions given in `over`. + """ + if input_component.dtype.kind != "b": + raise ValueError( + "Input to `any` must be a boolean array. " + f"Received {input_component.name} of dtype {input_component.dtype}" + ) + available_dims = set(input_component.dims).intersection(_to_str_list(over)) + return input_component.any(dim=available_dims, keep_attrs=True) + + +class Sum(HelperFunction): + """Apply a summation over dimension(s) in math expressions.""" + + NAME = "sum" + ALLOWED_IN: ClassVar[list[KIND_T]] = ["expression", "mask"] + + def as_math_string( # noqa: D102, override + self, array: str, *, over: str | list[str | xr.DataArray] + ) -> str: + overstring = r" \\ ".join(self._instr(i) for i in _to_str_list(over)) + return rf"\sum\limits_{{\substack{{{overstring}}}}} ({array})" + + def as_raw( + self, array: xr.DataArray, *, over: xr.DataArray | list[xr.DataArray] + ) -> xr.DataArray: + """ + Sum an expression array over the given dimension(s). + + Parameters + ---------- + array : xr.DataArray + Expression array. + over : xr.DataArray | list[xr.DataArray] + Dimension(s) over which to apply `sum`. + + Returns + ------- + xr.DataArray + Array with dimensions reduced by summing over the dimensions given + in `over`; dimensions not present in `array` are ignored. + """ + filtered_over = set(_to_str_list(over)).intersection(array.dims) + return array.sum(filtered_over) + + +class SelectFromLookupArrays(HelperFunction): + """N-dimensional vectorised indexing via lookup arrays.""" + + NAME = "select_from_lookup_arrays" + ALLOWED_IN: ClassVar[list[KIND_T]] = ["expression"] + + def as_math_string(self, array: str, **lookup_arrays: str) -> str: # noqa: D102, override + new_strings = { + (iterator := self._dim_iterator(dim)): rf"={array}[{iterator}]" + for dim, array in lookup_arrays.items() + } + return _update_iterator(array, new_strings, "add") + + def as_raw( + self, array: xr.DataArray, **lookup_arrays: xr.DataArray + ) -> xr.DataArray: + """ + Apply vectorised indexing on an arbitrary number of an input array's dimensions. + + Parameters + ---------- + array : xr.DataArray + Array on which to apply vectorised indexing. + **lookup_arrays : xr.DataArray + Keys are dimensions on which to apply vectorised indexing; values are + arrays whose values are either NaN or values from that dimension. + + Returns + ------- + xr.DataArray + `array` with rearranged values (coordinates remain unchanged). + Any NaN index coordinates in the lookup arrays will be NaN in the + returned array. + + Raises + ------ + ValueError + If `array` or any lookup array is not indexed over all the + dimensions given in the `lookup_arrays` keys. + """ + # Inspired by https://github.com/pydata/xarray/issues/1553#issuecomment-748491929 + # Reindex does not presently support vectorized lookups: https://github.com/pydata/xarray/issues/1553 + # Sel does (e.g. https://github.com/pydata/xarray/issues/4630) but can't handle missing keys + dims = set(lookup_arrays.keys()) + missing_dims_in_component = dims.difference(array.dims) + missing_dims_in_lookup_tables = any( + dim not in lookup.dims for dim in dims for lookup in lookup_arrays.values() + ) + if missing_dims_in_component: + raise ValueError( + f"Cannot select items from `{array.name}` on the dimensions {dims} " + f"since the array is not indexed over the dimensions {missing_dims_in_component}" + ) + if missing_dims_in_lookup_tables: + raise ValueError( + f"All lookup arrays used to select items from `{array.name}` " + f"must be indexed over the dimensions {dims}" + ) + + dim = "dim_0" + ixs = {} + masks = [] + + # Turn string lookup values to numeric ones. + # We stack the dimensions to handle multidimensional lookups + for index_dim, index in lookup_arrays.items(): + stacked_lookup = self._context.input_data[index.name].stack( + {dim: tuple(dims)} + ) + ix = array.indexes[index_dim].get_indexer(stacked_lookup) + if (ix == -1).all(): + received_lookup = ( + self._context.input_data[index.name].to_series().dropna() + ) + raise IndexError( + f"Trying to select items on the dimension {index_dim} from the " + f"{index.name} lookup array, but no matches found. Received: {received_lookup}" + ) + ixs[index_dim] = xr.DataArray( + np.fmax(0, ix), coords={dim: stacked_lookup[dim]} + ) + masks.append(ix >= 0) + + # Nullify any lookup values that are not given (i.e., are NaN in the lookup array) + mask = functools.reduce(lambda x, y: x & y, masks) + result = array[ixs] + if not mask.all(): + result[{dim: ~mask}] = np.nan + return result.drop_vars(dims).unstack(dim) + + +class GetValAtIndex(HelperFunction): + """Get the value of a dimension at a given integer index.""" + + NAME = "get_val_at_index" + ALLOWED_IN: ClassVar[list[KIND_T]] = ["expression", "mask"] + + def as_math_string(self, **dim_idx_mapping: str) -> str: # noqa: D102, override + dim, idx = self._mapping_to_dim_idx(**dim_idx_mapping) + return f"{dim}[{idx}]" + + def as_raw(self, **dim_idx_mapping: int) -> xr.DataArray: + """ + Get the value of a model dimension at a given integer index. + + This function is primarily useful for timeseries data, e.g. + `get_val_at_index(snapshot=0)` is the first snapshot. + + Parameters + ---------- + **dim_idx_mapping : int + A single mapping from a model dimension name to the (zero-indexed) + integer index of the value to extract. + + Returns + ------- + xr.DataArray + Dimensionless array containing one value. + """ + dim, idx = self._mapping_to_dim_idx(**dim_idx_mapping) + return self._context.input_data.coords[dim][int(idx)] + + @staticmethod + def _mapping_to_dim_idx(**dim_idx_mapping: Any) -> tuple[str, Any]: + if len(dim_idx_mapping) != 1: + raise ValueError("Supply one (and only one) dimension:index mapping") + return next(iter(dim_idx_mapping.items())) + + +class Roll(HelperFunction): + """Roll (a.k.a. shift) items along ordered dimensions.""" + + NAME = "roll" + ALLOWED_IN: ClassVar[list[KIND_T]] = ["expression"] + ignore_mask = True + + def as_math_string(self, array: str, **roll_kwargs: str) -> str: # noqa: D102, override + new_strings = { + self._dim_iterator(k): f"{-1 * int(v):+d}" for k, v in roll_kwargs.items() + } + return _update_iterator(array, new_strings, "add") + + def as_raw(self, array: xr.DataArray, **roll_kwargs: int) -> xr.DataArray: + """ + Roll the array along the given dimension(s) by the given number of places. + + Rolling keeps the array index labels in the same position, but moves the + data by the given number of places, e.g. `roll(storage_level, snapshot=1)` + aligns each snapshot with the previous snapshot's storage level. + + Parameters + ---------- + array : xr.DataArray + Array on which to roll data. + **roll_kwargs : int + Keys are dimension names on which to roll; values are the number of + places to roll data. + + Returns + ------- + xr.DataArray + `array` with rolled data. + """ + roll_kwargs_int: Mapping = {k: int(v) for k, v in roll_kwargs.items()} + return array.roll(roll_kwargs_int) + + +class Mask(HelperFunction): + """Apply a boolean condition to an array _within_ an expression string.""" + + NAME = "mask" + ALLOWED_IN: ClassVar[list[KIND_T]] = ["expression"] + + def as_math_string(self, array: str, condition: str) -> str: # noqa: D102, override + return rf"({array} \text{{if }} {condition} == True)" + + def as_raw(self, array: xr.DataArray, condition: xr.DataArray) -> xr.DataArray: + """ + Apply a mask condition to a math array within an expression string. + + For example, `sum(mask(flow_cap, node_grouping), over=node)` sums only + the group members flagged in a boolean `node_grouping` array. + + Parameters + ---------- + array : xr.DataArray + Math component array. + condition : xr.DataArray + Boolean mask array. If not `bool` dtype, NaN and 0 are taken as False + and all other values as True. + + Returns + ------- + xr.DataArray + The input array with the condition applied, broadcast across any new + dimensions provided by the condition. + """ + return array.where(condition.fillna(False).astype(bool)) + + +class GroupSum(HelperFunction): + """Apply a summation over an arbitrary array grouping.""" + + NAME = "group_sum" + ALLOWED_IN: ClassVar[list[KIND_T]] = ["expression"] + ignore_mask = True + + def as_math_string(self, array: str, groupby: str, group_dim: str) -> str: # noqa: D102, override + group_dim_singular = self._dim_iterator(group_dim) + sum_lim_string = rf"\text{{ if }} {groupby} = \text{{{group_dim_singular}}}" + over = [self._instr(i) for i in self._get_dims_from_iterators(groupby)] + foreach_string = r" \\ ".join([*over, sum_lim_string]) + return rf"\sum\limits_{{\substack{{{foreach_string}}}}} ({array})" + + def as_raw( + self, array: xr.DataArray, groupby: xr.DataArray, group_dim: xr.DataArray + ) -> xr.DataArray: + """ + Sum an array over the given groupings. + + For example, `group_sum(p * sign, bus, node)` sums the per-component + `p * sign` expression into per-`node` totals using the `bus` lookup. + + Parameters + ---------- + array : xr.DataArray + Expression array. + groupby : xr.DataArray + Array with which to group `array`; all dimensions over which it is + indexed are replaced by `group_dim` in the result. + group_dim : xr.DataArray + Dimension that the `groupby` values are members of. This becomes a + new dimension over which the result is indexed. + + Returns + ------- + xr.DataArray + Array with dimension(s) aggregated over the `groupby`. + + See Also + -------- + GroupDatetime : grouping over datetime periods without a separate `groupby` array. + """ + grouping_dims = groupby.dims + groups = array.stack(_stacked=grouping_dims).groupby( + groupby.rename(group_dim.name).stack(_stacked=grouping_dims) + ) + return groups.sum("_stacked") + + +class GroupDatetime(HelperFunction): + """Apply a summation over a datetime group on a datetime dimension.""" + + NAME = "group_datetime" + ALLOWED_IN: ClassVar[list[KIND_T]] = ["expression"] + ignore_mask = True + + def as_math_string(self, array: str, over: str, group: str) -> str: # noqa: D102, override + overstring = self._instr(over) + foreach_string = ( + rf"{overstring} \text{{ if }} \text{{{group}}}" + rf"(\text{{{self._dim_iterator(over)}}}) = \text{{{self._dim_iterator(group)}}}" + ) + return rf"\sum\limits_{{\substack{{{foreach_string}}}}} ({array})" + + def as_raw( + self, array: xr.DataArray, over: xr.DataArray, group: xr.DataArray + ) -> xr.DataArray: + """ + Sum an array over a datetime grouping of a datetime dimension. + + For example, `group_datetime(flow_in, snapshot, month) <= max_monthly` + constrains the monthly sum of a timeseries variable. + + Parameters + ---------- + array : xr.DataArray + Expression array. + over : xr.DataArray + Datetime dimension over which to group. + group : xr.DataArray + Datetime grouper dimension; any xarray/pandas datetime accessor name + ('date', 'dayofweek', 'month', ...). The `over` dimension is replaced + by this dimension in the result. + + Returns + ------- + xr.DataArray + Array with the datetime dimension aggregated over the grouper. + """ + group_name = str(group.name) + dtype = DTYPE_OPTIONS[self._context.math.dimensions[group_name].dtype] + group_sum_helper = GroupSum(self._context) + return group_sum_helper.as_raw( + array, getattr(array[str(over.name)].dt, group_name).astype(dtype), group + ) + + +class SumNextN(HelperFunction): + """ + Sum the current and next N items in an array. + + Works best for ordered arrays (datetime, integer) and is equivalent to a + summation over a rolling window. + """ + + NAME = "sum_next_n" + ALLOWED_IN: ClassVar[list[KIND_T]] = ["expression"] + + def as_math_string(self, array: str, over: str, N: int) -> str: # noqa: D102, override + over_singular = rf"\text{{{self._dim_iterator(over)}}}" + new_iterator = over[0] + updated_iterator_array = _update_iterator( + array, {self._dim_iterator(over): new_iterator}, "replace" + ) + return ( + rf"\sum\limits_{{\text{{{new_iterator}}}={over_singular}}}" + rf"^{{{over_singular}+{N}}} ({updated_iterator_array})" + ) + + def as_raw(self, array: xr.DataArray, over: xr.DataArray, N: int) -> xr.DataArray: + """ + Sum values from the current up to N-from-current position on a dimension. + + For example, `sum_next_n(flow_in, snapshot, 4) == sum_next_n(demand, snapshot, 4)` + requires flexible demand to be met within a 4-snapshot window. + + Parameters + ---------- + array : xr.DataArray + Math component array. + over : xr.DataArray + Dimension over which to sum. + N : int + Number of items beyond the current value to sum over. + + Returns + ------- + xr.DataArray + Array of rolling-window sums, indexed as `array`. + + Notes + ----- + The rolling window does not wrap around to the start of the dimension when + reaching the end, so the final N items sum over progressively shorter + windows. This over-constrains a model unless the constraint is limited + (using its `mask` string) to the first `len(over) - N` items, e.g. + `mask: snapshot<=get_val_at_index(snapshot=-4)` if N == 4. + """ + # We cannot use the xarray rolling window method as it doesn't like + # operating on Python objects, which our optimisation problem components are. + results: list[xr.DataArray] = [] + for i in range(len(over)): + results.append( + array.isel({str(over.name): slice(i, i + int(N))}).sum( + str(over.name), min_count=1 + ) + ) + return xr.concat(results, dim=over).broadcast_like(array) + + +BUILTIN_HELPERS: tuple[type[HelperFunction], ...] = ( + MaskAny, + Sum, + SelectFromLookupArrays, + GetValAtIndex, + Roll, + Mask, + GroupSum, + GroupDatetime, + SumNextN, +) +"""The helper functions shipped with linopy.""" + + +def build_registry( + extra: Iterable[type[HelperFunction]] = (), +) -> dict[KIND_T, dict[str, type[HelperFunction]]]: + """ + Return a helper-function registry of the built-in helpers plus any `extra` ones. + + Parameters + ---------- + extra : Iterable[type[HelperFunction]], optional + User-defined :class:`HelperFunction` subclasses to add to the registry. + + Returns + ------- + dict[Literal["mask", "expression"], dict[str, type[HelperFunction]]] + Per string type, a mapping from helper `NAME` to helper class. + + Raises + ------ + ValueError + If an entry is not a :class:`HelperFunction` subclass, is missing + `NAME`/`ALLOWED_IN`, or clashes with an already-registered name. + """ + registry: dict[KIND_T, dict[str, type[HelperFunction]]] = { + "mask": {}, + "expression": {}, + } + for cls in (*BUILTIN_HELPERS, *extra): + if not (isinstance(cls, type) and issubclass(cls, HelperFunction)): + raise ValueError( + "Helper function must be subclassed from " + f"linopy.declarative.helpers.HelperFunction: {cls}" + ) + name = getattr(cls, "NAME", None) + allowed_in = getattr(cls, "ALLOWED_IN", None) + if not isinstance(name, str) or not allowed_in: + raise ValueError( + f"Helper function {cls.__name__} must define `NAME` and `ALLOWED_IN`" + ) + for kind in allowed_in: + if name in registry[kind]: + raise ValueError( + f"`{kind}` string helper function `{name}` already exists" + ) + registry[kind][name] = cls + return registry diff --git a/linopy/declarative/latex.py b/linopy/declarative/latex.py new file mode 100644 index 00000000..c6e89798 --- /dev/null +++ b/linopy/declarative/latex.py @@ -0,0 +1,497 @@ +""" +Linopy declarative LaTeX math documentation module. + +This module builds a human-readable mathematical formulation document from a declarative math definition. +Every active math component is rendered to LaTeX (equations, mask conditions, foreach sets, bounds) without building an optimisation problem. +Metadata and cross-references are then combined together with the LaTeX rendering to produce a complete mathematical formulation document. + +This module is adapted from the calliope Apache-2.0 licensed latex backend module: +- https://github.com/calliope-project/calliope/blob/9916116a06ec8c1feaf3c2606bdb8941b916ce85/src/calliope/backend/latex_backend.py +""" + +from __future__ import annotations + +import re +from collections.abc import Iterable +from dataclasses import dataclass, field, replace +from typing import Literal + +import pandas as pd +import xarray as xr + +from linopy.declarative import parsing +from linopy.declarative.build import _DeclarativeBase +from linopy.declarative.helpers import HelperFunction, dim_iterator +from linopy.declarative.nodes import Component, Node, find_refs, latex_number +from linopy.declarative.schema import ( + ConstraintDef, + ExpressionDef, + LookupDef, + ObjectiveDef, + ParameterDef, + VariableDef, +) + +FORMAT_T = Literal["md", "rst", "tex"] + +_REPR_STYLES = { + "parameters": "textit", + "lookups": "textit", + "variables": "textbf", + "expressions": "textbf", +} + +DOCUMENTED_GROUPS_T = ( + ParameterDef + | LookupDef + | VariableDef + | ExpressionDef + | ConstraintDef + | ObjectiveDef +) +DOCUMENTED_LINOPY_OBJ_GROUPS_T = ( + VariableDef | ExpressionDef | ConstraintDef | ObjectiveDef +) +DOCUMENTED_EXPRESSION_GROUPS_T = ExpressionDef | ConstraintDef | ObjectiveDef +DOCUMENTED_GROUPS: dict[str, str] = { + "parameters": "Parameters", + "lookups": "Lookups", + "variables": "Variables", + "expressions": "Expressions", + "constraints": "Constraints", + "objectives": "Objectives", +} +"""Component groups documented in LaTeX math docs, with their section titles.""" + + +@dataclass +class RenderedComponent: + """One math component rendered to LaTeX, ready for document generation.""" + + group: str + """Component group (e.g. "constraints").""" + + name: str + """Component name.""" + + title: str = "" + """The component's long name.""" + + description: str = "" + """The component's verbose description.""" + + foreach: str = "" + """LaTeX `\\forall` line body over the component's sets ("" if adimensional).""" + + mask: str = "" + """LaTeX rendering of the component's top-level mask ("" if trivially true).""" + + equations: list[dict[str, str]] = field(default_factory=list) + """Per equation variant, its "expression" LaTeX and (possibly empty) "mask" LaTeX.""" + + uses: list[str] = field(default_factory=list) + """Names of the math components this component references.""" + + used_in: list[str] = field(default_factory=list) + """Names of the math components referencing this component (filled by `build`).""" + + extras: dict[str, str] = field(default_factory=dict) + """Additional metadata to document (unit, default, sense, ...).""" + + +class LatexModelBuilder(_DeclarativeBase): + """ + Builder turning a declarative math definition into LaTeX math documentation. + + The counterpart of :class:`linopy.declarative.build.DeclarativeModelBuilder` + on the LaTeX route: instead of adding components to a linopy model, it + renders each component's math strings via the same parsing machinery. + + Examples + -------- + ```python + doc = LatexModelBuilder(math_def, input_data).build().generate_math_doc("md") + ``` + """ + + def __init__( + self, + math_def: dict, + input_data: xr.Dataset | None = None, + config: dict | None = None, + helpers: Iterable[type[HelperFunction]] = (), + ) -> None: + """ + Validate the math definition, ready to `build()` the documentation. + + Parameters + ---------- + math_def : dict + Declarative math definition. + input_data : xr.Dataset, optional + Model input data; only used to derive the dimensions of parameters + and lookups for their LaTeX subscripts. + config : dict, optional + Build configuration options (available to `config.` mask references). + helpers : Iterable[type[HelperFunction]], optional + User-defined helper functions, in addition to the built-in ones. + """ + super().__init__(math_def, input_data, config, helpers) + self.components: dict[str, dict[str, RenderedComponent]] = {} + self._ctx = replace(self._ctx, math_reprs=self._build_math_reprs()) + + def _build_math_reprs(self) -> dict[str, str]: + r""" + Return the decorated LaTeX representation of every referenceable component. + + Parameters/lookups render as `\textit{name}` and variables/expressions as + `\textbf{name}`, subscripted with the iterators of the dimensions they are + indexed over (a variable's `foreach`; an input's dimensions in the data). + """ + reprs: dict[str, str] = {} + for group, style in _REPR_STYLES.items(): + for name, definition in getattr(self.math, group)._active.items(): + if isinstance(definition, DOCUMENTED_LINOPY_OBJ_GROUPS_T): + dims = definition.foreach + else: + dims = ( + list(self.input_data[name].dims) + if name in self.input_data + else [] + ) + + iterators = ",".join(dim_iterator(self.math, str(dim)) for dim in dims) + subscript = rf"_\text{{{iterators}}}" if iterators else "" + reprs[name] = rf"\{style}{{{name}}}{subscript}" + return reprs + + def _foreach_string(self, definition: DOCUMENTED_LINOPY_OBJ_GROUPS_T) -> str: + r"""Return the LaTeX `\forall` line body for a component's `foreach` sets.""" + if not (sets := definition.foreach): + return "" + instrs = ", ".join( + rf"\text{{{dim_iterator(self.math, dim)}}} \in \text{{{dim}}}" + for dim in sets + ) + return rf"\forall{{}} {instrs}" + + def _mask_string(self, mask_node: Node, name: str) -> str: + """Return the LaTeX rendering of a component's parsed top-level mask ("" if true).""" + rendered = mask_node.to_latex( + replace(self._ctx, mode="mask", equation_name=name) + ) + return "" if rendered == "true" else rendered + + def _render_metadata(self, definition: DOCUMENTED_GROUPS_T) -> dict[str, str]: + """Return the documentable metadata (unit, default, ...) of a definition.""" + extras: dict[str, str] = {} + if unit := getattr(definition, "unit", None): + extras["Unit"] = unit + if pd.notnull(default := getattr(definition, "default", None)): + extras["Default"] = str(default) + return extras + + def add_component( + self, group: str, name: str, definition: DOCUMENTED_GROUPS_T + ) -> None: + """Render one math component and store it under `self.components`.""" + rendered = RenderedComponent( + group=group, + name=name, + title=definition.title, + description=definition.description, + extras=self._render_metadata(definition), + ) + if isinstance(definition, DOCUMENTED_LINOPY_OBJ_GROUPS_T): + self._add_linopy_obj_component(group, name, rendered, definition) + self.components.setdefault(definition._group, {})[name] = rendered + + def _add_linopy_obj_component( + self, + group: str, + name: str, + rendered: RenderedComponent, + definition: DOCUMENTED_LINOPY_OBJ_GROUPS_T, + ) -> None: + parsed = self.parsed[group][name] + rendered.mask = self._mask_string(parsed.mask, f"{group}:{name}") + rendered.foreach = self._foreach_string(definition) + + uses = find_refs(parsed.mask, Component) + if isinstance(definition, DOCUMENTED_EXPRESSION_GROUPS_T): + self._add_expr_component(rendered, uses, parsed) + if isinstance(definition, VariableDef): + self._add_var_component(name, rendered, uses, definition) + if isinstance(definition, ObjectiveDef): + self._add_obj_component(rendered, definition) + # Only cross-reference documented components (not dimensions or strings). + rendered.uses = sorted(uses & set(self._ctx.math_reprs)) + + # Escape special characters in text-mode LaTeX so KaTeX can render names + # and coordinate values containing e.g. underscores. Done here, on the + # final display strings, so `math_reprs` stay unescaped for evaluation. + for rendered_eq in rendered.equations: + rendered_eq["mask"] = _escape_text_mode(rendered_eq["mask"]) + rendered_eq["expression"] = _escape_text_mode(rendered_eq["expression"]) + rendered.foreach = _escape_text_mode(rendered.foreach) + rendered.mask = _escape_text_mode(rendered.mask) + + def _add_expr_component( + self, + rendered: RenderedComponent, + uses: set[str], + parsed: parsing.ParsedComponent, + ) -> None: + for equation in parsed.equations: + equation_ctx = replace(self._ctx, equation_name=equation.name) + rendered.equations.append( + { + "mask": parsing.as_latex_mask(equation, equation_ctx), + "expression": parsing.as_latex_expression(equation, equation_ctx), + } + ) + uses |= equation.references() + + def _add_var_component( + self, + name: str, + rendered: RenderedComponent, + uses: set[str], + definition: VariableDef, + ) -> None: + rendered.extras["Domain"] = definition.domain + rendered.equations.append( + {"mask": "", "expression": self._bounds_string(name, definition)} + ) + uses |= { + bound + for bound in (definition.bounds.lower, definition.bounds.upper) + if isinstance(bound, str) + } + + def _bounds_string(self, name: str, definition: VariableDef) -> str: + """Return the LaTeX bounds equation of a decision variable.""" + bounds = definition.bounds + reprs = self._ctx.math_reprs + lower, upper = ( + reprs.get(bound, rf"\textit{{{bound}}}") + if isinstance(bound, str) + else latex_number(bound) + for bound in (bounds.lower, bounds.upper) + ) + return rf"{lower} \leq {reprs[name]} \leq {upper}" + + def _add_obj_component( + self, rendered: RenderedComponent, definition: ObjectiveDef + ) -> None: + rendered.extras["Sense"] = ( + "minimise" if definition.sense == "min" else "maximise" + ) + + def build(self) -> LatexModelBuilder: + """ + Render all active math components and resolve their cross-references. + + Returns + ------- + LatexModelBuilder + Itself, with `self.components` filled, so that document generation + can be chained (`builder.build().generate_math_doc()`). + """ + for group in DOCUMENTED_GROUPS: + for name, definition in getattr(self.math, group)._active.items(): + self.add_component(group, name, definition) + + used_in: dict[str, set[str]] = {} + for group_components in self.components.values(): + for component in group_components.values(): + for ref in component.uses: + used_in.setdefault(ref, set()).add(component.name) + for group_components in self.components.values(): + for component in group_components.values(): + component.used_in = sorted(used_in.get(component.name, set())) + return self + + def generate_math_doc(self, format: FORMAT_T = "md") -> str: + """ + Generate a math documentation string from the rendered components. + + Parameters + ---------- + format : Literal["md", "rst", "tex"], default: "md" + Output format: Markdown, reStructuredText, or LaTeX source. + + Returns + ------- + str + The full documentation document. + """ + if not self.components: + self.build() + blocks = [_heading(format, 1, "Math formulation"), ""] + for group, group_title in DOCUMENTED_GROUPS.items(): + group_components = self.components.get(group) + if not group_components: + continue + blocks.extend([_heading(format, 2, group_title), ""]) + for component in group_components.values(): + blocks.extend(_component_doc(format, component)) + return "\n".join(blocks).rstrip() + "\n" + + +def latex_math_doc( + math_def: dict, + input_data: xr.Dataset | None = None, + config: dict | None = None, + format: FORMAT_T = "md", + helpers: Iterable[type[HelperFunction]] = (), +) -> str: + """ + Generate LaTeX math documentation from a declarative math definition. + + Parameters + ---------- + math_def : dict + Declarative math definition (see + :class:`linopy.declarative.schema.MathModel` for the expected structure). + input_data : xr.Dataset, optional + Model input data; only used to derive the dimensions of parameters and + lookups for their LaTeX subscripts. + config : dict, optional + Build configuration options (available to `config.` mask references). + format : Literal["md", "rst", "tex"], default: "md" + Output format: Markdown, reStructuredText, or LaTeX source. + helpers : Iterable[type[HelperFunction]], optional + User-defined helper functions, in addition to the built-in ones. + + Returns + ------- + str + The full documentation document. + """ + builder = LatexModelBuilder(math_def, input_data, config, helpers) + return builder.build().generate_math_doc(format) + + +# --------------------------------------------------------------------------- +# Document assembly +# --------------------------------------------------------------------------- + + +def _tex_text(text: str) -> str: + """Escape underscores for LaTeX text mode.""" + return text.replace("_", r"\_") + + +# Matches a `\text`/`\textbf`/`\textit`/... command and captures its (brace-free) +# argument; the special characters within are escaped by `_escape_text_mode`. +_TEXT_CMD_RE = re.compile(r"(\\text(?:bf|it|rm|sf|tt|normal)?)\{([^{}]*)\}") +_TEXT_SPECIAL_RE = re.compile(r"(? str: + r""" + Escape LaTeX-special characters inside `\text*{...}` arguments. + + KaTeX rejects bare `_`, `^`, ... in text mode, so a component or coordinate + named e.g. `storage_units` renders as `\text{storage_units}` and raises a + parse error. Only the *contents* of text commands are escaped: subscript + operators such as the `_` in `\textbf{x}_\text{n}` sit outside the braces and + are preserved, and characters already escaped (`\_`) are left untouched. + """ + + def _escape(match: re.Match[str]) -> str: + cmd, content = match.group(1), match.group(2) + escaped = _TEXT_SPECIAL_RE.sub(r"\\\1", content) + return f"{cmd}{{{escaped}}}" + + return _TEXT_CMD_RE.sub(_escape, latex) + + +def _heading(format: FORMAT_T, level: int, text: str) -> str: + """Return a document heading at the given level.""" + if format == "md": + return f"{'#' * level} {text}" + if format == "rst": + return f"{text}\n{'=-^'[level - 1] * len(text)}" + tex_levels = {1: "section", 2: "subsection", 3: "paragraph"} + return rf"\{tex_levels[level]}{{{_tex_text(text)}}}" + + +def _metadata_lines(format: FORMAT_T, key: str, value: str) -> list[str]: + """Return the lines of a "key: value" metadata entry.""" + if format == "tex": + # A blank line so that each entry renders as its own paragraph. + return [rf"\textbf{{{key}}}: {_tex_text(value)}", ""] + return [f"- **{key}**: {value}"] + + +def _array_block(lines: list[str]) -> str: + """Return a LaTeX `array` environment of the given lines.""" + joined = " \\\\\n ".join(lines) + return f"\\begin{{array}}{{l}}\n {joined}\n\\end{{array}}" + + +def _cases_block(equations: list[dict[str, str]]) -> str: + r""" + Return a LaTeX `cases` environment, one row per equation variant. + + Each row is the variant's expression, with its own mask (if any) as the + row's `\text{if }` condition, so that variants sharing a component's + `foreach`/top-level mask render as sub-clauses at the same nesting level. + """ + rows = [] + for equation in equations: + row = equation["expression"] + if equation["mask"]: + row += rf" & \quad \text{{if }} {equation['mask']}" + rows.append(row) + joined = " \\\\\n ".join(rows) + return f"\\begin{{cases}}\n {joined}\n\\end{{cases}}" + + +def _math_block(format: FORMAT_T, inner: str) -> list[str]: + """Return a display-math block wrapping the given LaTeX body.""" + if format == "md": + return ["$$", inner, "$$", ""] + if format == "rst": + indented = "\n".join(f" {line}" for line in inner.split("\n")) + return [".. math::", "", indented, ""] + return [r"\begin{equation}", inner, r"\end{equation}", ""] + + +def _component_doc(format: FORMAT_T, component: RenderedComponent) -> list[str]: + """Return the documentation lines of one rendered component.""" + blocks = [_heading(format, 3, component.name), ""] + if component.title: + blocks.extend([component.title, ""]) + if component.description: + blocks.extend([component.description, ""]) + metadata = dict(component.extras) + if component.uses: + metadata["Uses"] = ", ".join(component.uses) + if component.used_in: + metadata["Used in"] = ", ".join(component.used_in) + for key, value in metadata.items(): + blocks.extend(_metadata_lines(format, key, value)) + if metadata: + blocks.append("") + if component.equations: + header = [] + if component.foreach: + header.append(component.foreach) + if component.mask: + header.append(rf"\text{{if }} {component.mask}") + + if len(component.equations) == 1: + equation = component.equations[0] + lines = [*header] + if equation["mask"]: + lines.append(rf"\text{{if }} {equation['mask']}") + lines.append(equation["expression"]) + inner = _array_block(lines) + else: + cases = _cases_block(component.equations) + inner = f"{_array_block(header)}\n{cases}" if header else cases + blocks.extend(_math_block(format, inner)) + return blocks diff --git a/linopy/declarative/nodes.py b/linopy/declarative/nodes.py new file mode 100644 index 00000000..fc5c4f46 --- /dev/null +++ b/linopy/declarative/nodes.py @@ -0,0 +1,838 @@ +""" +Linopy declarative math AST module. + +This module contains the AST node types produced when parsing declarative math strings. +For each node, this module contains the `pyparsing` parse action(s) that build it, the data evaluator, and the LaTeX math renderer. +The evaluation context and shared utilities live here too. +The pyparsing grammars that produce the nodes live in :mod:`linopy.declarative.grammar`. + +This module is adapted from the calliope Apache-2.0 licensed math parsers: +- https://github.com/calliope-project/calliope/blob/9916116a06ec8c1feaf3c2606bdb8941b916ce85/src/calliope/backend/expression_parser.py +- https://github.com/calliope-project/calliope/blob/9916116a06ec8c1feaf3c2606bdb8941b916ce85/src/calliope/backend/where_parser.py + +""" + +from __future__ import annotations + +import operator +import re +from abc import ABC, abstractmethod +from collections.abc import Callable +from dataclasses import dataclass, field, fields, replace +from typing import TYPE_CHECKING, Any, Literal + +import numpy as np +import pandas as pd +import pyparsing as pp +import xarray as xr + +from linopy.declarative.helpers import ( + KIND_T, + HelperFunction, + _update_iterator, + dim_iterator, +) +from linopy.declarative.schema import ConfigModel, MathModel +from linopy.expressions import LinearExpression +from linopy.variables import Variable + +if TYPE_CHECKING: + from linopy.model import Model + +TRUE_ARRAY = xr.DataArray(True) + +MODE_T = Literal["mask", "raw", "expr"] + +COMPONENT_CATEGORY_T = Literal["any", "dimension", "input", "result"] + +_INPUT_GROUPS = ("parameters", "lookups", "dimensions") + +_OPERATIONS = { + "**": operator.pow, + "*": operator.mul, + "/": operator.truediv, + "+": operator.add, + "-": operator.sub, + "and": operator.and_, + "or": operator.or_, + "<=": operator.le, + ">=": operator.ge, + "<": operator.lt, + ">": operator.gt, + "==": operator.eq, +} + +_LATEX_OPERATORS = { + "**": "{val}^{{{operand}}}", + "*": r"{val} \times {operand}", + "/": r"\frac{{ {val} }}{{ {operand} }}", + "+": "{val} + {operand}", + "-": "{val} - {operand}", + "and": r"{val} \land {operand}", + "or": r"{val} \lor {operand}", +} + +_LATEX_IDENTITIES = {"+": "0", "-": "0", "and": "true", "or": "true"} +"""Operands that add nothing to a LaTeX operator chain (e.g. `0 + flow` is just `flow`).""" + +_LATEX_EQUATION_OPERATORS = {"<=": r" \leq ", ">=": r" \geq ", "=": " = "} + +_LATEX_MASK_OPERATORS = { + "<=": r"\mathord{\leq}", + ">=": r"\mathord{\geq}", + "==": r"\mathord{==}", + "<": r"\mathord{<}", + ">": r"\mathord{>}", +} + + +@dataclass +class Context: + """ + Context against which a parsed math AST is evaluated. + + The first five fields describe the model being built and are fixed for a + whole build; the remainder are evaluation state, set per equation and only + ever modified on immutable copies (via :func:`dataclasses.replace`). + """ + + model: Model + """Linopy model containing any already-built variables and expressions.""" + + input_data: xr.Dataset = field(default_factory=xr.Dataset) + """Model input data (parameters, lookups, dimensions).""" + + math: MathModel = field(default_factory=MathModel) + """Declarative math definition.""" + + config: ConfigModel = field(default_factory=ConfigModel) + """Build configuration options.""" + + helpers: dict[KIND_T, dict[str, type[HelperFunction]]] = field(default_factory=dict) + """Helper-function registry (see :func:`linopy.declarative.helpers.build_registry`).""" + + equation_name: str = "" + """Name of the equation being evaluated (used in error messages).""" + + mode: MODE_T = "raw" + """Evaluation mode: "mask" (boolean-array route), "raw" (un-normalised + data), or "expr" (results coerced to linopy expressions).""" + + apply_mask: bool = True + """In mask mode, whether component references coerce to existence booleans.""" + + sub_expressions: dict[str, Node] = field(default_factory=dict) + """Resolved `$name` sub-expression AST per name.""" + + slices: dict[str, Node] = field(default_factory=dict) + """Resolved `$name` slicer AST per name.""" + + mask: xr.DataArray = field(default_factory=lambda: TRUE_ARRAY) + """Boolean array defining where the evaluated expression applies.""" + + math_reprs: dict[str, str] = field(default_factory=dict) + """Custom LaTeX representations per component name, taking precedence when + rendering math strings (used by :class:`linopy.declarative.latex.LatexModelBuilder`).""" + + def demote(self) -> Context: + """Copy in raw mode if in expr mode (helper args, slicers, list items).""" + return replace(self, mode="raw") if self.mode == "expr" else self + + def helper_kind(self) -> KIND_T: + """Return the helper-registry kind matching the evaluation mode.""" + return "mask" if self.mode == "mask" else "expression" + + +# --------------------------------------------------------------------------- +# Utilities +# --------------------------------------------------------------------------- + + +def error(ctx: Context, node: Node, message: str) -> ValueError: + """Return a ValueError with the equation name and a caret at the node's position.""" + marker = " " * (pp.col(node.loc, node.instring) - 1) + "^" + return ValueError( + f"({ctx.equation_name}) | {message}\n {node.instring}\n {marker}" + ) + + +def _unwrap(value: Any) -> Any: + """Extract the scalar from a dimensionless DataArray, e.g. for use in `.sel`/`.isin`.""" + return ( + value.item() if isinstance(value, xr.DataArray) and value.ndim == 0 else value + ) + + +def to_linexpr(obj: Any, model: Model) -> Any: + """ + Normalise an evaluated object to a linopy expression. + + `Variable` objects are converted via `to_linexpr` and `xr.DataArray` objects + are wrapped in a constant `LinearExpression`; everything else is returned + unchanged. This is the single place where expression-route coercion happens. + """ + if isinstance(obj, Variable): + return obj.to_linexpr() + if isinstance(obj, xr.DataArray): + return LinearExpression(obj, model) + return obj + + +def latex_number(value: float | int) -> str: + r"""Format a number for LaTeX, mapping infinities to `\infty`.""" + if value == float("inf"): + return r"\infty" + if value == float("-inf"): + return r"-\infty" + return re.sub( + r"([\d]+?)e([+-])([\d]+)", r"\1\\mathord{\\times}10^{\2\3}", f"{value:.6g}" + ) + + +# --------------------------------------------------------------------------- +# AST nodes +# --------------------------------------------------------------------------- + + +@dataclass(frozen=True, kw_only=True) +class Node(ABC): + """ + Base class of all declarative math AST nodes. + + Nodes are immutable data produced by the grammars in + :mod:`linopy.declarative.grammar`; each concrete node carries its own parse + action(s) (`from_tokens`), data evaluator (`evaluate`), and LaTeX renderer + (`to_latex`). + """ + + instring: str = field(repr=False, compare=False) + """The full source string this node was parsed from (used in error messages).""" + + loc: int = field(default=0, repr=False, compare=False) + """Character offset of this node in `instring` (used in error messages).""" + + @abstractmethod + def evaluate(self, ctx: Context) -> Any: + """Evaluate this node to data (an `xr.DataArray` or a linopy expression).""" + + @abstractmethod + def to_latex(self, ctx: Context) -> str: + """Render this node as a LaTeX math string.""" + + +@dataclass(frozen=True, kw_only=True) +class Constant(Node): + """A literal number (including `inf`), boolean, or generic string.""" + + value: float | bool | str + + @classmethod + def number(cls, instring: str, loc: int, tokens: pp.ParseResults) -> Constant: + """Parse action for numeric literals.""" + return cls(value=float(tokens[0]), instring=instring, loc=loc) + + @classmethod + def string(cls, instring: str, loc: int, tokens: pp.ParseResults) -> Constant: + """Parse action for generic string literals.""" + return cls(value=str(tokens[0]), instring=instring, loc=loc) + + @classmethod + def boolean(cls, instring: str, loc: int, tokens: pp.ParseResults) -> Constant: + """Parse action for boolean literals.""" + return cls(value=str(tokens[0]).lower() == "true", instring=instring, loc=loc) + + def evaluate(self, ctx: Context) -> Any: + """Evaluate to a dimensionless array (or the plain string for string literals).""" + if isinstance(self.value, bool): + return xr.DataArray(np.bool_(self.value)) + if isinstance(self.value, str): + return self.value + return xr.DataArray(float(self.value), name=float(self.value)) + + def to_latex(self, ctx: Context) -> str: + """Render the literal value.""" + if isinstance(self.value, bool): + return str(self.value).lower() + if isinstance(self.value, str): + return self.value + return latex_number(float(self.value)) + + +@dataclass(frozen=True, kw_only=True) +class ListNode(Node): + """A literal list of items, e.g. `[a, b, 1]`.""" + + items: tuple[Node, ...] + + @classmethod + def from_tokens(cls, instring: str, loc: int, tokens: pp.ParseResults) -> ListNode: + """Parse action for `[item, item, ...]` lists.""" + return cls(items=tuple(tokens), instring=instring, loc=loc) + + def evaluate(self, ctx: Context) -> list[Any]: + """Evaluate to a plain list of the evaluated items.""" + return [item.evaluate(ctx.demote()) for item in self.items] + + def to_latex(self, ctx: Context) -> str: + """Render as a bracketed plain-text list.""" + return "[" + ",".join(_plain_string(item, ctx) for item in self.items) + "]" + + +@dataclass(frozen=True, kw_only=True) +class Component(Node): + """ + A reference to a named math component (parameter, lookup, dimension, variable or expression). + + `category` records which name set matched during parsing: mask-string grammars + distinguish dimensions / inputs / results at parse time, while expression + grammars leave it as "any" (resolved from the math definition at evaluation). + """ + + name: str + category: COMPONENT_CATEGORY_T = "any" + + @classmethod + def from_tokens_as( + cls, category: COMPONENT_CATEGORY_T + ) -> Callable[[str, int, pp.ParseResults], Component]: + """Return a parse action building a component reference of `category`.""" + + def _action(instring: str, loc: int, tokens: pp.ParseResults) -> Component: + return cls( + name=str(tokens[0]), category=category, instring=instring, loc=loc + ) + + return _action + + def evaluate(self, ctx: Context) -> Any: + """Evaluate the reference according to its category and the evaluation mode.""" + name = self.name + if self.category == "dimension": + # The mask string should evaluate successfully even if a dimension isn't defined. + return ctx.input_data.get(name, xr.DataArray()) + if self.category == "input": + da = ctx.input_data.get(name, xr.DataArray(False)) + if ctx.apply_mask and da.dtype.kind != "b": + da = da.notnull() & (da != np.inf) & (da != -np.inf) + elif da.isnull().any() and pd.notnull( + default := ctx.math.find(name)["default"] + ): + da = da.fillna(default) + return da + if self.category == "result": + result: Any = ctx.model[name] + if ctx.apply_mask: + result = ~result.isnull() + return result + + # category "any": resolve the component group from the math definition. + math_def = ctx.math.find(name) + group = math_def._group + evaluated: Any + if group in _INPUT_GROUPS: + # A parameter/lookup/dimension defined in the math but absent from the + # input data resolves to its default (NaN if none is set). + evaluated = ctx.input_data.get(name, xr.DataArray(np.nan)) + if hasattr(math_def, "dims") and math_def.dims is not None: + evaluated = evaluated.reindex( + {dim: ctx.input_data.coords[dim] for dim in math_def.dims}, + fill_value=np.nan, + ) + else: + # Model entries (variables / expressions): a model entry that was never + # built (e.g. skipped because its mask was empty) resolves to a NaN + # expression rather than raising. In expr mode the object is + # normalised to a linopy expression; in raw mode it is returned unchanged. + try: + evaluated = getattr(ctx.model, group)[name] + except KeyError: + return LinearExpression(xr.DataArray(np.nan), ctx.model) + if ctx.mode != "expr": + return evaluated + evaluated = to_linexpr(evaluated, ctx.model) + if evaluated.isnull().any() and pd.notna(default := math_def["default"]): + evaluated = evaluated.fillna(default) + return evaluated + + def to_latex(self, ctx: Context) -> str: + r""" + Render the component reference as LaTeX. + + A representation registered in `ctx.math_reprs` takes precedence (this is how + :class:`linopy.declarative.latex.LatexModelBuilder` decorates references), + followed by a `math_repr` data attribute, then the bare component name. + """ + name = self.name + custom = ctx.math_reprs.get(name) + if self.category == "dimension": + return name + if self.category == "input": + if custom is None: + array = ctx.input_data.get(name) + attr = array.attrs.get("math_repr") if array is not None else None + custom = str(attr) if attr is not None else rf"\textit{{{name}}}" + if ctx.apply_mask: + custom = rf"\exists ({custom})" + return custom + if self.category == "result": + if custom is not None: + return rf"\exists ({custom})" if ctx.apply_mask else custom + try: + return str( + ctx.model[name].attrs.get( + "math_repr", rf"\exists (\textbf{{{name}}})" + ) + ) + except KeyError: + return rf"\exists (\textbf{{{name}}})" + if custom is not None: + return custom + evaluated = self.evaluate(ctx) + attrs = getattr(evaluated, "attrs", {}) + return str(attrs.get("math_repr", name)) + + +@dataclass(frozen=True, kw_only=True) +class Sliced(Node): + """A sliced component, e.g. `flow[node=a]` or `flow[node=$n]`.""" + + obj: Component + slices: dict[str, Node] + + @classmethod + def from_tokens(cls, instring: str, loc: int, tokens: pp.ParseResults) -> Sliced: + """Parse action for `name[dim=slicer, ...]`.""" + slices = {str(grp["set_name"][0]): grp["slicer"][0] for grp in tokens["slices"]} + return cls(obj=tokens["obj"], slices=slices, instring=instring, loc=loc) + + def evaluate(self, ctx: Context) -> Any: + """Evaluate the component and select the evaluated slices from it.""" + slicer_ctx = ctx.demote() + evaluated_slices = { + dim: [_unwrap(i) for i in vals] + if isinstance(vals := slicer.evaluate(slicer_ctx), list) + else vals + for dim, slicer in self.slices.items() + } + return self.obj.evaluate(ctx).sel(**evaluated_slices) + + def to_latex(self, ctx: Context) -> str: + r""" + Render the sliced component as LaTeX. + + If the component's LaTeX representation carries an iterator substring + (e.g. `\textbf{flow}_\text{n}` from a `math_repr`), the slices are + injected into it (`\textbf{flow}_\text{n=a}` when sliced with `node=a`); + otherwise the slices are appended as a subscript. + """ + slice_strings = { + dim_iterator(ctx.math, dim): slicer.to_latex(ctx) + for dim, slicer in self.slices.items() + } + obj_string = self.obj.to_latex(ctx) + if re.search(r"_\\text\{", obj_string): + return _update_iterator( + obj_string, {it: f"={v}" for it, v in slice_strings.items()}, "add" + ) + subscript = ",".join(f"{k}={v}" for k, v in slice_strings.items()) + return rf"{obj_string}_\text{{{subscript}}}" + + +@dataclass(frozen=True, kw_only=True) +class SliceRef(Node): + """A `$name` reference to a named slicer, valid only inside slice brackets.""" + + name: str + + @classmethod + def from_tokens(cls, instring: str, loc: int, tokens: pp.ParseResults) -> SliceRef: + """Parse action for `$name` slicer references.""" + return cls(name=str(tokens[0]), instring=instring, loc=loc) + + def evaluate(self, ctx: Context) -> Any: + """Evaluate the resolved slicer AST.""" + return ctx.slices[self.name].evaluate(ctx.demote()) + + def to_latex(self, ctx: Context) -> str: + """Render the resolved slicer AST.""" + return ctx.slices[self.name].to_latex(ctx) + + +@dataclass(frozen=True, kw_only=True) +class SubExprRef(Node): + """A `$name` reference to a named sub-expression.""" + + name: str + + @classmethod + def from_tokens( + cls, instring: str, loc: int, tokens: pp.ParseResults + ) -> SubExprRef: + """Parse action for `$name` sub-expression references.""" + return cls(name=str(tokens[0]), instring=instring, loc=loc) + + def evaluate(self, ctx: Context) -> Any: + """Evaluate the resolved sub-expression AST.""" + return ctx.sub_expressions[self.name].evaluate(ctx) + + def to_latex(self, ctx: Context) -> str: + """Render the resolved sub-expression AST.""" + return ctx.sub_expressions[self.name].to_latex(ctx) + + +@dataclass(frozen=True, kw_only=True) +class Call(Node): + """A helper-function call, e.g. `sum(flow, over=node)`.""" + + func: str + args: tuple[Node, ...] + kwargs: dict[str, Node] + + @classmethod + def from_tokens(cls, instring: str, loc: int, tokens: pp.ParseResults) -> Call: + """Parse action for `name(*args, **kwargs)` helper calls.""" + token_dict = tokens.as_dict() + args = tuple( + arg[0] if isinstance(arg, (pp.ParseResults, list)) else arg + for arg in token_dict.get("args", []) + ) + kwargs = { + name: val[0] if isinstance(val, (pp.ParseResults, list)) else val + for name, val in token_dict.get("kwargs", {}).items() + } + return cls( + func=token_dict["func"], + args=args, + kwargs=kwargs, + instring=instring, + loc=loc, + ) + + def _helper(self, ctx: Context) -> type[HelperFunction]: + """Return the helper class for this call, validating it exists in the registry.""" + helpers = ctx.helpers.get(ctx.helper_kind(), {}) + if self.func not in helpers: + raise error(ctx, self, f"Invalid helper function defined: {self.func}") + return helpers[self.func] + + def evaluate(self, ctx: Context) -> Any: + """ + Evaluate the helper-function call. + + The helper dispatches to `as_expr` in expr mode and `as_raw` otherwise. + Its arguments, however, are always evaluated in a demoted (raw) mode: + helpers must receive un-normalised inputs (`xr.DataArray` for parameters/ + lookups/dimensions and the raw model object for variables/expressions) + rather than values coerced to boolean masks or `LinearExpression`. + """ + helper_cls = self._helper(ctx) + helper = helper_cls(ctx) + if helper_cls.ignore_mask: + ctx = replace(ctx, mask=TRUE_ARRAY) + arg_ctx = ctx.demote() + args = [arg.evaluate(arg_ctx) for arg in self.args] + kwargs = {name: val.evaluate(arg_ctx) for name, val in self.kwargs.items()} + if ctx.mode == "expr": + return helper.as_expr(*args, **kwargs) + return helper.as_raw(*args, **kwargs) + + def to_latex(self, ctx: Context) -> str: + """Render the call via the helper's `as_math_string`.""" + helper = self._helper(ctx)(ctx) + args = [_call_arg_math_string(arg, ctx) for arg in self.args] + kwargs = { + name: _call_arg_math_string(val, ctx) for name, val in self.kwargs.items() + } + return helper.as_math_string(*args, **kwargs) + + +@dataclass(frozen=True, kw_only=True) +class Unary(Node): + """A unary operation: leading `+`/`-` sign or boolean `not`.""" + + op: str + operand: Node + + @classmethod + def from_tokens(cls, instring: str, loc: int, tokens: pp.ParseResults) -> Unary: + """Parse action for unary `+`/`-`/`not` operations.""" + op, operand = tokens[0] + return cls(op=str(op), operand=operand, instring=instring, loc=loc) + + def evaluate(self, ctx: Context) -> Any: + """Evaluate the operand and apply the unary operator.""" + if self.op == "not": + return ~self.operand.evaluate(ctx.demote()) + evaluated = self.operand.evaluate(ctx) + return -1 * evaluated if self.op == "-" else evaluated + + def to_latex(self, ctx: Context) -> str: + """Render the unary operation.""" + if self.op == "not": + return rf"\neg ({self.operand.to_latex(ctx)})" + return self.op + self.operand.to_latex(ctx) + + +@dataclass(frozen=True, kw_only=True) +class Arith(Node): + """ + A chain of same-precedence binary operations. + + Covers arithmetic (`**`, `*`, `/`, `+`, `-`) and boolean (`and`, `or`) + operator chains: `first OP operand OP operand ...`. + """ + + first: Node + rest: tuple[tuple[str, Node], ...] + + @classmethod + def from_tokens(cls, instring: str, loc: int, tokens: pp.ParseResults) -> Arith: + """Parse action for infix operator chains.""" + items = tokens[0] + rest = tuple( + (str(op), operand) + for op, operand in zip(items[1::2], items[2::2], strict=True) + ) + return cls(first=items[0], rest=rest, instring=instring, loc=loc) + + def evaluate(self, ctx: Context) -> Any: + """Evaluate the operands (masked, unless boolean) and fold the operator chain.""" + boolean = self.rest[0][0] in ("and", "or") + val = self.first.evaluate(ctx) + if not boolean: + val = val.where(ctx.mask) + for op, operand in self.rest: + evaluated = operand.evaluate(ctx) + if not boolean: + evaluated = evaluated.where(ctx.mask) + val = _OPERATIONS[op](val, evaluated) + return val + + def to_latex(self, ctx: Context) -> str: + """Render the operator chain, skipping identity operands.""" + val = self.first.to_latex(ctx) + for op, operand in self.rest: + evaluated = operand.to_latex(ctx) + # We ignore identity elements that do nothing (e.g. `0 + flow` is `flow`) + if evaluated == _LATEX_IDENTITIES.get(op): + continue + if isinstance(self.first, Arith): + val = f"({val})" + if isinstance(operand, Arith): + evaluated = f"({evaluated})" + if val == _LATEX_IDENTITIES.get(op): + val = evaluated + else: + val = _LATEX_OPERATORS[op].format(val=val, operand=evaluated) + return val + + +@dataclass(frozen=True, kw_only=True) +class Compare(Node): + """A comparison `lhs OP rhs`: an equation in expressions, a condition in masks.""" + + lhs: Node + op: str + rhs: Node + + @classmethod + def from_tokens(cls, instring: str, loc: int, tokens: pp.ParseResults) -> Compare: + """Parse action for `lhs OP rhs` comparisons.""" + lhs, op, rhs = tokens + return cls(lhs=lhs, op=str(op), rhs=rhs, instring=instring, loc=loc) + + def evaluate(self, ctx: Context) -> Any: + """ + Evaluate the comparison. + + In expr mode, return a masked `(lhs, sign, rhs)` tuple for constraint + assembly; otherwise return the boolean comparison array. + """ + if ctx.mode == "expr": + return self._evaluate_equation(ctx) + unmasked_ctx = replace(ctx, apply_mask=False) + comparison = _OPERATIONS[self.op]( + self.lhs.evaluate(unmasked_ctx), self.rhs.evaluate(unmasked_ctx) + ) + return xr.DataArray(comparison) + + def _evaluate_equation(self, ctx: Context) -> tuple[Any, xr.DataArray, Any]: + """Evaluate the equation to a masked `(lhs, sign, rhs)` tuple for constraint assembly.""" + lhs = self.lhs.evaluate(ctx) + rhs = self.rhs.evaluate(ctx) + for side, arr in {"left": lhs, "right": rhs}.items(): + extra_dims = set(arr.dims).difference(set(ctx.mask.dims) | {"_term"}) + if extra_dims: + raise error( + ctx, + self, + f"The {side}-hand side of the equation is indexed over " + f"dimensions not present in `foreach`: {extra_dims}", + ) + lhs_masked = to_linexpr(lhs.where(ctx.mask), ctx.model) + rhs_masked = to_linexpr(rhs.where(ctx.mask), ctx.model) + sign_masked = xr.DataArray(self.op).where(ctx.mask) + return lhs_masked, sign_masked, rhs_masked + + def to_latex(self, ctx: Context) -> str: + """Render the comparison with mask or equation operator tables per mode.""" + if ctx.mode == "mask": + unmasked_ctx = replace(ctx, apply_mask=False) + lhs_str = self.lhs.to_latex(unmasked_ctx) + rhs_str = self.rhs.to_latex(unmasked_ctx) + # Wrap plain-text tokens (coordinate labels, bare numbers, booleans) in `\text{}` so they render upright + if "\\" not in rhs_str: + rhs_str = rf"\text{{{rhs_str}}}" + return lhs_str + _LATEX_MASK_OPERATORS[self.op] + rhs_str + lhs_str = self.lhs.to_latex(ctx) + rhs_str = self.rhs.to_latex(ctx) + return lhs_str + _LATEX_EQUATION_OPERATORS[self.op] + rhs_str + + +@dataclass(frozen=True, kw_only=True) +class Subset(Node): + """A dimension subset condition, e.g. `[a, b] in node`.""" + + items: tuple[Node, ...] + dim: Node + + @classmethod + def from_tokens(cls, instring: str, loc: int, tokens: pp.ParseResults) -> Subset: + """Parse action for `[item, ...] in dim` subsets.""" + items, dim = tokens + return cls(items=tuple(items), dim=dim, instring=instring, loc=loc) + + def evaluate(self, ctx: Context) -> Any: + """Evaluate to a boolean array flagging dimension items in the subset.""" + subset = [_unwrap(item.evaluate(ctx)) for item in self.items] + dim_array = self.dim.evaluate(replace(ctx, apply_mask=False)) + return dim_array.isin(subset) + + def to_latex(self, ctx: Context) -> str: + """Render the subset condition.""" + subset = [_unwrap(item.evaluate(ctx)) for item in self.items] + # Subsets can range over lookups as well as dimensions; dim_iterator + # falls back to the plain name when there is no dimension iterator. + dim_name = ( + self.dim.name if isinstance(self.dim, Component) else self.dim.to_latex(ctx) + ) + iterator = dim_iterator(ctx.math, dim_name) + subset_string = "[" + ",".join(str(i) for i in subset) + "]" + return rf"\text{{{iterator}}} \in \text{{{subset_string}}}" + + +@dataclass(frozen=True, kw_only=True) +class ConfigRef(Node): + """A reference to a build-configuration option, e.g. `config.foo`.""" + + option: str + + @classmethod + def from_tokens(cls, instring: str, loc: int, tokens: pp.ParseResults) -> ConfigRef: + """Parse action for `config.option` references.""" + return cls(option=str(tokens[0]), instring=instring, loc=loc) + + def evaluate(self, ctx: Context) -> xr.DataArray: + """Evaluate the config option to a dimensionless array.""" + try: + config_val = getattr(ctx.config, self.option) + except AttributeError: + raise error( + ctx, self, f"Unknown configuration option: {self.option}" + ) from None + if not isinstance(config_val, int | float | str | bool | np.bool_): + raise error( + ctx, + self, + f"Configuration option resolves to invalid type " + f"`{type(config_val).__name__}`, expected a number, string, or boolean.", + ) + return xr.DataArray(config_val) + + def to_latex(self, ctx: Context) -> str: + """Render the config option reference.""" + return rf"\text{{config.{self.option}}}" + + +# --------------------------------------------------------------------------- +# Tree utilities +# --------------------------------------------------------------------------- + + +def iter_nodes(node: Node) -> list[Node]: + """ + Return `node` and all its descendant nodes, depth first. + + Parameters + ---------- + node : Node + Root of the (sub-)tree to walk. + """ + found = [node] + for f in fields(node): + val = getattr(node, f.name) + items: list = [] + if isinstance(val, Node): + items = [val] + elif isinstance(val, tuple): + items = [ + i + for pair in val + for i in (pair if isinstance(pair, tuple) else (pair,)) + ] + elif isinstance(val, dict): + items = list(val.values()) + for item in items: + if isinstance(item, Node): + found.extend(iter_nodes(item)) + return found + + +def find_refs(node: Node, of_type: type[Node]) -> set[str]: + """ + Return the names of all nodes of `of_type` found in the tree rooted at `node`. + + Parameters + ---------- + node : Node + Root of the (sub-)tree to search. + of_type : type[Node] + Node type with a `name` attribute to collect + (e.g. :class:`SubExprRef`, :class:`SliceRef`, :class:`Component`). + """ + return {n.name for n in iter_nodes(node) if isinstance(n, of_type)} # type: ignore[attr-defined] + + +# --------------------------------------------------------------------------- +# LaTeX helpers and functional walker wrappers +# --------------------------------------------------------------------------- + + +def _plain_string(item: Node, ctx: Context) -> str: + """Return a plain-text representation of a list item for LaTeX rendering.""" + evaluated = item.evaluate(ctx.demote()) + return ( + str(evaluated.name) if isinstance(evaluated, xr.DataArray) else str(evaluated) + ) + + +def _call_arg_math_string(arg: Node, ctx: Context) -> Any: + """ + Evaluate one helper-function argument for LaTeX rendering. + + List arguments are passed as raw item lists (so helpers can extract names); + all other arguments are passed as LaTeX strings. + """ + if isinstance(arg, ListNode): + return arg.evaluate(ctx) + return arg.to_latex(ctx) + + +def evaluate(node: Node, ctx: Context) -> Any: + """Evaluate a math AST node to data (see :meth:`Node.evaluate`).""" + return node.evaluate(ctx) + + +def to_math_string(node: Node, ctx: Context) -> str: + """Render a math AST node as a LaTeX math string (see :meth:`Node.to_latex`).""" + return node.to_latex(ctx) diff --git a/linopy/declarative/parsing.py b/linopy/declarative/parsing.py new file mode 100644 index 00000000..7f37f597 --- /dev/null +++ b/linopy/declarative/parsing.py @@ -0,0 +1,678 @@ +""" +Linopy declarative math parsing module. + +This module turns a validated math component definition into a list of :class:`Equation` objects. +The objects hold the parsed expression/mask ASTs with all `$name` sub-expression and slicer references resolved. +They also provide the typed entry points that evaluate an equation to a boolean mask array, a linopy expression, a constraint tuple, or a LaTeX math string. + +This module is adapted from the calliope Apache-2.0 licensed equation parsing module: +https://github.com/calliope-project/calliope/blob/9916116a06ec8c1feaf3c2606bdb8941b916ce85/src/calliope/backend/parsing.py +""" + +from __future__ import annotations + +import functools +import itertools +import logging +import operator +from dataclasses import dataclass, field, replace +from typing import Any, Literal + +import pyparsing as pp +import xarray as xr + +from linopy.declarative import grammar +from linopy.declarative.nodes import ( + MODE_T, + TRUE_ARRAY, + Component, + Context, + Node, + SliceRef, + SubExprRef, + find_refs, + to_linexpr, +) +from linopy.declarative.schema import ( + BUILD_ORDER, + COMPONENTS_T, + EQUATION_GROUP_T, + MATH_DEFS_T, + ConstraintDef, + ExpressionDef, + MathModel, + ObjectiveDef, + _Equations, +) +from linopy.expressions import LinearExpression + +LOGGER = logging.getLogger(__name__) + +EQUATION_DEFS_T = ConstraintDef | ExpressionDef | ObjectiveDef +"""Math component definitions that carry `equations`/`sub_expressions`/`slices` keys.""" + +_ERR_BULLET = " * " + + +@dataclass(frozen=True) +class Equation: + """ + One fully-resolved equation of a math component. + + Produced by :func:`parse_component`: each combination of sub-expression and + slicer variants referenced by a user-defined equation yields one `Equation`. + """ + + name: str + """Unique equation name, including the chosen sub-expression/slicer variants.""" + + sets: tuple[str, ...] + """The component's `foreach` dimensions.""" + + expression: Node + """Parsed expression AST.""" + + masks: tuple[Node, ...] + """Parsed mask ASTs: the equation's own mask plus those of the chosen variants.""" + + sub_expressions: dict[str, Node] = field(default_factory=dict) + """Resolved `$name` sub-expression AST per name.""" + + slices: dict[str, Node] = field(default_factory=dict) + """Resolved `$name` slicer AST per name.""" + + def references(self) -> set[str]: + """Return the names of all math components referenced by this equation.""" + trees = [ + self.expression, + *self.masks, + *self.sub_expressions.values(), + *self.slices.values(), + ] + return set().union(*(find_refs(tree, Component) for tree in trees)) + + +@dataclass(frozen=True) +class ParsedComponent: + """The parsed math strings of one active component.""" + + mask: Node + """Parsed top-level `mask` AST (the trivial "True" AST for objectives).""" + + equations: list[Equation] = field(default_factory=list) + """Fully-resolved equations (empty for variables, which carry none).""" + + +@dataclass(frozen=True) +class ParsedMath: + """All parsed math strings of a math definition's active components.""" + + components: dict[str, dict[str, ParsedComponent]] + """Per build group, the parsed strings of each active component by name.""" + + checks: dict[str, Node] + """Parsed mask AST of each active input-data check by name.""" + + def __getitem__(self, group: str) -> dict[str, ParsedComponent]: + """Return the parsed components of a build group.""" + return self.components[group] + + +def parse_math(math: MathModel) -> ParsedMath: + """ + Parse every math string of the active components, aggregating ALL failures. + + Walks the masks, equations, sub-expressions, and slices of active variables, + expressions, constraints, and objectives, plus active check masks. Inactive + components and groups without a build path (`piecewise_constraints`, + `postprocessed`) are deliberately not parsed: definitions there are + declarations of intent, and a stale string in them must never block a build. + + Parameters + ---------- + math : MathModel + The validated math definition. + + Returns + ------- + ParsedMath + Every math string of the walked components, parsed exactly once. + + Raises + ------ + ValueError + One error listing every parse failure (syntax errors and undefined + `$name` references), grouped by component. + """ + collector = _ParsingCollector() + mask_parser = _mask_grammar(math) + raw: dict[str, dict[str, tuple[Node | None, list[Equation]]]] = {} + for group in BUILD_ORDER: + raw[group] = {} + for name in getattr(math, group)._active: + definition = getattr(math, group)[name] + component_name = f"{group}:{name}" + mask_node = collector.parse( + mask_parser, definition.mask, component_name, "mask" + ) + equations = ( + parse_component(group, name, definition, math, collector) + if group != "variables" + else [] + ) + raw[group][name] = (mask_node, equations) + raw_checks = { + name: collector.parse( + mask_parser, math.checks[name].mask, f"checks:{name}", "mask" + ) + for name in math.checks._active + } + collector.raise_errors() + components = { + group: { + name: ParsedComponent(mask=mask_node, equations=equations) + for name, (mask_node, equations) in group_raw.items() + if mask_node is not None # always true after raise_errors + } + for group, group_raw in raw.items() + } + checks = {name: node for name, node in raw_checks.items() if node is not None} + return ParsedMath(components=components, checks=checks) + + +def _expression_names(math: MathModel) -> frozenset[str]: + """Return the valid component names for expression-string parsing.""" + return frozenset(set().union(*math.parsing_components["expression"].values())) + + +def _mask_grammar(math: MathModel) -> pp.ParserElement: + """Return the mask-string grammar for the math definition's component names.""" + names = math.parsing_components["mask"] + return grammar.mask_grammar( + frozenset(names["dimensions"]), + frozenset(names["inputs"]), + frozenset(names["results"]), + ) + + +class _ParsingCollector: + """ + Collect parsing results and errors per component. + + Collected errors can be raised as a single ValueError at the end. + """ + + def __init__(self) -> None: + self.errors: dict[str, list[str]] = {} + + def add(self, component_name: str, message: str) -> None: + """Store a plain (caret-free) error message against a component.""" + self.errors.setdefault(component_name, []).append(f"{_ERR_BULLET}{message}") + + def parse( + self, parser: pp.ParserElement, string: str, component_name: str, position: str + ) -> Node | None: + """ + Parse `string`, returning its AST root or None if parsing fails. + + Failures are stored against `component_name` with a caret marker pointing + at the parse position, for raising later via :meth:`raise_errors`. + """ + try: + return parser.parse_string(string, parse_all=True)[0] + except pp.ParseException as excinfo: + pointer = ( + f"{_ERR_BULLET}{position} (line {excinfo.lineno}, char {excinfo.col}): " + ) + marker_pos = " " * (len(pointer) + excinfo.col - 1) + self.errors.setdefault(component_name, []).append( + f"{pointer}{excinfo.line}\n{marker_pos}^" + ) + return None + + def raise_errors(self) -> None: + """Raise all collected errors as one ValueError, grouped by component.""" + if self.errors: + sections = "\n".join( + f"{component}:\n" + "\n".join(bullets) + for component, bullets in self.errors.items() + ) + raise ValueError(sections) + + +def parse_mask(mask_string: str, math: MathModel, name: str = "") -> Node: + """ + Parse a standalone mask string, raising on invalid syntax. + + Parameters + ---------- + mask_string : str + The mask ("where"-condition) string to parse. + math : MathModel + Math definition providing the valid component names. + name : str, optional + Name to identify the string by in error messages. + """ + collector = _ParsingCollector() + parsed = collector.parse(_mask_grammar(math), mask_string, name, "mask") + collector.raise_errors() + assert parsed is not None + return parsed + + +def _parse_variants( + collector: _ParsingCollector, + component_name: str, + parser: pp.ParserElement, + mask_parser: pp.ParserElement, + expression_list: _Equations, + sets: tuple[str, ...], + position: str, + name_prefix: str, +) -> list[Equation]: + """Parse a list of `{mask, expression}` items into one Equation per item.""" + equations = [] + for idx, item in enumerate(expression_list): + position_id = f"{position}[{idx}]" + mask = collector.parse( + mask_parser, item.mask, component_name, f"{position_id}.mask" + ) + expression = collector.parse( + parser, item.expression, component_name, f"{position_id}.expression" + ) + if expression is not None and mask is not None: + equations.append( + Equation( + name=f"{name_prefix}:{idx}", + sets=sets, + expression=expression, + masks=(mask,), + ) + ) + return equations + + +def _expand( + component_name: str, + equations: list[Equation], + candidates: dict[str, list[Equation]], + kind: Literal["sub_expressions", "slices"], + collector: _ParsingCollector, +) -> list[Equation]: + """ + Expand equations with all combinations of their referenced `$name` variants. + + Each `$name` reference maps to a list of `{mask, expression}` variants. + An equation referencing them is replaced by one equation per element of the + cartesian product of those variant lists. + The chosen variants' masks and ASTs are merged in. + Undefined `$name` references are collected (not raised) so they aggregate with any other parse failures. + """ + expanded = [] + for equation in equations: + ref_type = SubExprRef if kind == "sub_expressions" else SliceRef + trees = [equation.expression, *equation.sub_expressions.values()] + refs = set().union(*(find_refs(tree, ref_type) for tree in trees)) + if not refs: + expanded.append(equation) + continue + undefined = refs.difference(candidates.keys()) + if undefined: + collector.add( + component_name, + f"Undefined {kind} found in equation: {sorted(undefined)}", + ) + continue + for combination in itertools.product(*(candidates[ref] for ref in refs)): + new_name = "-".join([equation.name, *(v.name for v in combination)]) + new_masks = ( + *equation.masks, + *(mask for variant in combination for mask in variant.masks), + ) + resolved = { + variant.name.split(":")[0]: variant.expression + for variant in combination + } + if kind == "sub_expressions": + new_equation = replace( + equation, name=new_name, masks=new_masks, sub_expressions=resolved + ) + else: + new_equation = replace( + equation, name=new_name, masks=new_masks, slices=resolved + ) + expanded.append(new_equation) + return expanded + + +def parse_component( + group: EQUATION_GROUP_T, + name: str, + definition: EQUATION_DEFS_T, + math: MathModel, + collector: _ParsingCollector | None = None, +) -> list[Equation]: + """ + Parse a math component's equations into fully-resolved :class:`Equation` objects. + + All `expression` and `mask` strings of the component's equations, sub-expressions, and slicers are parsed. + Then every equation is expanded with the cartesian product of the sub-expression and slicer variants it references. + + Parameters + ---------- + group : EQUATION_GROUP_T + Component group the definition belongs to (defines the equation grammar: + comparisons for constraints, arithmetic otherwise). + name : str + Name of the math component. + definition : EQUATION_DEFS_T + The component's (already schema-validated) definition. + math : MathModel + The full math definition, providing valid component names. + collector : _ParsingCollector, optional + Shared error collector (used by :func:`parse_math` to aggregate failures across components). + If not given, all failures of this component are raised at the end of the call. + + Returns + ------- + list[Equation] + One equation per user-defined equation and referenced variant combination. + """ + component_name = f"{group}:{name}" + own_collector = collector is None + collector = collector or _ParsingCollector() + names = _expression_names(math) + equation_parser = ( + grammar.equation_grammar(names) + if group == "constraints" + else grammar.arithmetic_grammar(names) + ) + mask_parser = _mask_grammar(math) + sets = tuple(definition.foreach) + + equations = _parse_variants( + collector, + component_name, + equation_parser, + mask_parser, + definition.equations, + sets, + "equations", + component_name, + ) + sub_expressions = { + sub_name: _parse_variants( + collector, + component_name, + grammar.sub_expression_grammar(names), + mask_parser, + sub_list, + sets, + f"sub_expressions.{sub_name}", + sub_name, + ) + for sub_name, sub_list in definition.sub_expressions.root.items() + } + slices = { + slice_name: _parse_variants( + collector, + component_name, + grammar.slice_grammar(names), + mask_parser, + slice_list, + sets, + f"slices.{slice_name}", + slice_name, + ) + for slice_name, slice_list in definition.slices.root.items() + } + + equations = _expand( + component_name, equations, sub_expressions, "sub_expressions", collector + ) + equations = _expand(component_name, equations, slices, "slices", collector) + if own_collector: + collector.raise_errors() + return equations + + +# --------------------------------------------------------------------------- +# Component-level masking +# --------------------------------------------------------------------------- + + +def foreach_mask(sets: tuple[str, ...], input_data: xr.Dataset) -> xr.DataArray: + """ + Return the initial boolean array spanning a component's `foreach` dimensions. + + Parameters + ---------- + sets : tuple[str, ...] + The component's `foreach` dimensions. + input_data : xr.Dataset + Model input data providing the dimension coordinates. + """ + missing_sets = set(sets).difference(input_data.dims) + if missing_sets: + LOGGER.debug( + f"Math parsing | indexed over unidentified set names: `{missing_sets}`." + ) + return xr.DataArray(False) + if not sets: + return TRUE_ARRAY + exists_and_foreach = [input_data[i].notnull() for i in sets] + return functools.reduce(operator.and_, exists_and_foreach) + + +def drop_dims_not_in_foreach(mask: xr.DataArray, sets: tuple[str, ...]) -> xr.DataArray: + """ + Reduce a mask array to a component's `foreach` dimensions. + + Any dimension not in `sets` is reduced with a boolean any-operation, and the + result is transposed to the order given by `sets`. + """ + unwanted_dims = set(mask.dims).difference(sets) + return (mask.sum(unwanted_dims) > 0).astype(bool).transpose(*sets) + + +def _mask_is_empty(mask: xr.DataArray, name: str, reason: str) -> bool: + """Return True (with a debug log) if `mask` leaves no valid data point.""" + if not mask.any(): + LOGGER.debug(f"Math parsing | {name} | Component not added; {reason}.") + return True + return False + + +def component_mask( + group: COMPONENTS_T, + name: str, + definition: MATH_DEFS_T, + mask_node: Node, + ctx: Context, + *, + align_to_foreach_sets: bool = True, +) -> xr.DataArray: + """ + Evaluate a component's top-level mask over its `foreach` dimensions. + + Combines the `foreach` existence array with the component's pre-parsed top-level `mask` AST, breaking early if no valid element remains. + + Parameters + ---------- + group : COMPONENTS_T + Component group the definition belongs to. + name : str + Name of the math component. + definition : MATH_DEFS_T + The component's (already schema-validated) definition. + mask_node : Node + The component's pre-parsed top-level `mask` AST + (e.g. from :func:`parse_math` / :class:`ParsedMath`). + ctx : Context + Evaluation context. + align_to_foreach_sets : bool, default: True + If True, reduce the result to the `foreach` dimensions + (see :func:`drop_dims_not_in_foreach`). + """ + component_name = f"{group}:{name}" + sets = tuple(definition.foreach) + initial_mask = foreach_mask(sets, ctx.input_data) + if _mask_is_empty( + initial_mask, component_name, "'foreach' does not apply anywhere" + ): + return initial_mask + + mask_ctx = replace(ctx, mode="mask", equation_name=component_name) + mask = xr.DataArray(initial_mask & mask_node.evaluate(mask_ctx)) + if _mask_is_empty(mask, component_name, "'mask' does not apply anywhere"): + return mask + + if align_to_foreach_sets: + mask = drop_dims_not_in_foreach(mask, sets) + return mask + + +def _equation_ctx( + equation: Equation, + ctx: Context, + mode: MODE_T, + **kwargs: Any, +) -> Context: + """Return a context copy carrying the equation's name and resolved references.""" + return replace( + ctx, + equation_name=equation.name, + mode=mode, + sub_expressions=equation.sub_expressions, + slices=equation.slices, + **kwargs, + ) + + +def as_mask( + equation: Equation, ctx: Context, *, initial_mask: xr.DataArray = TRUE_ARRAY +) -> xr.DataArray: + """ + Evaluate an equation's mask strings to a boolean array. + + Parameters + ---------- + equation : Equation + Parsed equation. + ctx : Context + Evaluation context. + initial_mask : xr.DataArray, optional + Mask to combine (boolean AND) with the equation's own masks. + E.g., the component-level mask from :func:`component_mask`. + + Returns + ------- + xr.DataArray + Boolean array defining on which index items the equation applies. + """ + mask_ctx = _equation_ctx(equation, ctx, "mask") + evaluated = [mask.evaluate(mask_ctx) for mask in equation.masks] + mask = xr.DataArray(functools.reduce(operator.and_, [initial_mask, *evaluated])) + _mask_is_empty(mask, equation.name, "'mask' does not apply anywhere") + return mask + + +def as_expression( + equation: Equation, ctx: Context, *, mask: xr.DataArray = TRUE_ARRAY +) -> LinearExpression: + """ + Evaluate an equation's arithmetic expression to a linopy expression. + + Parameters + ---------- + equation : Equation + Parsed equation (from an `expressions`/`objectives` component). + ctx : Context + Evaluation context. + mask : xr.DataArray, optional + Boolean array with which to mask the produced arrays. + + Returns + ------- + LinearExpression + The evaluated expression; a pure-parameter expression (evaluated to an + `xr.DataArray`) is coerced to a `LinearExpression`. + """ + expr_ctx = _equation_ctx(equation, ctx, "expr", mask=mask) + return to_linexpr(equation.expression.evaluate(expr_ctx), ctx.model) + + +def as_constraint( + equation: Equation, ctx: Context, *, mask: xr.DataArray = TRUE_ARRAY +) -> tuple[LinearExpression, xr.DataArray, LinearExpression]: + """ + Evaluate an equation of the form `LHS OP RHS` to a constraint tuple. + + Parameters + ---------- + equation : Equation + Parsed equation (from a `constraints` component). + ctx : Context + Evaluation context. + mask : xr.DataArray, optional + Boolean array with which to mask the produced arrays. + + Returns + ------- + tuple[LinearExpression, xr.DataArray, LinearExpression] + `(lhs, sign, rhs)` for constraint assembly. + `lhs` and `rhs` are coerced to `LinearExpression`. + `sign` is an array of the comparison operator. + """ + expr_ctx = _equation_ctx(equation, ctx, "expr", mask=mask) + lhs, sign, rhs = equation.expression.evaluate(expr_ctx) + return lhs, sign, rhs + + +def as_latex_mask( + equation: Equation, + ctx: Context, +) -> str: + """ + Render an equation's mask as a LaTeX math string. + + Parameters + ---------- + equation : Equation + Parsed equation. + ctx : Context + Evaluation context. + what : Literal["expression", "mask"], default: "expression" + Whether to render the equation's expression (including an equation's comparison operator) or its combined mask conditions. + + Returns + ------- + str + A valid LaTeX math string. + """ + mask_ctx = _equation_ctx(equation, ctx, "mask") + strings = [mask.to_latex(mask_ctx) for mask in equation.masks] + return r"\land{}".join(f"({s})" for s in strings if s != "true") + + +def as_latex_expression( + equation: Equation, + ctx: Context, +) -> str: + """ + Render an equation's expression as a LaTeX math string. + + Parameters + ---------- + equation : Equation + Parsed equation. + ctx : Context + Evaluation context. + + Returns + ------- + str + A valid LaTeX math string. + """ + expr_ctx = _equation_ctx(equation, ctx, "raw") + return equation.expression.to_latex(expr_ctx) diff --git a/linopy/declarative/schema.py b/linopy/declarative/schema.py new file mode 100644 index 00000000..d5a6cfe1 --- /dev/null +++ b/linopy/declarative/schema.py @@ -0,0 +1,606 @@ +""" +Linopy declarative math schema module. + +This module contains the pydantic models that validate declarative math and build-configuration definitions. + +This module is adapted from the calliope Apache-2.0 licensed math schema module: +https://github.com/calliope-project/calliope/blob/9916116a06ec8c1feaf3c2606bdb8941b916ce85/src/calliope/schemas/math_schema.py +""" + +from __future__ import annotations + +import logging +from collections.abc import Hashable, Iterable +from functools import cached_property +from typing import Annotated, Any, ClassVar, Literal, Self, TypeVar, get_args + +import numpy as np +from annotated_types import Len +from pydantic import AfterValidator, BaseModel, Field, RootModel, model_validator +from pydantic_core import PydanticCustomError + +LOGGER = logging.getLogger(__name__) + +# Modified from https://github.com/pydantic/pydantic-core/pull/820#issuecomment-1670475909 +T = TypeVar("T", bound=Hashable | list) + +DTYPE_OPTIONS: dict[str, type] = { + "string": str, + "float": float, + "bool": bool, + "integer": int, + "datetime": np.datetime64, + "date": np.datetime64, +} +"""Mapping from math-schema dtype names to Python/numpy types.""" + +COMPONENTS_T = Literal[ + "dimensions", + "parameters", + "lookups", + "variables", + "expressions", + "constraints", + "piecewise_constraints", + "objectives", + "postprocessed", +] + +EQUATION_GROUP_T = Literal["expressions", "constraints", "objectives"] + +BUILD_ORDER_T = Literal["variables", "expressions", "constraints", "objectives"] +BUILD_ORDER: tuple[BUILD_ORDER_T, ...] = get_args(BUILD_ORDER_T) + +"""Component groups in the order they are built into a linopy model.""" + + +def _validate_unique_list(v: list) -> list: + try: + unique = set(v) + except TypeError: + unique = set([tuple(i) for i in v]) + if len(v) != len(unique): + raise PydanticCustomError("unique_list", "List must be unique") + return v + + +UniqueList = Annotated[ + list[T], + AfterValidator(_validate_unique_list), + Field(json_schema_extra={"uniqueItems": True}), +] +"""A list with no repeated values.""" +# == +NonEmptyList = Annotated[list[T], Len(min_length=1)] +"""A list with at least one value in it.""" +NonEmptyUniqueList = Annotated[UniqueList[T], Len(min_length=1)] +"""A list with at least one value in it and no repeated values.""" +AttrStr = Annotated[str, Field(pattern=r"^[^_^\d][\w]*$")] +"""Single word string in snake_case (e.g., wind_offshore).""" +NumericVal = int | Annotated[float, Field(allow_inf_nan=True)] +"""Numerical integer or float value. Can be `nan` or infinite (`float(inf)`).""" + + +class LinopyDictModel(RootModel): + """Pydantic model storing a dictionary of user-named component definitions.""" + + def __setitem__(self, *args: Any, **kwargs: Any) -> None: + """Do not allow direct item setting.""" + raise PydanticCustomError( + "no_extra_dict", + f"Cannot set a {self.__class__.__name__} item directly. Re-validate a new definition dictionary instead.", + ) + + def __getitem__(self, key: str) -> Any: + """Expose the root attribute when getting an item by key.""" + return self.root[key] + + def __repr__(self, *args: Any, **kwargs: Any) -> str: + """Show the __repr__ of the root attribute when requesting the __repr__ of the class.""" + return self.root.__repr__(*args, **kwargs) + + def __rich_repr__(self) -> Iterable: + """Prettyprint the __repr__ of the root attribute when requesting the prettyprint of the class.""" + yield from self.root.items() + + @cached_property + def _active(self) -> dict[str, BaseModel]: + """Return only active components.""" + return {k: v for k, v in self.root.items() if v.active} + + +class LinopyListModel(RootModel): + """Pydantic model storing a list of definitions.""" + + def __iter__(self) -> Any: + """Iterate over root attribute contents when iterating over class.""" + return iter(self.root) + + def __getitem__(self, item: int) -> Any: + """Expose the root attribute when getting an item by index value.""" + return self.root[item] + + def __repr__(self, *args: Any, **kwargs: Any) -> str: + """Show the __repr__ of the root attribute when requesting the __repr__ of the class.""" + return self.root.__repr__(*args, **kwargs) + + def __rich_repr__(self) -> Iterable: + """Prettyprint the __repr__ of the root attribute when requesting the prettyprint of the class.""" + yield from self.root + + +class LinopyBaseModel(BaseModel): + """Base class for declarative math pydantic models.""" + + model_config = { + "extra": "forbid", + "frozen": True, + "revalidate_instances": "always", + "use_attribute_docstrings": True, + } + + def __getitem__(self, item: str) -> Any: + """Allow attribute access via item lookup.""" + return getattr(self, item) + + +class _ExpressionItem(LinopyBaseModel): + """Schema for equations, _subexpressions and slices.""" + + mask: str = "True" + """Condition to determine whether the accompanying expression is built.""" + expression: str + """Expression for this component. + - _Equations: LHS OPERATOR RHS, where LHS and RHS are math expressions and OPERATOR is one of [==, <=, >=]. + - _Subexpressions: be one term or a combination of terms using the operators [+, -, *, /, **]. + - Slices: a list of set items or a call to a helper function. + """ + + +class _MathComponent(LinopyBaseModel): + """Generic math component class.""" + + title: str = "" + """The component long name, for use in visualisation.""" + description: str = "" + """A verbose description of the component.""" + active: bool = True + """If False, this component will be ignored during the build phase.""" + + _group: ClassVar[COMPONENTS_T] + """Return the component group this component belongs to.""" + + +class DimensionDef(_MathComponent): + """Schema for named dimension.""" + + dtype: Literal["string", "datetime", "date", "float", "integer"] = "string" + """The data type of this dimension's items.""" + ordered: bool = False + """If True, the order of the dimension items is meaningful (e.g. chronological time).""" + iterator: str = "NEEDS_ITERATOR" + """The name of the iterator to use in the LaTeX math formulation for this dimension.""" + + _group: ClassVar[COMPONENTS_T] = "dimensions" + + @property + def default(self) -> float: + """Dummy field to align with lookups and dims.""" + return float("nan") + + +class ParameterDef(_MathComponent): + """Schema for named parameter.""" + + default: float | int = float("nan") + """The default value for the parameter, if not set in the data.""" + unit: str = "" + """The unit of the parameter, e.g. 'kW', 'm', 'kg', 'energy', 'power', ...""" + dims: UniqueList[AttrStr] | None = Field(default=None) + """The dimensions over which the parameter can be defined. + + It is not necessary for a parameter to be defined over all of its dimensions. + + If undefined, it is assumed that the parameter can be indexed over all dimensions. + If an empty list, it is assumed that the parameter is a scalar and not indexed over any dimensions. + """ + + @property + def dtype(self) -> Literal["float"]: + """Dummy field to align with lookups and dims.""" + return "float" + + _group: ClassVar[COMPONENTS_T] = "parameters" + + +class LookupDef(_MathComponent): + """Schema for named lookup arrays.""" + + default: AttrStr | float | int | bool = float("nan") + """The default value for the lookup, if not set in the data.""" + dtype: Literal["float", "string", "bool", "datetime", "date"] = "string" + """The lookup data type.""" + dims: UniqueList[AttrStr] | None = Field(default=None) + """The dimensions over which the lookup can be defined. + + It is not necessary for a lookup to be defined over all of its dimensions. + + If undefined, it is assumed that the lookup can be indexed over all dimensions. + If an empty list, it is assumed that the lookup is a scalar and not indexed over any dimensions.""" + one_of: list | None = None + """If given, the lookup values must be one of these items.""" + + _group: ClassVar[COMPONENTS_T] = "lookups" + + +class _MathIndexedComponent(_MathComponent): + """Generic indexed component class.""" + + foreach: UniqueList[AttrStr] = Field(default_factory=list) + """Sets (a.k.a. dimensions) of the model over which the math formulation component + will be built.""" + mask: str = "True" + """Top-level condition to determine whether the component exists in this + optimisation problem. At all if `foreach` is not given, or for specific index items + within the product of the sets given by `foreach`.""" + + +class _Equations(LinopyListModel): + """List of equations that can be updated when a parent pydantic model is updated.""" + + root: list[_ExpressionItem] = Field(default_factory=list) + + +class _SubExpressions(LinopyDictModel): + """Dictionary of sub-expressions that can be updated when a parent pydantic model is updated.""" + + root: dict[AttrStr, _Equations] = Field(default_factory=dict) + + +class _MathEquationComponent(_MathComponent): + """Components necessary to generate math expressions.""" + + equations: _Equations = _Equations() + """Constraint math equations.""" + sub_expressions: _SubExpressions = _SubExpressions() + """Named sub-expressions.""" + slices: _SubExpressions = _SubExpressions() + """Named index slices.""" + + @model_validator(mode="after") + def must_have_equations_if_active(self) -> Self: + """Ensure that equations are defined if the component is active.""" + if self.active and not self.equations.root: + raise ValueError("Must have equations defined if component is active.") + return self + + +class ConstraintDef(_MathIndexedComponent, _MathEquationComponent): + """Schema for named constraints.""" + + _group: ClassVar[COMPONENTS_T] = "constraints" + + +class PiecewiseConstraintDef(_MathIndexedComponent): + """ + Schema for named piece-wise constraints. + + These link an `x`-axis decision variable with a `y`-axis decision variable with + values at specified breakpoints. + """ + + x_expression: str + """X variable name whose values are assigned at each breakpoint.""" + y_expression: str + """Y variable name whose values are assigned at each breakpoint.""" + x_values: str + """X parameter name containing data, indexed over the `breakpoints` dimension.""" + y_values: str + """Y parameter name containing data, indexed over the `breakpoints` dimension.""" + + _group: ClassVar[COMPONENTS_T] = "piecewise_constraints" + + +class ExpressionDef(_MathIndexedComponent, _MathEquationComponent): + """ + Schema for named expressions. + + Can be used to combine parameters and variables and then used in one or more + expressions elsewhere in the math formulation (i.e., in constraints, objectives, + and other expressions). + + NOTE: If expecting to use expression `A` in expression `B`, `A` must + be defined above `B`. + """ + + unit: str = "" + """Generalised unit of the component (e.g., length, time, quantity_per_hour, ...).""" + default: NumericVal = float("nan") + """If set, will be the default value for the expression.""" + equations: _Equations = _Equations() + """Expression math equations.""" + sub_expressions: _SubExpressions = _SubExpressions() + """Expression named sub-expressions.""" + slices: _SubExpressions = _SubExpressions() + """Expression named index slices.""" + order: int = 0 + """Order in which to apply this expression relative to all others, if different to its definition order.""" + + _group: ClassVar[COMPONENTS_T] = "expressions" + + +class _Bounds(LinopyBaseModel): + """ + Bounds of decision variables. + + Either derived per-index item from a multi-dimensional input parameter, or given as + a single value that is applied across all decision variable index items. + """ + + upper: AttrStr | NumericVal = float("inf") + """Decision variable upper bound, either as a reference to an input parameter or as a number.""" + lower: AttrStr | NumericVal = float("-inf") + """Decision variable lower bound, either as a reference to an input parameter or as a number.""" + + +class VariableDef(_MathIndexedComponent): + """ + Schema for optimisation problem variables. + + A decision variable must be referenced in at least one constraint or in the + objective for it to exist in the optimisation problem that is sent to the solver. + """ + + unit: str = "" + """Generalised unit of the component (e.g., length, time, quantity_per_hour, ...).""" + default: NumericVal = float("nan") + """If set, will be the default value for the variable.""" + domain: Literal["real", "integer"] = "real" + """Allowed values that the decision variable can take. + Either real (a.k.a. continuous) or integer.""" + bounds: _Bounds = _Bounds() + + _group: ClassVar[COMPONENTS_T] = "variables" + + +class ObjectiveDef(_MathEquationComponent): + """ + Schema for optimisation problem objectives. + + Only one objective, the one referenced in model configuration `build.objective` + will be activated for the optimisation problem. + """ + + sense: Literal["min", "max"] + """Whether the objective function should be minimised or maximised in the + optimisation.""" + + _group: ClassVar[COMPONENTS_T] = "objectives" + + @property + def foreach(self) -> UniqueList: + """Dummy field to align with other math components.""" + return [] + + @property + def mask(self) -> str: + """Dummy field to align with other math components.""" + return "True" + + +class PostprocessedExpressionDef(ExpressionDef): + """ + Schema for postprocessed expressions. + + Can be used to combine parameters, variables, and expressions into a single expression solving the model. + + NOTE: If expecting to use postprocessed array `A` in postprocessed array `B`, `A` must + be defined above `B`. + """ + + _group: ClassVar[COMPONENTS_T] = "postprocessed" + + +class CheckDef(LinopyBaseModel): + """Schema for input data checks.""" + + mask: str + """Top-level condition to check""" + message: str + """Message to display when the `mask` array returns True, if raising or warning on error.""" + errors: Literal["raise", "warn"] = "raise" + """How to respond to any instances in which the `mask` array returns True.""" + active: bool = True + """If False, this check will be ignored during the build phase.""" + + +class DimensionDefs(LinopyDictModel): + """Linopy model dimensions dictionary.""" + + root: dict[AttrStr, DimensionDef] = Field(default_factory=dict) + + +class ParameterDefs(LinopyDictModel): + """Linopy model parameters dictionary.""" + + root: dict[AttrStr, ParameterDef] = Field(default_factory=dict) + + +class LookupDefs(LinopyDictModel): + """Linopy model lookup dictionary.""" + + root: dict[AttrStr, LookupDef] = Field(default_factory=dict) + + +class VariableDefs(LinopyDictModel): + """Linopy model variables dictionary.""" + + root: dict[AttrStr, VariableDef] = Field(default_factory=dict) + + +class ExpressionDefs(LinopyDictModel): + """Linopy model expressions dictionary.""" + + root: dict[AttrStr, ExpressionDef] = Field(default_factory=dict) + + +class ConstraintDefs(LinopyDictModel): + """Linopy model constraints dictionary.""" + + root: dict[AttrStr, ConstraintDef] = Field(default_factory=dict) + + +class PiecewiseConstraintDefs(LinopyDictModel): + """Linopy model piecewise_constraints dictionary.""" + + root: dict[AttrStr, PiecewiseConstraintDef] = Field(default_factory=dict) + + +class ObjectiveDefs(LinopyDictModel): + """Linopy model objectives dictionary.""" + + root: dict[AttrStr, ObjectiveDef] = Field(default_factory=dict) + + +class PostprocessedExpressionDefs(LinopyDictModel): + """Linopy model postprocessed expressions dictionary.""" + + root: dict[AttrStr, PostprocessedExpressionDef] = Field(default_factory=dict) + + +class Checks(LinopyDictModel): + """Linopy math checks dictionary.""" + + root: dict[AttrStr, CheckDef] = Field(default_factory=dict) + + +class MathModel(LinopyBaseModel): + """ + Declarative definition of a linopy optimisation problem. + + Contains all mathematical programming components from which a linopy model + can be built. + """ + + model_config = {"title": "Model math schema"} + + dimensions: DimensionDefs = DimensionDefs() + """All dimensions to include in the optimisation problem.""" + parameters: ParameterDefs = ParameterDefs() + """All parameters to include in the optimisation problem.""" + lookups: LookupDefs = LookupDefs() + """All lookups to include in the optimisation problem.""" + variables: VariableDefs = VariableDefs() + """All decision variables to include in the optimisation problem.""" + expressions: ExpressionDefs = ExpressionDefs() + """All expressions that can be applied to the optimisation problem.""" + constraints: ConstraintDefs = ConstraintDefs() + """All constraints to apply to the optimisation problem.""" + piecewise_constraints: PiecewiseConstraintDefs = PiecewiseConstraintDefs() + """All _piecewise_ constraints to apply to the optimisation problem.""" + objectives: ObjectiveDefs = ObjectiveDefs() + """Possible objectives to apply to the optimisation problem.""" + postprocessed: PostprocessedExpressionDefs = PostprocessedExpressionDefs() + """All postprocessed expressions generated after math has completed.""" + checks: Checks = Checks() + """Checks to apply before building the optimisation problem.""" + + @model_validator(mode="after") + def unique_component_names(self) -> Self: + """Ensure all component names are unique.""" + groups = sorted( + ( + {name for name in getattr(self, field)._active} + for field in type(self).model_fields + ), + key=len, + ) + seen: set[str] = set() + duplicates: set[str] = set() + for field_names in groups: + duplicates |= field_names & seen + seen |= field_names + if duplicates: + raise ValueError( + f"Non-unique names in math components: {sorted(duplicates)}." + ) + + return self + + @model_validator(mode="after") + def validate_dimensions(self) -> Self: + """Ensure all dimensions referenced in the math schema are defined.""" + defined_dims = set(self.dimensions._active) + for field in type(self).model_fields: + if field == "dimensions": + continue + dim_field = "dims" if field in ("parameters", "lookups") else "foreach" + for component, _def in getattr(self, field)._active.items(): + if hasattr(_def, dim_field): + component_dims = set(getattr(_def, dim_field) or defined_dims) + missing = component_dims - defined_dims + if missing: + raise ValueError( + f"Component `{field}.{component}` references undefined dimensions: {sorted(missing)}." + ) + return self + + @cached_property + def parsing_components(self) -> dict[str, dict[str, set[str]]]: + """ + Return the valid component names available to each parser. + + Returns + ------- + dict[str, dict[str, set[str]]] + Per parser kind ("expression" / "mask"), the valid names grouped by + where they are defined in the math ("dimensions" / "inputs" / "results"). + """ + parsing_components = { + "dimensions": ["dimensions"], + "inputs": ["lookups", "parameters"], + "results": ["variables", "expressions"], + } + + def _names() -> dict[str, set[str]]: + return { + k: set().union(*[getattr(self, i)._active for i in v]) + for k, v in parsing_components.items() + } + + mask_names = _names() + all_active = mask_names["results"].union(mask_names["inputs"]) + for component in ["inputs", "results"]: + all_names = set().union( + *(getattr(self, k).root for k in parsing_components[component]) + ) + mask_names[component] |= all_names - all_active + all_components = {"expression": _names(), "mask": mask_names} + + return all_components + + def find( + self, component: str, subset: Iterable[COMPONENTS_T] | None = None + ) -> _MathComponent: + """Find a component in the math schema.""" + fields: Iterable = subset or (set(type(self).model_fields) - {"checks"}) + + found = {f for f in fields if component in getattr(self, f)._active} + if not found: + raise KeyError(f"Component name `{component}` not found in math schema.") + if len(found) > 1: + raise ValueError( + f"Component name `{component}` found in multiple places: {found}." + ) + return getattr(self, found.pop())[component] + + +MATH_DEFS_T = ( + ConstraintDef | VariableDef | ExpressionDef | ObjectiveDef | PiecewiseConstraintDef +) + + +class ConfigModel(LinopyBaseModel): + """Base configuration options used when building a Linopy optimisation problem.""" + + model_config = {"title": "Model build configuration", "extra": "allow"} diff --git a/linopy/expressions.py b/linopy/expressions.py index 21a4160e..7a7d5a8b 100644 --- a/linopy/expressions.py +++ b/linopy/expressions.py @@ -11,10 +11,26 @@ import logging import operator from abc import ABC, abstractmethod -from collections.abc import Callable, Hashable, Iterable, Iterator, Mapping, Sequence +from collections.abc import ( + Callable, + Hashable, + ItemsView, + Iterable, + Iterator, + Mapping, + Sequence, +) from dataclasses import dataclass, field from itertools import product, zip_longest -from typing import TYPE_CHECKING, Any, Self, TypeAlias, TypeVar, cast, overload +from typing import ( + TYPE_CHECKING, + Any, + Self, + TypeAlias, + TypeVar, + cast, + overload, +) from warnings import warn import numpy as np @@ -55,6 +71,7 @@ filter_nulls_polars, format_coord, format_single_expression, + format_string_as_variable_name, forward_as_properties, generate_indices_for_printout, get_dims_with_index_levels, @@ -64,6 +81,7 @@ is_constant, iterate_slices, maybe_group_terms_polars, + save_join, to_dataframe, to_polars, ) @@ -735,6 +753,7 @@ def __init__(self, data: Dataset | Any | None, model: Model) -> None: # TODO: add a warning here, routines should be safe against this data = data.drop_vars(drop_dims) + data = data.assign_attrs(name=None) self._model = model self._data = cast(Dataset, data) @@ -1235,6 +1254,13 @@ def loc(self) -> LocIndexer: def type(self) -> str: return "LinearExpression" + @property + def name(self) -> str: + """ + Return the name of the variable. + """ + return str(self.attrs["name"]) + @property def data(self) -> Dataset: return self._data @@ -2827,6 +2853,127 @@ def merge( return cls(ds, model) +@dataclass(repr=False) +class Expressions: + """ + An expressions container used for storing multiple expression arrays. + """ + + data: dict[str, LinearExpression | QuadraticExpression] + model: Model + + def _formatted_names(self) -> dict[str, str]: + """ + Get a dictionary of formatted names to the proper variable names. + This map enables a attribute like accession of variable names which + are not valid python variable names. + """ + return {format_string_as_variable_name(n): n for n in self} + + @overload + def __getitem__(self, names: str) -> LinearExpression | QuadraticExpression: ... + + @overload + def __getitem__(self, names: list[str]) -> Expressions: ... + + def __getitem__( + self, names: str | list[str] + ) -> LinearExpression | QuadraticExpression | Expressions: + if isinstance(names, str): + return self.data[names] + return Expressions({name: self.data[name] for name in names}, self.model) + + def __getattr__(self, name: str) -> LinearExpression | QuadraticExpression: + # If name is an attribute of self (including methods and properties), return that + if name in self.data: + return self.data[name] + else: + if name in (formatted_names := self._formatted_names()): + return self.data[formatted_names[name]] + raise AttributeError( + f"Expressions has no attribute `{name}` or the attribute is not accessible / raises an error." + ) + + def __getstate__(self) -> dict: + return self.__dict__ + + def __setstate__(self, d: dict) -> None: + self.__dict__.update(d) + + def __dir__(self) -> list[str]: + base_attributes = list(super().__dir__()) + formatted_names = [ + n for n in self._formatted_names() if n not in base_attributes + ] + return base_attributes + formatted_names + + def _format_items(self, exclude: set[str] | None = None) -> str: + """Format expression items, optionally excluding names in a group.""" + r = "" + count = 0 + for name, ds in self.items(): + if exclude and name in exclude: + continue + count += 1 + coords = ( + " (" + ", ".join(str(coord) for coord in ds.coords) + ")" + if ds.coords + else "" + ) + r += f" * {name}{coords}\n" + if count == 0: + r += "\n" + return r + + def __repr__(self) -> str: + """ + Return a string representation of the expressions container. + """ + r = "linopy.model.Expressions" + line = "-" * len(r) + r += f"\n{line}\n" + r += self._format_items() + return r + + def __len__(self) -> int: + return self.data.__len__() + + def __iter__(self) -> Iterator[str]: + return self.data.__iter__() + + def items(self) -> ItemsView[str, LinearExpression | QuadraticExpression]: + return self.data.items() + + def _ipython_key_completions_(self) -> list[str]: + """ + Provide method for the key-autocompletions in IPython. + + See + http://ipython.readthedocs.io/en/stable/config/integrating.html#tab-completion + For the details. + """ + return list(self) + + def add(self, expression: LinearExpression | QuadraticExpression) -> None: + """ + Add an expression to the expressions container. + """ + self.data[expression.name] = expression + + def remove(self, name: str) -> None: + """ + Remove variable `name` from the variables. + """ + self.data.pop(name) + + @property + def solution(self) -> Dataset: + """ + Get the solution of variables. + """ + return save_join(*[v.solution.rename(k) for k, v in self.items()]) + + class ScalarLinearExpression: """ A scalar linear expression container. diff --git a/linopy/io.py b/linopy/io.py index 462fa5b8..8fdb9338 100644 --- a/linopy/io.py +++ b/linopy/io.py @@ -958,6 +958,8 @@ def to_netcdf(m: Model, *args: Any, **kwargs: Any) -> None: def with_prefix(ds: xr.Dataset, prefix: str) -> xr.Dataset: to_rename = set([*ds.dims, *ds.coords, *ds]) ds = ds.rename({d: f"{prefix}-{d}" for d in to_rename}) + if any(v is None for v in ds.attrs.values()): + breakpoint() ds.attrs = {f"{prefix}-{k}": v for k, v in ds.attrs.items()} # Flatten multiindexes diff --git a/linopy/model.py b/linopy/model.py index 24594c9c..6ea2259a 100644 --- a/linopy/model.py +++ b/linopy/model.py @@ -57,6 +57,7 @@ ) from linopy.dualization import dualize from linopy.expressions import ( + Expressions, LinearExpression, QuadraticExpression, ScalarLinearExpression, @@ -133,6 +134,7 @@ class Model: _solver: solvers.Solver | None _variables: Variables + _expressions: Expressions _constraints: Constraints _objective: Objective _parameters: Dataset @@ -144,6 +146,7 @@ class Model: _cCounter: int _dtypes: dict[DtypeKey, type[np.signedinteger]] _varnameCounter: int + _exprnameCounter: int _connameCounter: int _pwlCounter: int _blocks: DataArray | None @@ -155,6 +158,7 @@ class Model: __slots__ = ( # containers "_variables", + "_expressions", "_constraints", "_objective", "_parameters", @@ -168,6 +172,7 @@ class Model: "_cCounter", "_dtypes", "_varnameCounter", + "_exprnameCounter", "_connameCounter", "_pwlCounter", "_blocks", @@ -258,6 +263,7 @@ def __init__( dtypes ) self._variables: Variables = Variables({}, model=self) + self._expressions: Expressions = Expressions({}, model=self) self._constraints: Constraints = Constraints({}, model=self) self._objective: Objective = Objective(LinearExpression(None, self), self) self._parameters: Dataset = Dataset() @@ -267,6 +273,7 @@ def __init__( self._xCounter: int = 0 self._cCounter: int = 0 self._varnameCounter: int = 0 + self._exprnameCounter: int = 0 self._connameCounter: int = 0 self._pwlCounter: int = 0 self._blocks: DataArray | None = None @@ -326,6 +333,13 @@ def variables(self) -> Variables: """ return self._variables + @property + def expressions(self) -> Expressions: + """ + Expressions assigned to the model. + """ + return self._expressions + @property def constraints(self) -> Constraints: """ @@ -572,6 +586,7 @@ def scalar_attrs(self) -> list[str]: "_xCounter", "_cCounter", "_varnameCounter", + "_exprnameCounter", "_connameCounter", "_pwlCounter", "force_dim_names", @@ -590,11 +605,13 @@ def __repr__(self) -> str: var_names, con_names = _get_piecewise_groups(self) var_string = self.variables._format_items(exclude=var_names) con_string = self.constraints._format_items(exclude=con_names) + expr_string = self.expressions._format_items() model_string = f"Linopy {self.type} model" return ( f"{model_string}\n{'=' * len(model_string)}\n\n" f"Variables:\n----------\n{var_string}\n" + f"Expressions:\n------------\n{expr_string}\n" f"Constraints:\n------------\n{con_string}" f"{pwl_repr_summary(self)}" f"\nStatus:\n-------\n{self.status}" @@ -913,6 +930,82 @@ def add_variables( self.variables.add(variable) return variable + def add_expressions( + self, + data: Variable + | LinearExpression + | QuadraticExpression + | Sequence[tuple[ConstantLike, Variable | str]], + name: str | None = None, + mask: MaskLike | None = None, + ) -> LinearExpression | QuadraticExpression: + """ + Assign a new, possibly multi-dimensional array of expressions to the + model. + + Parameters + ---------- + data : Variable, LinearExpression, QuadraticExpression, or Sequence of (constant, variable) tuples + The expression(s) to add. This can be a Variable or LinearExpression, a sequence of (constant, variable) tuples which will be summed up, or a callable that takes the model as input and returns either of the previous two. + coords : list/xarray.Coordinates, optional + The coords of the expression array. + The default is None. + name : str, optional + Reference name of the added expressions. The default None results in + a name like "expr1", "expr2" etc. + mask : array_like, optional + Boolean mask with False values for expressions which are skipped. + The shape of the mask has to match the shape the added expressions. + Default is None. + + Raises + ------ + ValueError + If neither lower bound and upper bound have coordinates, nor + `coords` are directly given. + + Returns + ------- + linopy.LinearExpression | linopy.QuadraticExpression + Expression which was added to the model. + + + Examples + -------- + >>> from linopy import Model + >>> import pandas as pd + >>> m = Model() + >>> time = pd.RangeIndex(10, name="Time") + >>> x = m.add_variables(lower=0, coords=[time], name="x") + >>> expr = m.add_expressions(x + 1, name="expr") + """ + if name is None: + name = f"expr{self._exprnameCounter}" + self._exprnameCounter += 1 + + if name in self.expressions: + raise ValueError(f"Expression '{name}' already assigned to model") + + expr: LinearExpression | QuadraticExpression + if isinstance(data, Variable): + expr = data.to_linexpr() + elif isinstance(data, Sequence): + expr = self.linexpr(*data) + else: + expr = data + self.check_force_dim_names(expr.data) + self._check_valid_dim_names(expr.data) + + if mask is not None: + mask = as_dataarray(mask, coords=expr.coords, dims=expr.dims).astype(bool) + expr = expr.where(mask) + if self.chunk: + expr = expr.chunk(self.chunk) + + expr.attrs["name"] = name + self.expressions.add(expr) + return expr + def add_sos_constraints( self, variable: Variable, diff --git a/linopy/objective.py b/linopy/objective.py index a51b2207..b729cefa 100644 --- a/linopy/objective.py +++ b/linopy/objective.py @@ -192,6 +192,9 @@ def expression( if (expr.const != 0.0) and not np.isnan(expr.const): raise ValueError("Constant values in objective function not supported.") + # TODO-871: If we want to track the objective name from declarative math IO then there should be the ability to set this dynamically. + # For now, we just set it to "objective" to avoid issues with the name being None (for .nc serialisation). + expr.attrs["name"] = "objective" self._expression = expr @property diff --git a/pyproject.toml b/pyproject.toml index a5186be1..4820ba74 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -39,6 +39,9 @@ dependencies = [ "tqdm", "deprecation", "packaging", + "pydantic>=2", + "pyparsing>=3", + "pyyaml", ] [project.urls] diff --git a/test/resources/inputs.nc b/test/resources/inputs.nc new file mode 100644 index 00000000..a5a5e9d6 Binary files /dev/null and b/test/resources/inputs.nc differ diff --git a/test/resources/math.yaml b/test/resources/math.yaml new file mode 100644 index 00000000..b073f2e6 --- /dev/null +++ b/test/resources/math.yaml @@ -0,0 +1,442 @@ +dimensions: + name: + dtype: string + ordered: false + snapshot: + dtype: datetime + ordered: true + node: + dtype: string + ordered: false + +parameters: + p_max_pu: + dims: [name, snapshot] + default: .inf + unit: pu + e_max_pu: + dims: [name, snapshot] + default: .inf + unit: pu + s_max_pu: + dims: [name, snapshot] + default: .inf + unit: pu + p_min_pu: + dims: [name, snapshot] + default: 0 + unit: pu + e_min_pu: + dims: [name, snapshot] + default: 0 + unit: pu + s_min_pu: + dims: [name, snapshot] + default: 0 + unit: pu + p_nom_min: + dims: [name] + default: 0 + unit: MW + e_nom_min: + dims: [name] + default: 0 + unit: MWh + s_nom_min: + dims: [name] + default: 0 + unit: MVA + p_nom_max: + dims: [name] + default: .inf + unit: MW + e_nom_max: + dims: [name] + default: .inf + unit: MWh + s_nom_max: + dims: [name] + default: .inf + unit: MVA + capital_cost: + dims: [name] + default: 0 + unit: $\frac{currency}{nominal_capacity}$ + overnight_cost: + dims: [name] + default: 0 + unit: $\frac{currency}{nominal_capacity}$ + fom_cost: + dims: [name] + default: 0 + unit: $\frac{currency}{nominal_capacity}$ + nyears: + dims: [name] + default: 1.0 + unit: years + discount_rate: + dims: [name] + default: 0 + unit: unitless + lifetime: + dims: [name] + default: .inf + unit: years + snapshot_weightings: + dims: [snapshot] + default: 1 + unit: hours + marginal_cost: + dims: [name, snapshot] + default: 0 + unit: $\frac{currency}{dispatch}$ + efficiency: + dims: [name, snapshot] + default: 1 + unit: unitless + rate1: + dims: [name, snapshot] + default: 1 + unit: unitless + sign: + dims: [name] + default: 1 + unit: unitless + p_set: + dims: [name, snapshot] + default: 0 + unit: MWh + efficiency_dispatch: + dims: [name, snapshot] + default: 1 + unit: unitless + efficiency_store: + dims: [name, snapshot] + default: 1 + unit: unitless + inflow: + dims: [name, snapshot] + default: 0 + unit: MWh + e_initial: + dims: [name] + default: 0 + unit: MWh + state_of_charge_initial: + dims: [name] + default: 0 + unit: MWh + standing_loss: + dims: [name, snapshot] + default: 0 + unit: $\frac{}{hour}$ + +lookups: + active: + dims: [name, snapshot] + dtype: bool + default: true + committable: + dims: [name] + dtype: bool + default: false + component: + dims: [name] + dtype: string + one_of: + - "buses" + - "global_constraints" + - "lines" + - "transformers" + - "links" + - "loads" + - "generators" + - "processes" + - "storage_units" + - "stores" + s_nom_extendable: + dims: [name] + dtype: bool + default: true + p_nom_extendable: + dims: [name] + dtype: bool + default: true + e_nom_extendable: + dims: [name] + dtype: bool + default: true + bus: + dims: [name] + dtype: string + bus0: + dims: [name] + dtype: string + bus1: + dims: [name] + dtype: string + cyclic_state_of_charge: + dims: [name] + dtype: bool + default: true + +variables: + p: + foreach: [name, snapshot] + mask: "[generators, links, loads, processes, stores] in component" + unit: MWh + default: 0 + bounds: + lower: -.inf + upper: .inf + p_dispatch: + foreach: [name] + mask: component == storage_units + unit: MWh + default: 0 + bounds: + lower: 0 + upper: .inf + p_store: + foreach: [name] + mask: component == storage_units + unit: MWh + default: 0 + bounds: + lower: 0 + upper: .inf + e: + foreach: [name, snapshot] + mask: component == stores + unit: MWh + default: 0 + bounds: + lower: 0 + upper: .inf + state_of_charge: + foreach: [name, snapshot] + mask: component == storage_units + unit: MWh + default: 0 + bounds: + lower: 0 + upper: .inf + s: + foreach: [name, snapshot] + mask: "[lines, transformers] in component" + unit: MVAh + default: 0 + bounds: + lower: 0 + upper: .inf + p_nom: + foreach: [name] + mask: "[generators, links, processes, storage_units, stores] in component" + unit: MW + default: .inf + bounds: + lower: 0 + upper: .inf + s_nom: + foreach: [name] + mask: "[lines, transformers] in component" + unit: MW + default: .inf + bounds: + lower: 0 + upper: .inf + e_nom: + foreach: [name] + mask: component == stores + unit: MW + default: .inf + bounds: + lower: 0 + upper: .inf + spill: + foreach: [name, snapshot] + mask: component == storage_units + unit: MWh + default: 0 + bounds: + lower: 0 + upper: .inf + +expressions: + + capex: + foreach: [name] + mask: capital_cost or overnight_cost or fom_cost + equations: + - expression: $capacity * (capital_cost + fom_cost + $cc) + sub_expressions: + capacity: + - mask: "[generators, links, processes, stores] in component" + expression: p_nom + - mask: component == stores + expression: e_nom + - mask: component == lines + expression: s_nom + cc: # capital cost derived from overnight cost + - mask: overnight_cost and discount_rate == 0 and lifetime == .inf + expression: "0" + - mask: overnight_cost and discount_rate == 0 and not lifetime == .inf + expression: 1 / lifetime + - mask: overnight_cost and discount_rate > 0 and not lifetime == .inf + expression: overnight_cost * discount_rate / (1 - 1 / (1 + discount_rate) ** lifetime) + opex: + foreach: [name, snapshot] + mask: marginal_cost + equations: + - expression: (p + s) * marginal_cost * snapshot_weightings + +constraints: + operation_upper: + description: "Upper limit on per-snapshot dispatch" + foreach: [name, snapshot] + mask: active and not committable + equations: + - mask: "[generators, links, processes] in component and not p_max_pu == 0" + expression: p <= p_max_pu * p_nom + - mask: "[generators, links, processes] in component and p_max_pu == 0" + expression: p <= 0 + - mask: "component == storage_units and not p_max_pu == 0" + expression: p_dispatch <= p_max_pu * p_nom + - mask: "component == storage_units and p_max_pu == 0" + expression: p_dispatch <= 0 + - mask: "component == stores and not e_max_pu == 0" + expression: e <= e_max_pu * e_nom + - mask: "component == stores and e_max_pu == 0" + expression: e <= 0 + - mask: "[lines, transformers] in component and not s_max_pu == 0" + expression: s <= s_max_pu * s_nom + - mask: "[lines, transformers] in component and s_max_pu == 0" + expression: s <= 0 + + operation_upper_store: + description: "Upper limit on per-snapshot storage unit charging" + foreach: [name, snapshot] + mask: active and not committable and component == storage_units + equations: + - mask: "not p_max_pu == 0" + expression: p_store <= p_max_pu * p_nom + - mask: "p_max_pu == 0" + expression: p_store <= 0 + + operation_lower: + description: "Lower limit on per-snapshot dispatch" + foreach: [name, snapshot] + mask: active and not committable + equations: + - mask: "[generators, links, processes] in component and not p_min_pu == 0" + expression: p >= p_min_pu * p_nom + - mask: "[generators, links, processes] in component and p_min_pu == 0" + expression: p >= 0 + - mask: "component == storage_units and not p_min_pu == 0" + expression: p_dispatch >= p_min_pu * p_nom + - mask: "component == storage_units and p_min_pu == 0" + expression: p_dispatch >= 0 + - mask: "component == stores and not e_min_pu == 0" + expression: e >= e_min_pu * e_nom + - mask: "component == stores and e_min_pu == 0" + expression: e >= 0 + - mask: "[lines, transformers] in component and not s_min_pu == 0" + expression: s >= s_min_pu * s_nom + - mask: "[lines, transformers] in component and s_min_pu == 0" + expression: s >= 0 + + operation_lower_store: + description: "Lower limit on per-snapshot storage unit charging" + foreach: [name, snapshot] + mask: active and not committable and component == storage_units + equations: + - mask: "not p_min_pu == 0" + expression: p_store >= p_min_pu * p_nom + - mask: "p_min_pu == 0" + expression: p_store >= 0 + + capacity_lower: + foreach: [name] + mask: any(active, over=snapshot) + equations: + - mask: "[generators, links, processes, storage_units] in component and p_nom_extendable==True" + expression: p_nom >= p_nom_min + - mask: "component == stores and e_nom_extendable==True" + expression: e_nom >= e_nom_min + - mask: "[lines, transformers] in component and s_nom_extendable==True" + expression: s_nom >= s_nom_min + + capacity_upper: + foreach: [name] + mask: any(active, over=snapshot) + equations: + - mask: "[generators, links, processes, storage_units] in component and p_nom_extendable==True and not p_nom_max == inf" + expression: p_nom <= p_nom_max + - mask: "component == stores and e_nom_extendable==True and not e_nom_max == inf" + expression: e_nom <= e_nom_max + - mask: "[lines, transformers] in component and s_nom_extendable==True and not s_nom_max == inf" + expression: s_nom <= s_nom_max + + fix_load: + description: "Fix load flow" + foreach: [name, snapshot] + mask: component == loads + equations: + - expression: p = p_set + + nodal_balance: + foreach: [node, snapshot] + equations: + - expression: >- + group_sum((p + p_dispatch - p_store) * sign, bus, node) - + group_sum(p * sign, bus0, node) + + group_sum(p * sign * (efficiency + rate1), bus1, node) + = 0 + + storage_balance: + foreach: [name, snapshot] + mask: active + equations: + - mask: not cyclic_state_of_charge==True and snapshot == get_val_at_index(snapshot=0) + expression: $soc + ($dispatch * snapshot_weightings) = $soc_init + - mask: (not cyclic_state_of_charge and not snapshot == get_val_at_index(snapshot=0)) or cyclic_state_of_charge + expression: >- + $soc + ($dispatch * snapshot_weightings) = + roll($soc * (standing_loss ** snapshot_weightings), snapshot=1) + sub_expressions: + soc: + - mask: component == stores + expression: e + - mask: component == storage_units + expression: state_of_charge + dispatch: + - mask: component == stores + expression: p + - mask: component == storage_units + expression: (1 / efficiency_dispatch) * p_dispatch - efficiency_store * p_store + spill - inflow + soc_init: + - mask: component == stores + expression: e_initial + - mask: component == storage_units + expression: state_of_charge_initial + +objectives: + cost_minimisation: + equations: + - expression: sum(opex, over=[snapshot]) + sum(capex, over=name) + sense: min + +checks: + overnight_capital_overlap: + mask: overnight_cost and capital_cost + message: can only define one of `overnight_cost` and `capital_cost` + errors: raise + active: false + rate_efficiency_overlap: + mask: efficiency and rate1 + message: cannot define `efficiency` and `rate1` for the same component + errors: raise + efficiency_range: + mask: efficiency > 1 or efficiency < 0 + message: must be between 0 and 1 + errors: raise diff --git a/test/test_declarative_parsing.py b/test/test_declarative_parsing.py new file mode 100644 index 00000000..bc64ada3 --- /dev/null +++ b/test/test_declarative_parsing.py @@ -0,0 +1,969 @@ +""" +Tests for the declarative math interface. + +Covers the grammar (string -> AST), the three evaluation routes (boolean mask +array / linopy expression / LaTeX math string), `$name` sub-expression and slicer +resolution, helper-function registration, and the model builder. +""" + +from __future__ import annotations + +import re +from pathlib import Path +from typing import Any, Literal + +import numpy as np +import pandas as pd +import pytest +import xarray as xr +import yaml + +from linopy.declarative import grammar, nodes, parsing +from linopy.declarative.build import DeclarativeModelBuilder, declarative_model +from linopy.declarative.helpers import HelperFunction, build_registry +from linopy.declarative.latex import ( + LatexModelBuilder, + _escape_text_mode, + latex_math_doc, +) +from linopy.declarative.schema import COMPONENTS_T, ConfigModel, MathModel +from linopy.expressions import LinearExpression +from linopy.variables import Variable + +NODES = ["a", "b", "c"] + + +@pytest.fixture +def math() -> dict: + """Minimal but representative math definition exercising every route.""" + return { + "dimensions": {"node": {"dtype": "string", "iterator": "n"}}, + "parameters": { + "cost": {"default": 0, "dims": ["node"]}, + "cap_max": {"default": float("inf")}, + "param_inactive": {"default": 0, "active": False}, + }, + "lookups": { + "active": {"default": True, "dtype": "bool", "dims": ["node"]}, + }, + "variables": { + "flow": { + "foreach": ["node"], + "bounds": {"lower": 0, "upper": float("inf")}, + }, + "flow_inactive": { + "foreach": ["node"], + "bounds": {"lower": 0, "upper": float("inf")}, + "active": False, + }, + }, + "expressions": { + "total_cost": { + "foreach": ["node"], + "equations": [{"expression": "flow * cost"}], + }, + # Pure-parameter expression: evaluates to a DataArray, must be coerced. + "cost_plus_one": { + "foreach": ["node"], + "equations": [{"expression": "cost + 1"}], + }, + # Sub-expression reference to a variable (regression for the expr route). + "sub_expr_test": { + "foreach": ["node"], + "equations": [{"expression": "$foo * cost"}], + "sub_expressions": {"foo": [{"expression": "flow"}]}, + }, + }, + "constraints": { + "cap": { + "foreach": ["node"], + "equations": [{"expression": "flow <= cap_max"}], + }, + }, + "objectives": { + "obj": { + "equations": [{"expression": "sum(total_cost, over=node)"}], + "sense": "min", + }, + }, + } + + +@pytest.fixture +def inputs() -> xr.Dataset: + return xr.Dataset( + { + "cost": ("node", [1.0, 2.0, 3.0]), + "cap_max": ("node", [10.0, 20.0, 30.0]), + }, + coords={"node": NODES}, + ) + + +def _ctx(builder: DeclarativeModelBuilder, **kwargs: Any) -> nodes.Context: + """Build a fresh evaluation context from a builder's validated components.""" + return nodes.Context( + model=builder.model, + input_data=builder.input_data, + math=builder.math, + config=kwargs.pop("config", builder.config), + helpers=kwargs.pop("helpers", build_registry()), + **kwargs, + ) + + +def _first_equation( + builder: DeclarativeModelBuilder, group: COMPONENTS_T, name: str +) -> tuple[parsing.Equation, xr.DataArray, nodes.Context]: + """Return a pre-parsed component's first equation, sub-mask, and context.""" + ctx = _ctx(builder) + definition = getattr(builder.math, group)[name] + parsed = builder.parsed[group][name] + mask = parsing.component_mask(group, name, definition, parsed.mask, ctx) + equation = parsed.equations[0] + sub_mask = parsing.as_mask(equation, ctx, initial_mask=mask) + sub_mask = parsing.drop_dims_not_in_foreach(sub_mask, equation.sets) + return equation, sub_mask, ctx + + +@pytest.fixture +def builder_with_flow(math: dict, inputs: xr.Dataset) -> DeclarativeModelBuilder: + """A builder with the `flow` variable already added to the model.""" + builder = DeclarativeModelBuilder(math, inputs, {}) + builder.add_variable("flow", builder.math.variables["flow"]) + return builder + + +class TestSchema: + @pytest.fixture + def validated_math(self, math: dict) -> MathModel: + """Return a validated MathModel for the given math dict.""" + return MathModel.model_validate(math) + + def test_duplicate_names(self, math: dict) -> None: + """Do not allow the same name to be defined between two component types.""" + math["parameters"]["flow"] = math["parameters"]["cost"].copy() + with pytest.raises(ValueError, match="Non-unique names in math components"): + MathModel.model_validate(math) + + def test_invalid_dimensions(self, math: dict) -> None: + """Do not allow a component to reference a dimension that is not defined.""" + math["parameters"]["cost"]["dims"] = ["node", "foo"] + with pytest.raises( + ValueError, match="`parameters.cost` references undefined dimensions" + ): + MathModel.model_validate(math) + + def test_valid_parsing_components(self, validated_math: MathModel) -> None: + """Get expected list of parsing components.""" + assert validated_math.parsing_components == { + "expression": { + "dimensions": {"node"}, + "inputs": {"cost", "active", "cap_max"}, + "results": {"total_cost", "cost_plus_one", "sub_expr_test", "flow"}, + }, + "mask": { + "dimensions": {"node"}, + "inputs": {"cost", "active", "cap_max", "param_inactive"}, + "results": { + "total_cost", + "cost_plus_one", + "sub_expr_test", + "flow", + "flow_inactive", + }, + }, + } + + def test_find(self, validated_math: MathModel) -> None: + """Find a component by name.""" + assert validated_math.find("flow") is validated_math.variables["flow"] + + def test_find_subset(self, validated_math: MathModel) -> None: + """Find a component by name and subset.""" + assert ( + validated_math.find("flow", ["variables"]) + is validated_math.variables["flow"] + ) + + def test_find_not_found(self, validated_math: MathModel) -> None: + """Raise on missing component.""" + with pytest.raises(KeyError, match="not found in math schema"): + validated_math.find("foo") + + def test_find_not_found_subset(self, validated_math: MathModel) -> None: + """Raise on missing component in subset.""" + with pytest.raises(KeyError, match="not found in math schema"): + validated_math.find("flow", ["expressions"]) + + +class TestGrammar: + """String -> AST parsing.""" + + NAMES = frozenset({"flow", "cost", "cap_max", "node"}) + + def test_equation_tree_shape(self) -> None: + tree = grammar.equation_grammar(self.NAMES).parse_string( + "flow <= cap_max", parse_all=True + )[0] + assert isinstance(tree, grammar.Compare) + assert tree.op == "<=" + assert isinstance(tree.lhs, grammar.Component) and tree.lhs.name == "flow" + + def test_equation_rejects_mask_only_operators(self) -> None: + import pyparsing as pp + + with pytest.raises(pp.ParseException): + grammar.equation_grammar(self.NAMES).parse_string( + "flow < cap_max", parse_all=True + ) + + def test_arithmetic_tree_shape(self) -> None: + tree = grammar.arithmetic_grammar(self.NAMES).parse_string( + "flow * cost + 1", parse_all=True + )[0] + assert isinstance(tree, grammar.Arith) + assert tree.rest[0][0] == "+" + + def test_sliced_component(self) -> None: + tree = grammar.arithmetic_grammar(self.NAMES).parse_string( + "flow[node=$n]", parse_all=True + )[0] + assert isinstance(tree, grammar.Sliced) + assert isinstance(tree.slices["node"], grammar.SliceRef) + + def test_sub_expression_grammar_rejects_refs(self) -> None: + import pyparsing as pp + + with pytest.raises(pp.ParseException): + grammar.sub_expression_grammar(self.NAMES).parse_string( + "$foo + 1", parse_all=True + ) + + def test_find_refs_in_call_kwargs(self) -> None: + tree = grammar.arithmetic_grammar(self.NAMES).parse_string( + "sum($foo, over=node) + flow[node=$n]", parse_all=True + )[0] + assert grammar.find_refs(tree, grammar.SubExprRef) == {"foo"} + assert grammar.find_refs(tree, grammar.SliceRef) == {"n"} + assert grammar.find_refs(tree, grammar.Component) == {"node", "flow"} + + def test_node_repr_is_clean(self) -> None: + tree = grammar.arithmetic_grammar(self.NAMES).parse_string( + "flow * cost + 1", parse_all=True + )[0] + rendered = repr(tree) + assert "instring=" not in rendered + assert "loc=" not in rendered + + def test_parse_error_carries_position_marker( + self, math: dict, inputs: xr.Dataset + ) -> None: + math["constraints"]["cap"]["equations"][0]["expression"] = "flow <= <=" + with pytest.raises(ValueError, match="equations\\[0\\].expression") as excinfo: + DeclarativeModelBuilder(math, inputs, {}) + message = str(excinfo.value) + assert "constraints:cap:" in message + assert "^" in message + + +class TestParseWalkthrough: + """Whole-dict parse walkthrough with aggregated errors.""" + + def test_errors_aggregate_across_components( + self, math: dict, inputs: xr.Dataset + ) -> None: + """Broken strings in two components raise as one grouped error.""" + math["constraints"]["cap"]["equations"][0]["expression"] = "flow <= <=" + math["expressions"]["total_cost"]["equations"][0]["expression"] = ( + "flow * * cost" + ) + with pytest.raises(ValueError) as excinfo: + DeclarativeModelBuilder(math, inputs, {}) + message = str(excinfo.value) + assert "constraints:cap:" in message + assert "expressions:total_cost:" in message + assert message.count("^") == 2 + + def test_undefined_ref_collected_alongside_syntax_errors( + self, math: dict, inputs: xr.Dataset + ) -> None: + """A syntax error does not short-circuit undefined `$ref` collection.""" + math["constraints"]["cap"]["equations"][0]["expression"] = "flow <= <=" + math["expressions"]["sub_expr_test"]["sub_expressions"] = { + "bar": [{"expression": "flow"}] + } + with pytest.raises(ValueError) as excinfo: + DeclarativeModelBuilder(math, inputs, {}) + message = str(excinfo.value) + assert "constraints:cap:" in message + assert "expressions:sub_expr_test:" in message + assert "Undefined sub_expressions" in message + + def test_inactive_components_are_skipped( + self, math: dict, inputs: xr.Dataset + ) -> None: + """Inactive components are neither parsed nor built.""" + math["expressions"]["broken"] = { + "active": False, + "foreach": ["node"], + "equations": [{"expression": "flow * * cost"}], + } + builder = DeclarativeModelBuilder(math, inputs, {}) + assert "broken" not in builder.parsed["expressions"] + model = builder.build() + assert "broken" not in model.expressions + + def test_check_masks_are_parsed(self, math: dict, inputs: xr.Dataset) -> None: + math["checks"] = {"bad": {"mask": "cost > >", "message": "boom"}} + with pytest.raises(ValueError, match="checks:bad"): + DeclarativeModelBuilder(math, inputs, {}) + + def test_inactive_check_masks_are_skipped( + self, math: dict, inputs: xr.Dataset + ) -> None: + math["checks"] = { + "bad": {"mask": "cost > >", "message": "boom", "active": False} + } + builder = DeclarativeModelBuilder(math, inputs, {}) + assert "bad" not in builder.parsed.checks + + def test_parsed_math_shape(self, math: dict, inputs: xr.Dataset) -> None: + builder = DeclarativeModelBuilder(math, inputs, {}) + assert set(builder.parsed.components) == { + "variables", + "expressions", + "constraints", + "objectives", + } + assert builder.parsed["variables"]["flow"].equations == [] + assert builder.parsed["constraints"]["cap"].equations + + +class TestMaskRoute: + """Mask strings evaluate to boolean arrays.""" + + def test_top_level_mask_returns_boolean_dataarray( + self, builder_with_flow: DeclarativeModelBuilder + ) -> None: + _, sub_mask, _ = _first_equation(builder_with_flow, "constraints", "cap") + assert isinstance(sub_mask, xr.DataArray) + assert sub_mask.dtype == bool + assert bool(sub_mask.all()) + + def test_mask_comparison_and_subset_and_helper_return_bool( + self, math: dict, inputs: xr.Dataset + ) -> None: + math["constraints"]["cap"]["equations"][0]["mask"] = ( + "cost > 1 and [a, b] in node and any(cap_max, over=node)" + ) + builder = DeclarativeModelBuilder(math, inputs, {}) + builder.add_variable("flow", builder.math.variables["flow"]) + _, sub_mask, _ = _first_equation(builder, "constraints", "cap") + assert sub_mask.dtype == bool + # `cost > 1` only holds for nodes b and c. + assert sub_mask.values.tolist() == [False, True, False] + + @pytest.mark.parametrize( + "mask_string", + ["True", "not cost > 1", "cost > 1 or cap_max <= 10", "config.foo == bar"], + ) + def test_mask_atoms_return_bool( + self, builder_with_flow: DeclarativeModelBuilder, mask_string: str + ) -> None: + node = parsing.parse_mask(mask_string, builder_with_flow.math) + config = ConfigModel.model_validate({"foo": "bar"}) + result = nodes.evaluate( + node, _ctx(builder_with_flow, mode="mask", config=config) + ) + assert isinstance(result, xr.DataArray) + assert result.dtype == bool + + def test_existence_coercion( + self, builder_with_flow: DeclarativeModelBuilder + ) -> None: + """Bare input references coerce to existence booleans on the mask route.""" + node = parsing.parse_mask("cap_max", builder_with_flow.math) + result = nodes.evaluate(node, _ctx(builder_with_flow, mode="mask")) + assert result.values.tolist() == [True, True, True] + + +class TestExpressionRoute: + """Expression strings evaluate to linopy expressions.""" + + def test_expression_with_variable_returns_linexpr( + self, builder_with_flow: DeclarativeModelBuilder + ) -> None: + equation, sub_mask, ctx = _first_equation( + builder_with_flow, "expressions", "total_cost" + ) + result = parsing.as_expression(equation, ctx, mask=sub_mask) + assert isinstance(result, LinearExpression) + + def test_pure_parameter_expression_coerced_to_linexpr( + self, builder_with_flow: DeclarativeModelBuilder + ) -> None: + equation, sub_mask, ctx = _first_equation( + builder_with_flow, "expressions", "cost_plus_one" + ) + result = parsing.as_expression(equation, ctx, mask=sub_mask) + # No decision variable is involved, but the contract is still LinearExpression. + assert isinstance(result, LinearExpression) + + def test_sub_expression_reference_returns_linexpr( + self, builder_with_flow: DeclarativeModelBuilder + ) -> None: + equation, sub_mask, ctx = _first_equation( + builder_with_flow, "expressions", "sub_expr_test" + ) + assert set(equation.sub_expressions) == {"foo"} + result = parsing.as_expression(equation, ctx, mask=sub_mask) + assert isinstance(result, LinearExpression) + + def test_sub_expression_variants_expand_to_cartesian_product( + self, math: dict, inputs: xr.Dataset + ) -> None: + """Two variants of one sub-expression yield two equations with merged masks.""" + math["expressions"]["sub_expr_test"]["sub_expressions"]["foo"] = [ + {"mask": "cost > 1", "expression": "flow"}, + {"mask": "not cost > 1", "expression": "flow * 2"}, + ] + builder = DeclarativeModelBuilder(math, inputs, {}) + builder.add_variable("flow", builder.math.variables["flow"]) + definition = builder.math.expressions["sub_expr_test"] + equations = parsing.parse_component( + "expressions", "sub_expr_test", definition, builder.math + ) + assert len(equations) == 2 + assert {eq.name for eq in equations} == { + "expressions:sub_expr_test:0-foo:0", + "expressions:sub_expr_test:0-foo:1", + } + # Each equation carries its own mask plus the chosen variant's mask. + assert all(len(eq.masks) == 2 for eq in equations) + ctx = _ctx(builder) + masks = [parsing.as_mask(eq, ctx) for eq in equations] + # The variant masks are complementary. + assert not (masks[0] & masks[1]).any() + assert (masks[0] | masks[1]).all() + # And the whole component still builds end-to-end. + builder.add_expression("sub_expr_test", definition) + assert "sub_expr_test" in builder.model.expressions + + def test_undefined_sub_expression_reference_raises( + self, math: dict, inputs: xr.Dataset + ) -> None: + math["expressions"]["sub_expr_test"]["sub_expressions"] = { + "bar": [{"expression": "flow"}] + } + with pytest.raises(ValueError, match="Undefined sub_expressions"): + DeclarativeModelBuilder(math, inputs, {}) + + def test_plain_and_list_slices( + self, inputs: xr.Dataset, builder_with_flow: DeclarativeModelBuilder + ) -> None: + ctx = _ctx( + builder_with_flow, + mode="expr", + mask=xr.full_like(inputs["cost"], True, bool), + ) + arith = grammar.arithmetic_grammar( + frozenset({"flow", "cost", "cap_max", "node"}) + ) + scalar_sliced = arith.parse_string("flow[node=a] * cost", parse_all=True)[0] + result = nodes.evaluate(scalar_sliced, ctx) + assert isinstance(result, LinearExpression) + + list_sliced = arith.parse_string("flow[node=[a, b]]", parse_all=True)[0] + result = nodes.evaluate(list_sliced, ctx) + assert result.data.sizes["node"] == 2 + + def test_slicer_reference(self, math: dict, inputs: xr.Dataset) -> None: + """`$name` slicer references resolve like sub-expressions (feature parity).""" + math["expressions"]["sliced"] = { + "foreach": ["node"], + "equations": [{"expression": "flow[node=$n] * cost"}], + "slices": {"n": [{"expression": "a"}]}, + } + builder = DeclarativeModelBuilder(math, inputs, {}) + builder.add_variable("flow", builder.math.variables["flow"]) + definition = builder.math.expressions["sliced"] + equations = parsing.parse_component( + "expressions", "sliced", definition, builder.math + ) + assert len(equations) == 1 + assert set(equations[0].slices) == {"n"} + builder.add_expression("sliced", definition) + assert "sliced" in builder.model.expressions + + +class TestConstraintRoute: + """Constraint equations evaluate to (lhs, sign, rhs) tuples.""" + + def test_equation_returns_lhs_sign_rhs_tuple( + self, builder_with_flow: DeclarativeModelBuilder + ) -> None: + equation, sub_mask, ctx = _first_equation( + builder_with_flow, "constraints", "cap" + ) + lhs, sign, rhs = parsing.as_constraint(equation, ctx, mask=sub_mask) + assert isinstance(lhs, LinearExpression) # decision variable side + assert isinstance(rhs, LinearExpression) # pure-parameter side, coerced + assert isinstance(sign, xr.DataArray) + assert set(np.unique(sign.values)) <= {"<="} + + def test_foreach_dim_mismatch_raises(self, math: dict, inputs: xr.Dataset) -> None: + # `sum` removed: the equation is indexed over `node` but foreach is empty. + math["constraints"]["cap"]["foreach"] = [] + builder = DeclarativeModelBuilder(math, inputs, {}) + builder.add_variable("flow", builder.math.variables["flow"]) + equation, sub_mask, ctx = _first_equation(builder, "constraints", "cap") + with pytest.raises(ValueError, match="not present in `foreach`"): + parsing.as_constraint(equation, ctx, mask=sub_mask) + + +class TestLatexRoute: + """Math strings render as LaTeX.""" + + def test_equation_latex(self, builder_with_flow: DeclarativeModelBuilder) -> None: + equation, _, ctx = _first_equation(builder_with_flow, "constraints", "cap") + assert parsing.as_latex_expression(equation, ctx) == r"flow \leq cap_max" + + def test_sum_latex(self, builder_with_flow: DeclarativeModelBuilder) -> None: + equation, _, ctx = _first_equation(builder_with_flow, "objectives", "obj") + assert ( + parsing.as_latex_expression(equation, ctx) + == r"\sum\limits_{\substack{\text{n} \in \text{node}}} (total_cost)" + ) + + def test_mask_latex(self, math: dict, inputs: xr.Dataset) -> None: + math["constraints"]["cap"]["equations"][0]["mask"] = ( + "cost > 1 and [a, b] in node" + ) + builder = DeclarativeModelBuilder(math, inputs, {}) + builder.add_variable("flow", builder.math.variables["flow"]) + equation, _, ctx = _first_equation(builder, "constraints", "cap") + assert parsing.as_latex_mask(equation, ctx) == ( + r"(\textit{cost}\mathord{>}\text{1} \land \text{n} \in \text{[a,b]})" + ) + + def test_mask_infinity_latex_not_wrapped_in_text( + self, math: dict, inputs: xr.Dataset + ) -> None: + # `\infty` (and other bare LaTeX commands) must stay in math mode; only + # plain-text tokens (numbers, coordinate labels, booleans) get `\text{}`. + math["constraints"]["cap"]["equations"][0]["mask"] = "cap_max == inf" + builder = DeclarativeModelBuilder(math, inputs, {}) + builder.add_variable("flow", builder.math.variables["flow"]) + equation, _, ctx = _first_equation(builder, "constraints", "cap") + rendered = parsing.as_latex_mask(equation, ctx) + assert r"\mathord{==}\infty" in rendered + assert r"\text{\infty}" not in rendered + + def test_sliced_component_latex( + self, builder_with_flow: DeclarativeModelBuilder + ) -> None: + ctx = _ctx(builder_with_flow) + arith = grammar.arithmetic_grammar(frozenset({"flow", "node"})) + tree = arith.parse_string("flow[node=a]", parse_all=True)[0] + assert nodes.to_math_string(tree, ctx) == r"flow_\text{n=a}" + + def test_identity_operands_are_skipped( + self, builder_with_flow: DeclarativeModelBuilder + ) -> None: + ctx = _ctx(builder_with_flow) + arith = grammar.arithmetic_grammar(frozenset({"flow"})) + tree = arith.parse_string("0 + flow", parse_all=True)[0] + assert nodes.to_math_string(tree, ctx) == "flow" + + +class _RecordArgs(HelperFunction): + """Test-only helper that records the types of the arguments it receives.""" + + NAME = "record_args" + ALLOWED_IN = ["expression"] + received: list[type] = [] + + def as_math_string(self, *args: Any, **kwargs: Any) -> str: # noqa: D102 + return "record_args" + + def as_raw(self, *args: Any, **kwargs: Any) -> LinearExpression | xr.DataArray: # noqa: D102 + type(self).received.extend(type(a) for a in args) + # Return the first argument so the enclosing expression stays valid. + return args[0] + + +class _Double(HelperFunction): + """Test-only helper doubling its argument.""" + + NAME = "double" + ALLOWED_IN = ["expression"] + + def as_math_string(self, array: Any) -> str: # noqa: D102 + return rf"2 \times {array}" + + def as_raw(self, array: Any) -> LinearExpression | xr.DataArray: # noqa: D102 + return 2 * array + + +class TestHelpers: + """Helper-function registration and argument evaluation.""" + + def test_helper_arguments_are_evaluated_raw( + self, math: dict, inputs: xr.Dataset + ) -> None: + """Helper args arrive un-normalised: raw Variable/DataArray, not LinearExpression.""" + math["expressions"]["total_cost"]["equations"][0]["expression"] = ( + "record_args(flow, cost)" + ) + builder = DeclarativeModelBuilder(math, inputs, {}) + builder.add_variable("flow", builder.math.variables["flow"]) + ctx = _ctx(builder, helpers=build_registry([_RecordArgs])) + definition = builder.math.expressions["total_cost"] + equation = parsing.parse_component( + "expressions", "total_cost", definition, builder.math + )[0] + _RecordArgs.received = [] + parsing.as_expression(equation, ctx) + assert _RecordArgs.received, "helper was not called" + assert Variable in _RecordArgs.received + assert xr.DataArray in _RecordArgs.received + assert LinearExpression not in _RecordArgs.received + + def test_non_subclass_rejected_by_registry(self) -> None: + with pytest.raises(ValueError, match="must be subclassed"): + build_registry([str]) # type: ignore[list-item] + + def test_unknown_helper_rejected(self, math: dict, inputs: xr.Dataset) -> None: + math["expressions"]["total_cost"]["equations"][0]["expression"] = ( + "unknown_helper(flow)" + ) + builder = DeclarativeModelBuilder(math, inputs, {}) + builder.add_variable("flow", builder.math.variables["flow"]) + definition = builder.math.expressions["total_cost"] + equation = parsing.parse_component( + "expressions", "total_cost", definition, builder.math + )[0] + with pytest.raises(ValueError, match="Invalid helper function"): + parsing.as_expression(equation, _ctx(builder)) + + def test_eval_error_carries_caret(self, math: dict, inputs: xr.Dataset) -> None: + """Evaluation errors point a caret at the failing node in the source string.""" + math["expressions"]["total_cost"]["equations"][0]["expression"] = ( + "unknown_helper(flow)" + ) + builder = DeclarativeModelBuilder(math, inputs, {}) + builder.add_variable("flow", builder.math.variables["flow"]) + definition = builder.math.expressions["total_cost"] + equation = parsing.parse_component( + "expressions", "total_cost", definition, builder.math + )[0] + with pytest.raises(ValueError) as excinfo: + parsing.as_expression(equation, _ctx(builder)) + message = str(excinfo.value) + assert "unknown_helper(flow)" in message + source_line, caret_line = message.splitlines()[-2:] + assert caret_line.strip() == "^" + assert caret_line.index("^") == source_line.index("unknown_helper") + + def test_duplicate_name_rejected(self) -> None: + class _ClashingSum(HelperFunction): + NAME = "sum" + ALLOWED_IN = ["expression"] + + def as_math_string(self, *args: Any, **kwargs: Any) -> str: # noqa: D102 + return "" + + def as_raw( + self, *args: Any, **kwargs: Any + ) -> LinearExpression | xr.DataArray: # noqa: D102 + return xr.DataArray() + + with pytest.raises(ValueError, match="already exists"): + build_registry([_ClashingSum]) + + def test_custom_helper_end_to_end(self, math: dict, inputs: xr.Dataset) -> None: + math["expressions"]["total_cost"]["equations"][0]["expression"] = ( + "double(flow) * cost" + ) + model = declarative_model(math, inputs, {}, helpers=[_Double]) + assert "total_cost" in model.expressions + + def test_get_val_at_index(self, builder_with_flow: DeclarativeModelBuilder) -> None: + ctx = _ctx(builder_with_flow) + arith = grammar.arithmetic_grammar(frozenset({"flow", "node"})) + tree = arith.parse_string("get_val_at_index(node=0)", parse_all=True)[0] + assert nodes.evaluate(tree, ctx).item() == "a" + + +class TestBuilder: + """Model assembly from parsed math.""" + + def test_overlapping_equation_masks_rejected( + self, math: dict, inputs: xr.Dataset + ) -> None: + math["constraints"]["cap"]["equations"] = [ + {"mask": "cost > 0", "expression": "flow <= cap_max"}, + {"mask": "cost > 1", "expression": "flow <= 2 * cap_max"}, + ] + with pytest.raises(ValueError, match="Overlapping 'mask' conditions"): + declarative_model(math, inputs, {}) + + def test_multiple_active_objectives_rejected( + self, math: dict, inputs: xr.Dataset + ) -> None: + math["objectives"]["obj2"] = math["objectives"]["obj"].copy() + with pytest.raises(ValueError, match="Only one active objective"): + declarative_model(math, inputs, {}) + + def test_references_are_sorted_lists(self, math: dict, inputs: xr.Dataset) -> None: + model = declarative_model(math, inputs, {}) + refs = model.constraints["cap"].attrs["references"] + assert refs == sorted(refs) + assert isinstance(refs, list) + assert set(refs) == {"cap_max", "flow"} + + def test_dtype_coercion(self, math: dict, inputs: xr.Dataset) -> None: + math["lookups"] = { + "flag": {"dtype": "bool", "default": False}, + "label": {"dtype": "string"}, + } + inputs["flag"] = ("node", [1.0, float("nan"), 0.0]) + inputs["label"] = ("node", ["x", "", "z"]) + builder = DeclarativeModelBuilder(math, inputs, {}) + assert builder.input_data["flag"].dtype == bool + assert builder.input_data["flag"].values.tolist() == [True, False, False] + # Empty strings are coerced to missing values. + assert builder.input_data["label"].isnull().sum() == 1 + + def test_checks_run_without_active_variable( + self, math: dict, inputs: xr.Dataset + ) -> None: + """Input checks must not require an `active` variable in the input data.""" + math["checks"] = { + "too_expensive": { + "mask": "cost > 100", + "message": "cost too high", + "errors": "raise", + } + } + # No `active` variable in the inputs, and the check does not trigger. + declarative_model(math, inputs, {}) + + def test_check_raises_when_triggered_without_active( + self, math: dict, inputs: xr.Dataset + ) -> None: + math["checks"] = { + "too_expensive": { + "mask": "cost > 0", + "message": "cost too high", + "errors": "raise", + } + } + with pytest.raises(ValueError, match="cost too high"): + declarative_model(math, inputs, {}) + + def test_check_warns( + self, caplog: pytest.LogCaptureFixture, math: dict, inputs: xr.Dataset + ) -> None: + math["checks"] = { + "pricey": {"mask": "cost > 0", "message": "prices!", "errors": "warn"} + } + with caplog.at_level("INFO", logger="linopy.declarative.build"): + declarative_model(math, inputs, {}) + assert "prices!" in caplog.text + + def test_input_dims_misalignment_build_level( + self, math: dict, inputs: xr.Dataset + ) -> None: + inputs["cost"] = inputs["cost"].expand_dims(wrong_dim=pd.Index([1, 2])) + with pytest.raises(ValueError, match="Parameter `cost` has dimensions"): + DeclarativeModelBuilder(math, inputs, {}) + + +class TestLatexDoc: + """LaTeX math documentation building.""" + + def test_components_render_with_decorated_reprs( + self, math: dict, inputs: xr.Dataset + ) -> None: + builder = LatexModelBuilder(math, inputs, {}).build() + cap = builder.components["constraints"]["cap"] + assert cap.foreach == r"\forall{} \text{n} \in \text{node}" + assert cap.equations == [ + { + "mask": "", + "expression": r"\textbf{flow}_\text{n} \leq \textit{cap\_max}_\text{n}", + } + ] + + def test_variable_bounds_equation(self, math: dict, inputs: xr.Dataset) -> None: + math["variables"]["flow"]["bounds"] = {"lower": 0, "upper": "cap_max"} + builder = LatexModelBuilder(math, inputs, {}).build() + flow = builder.components["variables"]["flow"] + assert flow.equations[0]["expression"] == ( + r"0 \leq \textbf{flow}_\text{n} \leq \textit{cap\_max}_\text{n}" + ) + assert flow.uses == ["cap_max"] + + def test_underscores_escaped_in_text_mode( + self, math: dict, inputs: xr.Dataset + ) -> None: + # `cap_max` is a parameter (rendered `\textit{...}`); its underscore must + # be escaped so KaTeX does not read it as a subscript operator. + builder = LatexModelBuilder(math, inputs, {}).build() + cap = builder.components["constraints"]["cap"] + assert r"\textit{cap\_max}" in cap.equations[0]["expression"] + # The subscript operator between the name and its dimension is preserved. + assert r"\textbf{flow}_\text{n}" in cap.equations[0]["expression"] + + def test_escape_text_mode_escapes_content_only(self) -> None: + # Underscores inside text commands (names and coordinate values) are + # escaped, while the subscript operator between them is left intact. + assert _escape_text_mode(r"\text{storage_units}") == r"\text{storage\_units}" + assert ( + _escape_text_mode(r"\textbf{p_nom}_\text{n}") == r"\textbf{p\_nom}_\text{n}" + ) + # Already-escaped underscores are not doubled up. + assert _escape_text_mode(r"\text{a\_b}") == r"\text{a\_b}" + + def test_no_unescaped_underscore_in_math_text( + self, math: dict, inputs: xr.Dataset + ) -> None: + # `cap_max` is a parameter whose underscore is rendered inside `\textit`; + # no `\text*{...}` argument in the document may hold an unescaped one. + doc = latex_math_doc(math, inputs, format="md") + for arg in re.findall(r"\\text(?:bf|it)?\{([^{}]*)\}", doc): + assert "_" not in arg.replace(r"\_", "") + + def test_cross_references(self, math: dict, inputs: xr.Dataset) -> None: + builder = LatexModelBuilder(math, inputs, {}).build() + cost = builder.components["parameters"]["cost"] + assert set(cost.used_in) == {"cost_plus_one", "sub_expr_test", "total_cost"} + # Dimensions are not cross-referenced. + obj = builder.components["objectives"]["obj"] + assert obj.uses == ["total_cost"] + assert obj.extras["Sense"] == "minimise" + + def test_equation_masks_render_as_if_conditions( + self, math: dict, inputs: xr.Dataset + ) -> None: + math["constraints"]["cap"]["equations"][0]["mask"] = "cost > 1" + doc = latex_math_doc(math, inputs, format="md") + assert r"\text{if } (\textit{cost}_\text{n}\mathord{>}\text{1})" in doc + + def test_infinity_not_wrapped_in_text(self, math: dict, inputs: xr.Dataset) -> None: + # `cap_max`'s default is `inf`; a mask comparing against it must render + # `\infty` as a bare math token, never `\text{\infty}` (which KaTeX + # would print as the literal string, not the symbol). + math["constraints"]["cap"]["equations"][0]["mask"] = "cap_max == inf" + doc = latex_math_doc(math, inputs, format="md") + assert r"\infty" in doc + assert r"\text{\infty}" not in doc + + def test_sub_expression_variants_produce_multiple_equations( + self, math: dict, inputs: xr.Dataset + ) -> None: + math["expressions"]["sub_expr_test"]["sub_expressions"]["foo"] = [ + {"mask": "cost > 1", "expression": "flow"}, + {"mask": "not cost > 1", "expression": "flow * 2"}, + ] + builder = LatexModelBuilder(math, inputs, {}).build() + equations = builder.components["expressions"]["sub_expr_test"].equations + assert len(equations) == 2 + assert equations[1]["expression"] == ( + r"\textbf{flow}_\text{n} \times 2 \times \textit{cost}_\text{n}" + ) + + def test_multiple_equations_render_as_one_block_with_cases( + self, math: dict, inputs: xr.Dataset + ) -> None: + # Sub-clauses of a component share one `foreach`/top-level mask and + # should render as `cases` rows at the same nesting level, not as + # separate top-level math blocks. + math["constraints"]["cap"]["equations"] = [ + {"mask": "cost > 1", "expression": "flow <= cap_max"}, + {"mask": "not cost > 1", "expression": "flow <= 0"}, + ] + doc = latex_math_doc(math, inputs, format="md") + section = doc.split("### cap\n")[1].split("### ")[0] + assert section.count(r"\begin{array}{l}") == 1 + assert section.count(r"\begin{cases}") == 1 + assert section.count("$$") == 2 # one opening, one closing delimiter + assert r"\text{if } (\textit{cost}_\text{n}\mathord{>}\text{1})" in section + assert ( + r"\text{if } (\neg (\textit{cost}_\text{n}\mathord{>}\text{1}))" in section + ) + + def test_single_equation_renders_inline_without_cases( + self, math: dict, inputs: xr.Dataset + ) -> None: + doc = latex_math_doc(math, inputs, format="md") + section = doc.split("### total_cost\n")[1].split("### ")[0] + assert r"\begin{cases}" not in section + assert r"\begin{array}{l}" in section + + def test_markdown_document_structure(self, math: dict, inputs: xr.Dataset) -> None: + doc = latex_math_doc(math, inputs, format="md") + assert doc.startswith("# Math formulation") + for heading in ("## Parameters", "## Variables", "## Constraints", "### cap"): + assert heading in doc + assert "$$" in doc + + def test_rst_document_structure(self, math: dict, inputs: xr.Dataset) -> None: + doc = latex_math_doc(math, inputs, format="rst") + assert ".. math::" in doc + assert "Math formulation\n================" in doc + + def test_tex_document_structure(self, math: dict, inputs: xr.Dataset) -> None: + doc = latex_math_doc(math, inputs, format="tex") + assert r"\section{Math formulation}" in doc + assert r"\begin{equation}" in doc + # Underscores are escaped in text-mode headings. + assert r"\paragraph{cap\_max}" in doc + + +class TestEndToEnd: + """Full builds from math definitions.""" + + def test_dummy_declarative_model_end_to_end( + self, math: dict, inputs: xr.Dataset + ) -> None: + model = declarative_model(math, inputs, {}) + assert "flow" in model.variables + assert "total_cost" in model.expressions + assert "cap" in model.constraints + assert model.objective is not None + # flow is indexed over the node dimension. + assert set(model.variables["flow"].dims) == {"node"} + + @pytest.fixture(scope="class") + @classmethod + def larger_math(cls) -> dict: + math_path = Path(__file__).parent / "resources" / "math.yaml" + math_dict = yaml.safe_load(math_path.read_text()) + return math_dict + + @pytest.fixture(scope="class") + @classmethod + def larger_inputs(cls) -> xr.Dataset: + inputs_path = Path(__file__).parent / "resources" / "inputs.nc" + return xr.load_dataset(inputs_path) + + def test_larger_declarative_model_end_to_end( + self, larger_math: dict, larger_inputs: xr.Dataset + ) -> None: + """The demo math.yaml in the resources directory validates, parses, and evaluates to a linopy model.""" + declarative_model(larger_math, larger_inputs, {}) + + @pytest.mark.parametrize("fmt", ("md", "rst", "tex")) + def test_larger_declarative_model_latex_doc_end_to_end( + self, + larger_math: dict, + larger_inputs: xr.Dataset, + fmt: Literal["md", "rst", "tex"], + ) -> None: + # And the full LaTeX math documentation generates in every format. + doc = latex_math_doc(larger_math, larger_inputs, format=fmt) + name = r"storage\_balance" if fmt == "tex" else "storage_balance" + assert name in doc