From e594709da2367c5cc0d15d51fd92f097c139a703 Mon Sep 17 00:00:00 2001 From: Bryn Pickering <17178478+brynpickering@users.noreply.github.com> Date: Wed, 10 Jun 2026 19:43:40 +0100 Subject: [PATCH 01/12] [WIP] port across Calliope YAML math parsing; update to work with linopy --- linopy/__init__.py | 1 + linopy/declarative.py | 0 linopy/declarative/__init__.py | 1 + linopy/declarative/build.py | 198 ++++ linopy/declarative/eval_attrs.py | 62 ++ linopy/declarative/expression_parser.py | 1312 +++++++++++++++++++++++ linopy/declarative/helper_functions.py | 951 ++++++++++++++++ linopy/declarative/mask_parser.py | 617 +++++++++++ linopy/declarative/parsing.py | 830 ++++++++++++++ linopy/declarative/schema.py | 705 ++++++++++++ pyproject.toml | 3 + 11 files changed, 4680 insertions(+) create mode 100644 linopy/declarative.py create mode 100644 linopy/declarative/__init__.py create mode 100644 linopy/declarative/build.py create mode 100644 linopy/declarative/eval_attrs.py create mode 100644 linopy/declarative/expression_parser.py create mode 100644 linopy/declarative/helper_functions.py create mode 100644 linopy/declarative/mask_parser.py create mode 100644 linopy/declarative/parsing.py create mode 100644 linopy/declarative/schema.py diff --git a/linopy/__init__.py b/linopy/__init__.py index e80e615da..d145a96e0 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 diff --git a/linopy/declarative.py b/linopy/declarative.py new file mode 100644 index 000000000..e69de29bb diff --git a/linopy/declarative/__init__.py b/linopy/declarative/__init__.py new file mode 100644 index 000000000..80bfc3f97 --- /dev/null +++ b/linopy/declarative/__init__.py @@ -0,0 +1 @@ +"""Linopy declarative text math interface.""" \ No newline at end of file diff --git a/linopy/declarative/build.py b/linopy/declarative/build.py new file mode 100644 index 000000000..5a9c3679c --- /dev/null +++ b/linopy/declarative/build.py @@ -0,0 +1,198 @@ +import time +import typing + +import xarray as xr + +from linopy import merge +from linopy.declarative import parsing +from linopy.declarative.schema import ( + LOGGER, + ConfigModel, + ConstraintDef, + MathModel, + ObjectiveDef, + VariableDef, +) +from linopy.expressions import LinearExpression +from linopy.model import Model + +ORDERED_COMPONENTS_T = typing.Literal[ + "variables", + # "global_expressions", + "constraints", + # "piecewise_constraints", + "objectives", +] + + +def declarative_model(math_def: dict, input_data: xr.Dataset, config: dict) -> Model: + """Build a Linopy Model from declarative math definitions and input data.""" + builder = DeclarativeModelBuilder(math_def, input_data, config) + return builder.build() + + +class DeclarativeModelBuilder: + def __init__(self, math_def: dict, input_data: xr.Dataset, config: dict): + self.model = Model() + self.math = MathModel.model_validate(math_def) + self.input_data = input_data + self.config = ConfigModel.model_validate(config) + + @staticmethod + def _sorted_by_order( + root: typing.Mapping[str, typing.Any], + ) -> list[tuple[str, typing.Any]]: + """Return (name, obj) pairs from a root mapping, sorted by obj.order.""" + return sorted(root.items(), key=lambda item: getattr(item[1], "order", 0)) + + def add_variable(self, name: str, definition: VariableDef): + references: set[str] = set() + parsed_component = parsing.ParsedBackendComponent( + "variables", name, definition, self.math.parsing_components + ) + mask = parsed_component.generate_top_level_mask( + self.input_data, + self.model, + self.math, + self.config, + align_to_foreach_sets=True, + break_early=True, + references=references, + ) + kwargs = { + "upper": definition.bounds.upper, + "lower": definition.bounds.lower, + "integer": definition.domain == "integer", + "binary": definition.domain == "binary", + } + if mask.any(): + self.model.add_variables(coords=mask.coords, name=name, mask=mask, **kwargs) + self.model.variables[name].attrs["references"] = references + else: + LOGGER.warning( + f"Optimisation Model | variables:{name} | No valid data points after applying 'where' condition. Variable not added to model." + ) + + def add_constraint(self, name: str, definition: ConstraintDef): + references: set[str] = set() + parsed_component = parsing.ParsedBackendComponent( + "constraints", name, definition, self.math.parsing_components + ) + mask = parsed_component.generate_top_level_mask( + self.input_data, + self.model, + self.math, + self.config, + align_to_foreach_sets=True, + break_early=True, + references=references, + ) + lhs = LinearExpression(float("nan"), self.model).where(mask) + sign = xr.DataArray().where(parsed_component.drop_dims_not_in_foreach(mask)) + rhs = LinearExpression(float("nan"), self.model).where(mask) + all_mask = mask.copy() + if mask.any(): + equations = parsed_component.parse_equations() + for equation in equations: + sub_mask = equation.evaluate_mask( + self.input_data, + self.model, + self.math, + self.config, + initial_mask=mask, + references=references, + ) + if not sub_mask.any(): + continue + sub_mask = parsed_component.drop_dims_not_in_foreach(sub_mask) + if (sign.notnull() & sub_mask).any(): + raise ValueError( + f"Optimisation Model | 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 = equation.evaluate_expression( + self.input_data, + self.model, + self.math, + mask=sub_mask, + references=references, + ) + all_mask = all_mask | sub_mask + if isinstance(lhs_to_fill, xr.DataArray): + lhs = lhs.fillna(lhs_to_fill) + else: + lhs = merge([lhs, lhs_to_fill]).where(all_mask) + sign = sign.fillna(sign_to_fill) + if isinstance(rhs_to_fill, xr.DataArray): + rhs = rhs.fillna(rhs_to_fill) + else: + rhs = merge([rhs, rhs_to_fill]).where(all_mask) + + self.model.add_constraints( + coords=all_mask.coords, + name=name, + lhs=lhs, + sign=sign.fillna( + "==" + ), # Default to equality to avoid errors; will be masked. + rhs=rhs, + mask=all_mask, + ) + self.model.constraints[name].attrs["references"] = references + + def add_objective(self, name: str, definition: ObjectiveDef): + references: set[str] = set() + parsed_component = parsing.ParsedBackendComponent( + "objectives", name, definition, self.math.parsing_components + ) + mask = parsed_component.generate_top_level_mask( + self.input_data, + self.model, + self.math, + self.config, + align_to_foreach_sets=True, + break_early=True, + references=references, + ) + expr = LinearExpression(float("nan"), self.model).where(mask) + if mask.any(): + equations = parsed_component.parse_equations() + for equation in equations: + sub_mask = equation.evaluate_mask( + self.input_data, + self.model, + self.math, + self.config, + initial_mask=mask, + references=references, + ) + if not sub_mask.any(): + continue + sub_mask = parsed_component.drop_dims_not_in_foreach(sub_mask) + if (~expr.isnull() & sub_mask).any(): + raise ValueError( + f"Optimisation Model | objectives:{name} | Overlapping 'mask' conditions between equations are not allowed. Please revise the 'mask' conditions to ensure they are mutually exclusive." + ) + expr_to_fill = equation.evaluate_expression( + self.input_data, + self.model, + self.math, + mask=sub_mask, + references=references, + ) + expr = expr_to_fill + self.model.add_objective(expr=expr, sense=definition.sense) + self.model.objective.attrs["references"] = references + + def build(self) -> Model: + for components in typing.get_args(ORDERED_COMPONENTS_T): + component = components.removesuffix("s") + ordered_items = self._sorted_by_order(self.math[components].root) + for name, definition in ordered_items: + start = time.time() + getattr(self, f"add_{component}")(name, definition) + end = time.time() - start + LOGGER.debug( + f"Optimisation Model | {components}:{name} | Built in {end:.4f}s" + ) + LOGGER.info(f"Optimisation Model | {components} | Generated.") + return self.model diff --git a/linopy/declarative/eval_attrs.py b/linopy/declarative/eval_attrs.py new file mode 100644 index 000000000..e525d5294 --- /dev/null +++ b/linopy/declarative/eval_attrs.py @@ -0,0 +1,62 @@ +# Copyright (C) since 2013 Calliope contributors listed in AUTHORS. +# Licensed under the Apache 2.0 License (see LICENSE file). +"""Parsing evaluation attributes.""" + +from __future__ import annotations + +import logging +from collections.abc import Callable +from dataclasses import dataclass, field +from typing import TYPE_CHECKING + +import xarray as xr + +from linopy.declarative.schema import ConfigModel, MathModel + +if TYPE_CHECKING: + from linopy.model import Model +TRUE_ARRAY = xr.DataArray(True) + +LOGGER = logging.getLogger(__name__) + + +@dataclass +class EvalAttrs: + """Attributes required for evaluating parsed expressions.""" + + model: Model + """Backend interface component dataset.""" + + equation_name: str = "" + """Name of the equation being evaluated.""" + + helper_functions: dict[str, Callable] = field(default_factory=dict) + """Helper functions available for evaluations.""" + + input_data: xr.Dataset = field(default_factory=xr.Dataset) + """Model input data.""" + + math: MathModel = field(default_factory=MathModel) + """Linopy math definitions.""" + + apply_mask: bool = True + """Whether to apply the 'where' condition.""" + + as_values: bool = False + """If True, return with the array contents evaluated to base Python objects. + If False, return with the array contents as they are in the backend dataset.""" + + config: ConfigModel = field(default_factory=ConfigModel) + """Build configuration options.""" + + references: set[str] = field(default_factory=set) + """References to dimensions/lookups/parameters/variables/global expressions used in the expression.""" + + slice_dict: dict = field(default_factory=dict) + """Dictionary to look up array slice expressions if referenced in the evaluated string.""" + + sub_expression_dict: dict = field(default_factory=dict) + """Dictionary to look up sub-expressions if referenced in the evaluated string.""" + + mask: xr.DataArray = field(default_factory=lambda: xr.DataArray(True)) + """Boolean array defining where the expression should be applied.""" diff --git a/linopy/declarative/expression_parser.py b/linopy/declarative/expression_parser.py new file mode 100644 index 000000000..650092af6 --- /dev/null +++ b/linopy/declarative/expression_parser.py @@ -0,0 +1,1312 @@ +# Copyright (C) since 2013 Calliope contributors listed in AUTHORS. +# Licensed under the Apache 2.0 License (see LICENSE file). + +## +# Part of the code in this file is adapted from +# https://github.com/pyparsing/pyparsing/blob/master/examples/eval_arith.py +# available under the MIT license +## +# Copyright 2009, 2011 Paul McGuire +# +# Permission is hereby granted, free of charge, to any person obtaining +# a copy of this software and associated documentation files (the +# "Software"), to deal in the Software without restriction, including +# without limitation the rights to use, copy, modify, merge, publish, +# distribute, sublicense, and/or sell copies of the Software, and to +# permit persons to whom the Software is furnished to do so, subject to +# the following conditions: + +# The above copyright notice and this permission notice shall be +# included in all copies or substantial portions of the Software. + +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, +# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF +# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. +# IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY +# CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, +# TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE +# SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. +## +"""Expression parsing functionality.""" + +from __future__ import annotations + +import re +from abc import ABC, abstractmethod +from collections.abc import Callable, Iterable, Iterator +from dataclasses import replace +from typing import Any, Literal, overload + +import numpy as np +import pandas as pd +import pyparsing as pp +import xarray as xr + +from linopy.declarative.eval_attrs import EvalAttrs +from linopy.declarative.helper_functions import ParsingHelperFunction +from linopy.expressions import LinearExpression +from linopy.variables import Variable + +pp.ParserElement.enable_packrat() + +SUB_EXPRESSION_CLASSIFIER = "$" + + +RETURN_T = Literal["array", "math_string"] + + +class EvalString(ABC): + """Parent class for all string evaluation classes - used in type hinting.""" + + name: str + eval_attrs: EvalAttrs + instring: str + + def __eq__(self, other): + """Functionality for '==' operations.""" + return self.__repr__() == other + + @abstractmethod + def __repr__(self) -> str: + """Return string representation of the parsed grammar.""" + + def error_msg(self, message: str) -> ValueError: + """Raise an error message with context.""" + return ValueError( + f"({self.eval_attrs.equation_name}, {self.instring}) | {message}" + ) + + +class EvalArrayOrMath(EvalString): + """Abstract class to evaluate expressions as either arrays or math strings.""" + + @abstractmethod + def as_math_string(self) -> str: + """Evaluate and return expression as LaTeX.""" + + @abstractmethod + def as_array(self) -> xr.DataArray | list[xr.DataArray]: + """ + Evaluate and return expression as a DataArray or list. + + If the evaluated expression returns a simple string or number, + this value will be assigned as both the `name` and the data of the returned DataArray. + The purpose of this is to be able to access the string/number value whether we query the array name or its data. + """ + + # Math strings evaluate to strings. + @overload + def eval( + self, return_type: Literal["math_string"], eval_attrs: EvalAttrs + ) -> str: ... + + # Arrays evaluate to arrays + @overload + def eval( + self, return_type: Literal["array"], eval_attrs: EvalAttrs + ) -> xr.DataArray | list[xr.DataArray]: ... + + def eval( + self, return_type: RETURN_T, eval_attrs: EvalAttrs + ) -> str | xr.DataArray | list[xr.DataArray]: + """ + Evaluate math string expression. + + Args: + return_type (Literal[math_string, input, array]): + Dictates how the expression should be evaluated (see `Returns` section). + eval_attrs (EvalAttrs): Evaluation attributes. + + Returns: + str | list[str | float] | xr.DataArray: + If `math_string` is desired, returns a valid LaTex math string. + If `array` is desired, returns xarray DataArray or a list of strings/numbers (if the expression represents a list). + """ + self.eval_attrs = eval_attrs + evaluated: str | list[str | float] | xr.DataArray | list[xr.DataArray] + if return_type == "array": + evaluated = self.as_array() + elif return_type == "math_string": + evaluated = self.as_math_string() + return evaluated + + +class EvalToCallable(EvalString): + """Parent class for callable functionality.""" + + @abstractmethod + def as_callable(self, return_type: RETURN_T) -> Callable: + """Callable processing.""" + ... + + def eval(self, return_type: RETURN_T, eval_attrs: EvalAttrs) -> Callable: + """ + Evaluate math string expression. + + Args: + return_type (str): Whether to return a math string or xarray DataArray. + eval_attrs (EvalAttrs): Evaluation attributes. + + Returns: + Callable: returns helper function. + """ + self.eval_attrs = eval_attrs + evaluated = self.as_callable(return_type) + return evaluated + + +class EvalOperatorOperand(EvalArrayOrMath): + """Evaluation of math operands.""" + + LATEX_OPERATOR_LOOKUP: dict[str, str] = { + "**": "{val}^{{{operand}}}", + "*": r"{val} \times {operand}", + "/": r"\frac{{ {val} }}{{ {operand} }}", + "+": "{val} + {operand}", + "-": "{val} - {operand}", + } + SKIP_IF: list[str] = ["+", "-"] + + def __init__(self, instring: str, loc: int, tokens: pp.ParseResults) -> None: + """ + Process successfully parsed expressions with operands separated by an operator. + + I.e.: OPERAND OPERATOR OPERAND OPERATOR OPERAND ... + + Args: + instring (str): String that was parsed (used in error message). + loc (int): + Location in parsed string where parsing error was logged. + This is not used; we include it as pyparsing injects it alongside `instring` when setting the parse action. + tokens (pp.ParseResults): + Contains a list of the form [operand (pp.ParseResults), operator (str), + operand (pp.ParseResults), operator (str), ...]. + """ + self.value: pp.ParseResults = tokens[0] + self.values = tokens + self.instring = instring + + def __repr__(self) -> str: + """Programming / official string representation.""" + first_operand = self.value[0].__repr__() + operand_operator_pairs = " ".join( + op + " " + val.__repr__() + for op, val in self._operator_operands(self.value[1:]) + ) + arithmetic_string = f"({first_operand} {operand_operator_pairs})" + return arithmetic_string + + def _operator_operands( + self, token_list: list + ) -> Iterator[tuple[str, pp.ParseResults]]: + """Generator to extract operators and operands in pairs.""" + it = iter(token_list) + while 1: + try: + yield (next(it), next(it)) + except StopIteration: + break + + def _apply_mask(self, evaluated: xr.DataArray) -> xr.DataArray: + """Util function to apply mask arrays to non-latex strings.""" + mask = self.eval_attrs.mask + try: + evaluated = evaluated.where(mask) + except AttributeError: + evaluated = evaluated.broadcast_like(mask).where(mask) + + return evaluated + + def _skip_component_on_conditional(self, component: str, operator_: str) -> bool: + """ + Conditional to skip adding to math string if element evaluates to zero. + + E.g., "0 + flow_cap" is better evaluated as simply "flow_cap". + """ + return component == "0" and operator_ in self.SKIP_IF + + @staticmethod + def _operate( + val: xr.DataArray, evaluated_operand: xr.DataArray, operator_: str + ) -> xr.DataArray: + """Apply evaluated operation on two DataArrays.""" + match operator_: + case "**": + val = val**evaluated_operand + case "*": + val = val * evaluated_operand + case "/": + val = val / evaluated_operand + case "+": + val = val + evaluated_operand + case "-": + val = val - evaluated_operand + return val + + def as_math_string(self) -> str: # noqa: D102, override + val = self.value[0].eval("math_string", self.eval_attrs) + + for operator_, operand in self._operator_operands(self.value[1:]): + evaluated_operand = operand.eval("math_string", self.eval_attrs) + # We ignore zeros that do nothing + if self._skip_component_on_conditional(evaluated_operand, operator_): + continue + if isinstance(self.value[0], type(self)): + val = "(" + val + ")" + if isinstance(operand, type(self)): + evaluated_operand = "(" + evaluated_operand + ")" + if self._skip_component_on_conditional(val, operator_): + val = evaluated_operand + else: + val = self.LATEX_OPERATOR_LOOKUP[operator_].format( + val=val, operand=evaluated_operand + ) + return val + + def as_array(self) -> xr.DataArray: # noqa: D102, override + val = self._apply_mask(self.value[0].eval("array", self.eval_attrs)) + + for operator_, operand in self._operator_operands(self.value[1:]): + evaluated_operand = self._apply_mask(operand.eval("array", self.eval_attrs)) + val = self._operate(val, evaluated_operand, operator_) + return val + + +class EvalSignOp(EvalArrayOrMath): + """Class for processing expressions with + or -.""" + + def __init__(self, instring: str, loc: int, tokens: pp.ParseResults) -> None: + """ + Parse action to process successfully parsed expressions with a leading + or - sign. + + Args: + instring (str): String that was parsed (used in error message). + loc (int): + Location in parsed string where parsing error was logged. + This is not used; we include it as pyparsing injects it alongside `instring` when setting the parse action. + tokens (pp.ParseResults): + Contains a list of the form [sign (str), operand (pp.ParseResults)]. + """ + self.sign, self.value = tokens[0] + self.values = tokens + self.instring = instring + + def __repr__(self) -> str: + """Programming / official string representation.""" + return str(f"({self.sign}){self.value.__repr__()}") + + # string return + @overload + def _eval(self, return_type: Literal["math_string"]) -> str: ... + + # array return + @overload + def _eval(self, return_type: Literal["array"]) -> xr.DataArray: ... + + def _eval(self, return_type: RETURN_T) -> xr.DataArray | str: + """Evaluate the element that will have the sign attached to it.""" + return self.value.eval(return_type, self.eval_attrs) + + def as_math_string(self) -> str: # noqa: D102 + return self.sign + self._eval("math_string") + + def as_array(self) -> xr.DataArray: # noqa: D102, override + evaluated = self._eval("array") + if self.sign == "-": + evaluated = -1 * evaluated + return evaluated + + +class EvalComparisonOp(EvalArrayOrMath): + """Class for processing comparison operations.""" + + OP_TRANSLATOR = {"<=": r" \leq ", ">=": r" \geq ", "==": " = "} + + def __init__(self, instring: str, loc: int, tokens: pp.ParseResults) -> None: + """ + Parse action to process successfully parsed equations of the form LHS OPERATOR RHS. + + Args: + instring (str): String that was parsed (used in error message). + loc (int): + Location in parsed string where parsing error was logged. + This is not used; we include it as pyparsing injects it alongside `instring` when setting the parse action. + tokens (pp.ParseResults): + Contains a list with an RHS (pp.ParseResults), operator (str), and LHS (pp.ParseResults). + """ + self.lhs, self.op, self.rhs = tokens + self.instring = instring + self.loc = loc + self.values = tokens + + def __repr__(self) -> str: + """Programming / official string representation.""" + return f"{self.lhs.__repr__()} {self.op} {self.rhs.__repr__()}" + + # string return + @overload + def _eval(self, return_type: Literal["math_string"]) -> tuple[str, str]: ... + + # array return + @overload + def _eval( + self, return_type: Literal["array"] + ) -> tuple[xr.DataArray, xr.DataArray]: ... + + def _eval( + self, return_type: RETURN_T + ) -> tuple[str, str] | tuple[xr.DataArray, xr.DataArray]: + """Evaluate the LHS and RHS of the comparison.""" + lhs = self.lhs.eval(return_type, self.eval_attrs) + rhs = self.rhs.eval(return_type, self.eval_attrs) + return lhs, rhs + + def as_math_string(self) -> str: # noqa: D102, override + lhs, rhs = self._eval("math_string") + return lhs + self.OP_TRANSLATOR[self.op] + rhs + + def as_array( + self, + ) -> tuple[LinearExpression, xr.DataArray, LinearExpression | xr.DataArray]: # noqa: D102, override: # noqa: D102, override + lhs, rhs = self._eval("array") + mask = self.eval_attrs.mask + for side, arr in {"left": lhs, "right": rhs}.items(): + extra_dims = set(arr.dims).difference(set(mask.dims) | {"_term"}) + if extra_dims: + raise self.error_msg( + f"The {side}-hand side of the equation is indexed over dimensions not present in `foreach`: {extra_dims}" + ) + lhs_masked = lhs.where(mask) + rhs_masked = rhs.where(mask) + if isinstance(lhs_masked, Variable): + lhs_masked = lhs_masked.to_linexpr() + if isinstance(rhs_masked, Variable): + rhs_masked = rhs_masked.to_linexpr() + sign_masked = xr.DataArray(self.op).where(mask) + return lhs_masked, sign_masked, rhs_masked + + +class EvalFunction(EvalArrayOrMath): + """Class to process parsed functions.""" + + def __init__(self, instring: str, loc: int, tokens: pp.ParseResults) -> None: + """ + Parse action to process successfully parsed helper function strings. + + Strings must be in the following form: helper_function_name(*args, **kwargs). + + Args: + instring (str): String that was parsed (used in error message). + loc (int): + Location in parsed string where parsing error was logged. + This is not used; we include it as pyparsing injects it alongside `instring` when setting the parse action. + tokens (pp.ParseResults): + Has a dictionary component with the parsed elements: + helper_function_name (pp.ParseResults), args (list), kwargs (dict). + """ + token_dict = tokens.as_dict() + self.func_name: pp.ParseResults = token_dict["helper_function_name"] + self.args: list = token_dict["args"] + self.kwargs: dict = token_dict["kwargs"] + self.values = tokens + self.instring = instring + + def __repr__(self) -> str: + """Programming / official string representation.""" + _kwargs = ", ".join(f"{k}={v}" for k, v in self.kwargs.items()) + return f"{str(self.func_name)}(args={self.args}, kwargs={{{_kwargs}}})" + + @overload + def _arg_eval(self, return_type: Literal["math_string"], arg: Any) -> str: ... + + @overload + def _arg_eval( + self, return_type: Literal["array"], arg: Any + ) -> xr.DataArray | list[str | float]: ... + + def _arg_eval( + self, return_type: RETURN_T, arg: Any + ) -> str | xr.DataArray | list[str | float]: + """Evaluate the arguments of the helper function.""" + if isinstance(arg, pp.ParseResults): + evaluated = arg[0].eval(return_type, self.eval_attrs) + elif isinstance(arg, list): + evaluated = [self._arg_eval(return_type, arg_) for arg_ in arg] + elif isinstance(arg, ListParser): + evaluated = arg.eval("array", self.eval_attrs) + else: + evaluated = arg.eval(return_type, self.eval_attrs) + if isinstance(evaluated, xr.DataArray) and isinstance(arg, EvalGenericString): + evaluated = evaluated.item() + return evaluated + + @overload + def _eval(self, return_type: Literal["math_string"]) -> str: ... + + @overload + def _eval(self, return_type: Literal["array"]) -> xr.DataArray: ... + + def _eval(self, return_type: RETURN_T) -> str | xr.DataArray: + """Pass evaluated arguments to evaluated helper function.""" + helper_function = self.func_name.eval(return_type, self.eval_attrs) + if helper_function.ignore_mask: + self.eval_attrs = replace(self.eval_attrs, mask=xr.DataArray(True)) + + args_ = [] + for arg in self.args: + args_.append(self._arg_eval(return_type, arg)) + + kwargs_ = {} + for kwarg_name, kwarg_val in self.kwargs.items(): + kwargs_[kwarg_name] = self._arg_eval(return_type, kwarg_val) + + evaluated = helper_function(*args_, **kwargs_) + return evaluated + + def as_math_string(self) -> str: # noqa: D102, override + return self._eval("math_string") + + def as_array(self) -> xr.DataArray: # noqa: D102, override + return self._eval("array") + + +class EvalHelperFuncName(EvalToCallable): + """For processing parsed helper function names.""" + + def __init__(self, instring: str, loc: int, tokens: pp.ParseResults) -> None: + """ + Parse action to process successfully parsed helper function names. + + This is a unique parse action so that we can catch invalid helper functions + most safely. + + Args: + instring (str): String that was parsed (used in error message). + loc (int): + Location in parsed string where parsing error was logged. + This is not used; we include it as pyparsing injects it alongside `instring` when setting the parse action. + tokens (pp.ParseResults): + Has one parsed element: helper_function_name (str). + """ + self.name = self.value = tokens[0] + self.instring = instring + self.loc = loc + self.values = tokens + + def __repr__(self) -> str: + """Programming / official string representation.""" + return str(self.name) + + def as_callable(self, return_type: RETURN_T) -> Callable: + """Evalluate and return the callable action of the helper function.""" + helper_functions = self.eval_attrs.helper_functions + if self.name not in helper_functions.keys(): + raise self.error_msg(f"Invalid helper function defined: {self.name}") + elif not isinstance(helper_functions[self.name], type(ParsingHelperFunction)): + raise self.error_msg( + f"Helper function must be subclassed from calliope.backend.helper_functions.ParsingHelperFunction: {self.name}" + ) + else: + return helper_functions[self.name](return_type, self.eval_attrs) + + +class EvalSlicedComponent(EvalArrayOrMath): + """For processing of sliced parameters / decision variables.""" + + def __init__(self, instring: str, loc: int, tokens: pp.ParseResults) -> None: + """ + Process successfully parsed sliced parameters or decision variables. + + In the form of param_or_var[*slices]. + + Args: + instring (str): String that was parsed (used in error message). + loc (int): + Location in parsed string where parsing error was logged. + This is not used; we include it as pyparsing injects it alongside `instring` when setting the parse action. + tokens (pp.ParseResults): + Has a dictionary component with the parsed elements: + param_or_var_name (str), slices (list of strings). + """ + token_dict = tokens.as_dict() + self.obj_name: pp.ParseResults = token_dict["param_or_var_name"] + + self.slices: dict[str, pp.ParseResults] = { + idx["set_name"][0]: idx["slicer"][0] for idx in token_dict["slices"] + } + self.values = tokens + self.instring = instring + + def __repr__(self) -> str: + """Programming / official string representation.""" + slices = ", ".join(f"{k}={v.__repr__()}" for k, v in self.slices.items()) + return f"SLICED_{self.obj_name}[{slices}]" + + @staticmethod + def _replace_rule(index_slices): + """ + String parsing rule to catch and replace dimension names with the names + their slices. + + E.g., `techs` -> `techs=pv`. + """ + + def __replace(term): + if len(term) == 1: + return term + else: + replacers = {k: f"{k}={v}" for k, v in index_slices.items()} + return ( + term[0] + + term[1] + + ",".join(replacers.get(k, k) for k in term[2]) + + term[3] + ) + + return __replace + + @overload + def _eval(self, return_type: Literal["math_string"]) -> tuple[str, dict]: ... + + @overload + def _eval(self, return_type: Literal["array"]) -> tuple[xr.DataArray, dict]: ... + + def _eval(self, return_type: RETURN_T) -> tuple[str | xr.DataArray, dict]: + """Evaluate the slice dim and vals of each slice element.""" + slices: dict[str, Any] = { + k: xr.concat(slice_, dim=k) + if isinstance(slice_ := v.eval(return_type, self.eval_attrs), list) + else slice_ + for k, v in self.slices.items() + } + + evaluated = self.obj_name.eval(return_type, self.eval_attrs) + return evaluated, slices + + def as_math_string(self) -> str: # noqa: D102, override + evaluated, slices = self._eval("math_string") + singular_slice_refs = { + self.eval_attrs.math.dimensions[k].iterator: v for k, v in slices.items() + } + id_ = pp.Combine( + pp.Word(pp.alphas, pp.alphanums) + + pp.ZeroOrMore("_" + pp.Word(pp.alphanums)) + + pp.Opt("_") + ) + id_formatted = pp.Combine("\\" + pp.Word(pp.alphas) + "{" + id_ + "}") + obj_parser = id_formatted + pp.Opt( + r"_\text{" + pp.Group(pp.DelimitedList(id_)) + "}" + ) + obj_parser.set_parse_action(self._replace_rule(singular_slice_refs)) + return obj_parser.parse_string(evaluated, parse_all=True)[0] + + def as_array(self) -> xr.DataArray: # noqa: D102, override + evaluated, slices = self._eval("array") + return evaluated.sel(**slices) + + +class EvalIndexSlice(EvalArrayOrMath): + """For processing `$slice` expressions.""" + + def __init__(self, instring: str, loc: int, tokens: pp.ParseResults) -> None: + """ + Process successfully parsed expression index `$slice` references. + + Args: + instring (str): String that was parsed (used in error message). + loc (int): + Location in parsed string where parsing error was logged. + This is not used; we include it as pyparsing injects it alongside `instring` when setting the parse action. + tokens (pp.ParseResults): + Has one parsed element containing the index slice name (str). + """ + self.name: str = tokens[0] + self.values = tokens + self.instring = instring + + def __repr__(self) -> str: + """Programming / official string representation.""" + return "REFERENCE:" + str(self.name) + + @overload + def _eval(self, return_type: Literal["math_string"], as_values: bool) -> str: ... + + @overload + def _eval( + self, return_type: Literal["array"], as_values: bool + ) -> xr.DataArray | list[xr.DataArray]: ... + + def _eval( + self, return_type: RETURN_T, as_values: bool + ) -> str | xr.DataArray | list[xr.DataArray]: + """Evaluate the referenced `slice`.""" + self.eval_attrs = replace(self.eval_attrs, as_values=as_values) + return self.eval_attrs.slice_dict[self.name][0].eval( + return_type, self.eval_attrs + ) + + def as_math_string(self) -> str: # noqa: D102, override + return self._eval("math_string", False) + + def as_array(self) -> xr.DataArray | list[xr.DataArray]: # noqa: D102, override + evaluated = self._eval("array", True) + return evaluated + + +class EvalSubExpressions(EvalArrayOrMath): + """For processing sub-expressions.""" + + def __init__(self, instring: str, loc: int, tokens: pp.ParseResults) -> None: + """ + Process successfully parsed `$sub_expressions`. + + Args: + instring (str): String that was parsed (used in error message). + loc (int): + Location in parsed string where parsing error was logged. + This is not used; we include it as pyparsing injects it alongside `instring` when setting the parse action. + tokens (pp.ParseResults): + Has one parsed element containing the sub_expression name (str). + """ + self.name: str = tokens[0] + self.values = tokens + self.instring = instring + + def __repr__(self) -> str: + """Programming / official string representation.""" + return "SUB_EXPRESSION:" + str(self.name) + + @overload + def _eval(self, return_type: Literal["math_string"]) -> str: ... + + @overload + def _eval(self, return_type: Literal["array"]) -> xr.DataArray: ... + + def _eval(self, return_type: RETURN_T) -> str | xr.DataArray: + """Evaluate the referenced sub_expression.""" + return self.eval_attrs.sub_expression_dict[self.name][0].eval( + return_type, self.eval_attrs + ) + + def as_math_string(self) -> str: # noqa: D102, override + return self._eval("math_string") + + def as_array(self) -> xr.DataArray: # noqa: D102, override + return self._eval("array") + + +class EvalNumber(EvalArrayOrMath): + """For processing numbers.""" + + def __init__(self, instring: str, loc: int, tokens: pp.ParseResults) -> None: + """ + Process successfully parsed numbers. + + Catches integers (1), floats (1.), and in scientific notation (1e1). + Also capture infinity (inf/.inf). + + Args: + instring (str): String that was parsed (used in error message). + loc (int): + Location in parsed string where parsing error was logged. + This is not used; we include it as pyparsing injects it alongside `instring` when setting the parse action. + tokens (pp.ParseResults): + Has one parsed element containing the number (str). + """ + self.value = tokens[0] + self.values = tokens + self.instring = instring + + def __repr__(self) -> str: + """Programming / official string representation.""" + return "NUM:" + str(self.value) + + def as_math_string(self) -> str: # noqa: D102, override + return re.sub( + r"([\d]+?)e([+-])([\d]+)", + r"\1\\mathord{\\times}10^{\2\3}", + f"{float(self.value):.6g}", + ) + + def as_array(self) -> xr.DataArray: # noqa: D102, override + return xr.DataArray(float(self.value), name=float(self.value)) + + +class ListParser(EvalArrayOrMath): + """For parsing lists.""" + + def __init__(self, instring: str, loc: int, tokens: pp.ParseResults) -> None: + """ + Process successfully parsed lists of generic strings. + + This is required since we call "eval()" on all elements of the where string, + so lists of strings need to be evaluatable as a whole "package". + + Args: + instring (str): String that was parsed (used in error message). + loc (int): + Location in parsed string where parsing error was logged. + This is not used; we include it as pyparsing injects it alongside `instring` when setting the parse action. + tokens (pp.ParseResults): a list of parsed string elements. + """ + self.val = tokens + self.instring = instring + + def __repr__(self) -> str: + """Programming / official string representation.""" + return f"{self.val}" + + def as_math_string(self) -> str: # noqa: D102, override + input_list = self.as_array() + return "[" + ",".join(str(i.name) for i in input_list) + "]" + + def as_array(self) -> list[xr.DataArray]: # noqa: D102, override + values = [val.eval("array", self.eval_attrs) for val in self.val] + # strings and numbers are returned as xarray arrays of size 1, + # so we extract those values. + return values + + +class EvalUnslicedComponent(EvalArrayOrMath): + """Evaluation of unsliced components.""" + + def __init__(self, instring: str, loc: int, tokens: pp.ParseResults) -> None: + """ + Parse action to process successfully parsed generic strings. + + This is required since we call "eval()" on all elements of the where string, + so even arbitrary strings (used in comparison operations) need to be evaluatable. + + Args: + instring (str): String that was parsed (used in error message). + loc (int): + Location in parsed string where parsing error was logged. + This is not used; we include it as pyparsing injects it alongside `instring` when setting the parse action. + tokens (pp.ParseResults): Has one parsed element: string name (str). + """ + self.val = tokens[0] + self.name = str(self.val) + self.values = tokens + self.instring = instring + + def __repr__(self) -> str: + """Programming / official string representation.""" + return f"COMPONENT:{self.name}" + + def as_math_string(self) -> str: # noqa: D102, override + self.eval_attrs = replace(self.eval_attrs, as_values=False) + evaluated = self.as_array() + self.eval_attrs.references.add(self.name) + + if "math_repr" in evaluated.attrs: + data_var_string = evaluated.attrs["math_repr"] + else: + data_var_string = self.name + + return data_var_string + + def as_array(self) -> xr.DataArray: # noqa: D102, override + if self.eval_attrs.math.find(self.name)._group in ["parameters", "lookups"]: + evaluated = self.eval_attrs.input_data[self.name] + elif self.eval_attrs.math.find(self.name)._group == "dimensions": + try: + evaluated = self.eval_attrs.input_data[self.name] + except KeyError: + evaluated = xr.DataArray(np.nan) + else: + evaluated = self.eval_attrs.model[self.name] + if evaluated.isnull().any() and pd.notna( + default := self.eval_attrs.math.find(self.name)["default"] + ): + evaluated = evaluated.fillna(default) + + self.eval_attrs.references.add(self.name) + return evaluated + + +class EvalGenericString(EvalArrayOrMath): + """For generic string parsing.""" + + def __init__(self, instring: str, loc: int, tokens: pp.ParseResults) -> None: + """ + Process successfully parsed generic strings. + + This is required since we call "eval()" on all elements of the where string, + so even arbitrary strings (used in comparison operations) need to be evaluatable. + + Args: + instring (str): String that was parsed (used in error message). + loc (int): + Location in parsed string where parsing error was logged. + This is not used; we include it as pyparsing injects it alongside `instring` when setting the parse action. + tokens (pp.ParseResults): Has one parsed element: string name (str). + """ + self.val = tokens[0] + self.instring = instring + + def __repr__(self) -> str: + """Programming / official string representation.""" + return f"STRING:{self.val}" + + def as_math_string(self): # noqa: D102, override + return str(self.val) + + def as_array(self) -> xr.DataArray: # noqa: D102, override + return xr.DataArray(str(self.val), name=str(self.val)) + + +def helper_function_parser( + *args: pp.ParserElement, + generic_identifier: pp.ParserElement, + allow_function_in_function: bool = False, +) -> pp.ParserElement: + """ + Process helper functions of the form `helper_function(*args, **kwargs)`. + + Helper functions can accept other parser elements as arguments, + i.e., components, parameters or variables, numbers, and other functions. + + Available helper functions are predefined in calliope.backend.helper_functions. + + Calling an unavailable helper will lead to a raised exception on evaluating the + parsed element. + + Based partially on: # https://stackoverflow.com/questions/61807705/pyparsing-generic-python-function-args-and-kwargs + + Args: + *args (pp.ParserElement): + Parser elements that can be arguments in the function (e.g., "number", "sliced_param_or_var"). + NOTE: the order of inclusion in the args list matters. The parser will parse based on first matches. + generic_identifier (pp.ParserElement): + Parser for valid python variables without leading underscore and not called "inf". + This parser has no parse action. + allow_function_in_function (bool, optional): + If True, allows functions to be defined inside functions. + Nested functions are evaluated from the greatest level of nesting up to the main helper function. + Defaults to True. + + Returns: + pp.ParserElement: + Parser for functions which will call the function with the specified + arguments on evaluation. + """ + helper_function = pp.Forward() + allowed_parser_elements_in_args = list(args) + lpar = pp.Suppress("(") + rpar = pp.Suppress(")") + + helper_function_name = generic_identifier.set_results_name("helper_function_name") + helper_function_name.set_parse_action(EvalHelperFuncName) + + if allow_function_in_function: + allowed_parser_elements_in_args.insert(0, helper_function) + + arg_values = pp.MatchFirst(allowed_parser_elements_in_args) + pp.NotAny("=") + + # define function arguments + arglist = pp.DelimitedList(arg_values.copy()) + args_ = pp.Group(arglist).set_results_name("args") + + # define function keyword arguments + key = generic_identifier + pp.Suppress("=") + kwarg_list = pp.DelimitedList(pp.dict_of(key, arg_values)) + kwargs_ = pp.Group(kwarg_list).set_results_name("kwargs") + + # build generic function + helper_func_args = args_ + pp.Suppress(",") + kwargs_ | pp.Opt( + args_, default=[] + ) + pp.Opt(kwargs_, default={}) + helper_function << ( + pp.Combine(helper_function_name + lpar) + helper_func_args + rpar + ) + + helper_function.set_parse_action(EvalFunction) + + return helper_function + + +def sliced_param_or_var_parser( + slicer: Iterable[pp.ParserElement], + generic_identifier: pp.ParserElement, + unsliced_object: pp.ParserElement, + allow_slice_references: bool = True, +) -> pp.ParserElement: + """ + Process strings representing sliced model parameters or variables. + + E.g. "source_use_max[node, tech]". + + If a parameter, must be a data variable in the Model.inputs xarray dataset. + + If a variable, must be an optimisation problem decision variable. + + The parser will not verify whether it has parsed a valid parameter or variable until + evaluation. + + Args: + slicer (Iterable[pp.ParserElement]): + List of parsers that can be used to define the slice of a parameter or variable. + E.g., "number", "evaluatable_identifier", "unsliced_param_or_var". + These elements will be parsed in the order they are given. + generic_identifier (pp.ParserElement): + Parser that evaluates to a string. + unsliced_object (pp.ParserElement): + Parser for valid backend objects. + On evaluation, this parser will access the backend object from the backend dataset. + allow_slice_references (bool): + If True, allow reference to `slice` expressions (e.g. `$bar` in `foo[bars=$bar]`). + Defaults to True. + + Returns: + pp.ParserElement: + Parser which returns a dictionary with name of parameter/variable and list + of index items as separate entries on. + """ + lspar = pp.Suppress("[") + rspar = pp.Suppress("]") + + direct_slicer = pp.MatchFirst(slicer) + if allow_slice_references: + slicer_ref = pp.Suppress(SUB_EXPRESSION_CLASSIFIER) + generic_identifier + slicer_ref.set_parse_action(EvalIndexSlice) + slicer = (slicer_ref | direct_slicer)("slicer") + else: + slicer = direct_slicer("slicer") + + slice = pp.Group(generic_identifier("set_name") + pp.Suppress("=") + slicer) + + slices = pp.Group(pp.DelimitedList(slice))("slices") + sliced_object_name = unsliced_object("param_or_var_name") + + sliced_param_or_var = pp.Combine(sliced_object_name + lspar) + slices + rspar + sliced_param_or_var.set_parse_action(EvalSlicedComponent) + + return sliced_param_or_var + + +def sub_expression_parser(generic_identifier: pp.ParserElement) -> pp.ParserElement: + """ + Parse strings prepended with the YAML constraint sub-expression classifier `$`. + + E.g. "$my_sub_expr" + + Args: + generic_identifier (pp.ParserElement): + Parser for valid python variables without leading underscore and not called "inf". + This parser has no parse action. + + Returns: + pp.ParserElement: + Parser which produces a dictionary of the form {"sub_expression": "my_sub_expression"} on evaluation. + """ + sub_expression = pp.Combine( + pp.Suppress(SUB_EXPRESSION_CLASSIFIER) + generic_identifier + ) + sub_expression.set_parse_action(EvalSubExpressions) + + return sub_expression + + +def unsliced_object_parser(valid_component_names: Iterable[str]) -> pp.ParserElement: + """ + Parse unsliced objects and identify their corresponding parse actions. + + Creates a copy of the generic identifier and sets a parse action to find the string in + the list of input parameters or optimisation decision variables. + + Args: + valid_component_names (Iterable[str]): A + All backend object names, to ensure they are captured by this parser function. + + Returns: + pp.ParserElement: + Copy of input parser with added parse action to lookup an unsliced + parameter/variable value + """ + unsliced_param_or_var = pp.one_of(valid_component_names, as_keyword=True) + unsliced_param_or_var.set_parse_action(EvalUnslicedComponent) + + return unsliced_param_or_var + + +def evaluatable_identifier_parser( + identifier: pp.ParserElement, valid_component_names: Iterable +) -> pp.ParserElement: + """ + Create an evaluatable copy of the generic identifier that will return a string or a model component as an array. + + Args: + identifier (pp.ParserElement): + Parser for valid python variables without leading underscore and not called "inf". + This parser has no parse action. + valid_component_names (Iterable[str]): A + All backend object names, to ensure they are *not* captured by this parser function. + + Returns: + pp.ParserElement: + Parser for valid python variables without leading underscore and not called "inf". + Evaluates to a string or an array (if it is a model component). + """ + evaluatable_identifier = ( + ~pp.one_of(valid_component_names, as_keyword=True) + identifier + ).set_parse_action(EvalGenericString) + + return evaluatable_identifier + + +def list_parser(*args: pp.ParserElement) -> pp.ParserElement: + """ + Parse strings which define a list of other strings or numbers. + + Lists are defined as anything wrapped in square brackets (`[]`). + + Args: + *args (pp.ParserElement): + Parser elements that can be list elements (e.g., "number", "evaluatable_identifier", "unsliced_param_or_var"). + These elements will be parsed in the order they are given. + + Returns: + pp.ParserElement: Parser for valid lists of strings and/or numbers. + """ + list_elements = pp.MatchFirst(args) + id_list = pp.Suppress("[") + pp.DelimitedList(list_elements) + pp.Suppress("]") + id_list.set_parse_action(ListParser) + return id_list + + +def setup_base_parser_elements() -> tuple[pp.ParserElement, pp.ParserElement]: + """ + Setup parser elements that will be components of other parsers. + + Returns: + tuple[pp.ParserElement, pp.ParserElement]: (number, generic_identifier) + number: parser for numbers (integer, float, scentific notation, "inf"/".inf"). + generic_identifier: parser for valid python variables without leading + underscore and not called "inf". This parser has no parse action. + """ + inf_kw = pp.Combine(pp.Opt(pp.Suppress(".")) + pp.Keyword("inf", caseless=True)) + number = pp.pyparsing_common.number | inf_kw + generic_identifier = ~inf_kw + pp.Word(pp.alphas, pp.alphanums + "_") + + number.set_parse_action(EvalNumber) + + return number, generic_identifier + + +def arithmetic_parser(*args, arithmetic: pp.Forward | None = None) -> pp.Forward: + """ + Parsing grammar to combine equation elements using basic arithmetic (+, -, *, /, **). + + Can handle the difference between a sign (e.g., -1,+1) and a addition/subtraction (0 - 1, 0 + 1). + Whitespace is ignored on parsing (i.e., "1+1+foo" is equivalent to "1 + 1 + foo"). + + Args: + *args: arguments in the form of a list. These can be: + helper_function (pp.ParserElement): parsing grammar to process helper functions + of the form `helper_function(*args, **kwargs)`. + sliced_param_or_var (pp.ParserElement): parser for sliced parameters or variables, e.g. "foo[bar]" + sub_expression (pp.ParserElement): parser for constraint sub expressions, e.g. "$foo" + unsliced_param_or_var (pp.ParserElement): parser for unsliced parameters or variables, e.g. "foo" + number (pp.ParserElement): parser for numbers (integer, float, scientific notation, "inf"/".inf"). + arithmetic (pp.Forward | None, optional): If given, add arithmetic rules to this + existing parsing rule (otherwise, arithmetic rules will be a newly generated rule). + Defaults to None. + + Returns: + pp.Forward: parser for strings which use arithmetic operations to combine other parser elements. + """ + signop = pp.one_of(["+", "-"]) + multop = pp.one_of(["*", "/"]) + expop = pp.Literal("**") + if arithmetic is None: + arithmetic = pp.Forward() + + arithmetic <<= pp.infixNotation( + # the order matters if two could capture the same string, e.g. "inf". + pp.MatchFirst(args), + [ + (signop, 1, pp.opAssoc.RIGHT, EvalSignOp), + (expop, 2, pp.opAssoc.LEFT, EvalOperatorOperand), + (multop, 2, pp.opAssoc.LEFT, EvalOperatorOperand), + (signop, 2, pp.opAssoc.LEFT, EvalOperatorOperand), + ], + ) + + return arithmetic + + +def equation_comparison_parser(arithmetic: pp.ParserElement) -> pp.ParserElement: + """ + Parsing grammar to combine equation elements either side of a comparison operator (<= >= ==). + + Whitespace is ignored on parsing (i.e., "1+foo==$bar" is equivalent to "1 + 1 == $bar"). + + Args: + arithmetic (pp.ParserElement): + Parser for arithmetic operations to combine other parser elements. + + Returns: + pp.ParserElement: + Parser for strings of the form "LHS OPERATOR RHS". + """ + comparison_operators = pp.one_of(["<=", ">=", "=="]) + equation_comparison = arithmetic + comparison_operators + arithmetic + equation_comparison.set_parse_action(EvalComparisonOp) + + return equation_comparison + + +def generate_slice_parser(valid_component_names: Iterable) -> pp.ParserElement: + """ + Create parser for index slice reference expressions. + + These expressions are linked to the equation expression by e.g. `$bar` in `foo[bars=$bar]`. + Unlike sub-expressions and equation expressions, these strings cannot contain arithmetic + nor references to sub expressions. + + Args: + valid_component_names (Iterable): + Allowed names for optimisation problem components (parameters, decision variables, expressions), + to allow the parser to separate these from generic strings. + + Returns: + pp.ParserElement: Parser for expression strings under the constraint key "slices". + """ + number, identifier = setup_base_parser_elements() + evaluatable_identifier = evaluatable_identifier_parser( + identifier, valid_component_names + ) + unsliced_param = unsliced_object_parser(valid_component_names) + helper_func_list = list_parser(number, evaluatable_identifier, unsliced_param) + slice_list = list_parser(number, evaluatable_identifier) + sliced_param = sliced_param_or_var_parser( + [number, evaluatable_identifier, slice_list], + identifier, + unsliced_param, + allow_slice_references=False, + ) + + helper_function = helper_function_parser( + sliced_param, + unsliced_param, + number, + helper_func_list, + evaluatable_identifier, + generic_identifier=identifier, + allow_function_in_function=True, + ) + + return ( + helper_function + | sliced_param + | unsliced_param + | number + | slice_list + | evaluatable_identifier + ) + + +def generate_sub_expression_parser(valid_component_names: Iterable) -> pp.Forward: + """ + Create parser for sub expressions. + + These expressions are linked to the equation expression by e.g. `$bar`. + This parser allows arbitrarily nested arithmetic and function calls (and arithmetic inside function calls) + and reference to index slice expressions. + + Args: + valid_component_names (Iterable): + Allowed names for optimisation problem components (parameters, decision variables, expressions), + to allow the parser to separate these from generic strings. + + Returns: + pp.ParserElement: Parser for expression strings under the constraint key "sub_expressions". + """ + number, identifier = setup_base_parser_elements() + evaluatable_identifier = evaluatable_identifier_parser( + identifier, valid_component_names + ) + unsliced_param = unsliced_object_parser(valid_component_names) + helper_func_list = list_parser(number, evaluatable_identifier, unsliced_param) + slice_list = list_parser(number, evaluatable_identifier) + sliced_param = sliced_param_or_var_parser( + [number, evaluatable_identifier, slice_list], identifier, unsliced_param + ) + + arithmetic = pp.Forward() + helper_function = helper_function_parser( + arithmetic, + helper_func_list, + evaluatable_identifier, + generic_identifier=identifier, + ) + arithmetic = arithmetic_parser( + helper_function, sliced_param, number, unsliced_param, arithmetic=arithmetic + ) + return arithmetic + + +def generate_arithmetic_parser(valid_component_names: Iterable) -> pp.ParserElement: + """ + Create parser for arithmetic expressions (+, -, /, *, **). + + This parser allows arbitrarily nested arithmetic and function calls (and arithmetic inside function calls) + and reference to sub-expressions and index slice expressions. + + Args: + valid_component_names (Iterable): + Allowed names for optimisation problem components (parameters, decision variables, global_expressions), + to allow the parser to separate these from generic strings. + + Returns: + pp.ParserElement: Partial parser for expression strings under the constraint key "equation/equations". + """ + number, identifier = setup_base_parser_elements() + evaluatable_identifier = evaluatable_identifier_parser( + identifier, valid_component_names + ) + unsliced_param = unsliced_object_parser(valid_component_names) + helper_func_list = list_parser(number, evaluatable_identifier, unsliced_param) + slice_list = list_parser(number, evaluatable_identifier) + sliced_param = sliced_param_or_var_parser( + [number, evaluatable_identifier, slice_list], identifier, unsliced_param + ) + sub_expression = sub_expression_parser(identifier) + + arithmetic = pp.Forward() + helper_function = helper_function_parser( + arithmetic, + helper_func_list, + evaluatable_identifier, + generic_identifier=identifier, + ) + arithmetic = arithmetic_parser( + helper_function, + sub_expression, + sliced_param, + number, + unsliced_param, + arithmetic=arithmetic, + ) + + return arithmetic + + +def generate_equation_parser(valid_component_names: Iterable) -> pp.ParserElement: + """ + Create parser for equation expressions of the form LHS OPERATOR RHS (e.g. `foo == 1 + bar`). + + This parser allows arbitrarily nested arithmetic and function calls (and arithmetic inside function calls) + and reference to sub-expressions and index slice expressions. + + Args: + valid_component_names (Iterable): + Allowed names for optimisation problem components (parameters, decision variables, global_expressions), + to allow the parser to separate these from generic strings. + + Returns: + pp.ParserElement: Parser for expression strings under the constraint key "equation/equations". + """ + arithmetic = generate_arithmetic_parser(valid_component_names) + equation_comparison = equation_comparison_parser(arithmetic) + + return equation_comparison diff --git a/linopy/declarative/helper_functions.py b/linopy/declarative/helper_functions.py new file mode 100644 index 000000000..e8f0c77e9 --- /dev/null +++ b/linopy/declarative/helper_functions.py @@ -0,0 +1,951 @@ +# Copyright (C) since 2013 Calliope contributors listed in AUTHORS. +# Licensed under the Apache 2.0 License (see LICENSE file). + +""" +Functions that can be used to process data in math `mask` and `expression` strings. + +`NAME` is the function name to use in the math strings. +""" + +import functools +import re +from abc import ABC, abstractmethod +from collections.abc import Mapping +from typing import Any, Literal, overload + +import numpy as np +import pandas as pd +import xarray as xr + +from linopy.declarative.eval_attrs import EvalAttrs + +DTYPE_OPTIONS = { + "string": str, + "float": float, + "bool": bool, + "datetime": np.datetime64, + "date": np.datetime64, + "integer": int, +} +_registry: dict[ + Literal["mask", "expression"], dict[str, type["ParsingHelperFunction"]] +] = {"mask": {}, "expression": {}} + + +class ParsingHelperFunction(ABC): + """Abstract base class for helper function parsing.""" + + def __init__( + self, return_type: Literal["array", "math_string"], attrs: "EvalAttrs" + ) -> None: + """ + Abstract helper function class, which all helper functions must subclass. + + The abstract properties and methods defined here must be defined by all helper functions. + """ + self._return_type = return_type + self._attrs = attrs + + @property + @abstractmethod + def ALLOWED_IN(self) -> list[Literal["mask", "expression"]]: + """List of parseable math strings that this function can be accessed from.""" + + @property + @abstractmethod + def NAME(self) -> str: + """Helper function name that is used in the math expression/mask string.""" + + @property + def ignore_mask(self) -> bool: + """If True, `mask` arrays will not be applied to the incoming data variables (valid for expression helpers).""" + return False + + @abstractmethod + def as_math_string(self, *args: Any, **kwargs: Any) -> str: + """ + Method to update LaTeX math strings to include the action applied by the helper function. + + This method is called when the class is initialised with ``return_type=math_string``. + """ + + @abstractmethod + def as_array(self, *args: Any, **kwargs: Any) -> xr.DataArray | Expression: + """ + Method to apply the helper function to provide an n-dimensional array output. + + This method is called when the class is initialised with ``return_type=array``. + """ + + def __call__(self, *args: Any, **kwargs: Any) -> Any: + """ + When a helper function is accessed by evaluating a parsing string, this method is called. + + The value of `return_type` on initialisation of the class defines whether this + method returns either: + - a string (``return_type=math_string``) + - :meth:xr.DataArray (``return_type=array``) + """ + if self._return_type == "math_string": + return self.as_math_string(*args, **kwargs) + elif self._return_type == "array": + return self.as_array(*args, **kwargs) + + def __init_subclass__(cls) -> None: + """ + Override subclass definition. + + 1. Do not allow new helper functions to have a name that is already defined (be it a built-in function or a custom function). + 2. Wrap helper function __call__ in a check for the function being allowed in specific parsing string types. + """ + super().__init_subclass__() + for allowed in cls.ALLOWED_IN: + if cls.NAME in _registry[allowed].keys(): + raise ValueError( + f"`{allowed}` string helper function `{cls.NAME}` already exists" + ) + for allowed in cls.ALLOWED_IN: + _registry[allowed][cls.NAME] = cls + + @staticmethod + def _update_iterator( + instring: str, + iterator_converter: dict[str, str], + method: Literal["add", "replace"], + ) -> str: + r""" + Utility function for generating latex strings in multiple helper functions. + + Find an iterator in the iterator substring of the component string + (anything wrapped in `_text{}`). Other parts of the iterator substring can be anything + except curly braces, e.g. the standalone `foo` will be found here and acted upon: + `\\textit{my_param}_\text{bar,foo,foo=bar,foo+1}` + + Args: + instring (str): String in which the iterator substring can be found. + iterator_converter (dict[str, str]): + key: the iterator to search for. + val: The new string to **append** to the iterator name (if method = add) or **replace**. + method (Literal[add, replace]): Whether to add to the iterator or replace it entirely + Returns: + str: `instring`, but with `iterator` replaced with `iterator + new_string` + """ + + def __replace_in_iterator(matched): + iterator_list = matched.group(2).split(",") + new_iterator_list = [] + for it in iterator_list: + if it in iterator_converter: + it = ( + it + iterator_converter[it] + if method == "add" + else iterator_converter[it] + ) + new_iterator_list.append(it) + + return matched.group(1) + ",".join(new_iterator_list) + matched.group(3) + + return re.sub(r"(_\\text{)([^{}]*?)(})", __replace_in_iterator, instring) + + def _get_dims_from_iterators(self, instring: str) -> list[str]: + """ + For a given math string describing a math component, extract the iterators and return the dimensions (a.k.a., sets) that they are members of. + + Args: + instring (str): string describing a math component. + + Returns: + list[str]: List of dimensions over which the math component is iterating. + + """ + + def __extract_dims(matched) -> str: + iterators = matched.group(2) + # Split on `,`, add 's' back in to singular iterators to refer to dimension names, + # then rejoin `,` as we must return a string. + return ",".join( + [ + dim_name + for i in iterators.split(",") + for dim_name, dim_math in self._attrs.math.dimensions.root.items() + if dim_math.iterator == i + ] + ) + + dims = re.sub(r"^.*(_\\text{)([^{}]*?)(})", __extract_dims, instring) + return dims.split(",") + + def _instr(self, dim: str) -> str: + """ + Utility function for generating latex strings in multiple helper functions. + + Args: + dim (str): Dimension suffixed with a "s" (e.g., "techs") + + Returns: + str: LaTeX string for iterator in a set (e.g., "tech in techs") + """ + iterator = self._dim_iterator(dim) + return rf"\text{{{iterator}}} \in \text{{{dim}}}" + + def _to_str_list( + self, vals: list[str | xr.DataArray] | str | xr.DataArray | list[xr.DataArray] + ) -> list[str]: + """ + Force a string to a list of length one if not already provided as a list. + + Args: + vals (list[str] | str): Values (or single value) to force to a list. + + Returns: + list[str]: Input forced to a list. + """ + if not isinstance(vals, list): + vals = [vals] + return [str(i.name) if isinstance(i, xr.DataArray) else i for i in vals] + + def _dim_iterator(self, dim: str) -> str: + return self._attrs.math.dimensions[dim].iterator + + +class MaskAny(ParsingHelperFunction): + """Apply `any` over a dimension in `mask` string.""" + + # Class name doesn't match NAME to avoid a clash with typing.Any + #: + NAME = "any" + #: + ALLOWED_IN = ["mask"] + + def as_math_string( # noqa: D102, override + self, array: str, *, over: str | list[str | xr.DataArray] + ) -> str: + over_list = self._to_str_list(over) + overstring = r" \\ ".join(self._instr(i) for i in over_list) + substack_overstring = rf"\substack{{{overstring}}}" + # Using bigvee for "collective-or" + return rf"\bigvee\limits_{{{substack_overstring}}} ({array})" + + def as_array( + self, input_component: xr.DataArray, *, over: xr.DataArray | list[xr.DataArray] + ) -> xr.DataArray: + """ + Reduce the boolean mask array of a model input by applying `any` over some dimension(s). + + If the component exists in the model, returns a boolean array with dimensions reduced + by applying a boolean OR operation along the dimensions given in `over`. + If the component does not exist, returns a dimensionless False array. + + Args: + input_component (str): Reference to a model input. + over (str | list[str]): dimension(s) over which to apply `any`. + + Returns: + xr.DataArray: resulting array. + """ + 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(self._to_str_list(over)) + + return input_component.any(dim=available_dims, keep_attrs=True) + + +class Defined(ParsingHelperFunction): + """Find all items of one dimension that are defined in an item of another dimension.""" + + #: + NAME = "defined" + #: + ALLOWED_IN = ["mask"] + + def as_math_string( # noqa: D102, override + self, *, within: xr.DataArray, how: Literal["all", "any"], **dims + ) -> str: + substrings = [] + for name, vals in dims.items(): + substrings.append(self._latex_substring(how, name, vals, str(within.name))) + if len(substrings) == 1: + return substrings[0] + else: + return rf"\bigwedge({', '.join(substrings)})" + + def as_array( + self, *, within: xr.DataArray, how: Literal["all", "any"], **dims: str + ) -> xr.DataArray: + """ + Find whether members of a model dimension are defined inside another. + + For instance, whether a node defines a specific tech (or group of techs). + Or, whether a tech defines a specific carrier. + + Args: + within (str): the model dimension to check. + how (Literal[all, any]): Whether to return True for `any` match of nested members or for `all` nested members. + **dims (str): + **key**: dimension whose members will be searched for as being defined under the primary dimension (`within`). + Must be one of the core model dimensions: [nodes, techs, carriers] + **value**: subset of the dimension members to find. + Transmission techs can be called using the base tech name (e.g., `ac_transmission`) and all link techs will be collected (e.g., [`ac_transmission:region1`, `ac_transmission:region2`]). + + + Returns: + xr.DataArray: + For each member of `within`, True if any/all member(s) in `dims` is nested within that member. + + Examples: + Check for any of a list of techs being defined at nodes. + Assuming a YAML definition of: + + ```yaml + nodes: + node1: + techs: + tech1: + tech3: + node2: + techs: + tech2: + tech3: + ``` + Then: + ``` + >>> defined(techs=[tech1, tech2], within=nodes, how=any) + [out] + array([ True, False]) + Coordinates: + * nodes (nodes) >> defined(techs=[tech1, tech2], within=nodes, how=all) + [out] + array([ False, False]) + Coordinates: + * nodes (nodes) set: + """ + From the definition matrix, get the dimensions that have not been defined. + + This includes dimensions not defined as keys of `dims` or as the value of `within`. + + Args: + dim_names (list[str]): Keys of `dims`. + within (str): dimension whose members are being checked. + + Raises: + ValueError: Can only define dimensions that exist in model.definition_matrix. + + Returns: + set: Undefined dimensions to remove from the definition matrix. + """ + definition_matrix = self._attrs.input_data.definition_matrix + missing_dims = set([*dim_names, within]).difference(definition_matrix.dims) + if missing_dims: + raise ValueError( + f"Unexpected model dimension referenced in `{self.NAME}` helper function. " + "Only dimensions given by `model.inputs.definition_matrix` can be used. " + f"Received: {missing_dims}" + ) + return set(definition_matrix.dims).difference([*dim_names, within]) + + def _latex_substring( + self, how: Literal["all", "any"], dim: str, vals: str | list[str], within: str + ) -> str: + if how == "all": + # Using wedge for "collective-and" + tex_how = "wedge" + elif how == "any": + # Using vee for "collective-or" + tex_how = "vee" + vals = self._to_str_list(vals) + within_iterator = self._dim_iterator(within) + dim_iterator = self._dim_iterator(dim) + selection = rf"\text{{{dim_iterator}}} \in \text{{[{','.join(vals)}]}}" + + return rf"\big{tex_how}\limits_{{\substack{{{selection}}}}}\text{{{dim_iterator} defined in {within_iterator}}}" + + +class Sum(ParsingHelperFunction): + """Apply a summation over dimension(s) in math expressions.""" + + NAME = "sum" + #: + ALLOWED_IN = ["expression", "mask"] + + def as_math_string( # noqa: D102, override + self, array: str, *, over: str | list[str | xr.DataArray] + ) -> str: + over_list = self._to_str_list(over) + overstring = r" \\ ".join(self._instr(i) for i in over_list) + substack_overstring = rf"\substack{{{overstring}}}" + return rf"\sum\limits_{{{substack_overstring}}} ({array})" + + def as_array( + self, array: xr.DataArray, *, over: xr.DataArray | list[xr.DataArray] + ) -> xr.DataArray: + """ + Sum an expression array over the given dimension(s). + + Args: + array (xr.DataArray): expression array + over (xr.DataArray | list[xr.DataArray]): + Dimension(s) over which to apply `sum`. + Array names will be extracted from the DataArray objects. + + Returns: + xr.DataArray: + Array with dimensions reduced by applying a summation over the dimensions given in `over`. + NaNs are ignored (xarray.DataArray.sum arg: `skipna: True`) and if all values along the dimension(s) are NaN, + the summation will lead to a NaN (xarray.DataArray.sum arg: `min_count=1`). + """ + filtered_over = set(self._to_str_list(over)).intersection(array.dims) + return array.sum(filtered_over) + + +class SelectFromLookupArrays(ParsingHelperFunction): + """N-dimensional indexing functionality.""" + + #: + NAME = "select_from_lookup_arrays" + #: + ALLOWED_IN = ["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() + } + array = self._update_iterator(array, new_strings, "add") + return array + + def as_array( + self, array: xr.DataArray, **lookup_arrays: xr.DataArray + ) -> xr.DataArray: + """ + Apply vectorised indexing on an arbitrary number of an input array's dimensions. + + Args: + array (xr.DataArray): Array on which to apply vectorised indexing. + **lookup_arrays (xr.DataArray): + key: dimension on which to apply vectorised indexing + value: array whose values are either NaN or values from the dimension given in the key. + + Raises: + ValueError: `array` must be indexed over the dimensions given in the `lookup_arrays` dict keys. + ValueError: All `lookup_arrays` must be indexed over all the dimensions given in the `lookup_arrays` dict keys. + + 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. + + Examples: + >>> coords = {"foo": ["A", "B", "C"]} + >>> array = xr.DataArray([1, 2, 3], coords=coords) + >>> lookup_array = xr.DataArray( + ... np.array(["B", "A", np.nan], dtype="O"), coords=coords, name="bar" + ... ) + >>> model_data = xr.Dataset({"bar": lookup_array}) + >>> select_from_lookup_arrays = SelectFromLookupArrays( + ... model_data=model_data + ... ) + >>> select_from_lookup_arrays(array, foo=lookup_array) + + array([ 2., 1., nan]) + Coordinates: + * foo (foo) object 'A' 'B' 'C' + + The lookup array assigns the value at "B" to "A" and vice versa. + "C" is masked since the lookup array value is NaN. + """ + # 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} 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}` 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._attrs.input_data[index.name].stack({dim: dims}) + ix = array.indexes[index_dim].get_indexer(stacked_lookup) + if (ix == -1).all(): + received_lookup = ( + self._attrs.input_data[index.name].to_series().dropna() + ) + raise IndexError( + f"Trying to select items on the dimension {index_dim} from the {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) + + # Create a mask to nullify any lookup values that are not given (i.e., are np.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 + unstacked_result = result.drop_vars(dims).unstack(dim) + return unstacked_result + + +class GetValAtIndex(ParsingHelperFunction): + """Getter functionality for obtaining values at specific integer indices.""" + + #: + NAME = "get_val_at_index" + #: + ALLOWED_IN = ["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_array(self, **dim_idx_mapping: int) -> xr.DataArray: + """ + Get value of a model dimension at a given integer index. + + This function is primarily useful for timeseries data. + + Args: + **dim_idx_mapping (int): kwargs with + key (str): Model dimension in which to extract value. + value (int): Integer index of the value to extract (assuming zero-indexing). + + Returns: + xr.DataArray: Dimensionless array containing one value. + + Examples: + >>> coords = { + ... "timesteps": [ + ... "2000-01-01 00:00", + ... "2000-01-01 01:00", + ... "2000-01-01 02:00", + ... ] + ... } + >>> model_data = xr.Dataset(coords=coords) + >>> get_val_at_index = GetValAtIndex(model_data=model_data) + >>> get_val_at_index(model_data)(timesteps=0) + + array('2000-01-01 00:00', dtype='>> get_val_at_index(model_data)(timesteps=-1) + + array('2000-01-01 00:00', dtype=' tuple[str, int]: ... + + # For as_math_string + @overload + @staticmethod + def _mapping_to_dim_idx(**dim_idx_mapping: str) -> tuple[str, str]: ... + + @staticmethod + def _mapping_to_dim_idx(**dim_idx_mapping) -> tuple[str, str | int]: + 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(ParsingHelperFunction): + """Roll (a.k.a. shift) items along ordered dimensions.""" + + #: + NAME = "roll" + #: + ALLOWED_IN = ["expression"] + + @property + def ignore_mask(self) -> bool: + """Whether or not to ignore `mask` functionality.""" + return 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() + } + component = self._update_iterator(array, new_strings, "add") + return component + + def as_array(self, array: xr.DataArray, **roll_kwargs: int) -> xr.DataArray: + """ + Roll (a.k.a., shift) 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. + + Args: + array (xr.DataArray): Array on which to roll data. + **roll_kwargs (int): kwargs with the following + key (str): name of dimension on which to roll. + value (int): number of places to roll data. + + Returns: + xr.DataArray: `array` with rolled data. + + Examples: + >>> array = xr.DataArray([1, 2, 3], coords={"foo": ["A", "B", "C"]}) + >>> model_data = xr.Dataset({"bar": array}) + >>> roll = Roll() + >>> roll("bar", foo=1) + + array([3, 1, 2]) + Coordinates: + * foo (foo) str: # noqa: D102, override + return rf"({array} \text{{if }} {condition} == True)" + + def as_array(self, array: xr.DataArray, condition: xr.DataArray) -> xr.DataArray: + """ + Apply a `mask` condition to a math array within an expression string. + + Args: + array (xr.DataArray): Math component array. + condition (xr.DataArray): + Boolean mask array. + If not `bool` type, NaNs and 0 will be assumed as False and all other values will be assumed as True. + + Returns: + xr.DataArray: + Returns the input array with the condition applied, + including having been broadcast across any new dimensions provided by the condition. + + Examples: + One common use-case is to introduce a new dimension to the variable which represents subsets of one of the main model dimensions. + In this case, each member of `cap_node_groups` is a subset of `nodes` and we want to sum `flow_cap` over each of those subsets and set a maximum value. + + input: + ```yaml + data_definitions: + node_grouping: + data: True + index: [[group_1, region1], [group_1, region1_1], [group_2, region1_2], [group_2, region1_3], [group_3, region2]] + dims: [cap_node_groups, nodes] + node_group_max: + data: [1, 2, 3] + index: [group_1, group_2, group_3] + dims: cap_node_groups + ``` + + math: + ```yaml + constraints: + my_new_constraint: + foreach: [techs, cap_node_groups] + equations: + - expression: sum(mask(flow_cap, node_grouping), over=nodes) <= node_group_max + ``` + """ + return array.where(condition.fillna(False).astype(bool)) + + +class GroupSum(ParsingHelperFunction): + """Apply a summation over an array grouping.""" + + #: + NAME = "group_sum" + #: + ALLOWED_IN = ["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]) + overstring = rf"\substack{{{foreach_string}}}" + return rf"\sum\limits_{{{overstring}}} ({array})" + + def as_array( + self, array: xr.DataArray, groupby: xr.DataArray, group_dim: xr.DataArray + ) -> xr.DataArray: + """ + Sum an array over the given groupings. + + Args: + array (xr.DataArray): expression array + groupby (xr.DataArray): Array with which to group the array. + group_dim (str): Name of dimension that the `groupby` values are members of. + This will become a new dimension over which the array is indexed once grouping is complete. + + Returns: + xr.DataArray: + Array with dimension(s) aggregated over the `groupby`. + + Note: + - The array is returned with all dimensions over which `groupby` is indexed replaced by a new dimension named by `group_dim`. + - To groupby datetime periods (weeks, months, dates, etc.), consider using `group_datetime` for convenience, as you do not need to define a separate `groupby` array. + + Examples: + To get the sum over an ad-hoc combination of techs at nodes, e.g. to limit their overall outflow in any given timestep, you would do the following: + + 1. Define an array linking node-tech combinations with a group: + ```yaml + data_definitions: + # You may prefer to define this in a CSV file or when referring to the techs within the `nodes` model definition. + power_plant_groups: + data: [low_emission_plant, low_emission_plant, high_emission_plant, high_emission_plant] + index: [ + [tech_1, node_1], + [tech_2, node_1], + [tech_1, node_2], + [tech_2, node_2], + ] + dims: [techs, nodes] + ``` + 2. Define a set of outflow limits: + ```yaml + data_definitions: + emission_limits: + data: [20, 10] + index: [low_emission_plant, high_emission_plant] + dims: [emission_groups] + ``` + 3. Define the math to link the two, using `group_sum`: + ```yaml + constraints: + node_tech_emission_group_max: + foreach: [emission_groups, carriers, timesteps] + mask: emission_limits + equations: + - expression: group_sum(flow_out, power_plant_groups, emission_groups) <= emission_limits + ``` + """ + # We can't apply typical xarray rolling window functionality + grouped: dict[str | int, xr.DataArray] = {} + + grouping_dims = groupby.dims + groups = array.stack(_stacked=grouping_dims).groupby( + groupby.stack(_stacked=grouping_dims) + ) + for group_name, array_subset in groups: + grouped[group_name] = array_subset.sum("_stacked", min_count=1, skipna=True) + + array = xr.concat( + grouped.values(), dim=pd.Index(grouped.keys(), name=group_dim.name) + ) + return array + + +class GroupDatetime(ParsingHelperFunction): + """Apply a summation over a datetime group on a datetime dimension in math expressions.""" + + NAME = "group_datetime" + #: + ALLOWED_IN = ["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}}}(\text{{{self._dim_iterator(over)}}}) = \text{{{self._dim_iterator(group)}}}" + overstring = rf"\substack{{{foreach_string}}}" + + return rf"\sum\limits_{{{overstring}}} ({array})" + + def as_array( + self, array: xr.DataArray, over: xr.DataArray, group: xr.DataArray + ) -> xr.DataArray: + """ + Sum an expression array over the given dimension(s). + + Args: + array (xr.DataArray): expression array + over (xr.DataArray): dimension name over which to group + group (xr.DataArray): datetime grouper. + Any xarray/pandas datetime grouper options + datetime grouper options include 'date', 'dayofweek', 'month', etc. + + + Returns: + xr.DataArray: + Array with datetime dimension aggregated over the grouper. + + Note: + - The array is returned with the `over` dimension replaced by the name of the grouper. + So, if you select to resample to monthly, the returned array will include the `month` dimension. + - the `date`/`time` groupers will return the date/time as a string in ISO8601 format (e.g. "2025-01-01"/"01:00:00"). + All other groupers will return integer values (e.g. month 1, 2, 3, etc.). + + Examples: + One common use-case is to allow demand to be met at any point on a given date. + For such a demand tech, the daily demand should be indexed over `date`, e.g.: + + sink_use_equals_daily.csv + ``` + date,sink_use_equals_daily + 2000-01-01,10 + 2000-01-02,15 + ... + ``` + + Then, to set the daily flow into the demand tech to those values: + ```yaml + constraints: + daily_demand: + foreach: [nodes, techs, carriers, date] + mask: sink_use_equals_daily + equations: + - expression: "group_datetime(flow_in, timesteps, date) == sink_use_equals_daily" + ``` + + Similarly, a monthly maximum resource to a supply technology might be used, to simulate e.g. biofuel feedstock availability: + + source_use_max_monthly.csv + ``` + month,source_use_max_monthly + 1,10 + 2,15 + ... + ``` + + Then, to set the daily flow into the demand tech to those values: + ```yaml + constraints: + daily_demand: + foreach: [nodes, techs, carriers, month] + mask: source_use_max_monthly + equations: + - expression: "group_datetime(flow_in, timesteps, month) <= source_use_max_monthly" + ``` + """ + dtype = DTYPE_OPTIONS[self._attrs.math.dimensions[group.name].dtype] + group_sum_helper = GroupSum(self._return_type, self._attrs) + array = group_sum_helper( + array, getattr(array[over.name].dt, group.name).astype(dtype), group + ) + + return array + + +class SumNextN(ParsingHelperFunction): + """ + Sum the 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 = ["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 = self._update_iterator( + array, {self._dim_iterator(over): new_iterator}, "replace" + ) + + return rf"\sum\limits_{{\text{{{new_iterator}}}={over_singular}}}^{{{over_singular}+{N}}} ({updated_iterator_array})" + + def as_array(self, array: xr.DataArray, over: xr.DataArray, N: int) -> xr.DataArray: + """ + Sum values from current up to N from current on the dimension `over`. + + Works best for ordered arrays (datetime, integer). + + + Args: + array (xr.DataArray): Math component array. + over (str): Dimension over which to sum + N (int): number of items beyond the current value to sum from + + Returns: + xr.DataArray: + Returns the input array with the condition applied, + including having been broadcast across any new dimensions provided by the condition. + + Note: + - The rolling window does not wrap around to the start of the set when reaching the end. + That is, if you have N = 4 then for a dimension of length T, at T - 1 it will sum over dimension positions (T - 1, T), not (T - 1, T, 0, 1). + - You will find that this over-constrains the model unless you limit the constraint (using the `mask` string) to only apply over `len(over) - N`. + This is linked to the abovementioned lack of wrapping. + E.g. `mask: timesteps<=get_val_at_index(timesteps=-24)` if N == 24. + - This function is based on an integer number of steps from the current step. + For datetime dimensions like `timesteps`, you will (a) need to be using a regular time frequency (e.g. hourly) and (b) update `N` to reflect the resolution of your time dimension + (N = 4 in if resample.timesteps=`1h` -> N = 2 if resample.timesteps=`2h`). + + Examples: + One common use-case is to collate N timesteps beyond a given timestep to apply a constraint to it + (e.g., demand must be less than X in the next 24 hours): + + For such a demand tech, the portion of its demand that is flexible should be separated from `sink_use_equals` to e.g., + a `sink_use_flexible` timeseries parameter which we will use in the DSR constraint: + + ```yaml + constraints: + 4hr_demand_side_response: + foreach: ["nodes", "techs", "carriers", "timesteps"] + mask: "carrier_in AND sink_use_flexible AND timesteps<=get_val_at_index(timesteps=-24)" + equations: + - expression: sum_next_n(flow_in, timesteps, 4) == sum_next_n(sink_use_flexible, timesteps, 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 + ) + ) + final_array = xr.concat(results, dim=over).broadcast_like(array) + return final_array diff --git a/linopy/declarative/mask_parser.py b/linopy/declarative/mask_parser.py new file mode 100644 index 000000000..623b67ca0 --- /dev/null +++ b/linopy/declarative/mask_parser.py @@ -0,0 +1,617 @@ +# Copyright (C) since 2013 Calliope contributors listed in AUTHORS. +# Licensed under the Apache 2.0 License (see LICENSE file). +"""Parsing for 'mask' statements.""" + +from __future__ import annotations + +import operator +from collections.abc import Iterable +from dataclasses import replace +from typing import Any + +import numpy as np +import pandas as pd +import pyparsing as pp +import xarray as xr + +from linopy.declarative import expression_parser + +pp.ParserElement.enable_packrat() +BOOLEANTYPE = np.bool_ | np.typing.NDArray[np.bool_] + + +def get_dot_attr(var: Any, attr: str) -> Any: + """ + Get nested attributes in dot notation. + + Works for nested objects (e.g., dictionaries, pydantic models). + + Args: + var (Any): Object to extract nested attributes from. + attr (str): Name of the attribute (e.g., "foo.bar"). + + Returns: + Any: Value at the given location. + """ + levels = attr.split(".", 1) + + if isinstance(var, dict): + value = var[levels[0]] + else: + value = getattr(var, levels[0]) + + if len(levels) > 1: + value = get_dot_attr(value, levels[1]) + return value + + +class EvalNot(expression_parser.EvalSignOp, expression_parser.EvalArrayOrMath): + """Parse action to process successfully parsed expressions with a leading `not`.""" + + def as_math_string(self) -> str: # noqa: D102, override + evaluated = self.value.eval("math_string", self.eval_attrs) + return rf"\neg ({evaluated})" + + def as_array(self) -> xr.DataArray: # noqa: D102, override + evaluated = self.value.eval("array", self.eval_attrs) + return ~evaluated + + +class EvalAndOr( + expression_parser.EvalOperatorOperand, expression_parser.EvalArrayOrMath +): + """ + Processing of successfully parsed expressions with and/or operators. + + E.g., "OPERAND OPERATOR OPERAND OPERATOR OPERAND ..." + """ + + LATEX_OPERATOR_LOOKUP: dict[str, str] = { + "and": r"{val} \land {operand}", + "or": r"{val} \lor {operand}", + } + SKIP_IF = ["and", "or"] + + def _skip_component_on_conditional(self, component: str, operator_: str) -> bool: + return component == "true" and operator_ in self.SKIP_IF + + @staticmethod + def _operate( + val: xr.DataArray, evaluated_operand: xr.DataArray, operator_: str + ) -> xr.DataArray: + """Apply bitwise comparison between boolean xarray dataarrays.""" + match operator_: + case "and": + val = operator.and_(val, evaluated_operand) + case "or": + val = operator.or_(val, evaluated_operand) + return val + + def _apply_mask(self, evaluated: xr.DataArray) -> xr.DataArray: + """Override func from parent class to effectively do nothing.""" + return evaluated + + def as_math_string(self) -> str: # noqa: D102, override + return super().as_math_string() + + def as_array(self) -> xr.DataArray: # noqa: D102, override + return super().as_array() + + +class ConfigOptionParser(expression_parser.EvalArrayOrMath): + """Parsing of configuration options.""" + + def __init__(self, instring: str, loc: int, tokens: pp.ParseResults) -> None: + """ + Parse action to process successfully parsed configuration option names. + + Args: + instring (str): String that was parsed (used in error message). + loc (int): + Location in parsed string mask parsing error was logged. + This is not used; we include it as pyparsing injects it alongside `instring` when setting the parse action. + tokens (pp.ParseResults): + Has two parsed elements: config group name (str) and config option (str). + """ + self.config_option = tokens[0] + self.instring = instring + self.loc = loc + + def __repr__(self): + """Programming / official string representation.""" + return f"CONFIG:{self.config_option}" + + def as_math_string(self) -> str: # noqa: D102, override + return rf"\text{{config.{self.config_option}}}" + + def as_array(self) -> xr.DataArray: # noqa: D102, override + config_val = get_dot_attr(self.eval_attrs.config, self.config_option) + + if not isinstance(config_val, int | float | str | bool | np.bool_): + raise self.error_msg( + f"mask string | Configuration option resolves to invalid " + f"type `{type(config_val).__name__}`, expected a number, string, or boolean." + ) + else: + return xr.DataArray(config_val) + + +class ResultArrayParser(expression_parser.EvalArrayOrMath): + """Variable/Expression array processing.""" + + def __init__(self, instring: str, loc: int, tokens: pp.ParseResults) -> None: + """ + Parse action to process successfully parsed model variable/global expression names. + + Args: + instring (str): String that was parsed (used in error message). + loc (int): + Location in parsed string mask parsing error was logged. + This is not used; we include it as pyparsing injects it alongside `instring` when setting the parse action. + tokens (pp.ParseResults): + Has one parsed element: model data variable name (str). + """ + self.array_name = tokens[0] + self.instring = instring + self.loc = loc + + def __repr__(self): + """Programming / official string representation.""" + return f"RESULT:{self.array_name}" + + def as_math_string(self) -> str: # noqa: D102, override + self.eval_attrs.references.add(self.array_name) + math_repr = self.eval_attrs.model[self.array_name].attrs.get( + "math_repr", rf"\exists (\textbf{{{self.array_name}}})" + ) + + return math_repr + + def as_array(self) -> xr.DataArray: # noqa: D102, override + self.eval_attrs.references.add(self.array_name) + da = self.eval_attrs.model[self.array_name] + if self.eval_attrs.apply_mask: + da = ~da.isnull() + return da + + +class InputArrayParser(expression_parser.EvalArrayOrMath): + """Input array processing.""" + + def __init__(self, instring: str, loc: int, tokens: pp.ParseResults) -> None: + """ + Parse action to process successfully parsed model input array names. + + Args: + instring (str): String that was parsed (used in error message). + loc (int): + Location in parsed string mask parsing error was logged. + This is not used; we include it as pyparsing injects it alongside `instring` when setting the parse action. + tokens (pp.ParseResults): + Has one parsed element: model data variable name (str). + """ + self.array_name = tokens[0] + self.instring = instring + self.loc = loc + + def __repr__(self): + """Programming / official string representation.""" + return f"INPUT:{self.array_name}" + + def as_math_string(self) -> str: # noqa: D102, override + self.eval_attrs.references.add(self.array_name) + + math_repr = self.eval_attrs.input_data[self.array_name].attrs.get( + "math_repr", rf"\textit{{{self.array_name}}}" + ) + if self.eval_attrs.apply_mask: + math_repr = rf"\exists ({math_repr})" + return math_repr + + def as_array(self) -> xr.DataArray: # noqa: D102, override + self.eval_attrs.references.add(self.array_name) + da = self.eval_attrs.input_data.get(self.array_name, xr.DataArray(False)) + if self.eval_attrs.apply_mask and da.dtype.kind != "b": + da = da.notnull() & (da != np.inf) & (da != -np.inf) + elif da.isnull().any() and pd.notnull( + default := self.eval_attrs.math.find(self.array_name).default + ): + da = da.fillna(default) + return da + + +class DimensionArrayParser(expression_parser.EvalArrayOrMath): + """Dimension array processing.""" + + def __init__(self, instring: str, loc: int, tokens: pp.ParseResults) -> None: + """ + Parse action to process successfully parsed model dimension names. + + Args: + instring (str): String that was parsed (used in error message). + loc (int): + Location in parsed string mask parsing error was logged. + This is not used; we include it as pyparsing injects it alongside `instring` when setting the parse action. + tokens (pp.ParseResults): + Has one parsed element: model data variable name (str). + """ + self.array_name = tokens[0] + self.instring = instring + self.loc = loc + + def __repr__(self): + """Programming / official string representation.""" + return f"DIM:{self.array_name}" + + def as_math_string(self) -> str: # noqa: D102, override + return self.array_name + + def as_array(self) -> xr.DataArray: # noqa: D102, override + # We want the mask string to evaluate successfully even if a dimension hasn't been defined. + da = self.eval_attrs.input_data.get(self.array_name, xr.DataArray()) + return da + + +class ComparisonParser( + expression_parser.EvalComparisonOp, expression_parser.EvalArrayOrMath +): + """Parse action to process successfully parsed strings of the form x=y.""" + + OP_TRANSLATOR = { + "<=": r"\mathord{\leq}", + ">=": r"\mathord{\geq}", + "==": r"\mathord{==}", + "<": r"\mathord{<}", + ">": r"\mathord{>}", + } + + def __repr__(self): + """Return string representation of the parsed grammar.""" + return f"{self.lhs}{self.op}{self.rhs}" + + def as_math_string(self) -> str: # noqa: D102, override + self.eval_attrs = replace(self.eval_attrs, apply_mask=False) + lhs, rhs = self._eval("math_string") + if r"\text" not in rhs: + rhs = rf"\text{{{rhs}}}" + return lhs + self.OP_TRANSLATOR[self.op] + rhs + + def as_array(self) -> xr.DataArray: # noqa: D102, override + self.eval_attrs = replace(self.eval_attrs, apply_mask=False) + lhs, rhs = self._eval("array") + match self.op: + case "<=": + comparison = lhs <= rhs + case ">=": + comparison = lhs >= rhs + case "<": + comparison = lhs < rhs + case ">": + comparison = lhs > rhs + case "==": + comparison = lhs == rhs + return xr.DataArray(comparison) + + +class SubsetParser(expression_parser.EvalArrayOrMath): + """Dimension subset parsing.""" + + def __init__(self, instring: str, loc: int, tokens: pp.ParseResults) -> None: + """ + Parse action to process successfully parsed dimension subsetting. + + Args: + instring (str): String that was parsed (used in error message). + loc (int): + Location in parsed string mask parsing error was logged. + This is not used; we include it as pyparsing injects it alongside `instring` when setting the parse action. + tokens (pp.ParseResults): + Has two parsed elements: model set name (str), set items (Any). + """ + self.val, self.set_name = tokens + self.instring = instring + self.loc = loc + + def __repr__(self): + """Return string representation of the parsed grammar.""" + return f"SUBSET:{self.set_name}{self.val}" + + def _eval(self) -> list[str | float]: + """Evaluate each element of the subset list.""" + values = [val.eval("array", self.eval_attrs) for val in self.val] + return [val.item() if isinstance(val, xr.DataArray) else val for val in values] + + def as_math_string(self) -> str: # noqa: D102, override + subset = self._eval() + dim = self.set_name.eval("math_string", self.eval_attrs) + iterator = self.eval_attrs.math.dimensions[dim].iterator + subset_string = "[" + ",".join(str(i) for i in subset) + "]" + return rf"\text{{{iterator}}} \in \text{{{subset_string}}}" + + def as_array(self) -> xr.DataArray: # noqa: D102, override + subset = self._eval() + set_item_in_subset = self.set_name.eval("array", self.eval_attrs).isin(subset) + return set_item_in_subset + + +class BoolOperandParser(expression_parser.EvalArrayOrMath): + """Boolean operand parsing.""" + + def __init__(self, instring: str, loc: int, tokens: pp.ParseResults) -> None: + """ + Parse action to process successfully parsed boolean strings. + + Args: + instring (str): String that was parsed (used in error message). + loc (int): + Location in parsed string mask parsing error was logged. + This is not used; we include it as pyparsing injects it alongside `instring` when setting the parse action. + tokens (pp.ParseResults): Has one parsed element: boolean (str). + """ + self.val = tokens[0].lower() + self.instring = instring + + def __repr__(self): + """Programming / official string representation.""" + return f"BOOL:{self.val}" + + def as_math_string(self): # noqa: D102, override + return self.val + + def as_array(self) -> xr.DataArray: # noqa: D102, override + if self.val == "true": + bool_val = xr.DataArray(np.True_) + elif self.val == "false": + bool_val = xr.DataArray(np.False_) + return bool_val + + +class GenericStringParser(expression_parser.EvalString): + """Parsing of generic strings.""" + + def __init__(self, instring: str, loc: int, tokens: pp.ParseResults) -> None: + """ + Parse action to process successfully parsed generic strings. + + This is required since we call "eval()" on all elements of the mask string, + so even arbitrary strings (used in comparison operations) need to be evaluatable. + + Args: + instring (str): String that was parsed (used in error message). + loc (int): + Location in parsed string mask parsing error was logged. + This is not used; we include it as pyparsing injects it alongside `instring` when setting the parse action. + tokens (pp.ParseResults): Has one parsed element: string name (str). + """ + self.val = tokens[0] + self.instring = instring + + def __repr__(self) -> str: + """Return string representation of the parsed grammar.""" + return f"STRING:{self.val}" + + def eval(self, *args, **kwargs) -> str: + """Evaluation just returns the string of values.""" + return str(self.val) + + +def data_var_parser( + names: Iterable, parse_action: type[expression_parser.EvalArrayOrMath] +) -> pp.ParserElement: + """ + Process model data variables which can be any valid python identifier (string + "_"). + + Args: + names (Iterable): List of valid component names. + parse_action (type[expression_parser.EvalArrayOrMath]): Parse action to evaluate the parsed string. + + Returns: + pp.ParserElement: parser for model data variables which will access the data + variable from the Calliope model dataset. + """ + data_var = pp.one_of(names, as_keyword=True) + data_var.set_parse_action(parse_action) + + return data_var + + +def config_option_parser(generic_identifier: pp.ParserElement) -> pp.ParserElement: + """ + Parsing grammar to process model configuration option key names of the form "x.y.z". + + Args: + generic_identifier (pp.ParserElement): + Parser for valid python variables without leading underscore and not called "inf". + This parser has no parse action. + + Returns: + pp.ParserElement: + Parser for configuration options which will be accessed from the configuration + dictionary attached to the attributes of the Calliope model dataset. + """ + data_var = pp.Suppress("config.") + generic_identifier + data_var.set_parse_action(ConfigOptionParser) + + return data_var + + +def bool_parser() -> pp.ParserElement: + """Parsing grammar for True/False (any case), which will evaluate to np.bool_.""" + TRUE = pp.Keyword("True", caseless=True) + FALSE = pp.Keyword("False", caseless=True) + bool_operand = TRUE | FALSE + bool_operand.set_parse_action(BoolOperandParser) + + return bool_operand + + +def evaluatable_string_parser( + generic_identifier: pp.ParserElement, valid_components: Iterable +) -> pp.ParserElement: + """Parsing grammar to make generic strings used in comparison operations evaluatable.""" + evaluatable_identifier = ( + ~pp.one_of(valid_components, as_keyword=True) + generic_identifier + ) + evaluatable_identifier.set_parse_action(GenericStringParser) + + return evaluatable_identifier + + +def comparison_parser( + lhs: list[pp.ParserElement], rhs: list[pp.ParserElement] +) -> pp.ParserElement: + """ + Parsing grammar to process comparisons of the form `variable_or_config=comparator`. + + Args: + lhs (list[pp.ParserElement]): + Parsers that can be included on the left-hand side of the comparison; will be matched in the order provided. + rhs (list[pp.ParserElement]): + Parsers that can be included on the right-hand side of the comparison; will be matched in the order provided. + + Returns: + pp.ParserElement: + Parser which will return a bool/boolean array as a result of the comparison. + """ + comparison_operators = pp.oneOf(["<", ">", "==", ">=", "<="]) + comparison_expression = ( + pp.MatchFirst(lhs) + comparison_operators + pp.MatchFirst(rhs) + ) + comparison_expression.set_parse_action(ComparisonParser) + + return comparison_expression + + +def subset_parser( + data_var: pp.ParserElement, *subset_items: pp.ParserElement +) -> pp.ParserElement: + """ + Parsing grammar to process subsets. + + Args: + data_var (pp.ParserElement): data variable parser + *subset_items (pp.ParserElement): parsers that can be included in the subset list; will be matched in the order provided. + + Returns: + pp.ParserElement: subset parser. + """ + subset = pp.Group(pp.delimited_list(pp.MatchFirst(subset_items))) + subset_expression = ( + pp.Suppress("[") + + subset + + pp.Suppress("]") + + pp.Suppress(pp.White(" ", min=1)) + + pp.Suppress("in") + + pp.Suppress(pp.White(" ", min=1)) + + data_var + ) + subset_expression.set_parse_action(SubsetParser) + + return subset_expression + + +def mask_parser(*args: pp.ParserElement) -> pp.ParserElement: + """ + Parser for strings which use AND/OR/NOT operators to combine other parser elements. + + Args: + *args (pp.ParserElement): + parsers that can be included in the mask string; will be matched in the order provided. + + Returns: + pp.ParserElement: mask parser. + """ + notop = pp.Keyword("not", caseless=True) + andorop = pp.Keyword("and", caseless=True) | pp.Keyword("or", caseless=True) + + mask_rules = pp.infixNotation( + pp.MatchFirst(args), + [ + (notop, 1, pp.opAssoc.RIGHT, EvalNot), + (andorop, 2, pp.opAssoc.LEFT, EvalAndOr), + ], + ) + + return mask_rules + + +def generate_mask_string_parser( + dimensions: Iterable, + inputs: Iterable, + results: Iterable, + postprocessed: Iterable | None = None, +) -> pp.ParserElement: + """ + Creates and executes the mask parser. + + Args: + dimensions (Iterable): List of valid dimension names. + inputs (Iterable): List of valid input names. + results (Iterable): List of valid variable/global expression names. + postprocessed (Iterable | None): List of valid postprocessed expression names. Defaults to None. + + Returns: + pp.ParseResults: evaluatable to a bool/boolean array. + """ + postprocessed = postprocessed if postprocessed is not None else set() + number, generic_identifier = expression_parser.setup_base_parser_elements() + dimensions_parser = data_var_parser(dimensions, DimensionArrayParser) + inputs_parser = data_var_parser(inputs, InputArrayParser) + results_parser = data_var_parser(results | postprocessed, ResultArrayParser) + config_option = config_option_parser(generic_identifier) + bool_operand = bool_parser() + unique_evaluatable_string = evaluatable_string_parser( + generic_identifier, set().union(dimensions, inputs, results, postprocessed) + ) + general_evaluatable_string = evaluatable_string_parser(generic_identifier, []) + id_list = expression_parser.list_parser( + number, unique_evaluatable_string, dimensions_parser + ) + subset = subset_parser( + dimensions_parser, config_option, number, general_evaluatable_string + ) + + arithmetic = pp.Forward() + comparison_helper_function = expression_parser.helper_function_parser( + unique_evaluatable_string, + number, + id_list, + arithmetic, + generic_identifier=generic_identifier, + ) + arithmetic_elements = [ + comparison_helper_function, + number, + dimensions_parser, + inputs_parser, + config_option, + ] + if postprocessed: + arithmetic_elements.insert(-2, results_parser) + + comparison_arithmetic = expression_parser.arithmetic_parser( + *arithmetic_elements, arithmetic=arithmetic + ) + comparison = comparison_parser( + lhs=[comparison_arithmetic], + rhs=[ + comparison_helper_function, + bool_operand, + number, + general_evaluatable_string, + ], + ) + + helper_function = expression_parser.helper_function_parser( + unique_evaluatable_string, + number, + id_list, + dimensions_parser, + inputs_parser, + results_parser, + config_option, + generic_identifier=generic_identifier, + ) + return mask_parser( + bool_operand, comparison, helper_function, subset, inputs_parser, results_parser + ) diff --git a/linopy/declarative/parsing.py b/linopy/declarative/parsing.py new file mode 100644 index 000000000..e1c7760e7 --- /dev/null +++ b/linopy/declarative/parsing.py @@ -0,0 +1,830 @@ +# Copyright (C) since 2013 Calliope contributors listed in AUTHORS. +# Licensed under the Apache 2.0 License (see LICENSE file). +"""Methods for math syntax parsing.""" + +from __future__ import annotations + +import functools +import itertools +import logging +import operator +from collections.abc import Callable, Iterable +from typing import TYPE_CHECKING, Literal, overload + +import pyparsing as pp +import xarray as xr + +from linopy.declarative import ( + eval_attrs, + expression_parser, + helper_functions, + mask_parser, +) +from linopy.declarative.schema import MATH_DEFS_T, ConfigModel, MathModel, _Equations + +if TYPE_CHECKING: + from linopy.model import Model +TRUE_ARRAY = xr.DataArray(True) + +LOGGER = logging.getLogger(__name__) + + +class ParsedBackendEquation: + """Backend equation parser.""" + + def __init__( + self, + equation_name: str, + sets: list[str], + expression: pp.ParseResults, + mask_list: list[pp.ParseResults], + sub_expressions: dict[str, pp.ParseResults] | None = None, + slices: dict[str, pp.ParseResults] | None = None, + ) -> None: + """ + For parsing equation expressions and corresponding "mask" strings. + + Args: + equation_name (str): Name of equation. + sets (list[str]): + Model data sets with which to create the initial multi-dimensional masking array + of the evaluated "mask" string. + expression (pp.ParseResults): + Parsed arithmetic/equation expression. + mask_list (list[pp.ParseResults]): + List of parsed mask strings. + sub_expressions (dict[str, pp.ParseResults] | None, optional): + Dictionary of parsed sub-expressions with which to replace sub-expression references + on evaluation of the parsed expression. Defaults to None. + slices (dict[str, pp.ParseResults] | None, optional): + Dictionary of parsed array slices with which to replace slice references + on evaluation of the parsed expression / sub-expression. Defaults to None. + """ + self.name = equation_name + self.mask = mask_list + self.expression = expression + self.sub_expressions = ( + sub_expressions if sub_expressions is not None else dict() + ) + self.slices = slices if slices is not None else dict() + self.sets = sets + + def find_sub_expressions(self) -> set[str]: + """ + Identify all the references to sub_expressions in the parsed expression. + + Returns: + set[str]: Unique sub-expression references. + """ + valid_eval_classes: tuple = ( + expression_parser.EvalOperatorOperand, + expression_parser.EvalFunction, + ) + to_find = expression_parser.EvalSubExpressions + elements: list + if isinstance(self.expression[0], to_find): + elements = [self.expression[0]] + else: + elements = [self.expression[0].values] + + return self._find_items_in_expression(elements, to_find, valid_eval_classes) + + def find_slices(self) -> set[str]: + """ + Finds all references to array slices in the expression and sub-expressions. + + Returns: + set[str]: Unique slice references. + """ + valid_eval_classes = tuple( + [ + expression_parser.EvalOperatorOperand, + expression_parser.EvalFunction, + expression_parser.EvalSlicedComponent, + ] + ) + to_find = expression_parser.EvalIndexSlice + elements: list = [ + self.expression[0].values, + *list(self.sub_expressions.values()), + ] + + return self._find_items_in_expression(elements, to_find, valid_eval_classes) + + @staticmethod + def _find_items_in_expression( + parser_elements: list | pp.ParseResults, + to_find: type[expression_parser.EvalString], + valid_eval_classes: tuple[type[expression_parser.EvalString], ...], + ) -> set[str]: + """ + Recursively find sub-expressions / index items defined in an equation expression. + + Args: + parser_elements (list | pp.ParseResults): list of parser elements to check. + to_find (type[expression_parser.EvalString]): type of equation element to search for. + valid_eval_classes (tuple[type[expression_parser.EvalString], ...]): Other expression + elements that can be recursively searched + + Returns: + set[str]: All unique component / index item names. + """ + items: list = [] + recursive_func = functools.partial( + ParsedBackendEquation._find_items_in_expression, + to_find=to_find, + valid_eval_classes=valid_eval_classes, + ) + for parser_element in parser_elements: + if isinstance(parser_element, to_find): + items.append(parser_element.name) + + elif isinstance(parser_element, pp.ParseResults | list): + items.extend(recursive_func(parser_elements=parser_element)) + + elif isinstance(parser_element, valid_eval_classes): + items.extend(recursive_func(parser_elements=parser_element.values)) + return set(items) + + def add_expression_group_combination( + self, + expression_group_name: Literal["sub_expressions", "slices"], + expression_group_combination: Iterable[ParsedBackendEquation], + ) -> ParsedBackendEquation: + """ + Add parsed sub-expressions/index slices to a copy of self with updated names and mask lists. + + Args: + expression_group_name (Literal[sub_expressions, slices]): + Which of `sub-expressions`/`index slices` is being added. + expression_group_combination (Iterable[ParsedBackendEquation]): + All items of expression_group_name to be added. + + Returns: + ParsedBackendEquation: Copy of self with added sub-expressions/index slice dictionary and updated name + and mask list to include those corresponding to the dictionary entries. + """ + new_mask_list = [*self.mask] + for expr in expression_group_combination: + new_mask_list.extend(expr.mask) + new_name = f"{self.name}-{'-'.join([expr.name for expr in expression_group_combination])}" + expression_group_dict = { + expression_group_name: { + expr.name.split(":")[0]: expr.expression + for expr in expression_group_combination + } + } + return ParsedBackendEquation( + equation_name=new_name, + sets=self.sets, + expression=self.expression, + mask_list=new_mask_list, + **{ + "sub_expressions": self.sub_expressions, + "slices": self.slices, + **expression_group_dict, # type: ignore + }, + ) + + # Expecting array if not requesting latex string + @overload + def evaluate_mask( + self, + input_data: xr.Dataset, + model: Model, + math: MathModel, + config: ConfigModel, + *, + return_type: Literal["array"] = "array", + references: set | None = None, + initial_mask: xr.DataArray = TRUE_ARRAY, + ) -> xr.DataArray: ... + + # Expecting string if requesting latex string. + @overload + def evaluate_mask( + self, + input_data: xr.Dataset, + model: Model, + math: MathModel, + config: ConfigModel, + *, + return_type: Literal["math_string"], + references: set | None = None, + ) -> str: ... + + def evaluate_mask( + self, + input_data: xr.Dataset, + model: Model, + math: MathModel, + config: ConfigModel, + *, + return_type: str = "array", + references: set | None = None, + initial_mask: xr.DataArray = TRUE_ARRAY, + ) -> xr.DataArray | str: + """ + Evaluate parsed backend object dictionary `mask` string. + + Args: + input_data (xr.Dataset): Model input data. + model (Model): Linopy model. + math (MathModel): Calliope math definitions. + config (ConfigModel): Build configuration options. + return_type (str, optional): If "array", return xarray.DataArray. + If "math_string", return LaTex math string. + Defaults to "array". + references (set | None, optional): List of references to use in evaluation. + Defaults to None. + initial_mask (xr.DataArray, optional): If given, the mask array resulting + from evaluation will be further masked by this array. + Defaults to xr.DataArray(True) (i.e., no effect). + + Returns: + xr.DataArray | str: + If return_type == `array`: Boolean array defining on which index items a parsed component should be built. + If return_type == `math_string`: Valid LaTeX math string defining the "mask" conditions using logic notation. + """ + eval_attrs_ = { + "equation_name": self.name, + "helper_functions": helper_functions._registry["mask"], + "input_data": input_data, + "model": model, + "math": math, + "config": config, + } + if references is not None: + eval_attrs_["references"] = references + + evaluated_masks = [ + mask[0].eval(return_type, eval_attrs.EvalAttrs(**eval_attrs_)) + for mask in self.mask + ] + if return_type == "math_string": + return r"\land{}".join(f"({i})" for i in evaluated_masks if i != "true") + else: + mask = xr.DataArray( + functools.reduce(operator.and_, [initial_mask, *evaluated_masks]) + ) + if not mask.any(): + self.log_not_added("'mask' does not apply anywhere.") + return mask + + def drop_dims_not_in_foreach(self, mask: xr.DataArray) -> xr.DataArray: + """ + Remove all dimensions not included in "foreach" from the input array. + + Args: + mask (xr.DataArray): Array with potentially unwanted dimensions + + Returns: + xr.DataArray: + Array with same dimensions as the user-defined foreach sets. + Dimensions are ordered to match the order given by the sets. + """ + unwanted_dims = set(mask.dims).difference(self.sets) + return (mask.sum(unwanted_dims) > 0).astype(bool).transpose(*self.sets) + + # Expecting anything (most likely an array) if not requesting latex string. + @overload + def evaluate_expression( + self, + input_data: xr.Dataset, + model: Model, + math: MathModel, + *, + return_type: Literal["array"] = "array", + references: set | None = None, + mask: xr.DataArray = TRUE_ARRAY, + ) -> xr.DataArray: ... + + # Expecting string if requesting latex string. + @overload + def evaluate_expression( + self, + input_data: xr.Dataset, + model: Model, + math: MathModel, + *, + return_type: Literal["math_string"], + references: set | None = None, + ) -> str: ... + + def evaluate_expression( + self, + input_data: xr.Dataset, + model: Model, + math: MathModel, + *, + return_type: Literal["array", "math_string"] = "array", + references: set | None = None, + mask: xr.DataArray = TRUE_ARRAY, + ) -> xr.DataArray | str: + """ + Evaluate a math string to produce an array backend objects or a LaTex math string. + + Args: + input_data (xr.Dataset): Model input data. + model (xr.Dataset): Backend interface component dataset. + math (MathModel): Calliope math definitions. + + Keyword Args: + return_type (str, optional): + If "array", return xarray.DataArray. If "math_string", return LaTex math string. + Defaults to "array". + references (set | None, optional): + If given, any references in the math string to other model components + will be logged here. Defaults to None. + mask (xr.DataArray, optional): + If given, should be a boolean array with which to mask any produced arrays. + Defaults to xr.DataArray(True). + + Returns: + xr.DataArray | str: + If return_type == `array`: array of backend expression objects. + If return_type == `math_string`: Valid LaTeX math string defining the + "mask" conditions using logic notation. + """ + eval_attrs_ = { + "equation_name": self.name, + "slice_dict": self.slices, + "sub_expression_dict": self.sub_expressions, + "input_data": input_data, + "model": model, + "math": math, + "mask": mask, + "helper_functions": helper_functions._registry["expression"], + } + if references is not None: + eval_attrs_["references"] = references + evaluated = self.expression[0].eval( + return_type, eval_attrs.EvalAttrs(**eval_attrs_) + ) + return evaluated + + def raise_error_on_mask_expr_mismatch( + self, expression: xr.DataArray, mask: xr.DataArray + ) -> None: + """ + Checks if an evaluated expression is consistent with the `mask` array. + + Args: + expression (xr.DataArray): array of linear expressions or one side of a constraint equation. + mask (xr.DataArray): mask array; there should be a valid expression value for all True elements. + + Raises: + BackendError: + Raised if there is a dimension in the expression that is not in the mask. + BackendError: + Raised if the expression has any NaN mask the mask applies. + """ + broadcast_dims_mask = set(expression.dims).difference(set(mask.dims)) + if broadcast_dims_mask: + raise ValueError( + f"{self.name} | The linear expression array is indexed over dimensions not present in `foreach`: {broadcast_dims_mask}" + ) + # Check whether expression has NaN values in elements mask the expression should be valid. + incomplete_constraints = expression.isnull() & mask + if incomplete_constraints.any(): + raise ValueError( + f"{self.name} | Missing a linear expression for some coordinates selected by 'mask'. Adapting 'mask' might help." + ) + + def log_not_added( + self, + message: str, + level: Literal["info", "warning", "debug", "error", "critical"] = "debug", + ): + """ + Log to module-level logger with some prettification of the message. + + Args: + message (str): Message to log. + level (Literal["info", "warning", "debug", "error", "critical"], optional): + Log level. Defaults to "debug". + """ + getattr(LOGGER, level)( + f"Math parsing | {self.name} | Component not added; {message}" + ) + + +class ParsedBackendComponent(ParsedBackendEquation): + """Backend component parser.""" + + _ERR_BULLET: str = " * " + _ERR_STRING_ORDER: list[str] = ["expression_group", "id", "expr_or_mask"] + PARSERS: dict[str, Callable] = { + "constraints": expression_parser.generate_equation_parser, + "global_expressions": expression_parser.generate_arithmetic_parser, + "postprocessed": expression_parser.generate_arithmetic_parser, + "objectives": expression_parser.generate_arithmetic_parser, + "piecewise_constraints": expression_parser.generate_arithmetic_parser, + } + + def __init__( + self, + group: Literal[ + "variables", + "global_expressions", + "constraints", + "piecewise_constraints", + "objectives", + "postprocessed", + ], + name: str, + unparsed_data: MATH_DEFS_T, + parsing_components: dict[str, dict[str, set[str]]], + ) -> None: + """ + Parse an optimisation problem configuration. + + Defined in a dictionary of strings loaded from YAML into a series of Python + objects that can be passed onto a solver interface like Pyomo or Gurobipy. + + Args: + group (Literal["variables", "global_expressions", "constraints", "objectives"]): + Optimisation problem component group to which the unparsed data belongs. + name (str): Name of the optimisation problem component + unparsed_data (T): Unparsed math formulation. Expected structure depends on + the group to which the optimisation problem component belongs. + parsing_components (dict[str, dict[str, Iterable[str]]]): + Dictionary of valid component names for different categories of model data to use in parsing `mask` and `expression` strings. + """ + self.name = f"{group}:{name}" + self.group = group + self._unparsed = unparsed_data + self._mask_components = parsing_components["mask"] + self._expression_components = set().union( + *parsing_components["expression"].values() + ) + self.mask: list[pp.ParseResults] = [] + self.equations: list[ParsedBackendEquation] = [] + self.equation_expression_parser: Callable = self.PARSERS.get( + group, lambda x: None + ) + + # capture errors to dump after processing, + # to make it easier for a user to fix the constraint YAML. + self._errors: list = [] + self._tracker = self._init_tracker() + + # Initialise switches + self._is_valid: bool = True + + # Add objects that are used by shared functions + self.sets: set[str] = set(unparsed_data.foreach) + + def get_parsing_position(self): + """Create "." separated list from tracked strings.""" + return ".".join( + filter(None, [self._tracker[i] for i in self._ERR_STRING_ORDER]) + ) + + def reset_tracker(self): + """Re-initialise error string tracking.""" + self._tracker = self._init_tracker() + + def _init_tracker(self): + """Initialise error string tracking as dictionary of `key: None`.""" + return {i: None for i in self._ERR_STRING_ORDER} + + def parse_top_level_mask( + self, errors: Literal["raise", "ignore"] = "raise" + ) -> None: + """ + Parse the "mask" string that is (optionally) given as a top-level key of the math component dictionary. + + Args: + errors (Literal["raise", "ignore"], optional): + Collected parsing errors can be raised directly or ignored. + If errors exist and are ignored, the parsed component cannot be successfully evaluated. Defaults to "raise". + """ + top_level_mask = self.parse_mask_string(self._unparsed.mask) + + if errors == "raise": + self.raise_caught_errors() + + if self._is_valid: + self.mask = [top_level_mask] + + def parse_equations( + self, errors: Literal["raise", "ignore"] = "raise" + ) -> list[ParsedBackendEquation]: + """ + Parse `expression` and `mask` strings of math component dictionary. + + Args: + errors (Literal["raise", "ignore"], optional): + Collected parsing errors can be raised directly or ignored. + If errors exist and are ignored, the parsed component cannot be successfully evaluated. Defaults to "raise". + + Returns: + list[ParsedBackendEquation]: + List of parsed equations ready to be evaluated. + The length of the list depends on the product of provided equations and sub-expression/slice references. + """ + equations = self.generate_expression_list( + expression_parser=self.equation_expression_parser( + self._expression_components + ), + expression_list=self._unparsed.equations, + expression_group="equations", + id_prefix=self.name, + ) + + sub_expression_dict = { + c_name: self.generate_expression_list( + expression_parser=expression_parser.generate_sub_expression_parser( + self._expression_components + ), + expression_list=c_list, + expression_group="sub_expressions", + id_prefix=c_name, + ) + for c_name, c_list in self._unparsed.sub_expressions.root.items() + } + slice_dict = { + idx_name: self.generate_expression_list( + expression_parser=expression_parser.generate_slice_parser( + self._expression_components + ), + expression_list=idx_list, + expression_group="slices", + id_prefix=idx_name, + ) + for idx_name, idx_list in self._unparsed.slices.root.items() + } + + if errors == "raise": + self.raise_caught_errors() + + equations_with_sub_expressions = [] + for equation in equations: + equations_with_sub_expressions.extend( + self.extend_equation_list_with_expression_group( + equation, sub_expression_dict, "sub_expressions" + ) + ) + equations_with_sub_expressions_and_slices: list[ParsedBackendEquation] = [] + for equation in equations_with_sub_expressions: + equations_with_sub_expressions_and_slices.extend( + self.extend_equation_list_with_expression_group( + equation, slice_dict, "slices" + ) + ) + + return equations_with_sub_expressions_and_slices + + def _parse_string( + self, parser: pp.ParserElement, parse_string: str + ) -> pp.ParseResults: + """ + Parse equation string according to predefined parsing grammar. + + Args: + parser (pp.ParserElement): Parsing grammar. + parse_string (str): String to parse according to parser grammar. + + Returns: + Optional[pp.ParseResults]: + Parsed string. If any parsing errors are caught, + they will be logged to `self._errors` to raise later. + """ + try: + parsed = parser.parse_string(parse_string, parse_all=True) + except pp.ParseException as excinfo: + parsed = pp.ParseResults([]) + self._is_valid = False + pointer = f"{self.get_parsing_position()} (line {excinfo.lineno}, char {excinfo.col}): " + marker_pos = " " * ( + len(pointer) + 2 * len(self._ERR_BULLET) + excinfo.col - 1 + ) + self._errors.append(f"{pointer}{excinfo.line}\n{marker_pos}^") + + return parsed + + def parse_mask_string(self, mask_string: str = "True") -> pp.ParseResults: + """ + Parse a "mask" string of the form "CONDITION OPERATOR CONDITION". + + The operator can be "and"/"or"/"not and"/"not or". + + Args: + mask_string (str): + string value from a math dictionary "mask" key. + Defaults to "True", to have no effect on the subsequent subsetting. + + Returns: + pp.ParseResults: Parsed string. If any parsing errors are caught, + they will be logged to `self._errors` to raise later. + """ + parser = mask_parser.generate_mask_string_parser(**self._mask_components) + self._tracker["expr_or_mask"] = "mask" + return self._parse_string(parser, mask_string) + + def generate_expression_list( + self, + expression_parser: pp.ParserElement, + expression_list: _Equations, + expression_group: Literal["equations", "sub_expressions", "slices"], + id_prefix: str = "", + ) -> list[ParsedBackendEquation]: + """ + Align user-defined constraint equations/sub-expressions. + + Achieved by parsing expressions, specifying a default "mask" string if not + defined, and providing an ID to enable returning to the initial dictionary. + + Args: + expression_parser (pp.ParserElement): parser to use. + expression_list (list[UnparsedEquation]): list of constraint equations + or sub-expressions with arithmetic expression string and optional + mask string. + expression_group (Literal["equations", "sub_expressions", "slices"]): + For error reporting, the constraint dict key corresponding to the parse_string. + id_prefix (str, optional): Extends the ID from a number corresponding to the + expression_list position `idx` to a tuple of the form (id_prefix, idx). + Defaults to "". + + Returns: + list[ParsedBackendEquation]: Aligned expression dictionaries with parsed + expression strings. + """ + parsed_equation_list = [] + + if expression_group == "equations": + to_track = {"expression_group": f"{expression_group}[{{id}}]"} + else: + to_track = { + "expression_group": expression_group, + "id": f"{id_prefix}[{{id}}]", + } + + for idx, expression_data in enumerate(expression_list): + self._tracker.update({k: v.format(id=idx) for k, v in to_track.items()}) + + parsed_mask = self.parse_mask_string(expression_data.mask) + + self._tracker["expr_or_mask"] = "expression" + parsed_expression = self._parse_string( + expression_parser, expression_data.expression + ) + if len(parsed_expression) > 0: + parsed_equation_list.append( + ParsedBackendEquation( + equation_name=":".join(filter(None, [id_prefix, str(idx)])), + sets=self.sets, + mask_list=[parsed_mask], + expression=parsed_expression, + ) + ) + self.reset_tracker() + + return parsed_equation_list + + def extend_equation_list_with_expression_group( + self, + parsed_equation: ParsedBackendEquation, + parsed_items: dict[str, list[ParsedBackendEquation]], + expression_group: Literal["sub_expressions", "slices"], + ) -> list[ParsedBackendEquation]: + """ + Extend equation expressions with sub-expression data. + + Finds all sub-expressions referenced in an equation expression and returns a + product of the sub-expression data. + + Args: + parsed_equation (ParsedBackendEquation): Equation data dictionary. + parsed_items (dict[str, list[ParsedBackendEquation]]): + Dictionary of expressions to replace within the equation data dictionary. + expression_group (Literal["sub_expressions", "slices"]): + Name of expression group that the parsed_items dict is referencing. + + Returns: + list[ParsedBackendEquation]: Expanded list of parsed equations with the + product of all references to items from the `expression_group` + producing a new equation object. E.g., if the input equation object has + a reference to an slice which itself has two expression options, two + equation objects will be added to the return list. + """ + if expression_group == "sub_expressions": + equation_items = parsed_equation.find_sub_expressions() + elif expression_group == "slices": + equation_items = parsed_equation.find_slices() + if not equation_items: + return [parsed_equation] + + invalid_items = equation_items.difference(parsed_items.keys()) + if invalid_items: + raise KeyError( + f"{self.name}: Undefined {expression_group} found in equation: {invalid_items}" + ) + + parsed_item_product = itertools.product( + *[parsed_items[k] for k in equation_items] + ) + + return [ + parsed_equation.add_expression_group_combination( + expression_group, parsed_item_combination + ) + for parsed_item_combination in parsed_item_product + ] + + def foreach_matrix(self, input_data: xr.Dataset) -> xr.DataArray: + """ + Generate a multi-dimensional array mask a constraint will be built. + + The multi-dimensional boolean array is based on the sets over which the + constraint is to be built (`foreach`) and the model `exists` array. + + Args: + input_data (xr.Dataset): Calliope model dataset. + + Returns: + xr.DataArray: boolean array indexed over ["nodes", "techs", "carriers"] + + any additional dimensions provided by `foreach`. + """ + if self.sets.difference(input_data.dims): + self.log_not_added( + f"indexed over unidentified set names: `{self.sets.difference(input_data.dims)}`." + ) + return xr.DataArray(False) + if not self.sets: + return xr.DataArray(True) + else: + exists_and_foreach = [input_data[i].notnull() for i in self.sets] + return functools.reduce(operator.and_, exists_and_foreach) + + def generate_top_level_mask( + self, + input_data: xr.Dataset, + model: Model, + math: MathModel, + config: ConfigModel, + *, + align_to_foreach_sets: bool = True, + break_early: bool = True, + references: set | None = None, + ) -> xr.DataArray: + """ + Generate a multi-dimentional "mask" array. + + The multi-dimensional array is created using model inputs and component sets + defined in foreach. The component top-level "mask" is then applied to the + array. + + Args: + input_data (xr.Dataset): Model input data. + model (xr.Dataset): Backend interface component dataset. + math (MathModel): Calliope math definitions. + config (ConfigModel): Build configuration options. + align_to_foreach_sets (bool, optional): + By default, all foreach arrays have the dimensions ("nodes", "techs", "carriers") + as well as any additional dimensions provided by the component's "foreach" key. + If this argument is True, the dimensions not included in "foreach" are removed from the array. + Defaults to True. + break_early (bool, optional): + If any intermediate array has no valid elements (i.e. all are False), + the function will return that array rather than continuing - saving + time and memory on large models. Defaults to True. + references (set | None, optional): references to use during evaluation. Defaults to None. + + Returns: + xr.DataArray: Boolean array defining on which index items a parsed component should be built. + """ + foreach_mask = self.foreach_matrix(input_data) + + if not foreach_mask.any(): + self.log_not_added("'foreach' does not apply anywhere.") + + if break_early and not foreach_mask.any(): + return foreach_mask + + self.parse_top_level_mask() + mask = self.evaluate_mask( + input_data, + model, + math, + config, + initial_mask=foreach_mask, + references=references if references is not None else set(), + ) + if break_early and not mask.any(): + return mask + + if align_to_foreach_sets: + mask = self.drop_dims_not_in_foreach(mask) + return mask + + def raise_caught_errors(self): + """Pipe parsing errors to the ModelError bullet point list generator.""" + errors = [] + if not self._is_valid: + errors.append({f"{self.name}": self._errors}) + if errors: + raise ValueError( + "\n".join(f"- {k}: {v}" for err in errors for k, v in err.items()) + ) diff --git a/linopy/declarative/schema.py b/linopy/declarative/schema.py new file mode 100644 index 000000000..14edbf99a --- /dev/null +++ b/linopy/declarative/schema.py @@ -0,0 +1,705 @@ +# Copyright (C) since 2013 Calliope contributors listed in AUTHORS. +# Licensed under the Apache 2.0 License (see LICENSE file). +"""Schema for Calliope mathematical definition.""" + +import logging +from collections.abc import Hashable, Iterable +from functools import cached_property +from typing import Annotated, ClassVar, Literal, Self, TypeVar + +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) + +COMPONENTS_T = Literal[ + "dimensions", + "parameters", + "lookups", + "variables", + "global_expressions", + "constraints", + "piecewise_constraints", + "objectives", + "postprocessed", +] + + +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 that is used to store dictionaries with user-defined keys and Calliope pydantic model values.""" + + def __setitem__(self, *args, **kwargs) -> None: + """Do not allow direct item setting.""" + raise PydanticCustomError( + "no_extra_dict", + f"Cannot set a {self.__class__.__name__} directly. Use the `update` method instead, which will return a copy.", + ) + + def __getitem__(self, key): + """Expose the root attribute when getting an item by key.""" + return self.root[key] + + def __repr__(self, *args, **kwargs): + """Show the __repr__ of the root attribute when requesting the __repr__ of the class.""" + return self.root.__repr__(*args, **kwargs) + + def __rich_repr__(self): + """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} + + def update( + self, update_def: dict | BaseModel, deep: bool = False, overwrite: bool = True + ) -> Self: + """ + Return a new iteration of the model with updated fields. + + Args: + update_def (dict | BaseModel): Dictionary or pydantic model with which to update the base model. + deep (bool, optional): Set to True to make a deep copy of the model. Defaults to False. + overwrite (bool, optional): Set to False to only update fields that are not already set in the base model. Defaults to True. + + Returns: + BaseModel: New model instance. + """ + update_dict: dict = ( + update_def.model_dump(exclude_unset=True) + if isinstance(update_def, BaseModel) + else update_def + ) + new_dict = dict() + # Iterate through dict to be updated and convert any sub-dicts into their respective pydantic model objects. + for key, val in update_dict.items(): + key_class = self.root.get(key, None) + if isinstance(key_class, LinopyBaseModel): + new_dict[key] = key_class.update(val, deep=deep, overwrite=overwrite) + elif isinstance(key_class, LinopyListModel): + if overwrite: + new_dict[key] = key_class.update(val) + else: + continue + elif key_class == val: + continue + else: + if key not in self.root or overwrite: + LOGGER.debug(f"Adding {self.__class__.__name__} entry: `{key}`") + new_dict[key] = self.model_validate({key: val})[key] + + return self.model_validate(self.root | new_dict) + + +class LinopyListModel(RootModel): + """Pydantic Model that is used to store lists of Linopy pydantic models.""" + + def __iter__(self): + """Iterate over root attribute contents when iterating over class.""" + return iter(self.root) + + def __getitem__(self, item: int): + """Expose the root attribute when getting an item by index value.""" + return self.root[item] + + def __repr__(self, *args, **kwargs): + """Show the __repr__ of the root attribute when requesting the __repr__ of the class.""" + return self.root.__repr__(*args, **kwargs) + + def __rich_repr__(self): + """Prettyprint the __repr__ of the root attribute when requesting the prettyprint of the class.""" + yield from self.root + + def update(self, update_list: list) -> Self: + """ + Return a new iteration of the model fields entirely replaced. + + We do not allow updating individual items in the list as it's hard to guarantee the order of items in the list. + + Args: + update_list (list): List with which to update the base model. + + Returns: + BaseModel: New model instance. + """ + return self.model_validate(update_list) + + +class LinopyBaseModel(BaseModel): + """A base class for creating pydantic models for Linopy models.""" + + model_config = { + "extra": "forbid", + "frozen": True, + "revalidate_instances": "always", + "use_attribute_docstrings": True, + } + + def __getitem__(self, item): + """Allow attribute access via item lookup.""" + return getattr(self, item) + + def update( + self, + update_def: dict | BaseModel, + deep: bool = False, + overwrite: bool = True, + _suppress_log: bool = False, + ) -> Self: + """ + Return a new iteration of the model with updated fields. + + Args: + update_def (dict | BaseModel): Dictionary or pydantic model with which to update the base model. + deep (bool, optional): Set to True to make a deep copy of the model. Defaults to False. + overwrite (bool, optional): Set to False to only update fields that are not already set in the base model. Defaults to True. + _suppress_log (bool, optional): + Set to True to suppress logging of updated fields. + This is an internal method argument used to avoid logging updates when the update method is called recursively. + Defaults to False. + + Returns: + BaseModel: New model instance. + """ + new_dict = dict() + # Iterate through dict to be updated and convert any sub-dicts into their respective pydantic model objects. + # Wrapped in `AttrDict` to allow users to define dot notation nested configuration. + # We revert to dict format to avoid issues with the `model_copy` method later. + update_dict = ( + update_def.model_dump(exclude_unset=True) + if isinstance(update_def, BaseModel) + else update_def + ) + for key, val in update_dict.items(): + key_class = getattr(self, key, None) + if isinstance(key_class, LinopyBaseModel | LinopyDictModel): + new_dict[key] = key_class.update(val, deep=deep, overwrite=overwrite) + elif isinstance(key_class, LinopyListModel): + if overwrite: + new_dict[key] = key_class.update(val) + else: + continue + elif key_class == val: + continue + else: + if not _suppress_log and ( + key not in self.model_fields_set + or (key in self.model_fields_set and overwrite) + ): + LOGGER.debug( + f"Updating {self.__class__.__name__} `{key}`: {key_class} -> {val}" + ) + new_dict[key] = val + updated = super().model_copy(update=new_dict, deep=deep) + if not overwrite: + extra_update = super().model_dump(exclude_unset=True, serialize_as_any=True) + updated = updated.update(extra_update, deep=deep, _suppress_log=True) + return updated.model_validate( + updated.model_dump(exclude_unset=True, serialize_as_any=True) + ) + + +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 variable 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.""" + resample_method: Literal["mean", "sum", "first"] = "first" + """If resampling is applied over any of the parameter's dimensions, the method to use to aggregate the data.""" + unit: str = "" + """The unit of the parameter, e.g. 'kW', 'm', 'kg', 'energy', 'power', ...""" + + @property + def dtype(self) -> Literal["float"]: + """Dummy variable 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.""" + resample_method: Literal["mean", "sum", "first"] = "first" + """If resampling is applied over any of the lookup's dimensions, the method to use to aggregate the data.""" + one_of: list | None = None + """If given, the lookup values must be one of these items.""" + pivot_values_to_dim: str | None = None + """If given, the lookup will be pivoted such that its values become the index of a new dimension and its new values are boolean, True where the index values match the old values. + For instance, if the lookup starts out indexed over `techs` with values of `[electricity, gas]` and `pivot_values_to_dim: carriers`, + then the lookup will be converted to a boolean array with the dimensions ['techs', 'carriers']. + """ + + _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.""" + + @property + def equations(self) -> _Equations: + """Dummy property to satisfy type hinting.""" + return _Equations() + + @property + def sub_expressions(self) -> _SubExpressions: + """Dummy property to satisfy type hinting.""" + return _SubExpressions() + + @property + def slices(self) -> _SubExpressions: + """Dummy property to satisfy type hinting.""" + return _SubExpressions() + + _group: ClassVar[COMPONENTS_T] = "piecewise_constraints" + + +class LinearExpressionDef(_MathIndexedComponent, _MathEquationComponent): + """ + Schema for named global 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 global expressions). + + NOTE: If expecting to use global expression `A` in global 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() + """Global expression math equations.""" + sub_expressions: _SubExpressions = _SubExpressions() + """Global expression named sub-expressions.""" + slices: _SubExpressions = _SubExpressions() + """Global expression named index slices.""" + order: int = 0 + """Order in which to apply this global expression relative to all others, if different to its definition order.""" + + _group: ClassVar[COMPONENTS_T] = "global_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() + + @property + def equations(self) -> _Equations: + """Dummy property to satisfy type hinting.""" + return _Equations() + + @property + def sub_expressions(self) -> _SubExpressions: + """Dummy property to satisfy type hinting.""" + return _SubExpressions() + + @property + def slices(self) -> _SubExpressions: + """Dummy property to satisfy type hinting.""" + return _SubExpressions() + + _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.""" + + @property + def foreach(self) -> UniqueList[AttrStr]: + """Objectives are always adimensional.""" + return [] + + @property + def mask(self) -> str: + """Dummy property to satisfy type hinting.""" + return "True" + + _group: ClassVar[COMPONENTS_T] = "objectives" + + +class PostprocessedExpressionDef(LinearExpressionDef): + """ + Schema for postprocessed expressions. + + Can be used to combine parameters, variables, and global 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 LinearExpressionDefs(LinopyDictModel): + """Linopy model global_expressions dictionary.""" + + root: dict[AttrStr, LinearExpressionDef] = 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): + """ + Mathematical definition of Calliope math. + + Contains mathematical programming components available for optimising with Calliope. + Can contain partial definitions if they are meant to be layered on top of another. + E.g.: layering 'base' and 'operate' math. + """ + + 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.""" + global_expressions: LinearExpressionDefs = LinearExpressionDefs() + """All global 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): + """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() + duplicates = 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 + + @cached_property + def parsing_components(self) -> dict[str, dict[str, set[str]]]: + """ + Return a set of valid component names in the model to use in `mask` string parsing. + + Returns: + dict[Literal["dimension_names", "input_names", "result_names"], set[str]]: + Set of valid names grouped by location in the math in which they are defined. + """ + parsing_components = { + "dimensions": ["dimensions"], + "inputs": ["lookups", "parameters"], + "results": ["variables", "global_expressions"], + } + + def _names(): + 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 + | LinearExpressionDef + | ObjectiveDef + | PiecewiseConstraintDef +) + + +class ConfigModel(LinopyBaseModel): + """Base configuration options used when building a Linopy optimisation problem.""" + + model_config = {"title": "Model build configuration"} + + foo: str = "bar" + """A dummy variable to test accessing the config items in declarative math.""" diff --git a/pyproject.toml b/pyproject.toml index 19d0abb39..6f3676ec1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -39,6 +39,9 @@ dependencies = [ "tqdm", "deprecation", "packaging", + "pydantic>=2", + "pyparsing>=3", + "pyyaml", ] [project.urls] From 06e4a6a289bd9747f392849812a536022f08cb81 Mon Sep 17 00:00:00 2001 From: Bryn Pickering <17178478+brynpickering@users.noreply.github.com> Date: Wed, 10 Jun 2026 19:45:49 +0100 Subject: [PATCH 02/12] ruff fix --- linopy/__init__.py | 1 + 1 file changed, 1 insertion(+) diff --git a/linopy/__init__.py b/linopy/__init__.py index d145a96e0..6bb10b8da 100644 --- a/linopy/__init__.py +++ b/linopy/__init__.py @@ -77,4 +77,5 @@ "read_netcdf", "segments", "tangent_lines", + "declarative_model" ) From b610fa3be5a0ada1c166b488acd2404fcfedd785 Mon Sep 17 00:00:00 2001 From: Bryn Pickering <17178478+brynpickering@users.noreply.github.com> Date: Tue, 7 Jul 2026 09:53:50 +0100 Subject: [PATCH 03/12] [WIP] Various messy fixes for demo --- linopy/declarative.py | 0 linopy/declarative/build.py | 324 +++++++++++++++++--- linopy/declarative/expression_parser.py | 381 ++++++++++++++++++------ linopy/declarative/helper_functions.py | 16 +- linopy/declarative/mask_parser.py | 10 +- linopy/declarative/parsing.py | 16 +- linopy/declarative/schema.py | 50 ++-- linopy/expressions.py | 131 +++++++- linopy/model.py | 92 ++++++ 9 files changed, 833 insertions(+), 187 deletions(-) delete mode 100644 linopy/declarative.py diff --git a/linopy/declarative.py b/linopy/declarative.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/linopy/declarative/build.py b/linopy/declarative/build.py index 5a9c3679c..cc44e4e01 100644 --- a/linopy/declarative/build.py +++ b/linopy/declarative/build.py @@ -1,10 +1,13 @@ +import textwrap import time import typing +from requests import exceptions +from tqdm.asyncio import tqdm import xarray as xr -from linopy import merge -from linopy.declarative import parsing +from linopy.expressions import merge +from linopy.declarative import eval_attrs, helper_functions, parsing from linopy.declarative.schema import ( LOGGER, ConfigModel, @@ -12,18 +15,33 @@ MathModel, ObjectiveDef, VariableDef, + ExpressionDef, ) from linopy.expressions import LinearExpression +from linopy.io import TQDM_COLOR +from linopy.io import TQDM_COLOR from linopy.model import Model - +import numpy as np ORDERED_COMPONENTS_T = typing.Literal[ "variables", - # "global_expressions", + "expressions", "constraints", # "piecewise_constraints", "objectives", ] +DTYPE_OPTIONS = { + "string": str, + "float": float, + "bool": bool, + "datetime": np.datetime64, + "date": np.datetime64, + "integer": int, +} + +DATETIME_DTYPE = "M" +"""Numpy type kind for datetime arrays""" + def declarative_model(math_def: dict, input_data: xr.Dataset, config: dict) -> Model: """Build a Linopy Model from declarative math definitions and input data.""" @@ -35,9 +53,66 @@ class DeclarativeModelBuilder: def __init__(self, math_def: dict, input_data: xr.Dataset, config: dict): self.model = Model() self.math = MathModel.model_validate(math_def) - self.input_data = input_data + self.input_data = self._update_dtypes(input_data) self.config = ConfigModel.model_validate(config) + self._check_inputs() + + def _update_dtypes(self, ds: xr.Dataset, id_: str = "") -> xr.Dataset: + """Update data types of coordinates or data variables in the dataset. + + Args: + ds (xr.Dataset): Dataset to update. + math (math_schema.CalliopeBuildMath): Model math definition. + id_ (str, optional): ID of the dataset being updated, for logging purposes. Defaults to an empty string. + + Raises: + ValueError: If there is a mismatch between the provided variable and its definition in the model math. + + Returns: + xr.Dataset: `ds` with data types updated. + """ + prefix = f"{id_} | " if id_ else "" + for var_name, var_data in ds.items(): + try: + math_def = self.math.find( + 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 = math_def.dtype # type: ignore + 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 "datetime": + updated_var = time._datetime_index( + var_data.to_series(), self.config.datetime_format + ).to_xarray() + case "date": + updated_var = ( + time._datetime_index( + var_data.to_series(), self.config.date_format + ) + .to_xarray() + .assign_attrs(var_data.attrs) + ) + case "bool": + updated_var = var_data.fillna(False).astype(dtype) + case _: + updated_var = var_data.astype(dtype) + + ds[var_name] = updated_var + return ds + @staticmethod def _sorted_by_order( root: typing.Mapping[str, typing.Any], @@ -45,7 +120,32 @@ def _sorted_by_order( """Return (name, obj) pairs from a root mapping, sorted by obj.order.""" return sorted(root.items(), key=lambda item: getattr(item[1], "order", 0)) - def add_variable(self, name: str, definition: VariableDef): + def _check_inputs(self) -> None: + data_checks = self.math.checks + check_results = {"raise": [], "warn": []} + parser_ = parsing.mask_parser.generate_mask_string_parser( + **self.math.parsing_components["mask"] + ) + eval_kwargs = { + "model": self.model, + "math": self.math, + "input_data": self.input_data, + "config": self.config, + "helper_functions": helper_functions._registry["mask"], + } + for name, check in data_checks.root.items(): + if check.active: + parsed_ = parser_.parse_string(check.mask, parse_all=True) + eval_attrs_ = eval_attrs.EvalAttrs(equation_name=name, **eval_kwargs) + evaluated = parsed_[0].eval("array", eval_attrs_) + if evaluated.any() and (evaluated & self.input_data.active).any(): + check_results[check.errors].append(check.message) + + print_warnings_and_raise_errors( + check_results["warn"], check_results["raise"], during="model input data checks" + ) + + def add_variable(self, name: str, definition: VariableDef) -> None: references: set[str] = set() parsed_component = parsing.ParsedBackendComponent( "variables", name, definition, self.math.parsing_components @@ -69,14 +169,14 @@ def add_variable(self, name: str, definition: VariableDef): self.model.add_variables(coords=mask.coords, name=name, mask=mask, **kwargs) self.model.variables[name].attrs["references"] = references else: - LOGGER.warning( - f"Optimisation Model | variables:{name} | No valid data points after applying 'where' condition. Variable not added to model." + LOGGER.info( + f"variables:{name} | No valid data points after applying mask. Variable not added to model." ) - def add_constraint(self, name: str, definition: ConstraintDef): + def add_expression(self, name: str, definition: ExpressionDef) -> None: references: set[str] = set() parsed_component = parsing.ParsedBackendComponent( - "constraints", name, definition, self.math.parsing_components + "expressions", name, definition, self.math.parsing_components ) mask = parsed_component.generate_top_level_mask( self.input_data, @@ -87,9 +187,7 @@ def add_constraint(self, name: str, definition: ConstraintDef): break_early=True, references=references, ) - lhs = LinearExpression(float("nan"), self.model).where(mask) - sign = xr.DataArray().where(parsed_component.drop_dims_not_in_foreach(mask)) - rhs = LinearExpression(float("nan"), self.model).where(mask) + expr = LinearExpression(float("nan"), self.model).where(mask) all_mask = mask.copy() if mask.any(): equations = parsed_component.parse_equations() @@ -105,41 +203,110 @@ def add_constraint(self, name: str, definition: ConstraintDef): if not sub_mask.any(): continue sub_mask = parsed_component.drop_dims_not_in_foreach(sub_mask) - if (sign.notnull() & sub_mask).any(): + if (~expr.isnull() & sub_mask).any(): raise ValueError( - f"Optimisation Model | constraints:{name} | Overlapping 'mask' conditions between equations are not allowed. Please revise the 'mask' conditions to ensure they are mutually exclusive." + f"expressions:{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 = equation.evaluate_expression( + expr_to_fill = equation.evaluate_expression( self.input_data, self.model, self.math, mask=sub_mask, references=references, ) - all_mask = all_mask | sub_mask - if isinstance(lhs_to_fill, xr.DataArray): - lhs = lhs.fillna(lhs_to_fill) + if isinstance(expr_to_fill, xr.DataArray): + expr = expr.fillna(expr_to_fill) else: - lhs = merge([lhs, lhs_to_fill]).where(all_mask) - sign = sign.fillna(sign_to_fill) - if isinstance(rhs_to_fill, xr.DataArray): - rhs = rhs.fillna(rhs_to_fill) - else: - rhs = merge([rhs, rhs_to_fill]).where(all_mask) - - self.model.add_constraints( - coords=all_mask.coords, - name=name, - lhs=lhs, - sign=sign.fillna( - "==" - ), # Default to equality to avoid errors; will be masked. - rhs=rhs, - mask=all_mask, + expr = merge([expr, expr_to_fill.where(sub_mask)]) + if not expr.isnull().all(): + self.model.add_expressions(name=name, data=expr, mask=all_mask) + self.model.expressions[name].attrs["references"] = references + else: + LOGGER.info( + f"expressions:{name} | No valid data points after applying mask. Expression not added to model." + ) + else: + LOGGER.info( + f"expressions:{name} | No valid data points after applying mask. Expression not added to model." ) - self.model.constraints[name].attrs["references"] = references - def add_objective(self, name: str, definition: ObjectiveDef): + def add_constraint(self, name: str, definition: ConstraintDef) -> None: + references: set[str] = set() + parsed_component = parsing.ParsedBackendComponent( + "constraints", name, definition, self.math.parsing_components + ) + mask = parsed_component.generate_top_level_mask( + self.input_data, + self.model, + self.math, + self.config, + align_to_foreach_sets=True, + break_early=True, + references=references, + ) + lhs = LinearExpression(float("nan"), self.model).where(mask) + sign = xr.DataArray().where(parsed_component.drop_dims_not_in_foreach(mask)) + rhs = LinearExpression(float("nan"), self.model).where(mask) + all_mask = mask.copy() + if not mask.any(): + LOGGER.info(f"constraints:{name} | No valid data points after applying mask. Constraint not added to model.") + return None + + equations = parsed_component.parse_equations() + for equation in equations: + sub_mask = equation.evaluate_mask( + self.input_data, + self.model, + self.math, + self.config, + initial_mask=mask, + references=references, + ) + if not sub_mask.any(): + LOGGER.info(f"constraints:{equation.name} | No valid data points after applying mask. Constraint not added to model.") + continue + sub_mask = parsed_component.drop_dims_not_in_foreach(sub_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 = equation.evaluate_expression( + self.input_data, + self.model, + self.math, + mask=sub_mask, + references=references, + ) + if isinstance(lhs_to_fill, xr.DataArray): + lhs_to_fill = LinearExpression(lhs_to_fill, self.model) + lhs = merge([lhs, lhs_to_fill]) + + if isinstance(rhs_to_fill, xr.DataArray): + rhs_to_fill = LinearExpression(rhs_to_fill, self.model) + rhs = merge([rhs, rhs_to_fill]) + + sign = sign.fillna(sign_to_fill) + + if sign.isnull().all(): + LOGGER.info(f"constraints:{name} | No valid data points after applying mask. Constraint not added to model.") + return None + + self.model.add_constraints( + coords=all_mask.coords, + name=name, + lhs=lhs, + sign=sign.fillna( + "==" + ), # Default to equality to avoid errors; will be masked. + rhs=rhs, + mask=all_mask, + ) + self.model.constraints[name].attrs["references"] = references + + def add_objective(self, name: str, definition: ObjectiveDef) -> None: references: set[str] = set() parsed_component = parsing.ParsedBackendComponent( "objectives", name, definition, self.math.parsing_components @@ -170,7 +337,8 @@ def add_objective(self, name: str, definition: ObjectiveDef): sub_mask = parsed_component.drop_dims_not_in_foreach(sub_mask) if (~expr.isnull() & sub_mask).any(): raise ValueError( - f"Optimisation Model | objectives:{name} | Overlapping 'mask' conditions between equations are not allowed. Please revise the 'mask' conditions to ensure they are mutually exclusive." + f"objectives:{name} | Overlapping 'mask' conditions between equations are not allowed. " + "Please revise the 'mask' conditions to ensure they are mutually exclusive." ) expr_to_fill = equation.evaluate_expression( self.input_data, @@ -187,12 +355,86 @@ def build(self) -> Model: for components in typing.get_args(ORDERED_COMPONENTS_T): component = components.removesuffix("s") ordered_items = self._sorted_by_order(self.math[components].root) - for name, definition in ordered_items: + ordered_items_tqdm = tqdm( + ordered_items, + desc=f"Building {components}.", + colour=TQDM_COLOR, + ) + for name, definition in ordered_items_tqdm: start = time.time() getattr(self, f"add_{component}")(name, definition) end = time.time() - start LOGGER.debug( - f"Optimisation Model | {components}:{name} | Built in {end:.4f}s" + f"{components}:{name} | Built in {end:.4f}s" ) - LOGGER.info(f"Optimisation Model | {components} | Generated.") + LOGGER.info(f"{components} | Generated.") return self.model + + +def print_warnings_and_raise_errors( + warnings: list[str] | dict[str, list[str]] | None = None, + errors: list[str] | dict[str, list[str]] | None = None, + during: str = "model processing", + bullet: str = " * ", +) -> None: + """Process collections of warnings/errors. + + Prints warnings / raises errors with a bullet point list of the concatenated + collections. + + Lists will return simple bullet lists: + E.g. warnings=["foo", "bar"] becomes: + + Possible issues found during model processing: + * foo + * bar + + Dicts of lists will return nested bullet lists: + E.g. errors={"foo": ["foobar", "foobaz"]} becomes: + + Errors during model processing: + * foo + * foobar + * foobaz + + Args: + warnings (list[str] | dict[str, list[str]] | None, optional): + List of warning strings or dictionary of warning strings. + If None or an empty list, no warnings will be printed. + Defaults to None. + errors (list[str] | dict[str, list[str]] | None, optional): + List of error strings or dictionary of error strings. + If None or an empty list, no errors will be raised. + Defaults to None. + during (str, optional): + Substring that will be placed at the top of the concatenated list of warnings/errors to point to during which phase of data processing they occurred. + Defaults to "model processing". + bullet (str, optional): Type of bullet points to use. Defaults to " * ". + + Raises: + ModelError: If errors is not None or is a non-empty list/dict + + """ + spacer = " " * len(bullet) + + def _sort_strings(stringlist: list[str]) -> list[str]: + return sorted(list(set(stringlist))) + + def _predicate(string_: str) -> bool: + return not string_.startswith((bullet, spacer)) + + def _indenter(strings: list[str] | dict[str, list[str]]) -> str: + if isinstance(strings, dict): + sorted_strings = [] + for k, v in strings.items(): + sorted_strings.append(str(k) + ":") + sorted_strings.extend(_sort_strings([spacer + bullet + i for i in v])) + else: + sorted_strings = _sort_strings(strings) + return textwrap.indent("\n".join(sorted_strings), bullet, predicate=_predicate) + + if warnings: + LOGGER.info(f"Possible issues found during {during}:\n" + _indenter(warnings)) + + if errors: + raise ValueError(f"Errors during {during}:\n" + _indenter(errors)) diff --git a/linopy/declarative/expression_parser.py b/linopy/declarative/expression_parser.py index 650092af6..7891d0641 100644 --- a/linopy/declarative/expression_parser.py +++ b/linopy/declarative/expression_parser.py @@ -35,7 +35,7 @@ from abc import ABC, abstractmethod from collections.abc import Callable, Iterable, Iterator from dataclasses import replace -from typing import Any, Literal, overload +from typing import Any, Literal, overload, TypeVar import numpy as np import pandas as pd @@ -44,15 +44,17 @@ from linopy.declarative.eval_attrs import EvalAttrs from linopy.declarative.helper_functions import ParsingHelperFunction -from linopy.expressions import LinearExpression +from linopy.expressions import LinearExpression, QuadraticExpression from linopy.variables import Variable pp.ParserElement.enable_packrat() SUB_EXPRESSION_CLASSIFIER = "$" - -RETURN_T = Literal["array", "math_string"] +EXPR_T = LinearExpression | QuadraticExpression +EXPRVAR_T = TypeVar("EXPRVAR_T", LinearExpression, QuadraticExpression) +ARRAY_T = TypeVar("ARRAY_T", LinearExpression, QuadraticExpression, xr.DataArray) +RETURN_T = Literal["expr", "array", "math_string"] class EvalString(ABC): @@ -62,7 +64,7 @@ class EvalString(ABC): eval_attrs: EvalAttrs instring: str - def __eq__(self, other): + def __eq__(self, other: Any) -> bool: """Functionality for '==' operations.""" return self.__repr__() == other @@ -94,6 +96,12 @@ def as_array(self) -> xr.DataArray | list[xr.DataArray]: The purpose of this is to be able to access the string/number value whether we query the array name or its data. """ + def as_expr(self) -> xr.DataArray | list[xr.DataArray]: + """ + Evaluate and return expression as a LinearExpression or QuadraticExpression. + """ + return self.as_array() + # Math strings evaluate to strings. @overload def eval( @@ -108,7 +116,69 @@ def eval( def eval( self, return_type: RETURN_T, eval_attrs: EvalAttrs - ) -> str | xr.DataArray | list[xr.DataArray]: + ) -> str | xr.DataArray | list[xr.DataArray] | EXPR_T: + """ + Evaluate math string expression. + + Args: + return_type (Literal[math_string, input, array]): + Dictates how the expression should be evaluated (see `Returns` section). + eval_attrs (EvalAttrs): Evaluation attributes. + + Returns: + str | list[str | float] | xr.DataArray: + If `math_string` is desired, returns a valid LaTex math string. + If `array` is desired, returns xarray DataArray or a list of strings/numbers (if the expression represents a list). + """ + self.eval_attrs = eval_attrs + evaluated: str | list[str | float] | xr.DataArray | list[xr.DataArray] | EXPR_T + if return_type == "array": + evaluated = self.as_array() + elif return_type == "math_string": + evaluated = self.as_math_string() + elif return_type == "expr": + evaluated = self.as_expr() + + return evaluated + +class EvalArrayOrMathExpr(EvalString): + """Abstract class to evaluate expressions as either arrays or math strings.""" + + @abstractmethod + def as_math_string(self) -> str: + """Evaluate and return expression as LaTeX.""" + + @abstractmethod + def as_array(self) -> xr.DataArray | list[xr.DataArray]: + """ + Evaluate and return expression as a DataArray or list. + + If the evaluated expression returns a simple string or number, + this value will be assigned as both the `name` and the data of the returned DataArray. + The purpose of this is to be able to access the string/number value whether we query the array name or its data. + """ + + @abstractmethod + def as_expr(self) -> EXPR_T: + """ + Evaluate and return expression as a LinearExpression or QuadraticExpression. + """ + + # Math strings evaluate to strings. + @overload + def eval( + self, return_type: Literal["math_string"], eval_attrs: EvalAttrs + ) -> str: ... + + # Arrays evaluate to arrays + @overload + def eval( + self, return_type: Literal["array"], eval_attrs: EvalAttrs + ) -> xr.DataArray | list[xr.DataArray]: ... + + def eval( + self, return_type: RETURN_T, eval_attrs: EvalAttrs + ) -> str | xr.DataArray | list[xr.DataArray] | EXPR_T: """ Evaluate math string expression. @@ -123,14 +193,128 @@ def eval( If `array` is desired, returns xarray DataArray or a list of strings/numbers (if the expression represents a list). """ self.eval_attrs = eval_attrs - evaluated: str | list[str | float] | xr.DataArray | list[xr.DataArray] + evaluated: str | list[str | float] | xr.DataArray | list[xr.DataArray] | EXPR_T if return_type == "array": evaluated = self.as_array() elif return_type == "math_string": evaluated = self.as_math_string() + elif return_type == "expr": + evaluated = self.as_expr() + return evaluated +class EvalComparisonOp(EvalString): + """Class for processing comparison operations.""" + + OP_TRANSLATOR = {"<=": r" \leq ", ">=": r" \geq ", "==": " = ", "=": " = "} + + def __init__(self, instring: str, loc: int, tokens: pp.ParseResults) -> None: + """ + Parse action to process successfully parsed equations of the form LHS OPERATOR RHS. + + Args: + instring (str): String that was parsed (used in error message). + loc (int): + Location in parsed string where parsing error was logged. + This is not used; we include it as pyparsing injects it alongside `instring` when setting the parse action. + tokens (pp.ParseResults): + Contains a list with an RHS (pp.ParseResults), operator (str), and LHS (pp.ParseResults). + """ + self.lhs, self.op, self.rhs = tokens + self.instring = instring + self.loc = loc + self.values = tokens + + def __repr__(self) -> str: + """Programming / official string representation.""" + return f"{self.lhs.__repr__()} {self.op} {self.rhs.__repr__()}" + + # string return + @overload + def _eval(self, return_type: Literal["math_string"]) -> tuple[str, str]: ... + + # array return + @overload + def _eval( + self, return_type: Literal["array"] + ) -> tuple[xr.DataArray, xr.DataArray]: ... + + def _eval( + self, return_type: RETURN_T + ) -> tuple[str, str] | tuple[xr.DataArray, xr.DataArray]: + """Evaluate the LHS and RHS of the comparison.""" + lhs = self.lhs.eval(return_type, self.eval_attrs) + rhs = self.rhs.eval(return_type, self.eval_attrs) + return lhs, rhs + + def as_math_string(self) -> str: # noqa: D102, override + lhs, rhs = self._eval("math_string") + return lhs + self.OP_TRANSLATOR[self.op] + rhs + + def as_array(self) -> xr.DataArray: # noqa: D102, override + self.eval_attrs = replace(self.eval_attrs, apply_mask=False) + lhs, rhs = self._eval("array") + match self.op: + case "<=": + comparison = lhs <= rhs + case ">=": + comparison = lhs >= rhs + case "<": + comparison = lhs < rhs + case ">": + comparison = lhs > rhs + case "==": + comparison = lhs == rhs + return xr.DataArray(comparison) + + def as_expr( + self, + ) -> tuple[LinearExpression, xr.DataArray, LinearExpression | xr.DataArray]: + lhs, rhs = self._eval("expr") + mask = self.eval_attrs.mask + for side, arr in {"left": lhs, "right": rhs}.items(): + extra_dims = set(arr.dims).difference(set(mask.dims) | {"_term"}) + if extra_dims: + raise self.error_msg( + f"The {side}-hand side of the equation is indexed over dimensions not present in `foreach`: {extra_dims}" + ) + lhs_masked = lhs.where(mask) + rhs_masked = rhs.where(mask) + if isinstance(lhs_masked, Variable): + lhs_masked = lhs_masked.to_linexpr() + if isinstance(rhs_masked, Variable): + rhs_masked = rhs_masked.to_linexpr() + sign_masked = xr.DataArray(self.op).where(mask) + return lhs_masked, sign_masked, rhs_masked + + def eval( + self, return_type: RETURN_T, eval_attrs: EvalAttrs + ) -> tuple[LinearExpression, xr.DataArray, LinearExpression | xr.DataArray] | str: + """ + Evaluate math string expression. + + Args: + return_type (Literal[math_string, input, array]): + Dictates how the expression should be evaluated (see `Returns` section). + eval_attrs (EvalAttrs): Evaluation attributes. + + Returns: + str | list[str | float] | xr.DataArray: + If `math_string` is desired, returns a valid LaTex math string. + If `array` is desired, returns xarray DataArray or a list of strings/numbers (if the expression represents a list). + """ + self.eval_attrs = eval_attrs + evaluated: str | tuple[LinearExpression, xr.DataArray, LinearExpression | xr.DataArray] + if return_type == "array": + evaluated = self.as_array() + elif return_type == "math_string": + evaluated = self.as_math_string() + elif return_type == "expr": + evaluated = self.as_expr() + + return evaluated + class EvalToCallable(EvalString): """Parent class for callable functionality.""" @@ -155,7 +339,7 @@ def eval(self, return_type: RETURN_T, eval_attrs: EvalAttrs) -> Callable: return evaluated -class EvalOperatorOperand(EvalArrayOrMath): +class EvalOperatorOperand(EvalArrayOrMathExpr): """Evaluation of math operands.""" LATEX_OPERATOR_LOOKUP: dict[str, str] = { @@ -187,7 +371,7 @@ def __init__(self, instring: str, loc: int, tokens: pp.ParseResults) -> None: self.instring = instring def __repr__(self) -> str: - """Programming / official string representation.""" + """String representation.""" first_operand = self.value[0].__repr__() operand_operator_pairs = " ".join( op + " " + val.__repr__() @@ -207,7 +391,7 @@ def _operator_operands( except StopIteration: break - def _apply_mask(self, evaluated: xr.DataArray) -> xr.DataArray: + def _apply_mask(self, evaluated: ARRAY_T) -> ARRAY_T: """Util function to apply mask arrays to non-latex strings.""" mask = self.eval_attrs.mask try: @@ -225,10 +409,28 @@ def _skip_component_on_conditional(self, component: str, operator_: str) -> bool """ return component == "0" and operator_ in self.SKIP_IF + @overload @staticmethod def _operate( val: xr.DataArray, evaluated_operand: xr.DataArray, operator_: str - ) -> xr.DataArray: + ) -> xr.DataArray: ... + + @overload + @staticmethod + def _operate( + val: xr.DataArray, evaluated_operand: EXPRVAR_T, operator_: str + ) -> EXPRVAR_T: ... + + @overload + @staticmethod + def _operate( + val: EXPRVAR_T, evaluated_operand: xr.DataArray, operator_: str + ) -> EXPRVAR_T: ... + + @staticmethod + def _operate( + val: xr.DataArray | EXPRVAR_T, evaluated_operand: xr.DataArray | EXPRVAR_T, operator_: str + ) -> xr.DataArray | EXPRVAR_T: """Apply evaluated operation on two DataArrays.""" match operator_: case "**": @@ -271,8 +473,17 @@ def as_array(self) -> xr.DataArray: # noqa: D102, override val = self._operate(val, evaluated_operand, operator_) return val + def as_expr(self) -> EXPR_T: # noqa: D102, override + val = self._apply_mask(self.value[0].eval("expr", self.eval_attrs)) + + for operator_, operand in self._operator_operands(self.value[1:]): + evaluated_operand = self._apply_mask(operand.eval("expr", self.eval_attrs)) + val = self._operate(val, evaluated_operand, operator_) + return val + + -class EvalSignOp(EvalArrayOrMath): +class EvalSignOp(EvalArrayOrMathExpr): """Class for processing expressions with + or -.""" def __init__(self, instring: str, loc: int, tokens: pp.ParseResults) -> None: @@ -303,7 +514,11 @@ def _eval(self, return_type: Literal["math_string"]) -> str: ... @overload def _eval(self, return_type: Literal["array"]) -> xr.DataArray: ... - def _eval(self, return_type: RETURN_T) -> xr.DataArray | str: + # expression return + @overload + def _eval(self, return_type: Literal["expr"]) -> EXPR_T: ... + + def _eval(self, return_type: RETURN_T) -> xr.DataArray | EXPR_T | str: """Evaluate the element that will have the sign attached to it.""" return self.value.eval(return_type, self.eval_attrs) @@ -316,77 +531,14 @@ def as_array(self) -> xr.DataArray: # noqa: D102, override evaluated = -1 * evaluated return evaluated - -class EvalComparisonOp(EvalArrayOrMath): - """Class for processing comparison operations.""" - - OP_TRANSLATOR = {"<=": r" \leq ", ">=": r" \geq ", "==": " = "} - - def __init__(self, instring: str, loc: int, tokens: pp.ParseResults) -> None: - """ - Parse action to process successfully parsed equations of the form LHS OPERATOR RHS. - - Args: - instring (str): String that was parsed (used in error message). - loc (int): - Location in parsed string where parsing error was logged. - This is not used; we include it as pyparsing injects it alongside `instring` when setting the parse action. - tokens (pp.ParseResults): - Contains a list with an RHS (pp.ParseResults), operator (str), and LHS (pp.ParseResults). - """ - self.lhs, self.op, self.rhs = tokens - self.instring = instring - self.loc = loc - self.values = tokens - - def __repr__(self) -> str: - """Programming / official string representation.""" - return f"{self.lhs.__repr__()} {self.op} {self.rhs.__repr__()}" - - # string return - @overload - def _eval(self, return_type: Literal["math_string"]) -> tuple[str, str]: ... - - # array return - @overload - def _eval( - self, return_type: Literal["array"] - ) -> tuple[xr.DataArray, xr.DataArray]: ... - - def _eval( - self, return_type: RETURN_T - ) -> tuple[str, str] | tuple[xr.DataArray, xr.DataArray]: - """Evaluate the LHS and RHS of the comparison.""" - lhs = self.lhs.eval(return_type, self.eval_attrs) - rhs = self.rhs.eval(return_type, self.eval_attrs) - return lhs, rhs - - def as_math_string(self) -> str: # noqa: D102, override - lhs, rhs = self._eval("math_string") - return lhs + self.OP_TRANSLATOR[self.op] + rhs - - def as_array( - self, - ) -> tuple[LinearExpression, xr.DataArray, LinearExpression | xr.DataArray]: # noqa: D102, override: # noqa: D102, override - lhs, rhs = self._eval("array") - mask = self.eval_attrs.mask - for side, arr in {"left": lhs, "right": rhs}.items(): - extra_dims = set(arr.dims).difference(set(mask.dims) | {"_term"}) - if extra_dims: - raise self.error_msg( - f"The {side}-hand side of the equation is indexed over dimensions not present in `foreach`: {extra_dims}" - ) - lhs_masked = lhs.where(mask) - rhs_masked = rhs.where(mask) - if isinstance(lhs_masked, Variable): - lhs_masked = lhs_masked.to_linexpr() - if isinstance(rhs_masked, Variable): - rhs_masked = rhs_masked.to_linexpr() - sign_masked = xr.DataArray(self.op).where(mask) - return lhs_masked, sign_masked, rhs_masked + def as_expr(self) -> EXPR_T: # noqa: D102, override + evaluated = self._eval("expr") + if self.sign == "-": + evaluated = -1 * evaluated + return evaluated -class EvalFunction(EvalArrayOrMath): +class EvalFunction(EvalArrayOrMathExpr): """Class to process parsed functions.""" def __init__(self, instring: str, loc: int, tokens: pp.ParseResults) -> None: @@ -446,7 +598,7 @@ def _eval(self, return_type: Literal["math_string"]) -> str: ... @overload def _eval(self, return_type: Literal["array"]) -> xr.DataArray: ... - def _eval(self, return_type: RETURN_T) -> str | xr.DataArray: + def _eval(self, return_type: RETURN_T) -> str | xr.DataArray | EXPR_T: """Pass evaluated arguments to evaluated helper function.""" helper_function = self.func_name.eval(return_type, self.eval_attrs) if helper_function.ignore_mask: @@ -469,6 +621,9 @@ def as_math_string(self) -> str: # noqa: D102, override def as_array(self) -> xr.DataArray: # noqa: D102, override return self._eval("array") + def as_expr(self) -> EXPR_T: # noqa: D102, override + return self._eval("expr") + class EvalHelperFuncName(EvalToCallable): """For processing parsed helper function names.""" @@ -510,7 +665,7 @@ def as_callable(self, return_type: RETURN_T) -> Callable: return helper_functions[self.name](return_type, self.eval_attrs) -class EvalSlicedComponent(EvalArrayOrMath): +class EvalSlicedComponent(EvalArrayOrMathExpr): """For processing of sliced parameters / decision variables.""" def __init__(self, instring: str, loc: int, tokens: pp.ParseResults) -> None: @@ -543,14 +698,14 @@ def __repr__(self) -> str: return f"SLICED_{self.obj_name}[{slices}]" @staticmethod - def _replace_rule(index_slices): + def _replace_rule(index_slices: dict) -> Callable: """ String parsing rule to catch and replace dimension names with the names + their slices. E.g., `techs` -> `techs=pv`. """ - def __replace(term): + def __replace(term: pp.ParseResults) -> str: if len(term) == 1: return term else: @@ -570,7 +725,10 @@ def _eval(self, return_type: Literal["math_string"]) -> tuple[str, dict]: ... @overload def _eval(self, return_type: Literal["array"]) -> tuple[xr.DataArray, dict]: ... - def _eval(self, return_type: RETURN_T) -> tuple[str | xr.DataArray, dict]: + @overload + def _eval(self, return_type: Literal["expr"]) -> tuple[EXPR_T, dict]: ... + + def _eval(self, return_type: RETURN_T) -> tuple[str | xr.DataArray | EXPR_T, dict]: """Evaluate the slice dim and vals of each slice element.""" slices: dict[str, Any] = { k: xr.concat(slice_, dim=k) @@ -603,6 +761,10 @@ def as_array(self) -> xr.DataArray: # noqa: D102, override evaluated, slices = self._eval("array") return evaluated.sel(**slices) + def as_expr(self) -> EXPR_T: + evaluated, slices = self._eval("expr") + return evaluated.sel(**slices) + class EvalIndexSlice(EvalArrayOrMath): """For processing `$slice` expressions.""" @@ -652,7 +814,8 @@ def as_array(self) -> xr.DataArray | list[xr.DataArray]: # noqa: D102, override return evaluated -class EvalSubExpressions(EvalArrayOrMath): + +class EvalSubExpressions(EvalArrayOrMathExpr): """For processing sub-expressions.""" def __init__(self, instring: str, loc: int, tokens: pp.ParseResults) -> None: @@ -693,6 +856,10 @@ def as_math_string(self) -> str: # noqa: D102, override def as_array(self) -> xr.DataArray: # noqa: D102, override return self._eval("array") + def as_expr(self) -> EXPR_T: + return self._eval("expr") + + class EvalNumber(EvalArrayOrMath): """For processing numbers.""" @@ -766,7 +933,7 @@ def as_array(self) -> list[xr.DataArray]: # noqa: D102, override return values -class EvalUnslicedComponent(EvalArrayOrMath): +class EvalUnslicedComponent(EvalArrayOrMathExpr): """Evaluation of unsliced components.""" def __init__(self, instring: str, loc: int, tokens: pp.ParseResults) -> None: @@ -805,15 +972,16 @@ def as_math_string(self) -> str: # noqa: D102, override return data_var_string def as_array(self) -> xr.DataArray: # noqa: D102, override - if self.eval_attrs.math.find(self.name)._group in ["parameters", "lookups"]: + group = self.eval_attrs.math.find(self.name)._group + if group in ["parameters", "lookups"]: evaluated = self.eval_attrs.input_data[self.name] - elif self.eval_attrs.math.find(self.name)._group == "dimensions": + elif group == "dimensions": try: evaluated = self.eval_attrs.input_data[self.name] except KeyError: evaluated = xr.DataArray(np.nan) else: - evaluated = self.eval_attrs.model[self.name] + evaluated = getattr(self.eval_attrs.model, group)[self.name] if evaluated.isnull().any() and pd.notna( default := self.eval_attrs.math.find(self.name)["default"] ): @@ -822,6 +990,29 @@ def as_array(self) -> xr.DataArray: # noqa: D102, override self.eval_attrs.references.add(self.name) return evaluated + def as_expr(self) -> xr.DataArray | EXPR_T: # noqa: D102, override + group = self.eval_attrs.math.find(self.name)._group + if group in ["parameters", "lookups", "dimensions"]: + try: + evaluated = self.eval_attrs.input_data[self.name] + except KeyError: + evaluated = xr.DataArray(np.nan) + else: + try: + evaluated = getattr(self.eval_attrs.model, group)[self.name] + except KeyError: + evaluated = LinearExpression(xr.DataArray(np.nan), self.eval_attrs.model) + if evaluated.isnull().any() and pd.notna( + default := self.eval_attrs.math.find(self.name)["default"] + ): + if isinstance(evaluated, Variable): + evaluated = evaluated.to_linexpr() + evaluated = evaluated.fillna(default) + + self.eval_attrs.references.add(self.name) + return evaluated + + class EvalGenericString(EvalArrayOrMath): """For generic string parsing.""" @@ -1148,7 +1339,7 @@ def equation_comparison_parser(arithmetic: pp.ParserElement) -> pp.ParserElement pp.ParserElement: Parser for strings of the form "LHS OPERATOR RHS". """ - comparison_operators = pp.one_of(["<=", ">=", "=="]) + comparison_operators = pp.one_of(["<=", ">=", "="]) equation_comparison = arithmetic + comparison_operators + arithmetic equation_comparison.set_parse_action(EvalComparisonOp) @@ -1254,7 +1445,7 @@ def generate_arithmetic_parser(valid_component_names: Iterable) -> pp.ParserElem Args: valid_component_names (Iterable): - Allowed names for optimisation problem components (parameters, decision variables, global_expressions), + Allowed names for optimisation problem components (parameters, decision variables, expressions), to allow the parser to separate these from generic strings. Returns: @@ -1300,7 +1491,7 @@ def generate_equation_parser(valid_component_names: Iterable) -> pp.ParserElemen Args: valid_component_names (Iterable): - Allowed names for optimisation problem components (parameters, decision variables, global_expressions), + Allowed names for optimisation problem components (parameters, decision variables, expressions), to allow the parser to separate these from generic strings. Returns: diff --git a/linopy/declarative/helper_functions.py b/linopy/declarative/helper_functions.py index e8f0c77e9..b876ffb68 100644 --- a/linopy/declarative/helper_functions.py +++ b/linopy/declarative/helper_functions.py @@ -18,6 +18,7 @@ import xarray as xr from linopy.declarative.eval_attrs import EvalAttrs +from linopy.expressions import LinearExpression DTYPE_OPTIONS = { "string": str, @@ -70,7 +71,7 @@ def as_math_string(self, *args: Any, **kwargs: Any) -> str: """ @abstractmethod - def as_array(self, *args: Any, **kwargs: Any) -> xr.DataArray | Expression: + def as_array(self, *args: Any, **kwargs: Any) -> LinearExpression | xr.DataArray: """ Method to apply the helper function to provide an n-dimensional array output. @@ -90,6 +91,8 @@ def __call__(self, *args: Any, **kwargs: Any) -> Any: return self.as_math_string(*args, **kwargs) elif self._return_type == "array": return self.as_array(*args, **kwargs) + elif self._return_type == "expr": + return self.as_array(*args, **kwargs) def __init_subclass__(cls) -> None: """ @@ -769,19 +772,14 @@ def as_array( ``` """ # We can't apply typical xarray rolling window functionality - grouped: dict[str | int, xr.DataArray] = {} grouping_dims = groupby.dims groups = array.stack(_stacked=grouping_dims).groupby( - groupby.stack(_stacked=grouping_dims) + groupby.rename(group_dim.name).stack(_stacked=grouping_dims) ) - for group_name, array_subset in groups: - grouped[group_name] = array_subset.sum("_stacked", min_count=1, skipna=True) - array = xr.concat( - grouped.values(), dim=pd.Index(grouped.keys(), name=group_dim.name) - ) - return array + grouped = groups.sum("_stacked") + return grouped class GroupDatetime(ParsingHelperFunction): diff --git a/linopy/declarative/mask_parser.py b/linopy/declarative/mask_parser.py index 623b67ca0..25042c2ad 100644 --- a/linopy/declarative/mask_parser.py +++ b/linopy/declarative/mask_parser.py @@ -330,7 +330,9 @@ def as_math_string(self) -> str: # noqa: D102, override def as_array(self) -> xr.DataArray: # noqa: D102, override subset = self._eval() - set_item_in_subset = self.set_name.eval("array", self.eval_attrs).isin(subset) + self.eval_attrs = replace(self.eval_attrs, apply_mask=False) + da = self.set_name.eval("array", replace(self.eval_attrs, apply_mask=False)) + set_item_in_subset = da.isin(subset) return set_item_in_subset @@ -483,7 +485,7 @@ def comparison_parser( def subset_parser( - data_var: pp.ParserElement, *subset_items: pp.ParserElement + data_vars: list[pp.ParserElement], *subset_items: pp.ParserElement ) -> pp.ParserElement: """ Parsing grammar to process subsets. @@ -503,7 +505,7 @@ def subset_parser( + pp.Suppress(pp.White(" ", min=1)) + pp.Suppress("in") + pp.Suppress(pp.White(" ", min=1)) - + data_var + + pp.MatchFirst(data_vars) ) subset_expression.set_parse_action(SubsetParser) @@ -568,7 +570,7 @@ def generate_mask_string_parser( number, unique_evaluatable_string, dimensions_parser ) subset = subset_parser( - dimensions_parser, config_option, number, general_evaluatable_string + [dimensions_parser, inputs_parser], config_option, number, general_evaluatable_string ) arithmetic = pp.Forward() diff --git a/linopy/declarative/parsing.py b/linopy/declarative/parsing.py index e1c7760e7..00246fdc4 100644 --- a/linopy/declarative/parsing.py +++ b/linopy/declarative/parsing.py @@ -21,7 +21,7 @@ mask_parser, ) from linopy.declarative.schema import MATH_DEFS_T, ConfigModel, MathModel, _Equations - +from linopy.expressions import LinearExpression if TYPE_CHECKING: from linopy.model import Model TRUE_ARRAY = xr.DataArray(True) @@ -294,10 +294,10 @@ def evaluate_expression( model: Model, math: MathModel, *, - return_type: Literal["array"] = "array", + return_type: Literal["expr"] = "expr", references: set | None = None, mask: xr.DataArray = TRUE_ARRAY, - ) -> xr.DataArray: ... + ) -> LinearExpression: ... # Expecting string if requesting latex string. @overload @@ -317,10 +317,10 @@ def evaluate_expression( model: Model, math: MathModel, *, - return_type: Literal["array", "math_string"] = "array", + return_type: Literal["expr", "math_string"] = "expr", references: set | None = None, mask: xr.DataArray = TRUE_ARRAY, - ) -> xr.DataArray | str: + ) -> LinearExpression | str: """ Evaluate a math string to produce an array backend objects or a LaTex math string. @@ -416,7 +416,7 @@ class ParsedBackendComponent(ParsedBackendEquation): _ERR_STRING_ORDER: list[str] = ["expression_group", "id", "expr_or_mask"] PARSERS: dict[str, Callable] = { "constraints": expression_parser.generate_equation_parser, - "global_expressions": expression_parser.generate_arithmetic_parser, + "expressions": expression_parser.generate_arithmetic_parser, "postprocessed": expression_parser.generate_arithmetic_parser, "objectives": expression_parser.generate_arithmetic_parser, "piecewise_constraints": expression_parser.generate_arithmetic_parser, @@ -426,7 +426,7 @@ def __init__( self, group: Literal[ "variables", - "global_expressions", + "expressions", "constraints", "piecewise_constraints", "objectives", @@ -443,7 +443,7 @@ def __init__( objects that can be passed onto a solver interface like Pyomo or Gurobipy. Args: - group (Literal["variables", "global_expressions", "constraints", "objectives"]): + group (Literal["variables", "expressions", "constraints", "objectives"]): Optimisation problem component group to which the unparsed data belongs. name (str): Name of the optimisation problem component unparsed_data (T): Unparsed math formulation. Expected structure depends on diff --git a/linopy/declarative/schema.py b/linopy/declarative/schema.py index 14edbf99a..05c8a44a5 100644 --- a/linopy/declarative/schema.py +++ b/linopy/declarative/schema.py @@ -12,6 +12,7 @@ from pydantic_core import PydanticCustomError LOGGER = logging.getLogger(__name__) +LOGGER.setLevel(logging.INFO) # == # Modified from https://github.com/pydantic/pydantic-core/pull/820#issuecomment-1670475909 T = TypeVar("T", bound=Hashable | list) @@ -21,7 +22,7 @@ "parameters", "lookups", "variables", - "global_expressions", + "expressions", "constraints", "piecewise_constraints", "objectives", @@ -281,8 +282,6 @@ class ParameterDef(_MathComponent): default: float | int = float("nan") """The default value for the parameter, if not set in the data.""" - resample_method: Literal["mean", "sum", "first"] = "first" - """If resampling is applied over any of the parameter's dimensions, the method to use to aggregate the data.""" unit: str = "" """The unit of the parameter, e.g. 'kW', 'm', 'kg', 'energy', 'power', ...""" @@ -301,15 +300,8 @@ class LookupDef(_MathComponent): """The default value for the lookup, if not set in the data.""" dtype: Literal["float", "string", "bool", "datetime", "date"] = "string" """The lookup data type.""" - resample_method: Literal["mean", "sum", "first"] = "first" - """If resampling is applied over any of the lookup's dimensions, the method to use to aggregate the data.""" one_of: list | None = None """If given, the lookup values must be one of these items.""" - pivot_values_to_dim: str | None = None - """If given, the lookup will be pivoted such that its values become the index of a new dimension and its new values are boolean, True where the index values match the old values. - For instance, if the lookup starts out indexed over `techs` with values of `[electricity, gas]` and `pivot_values_to_dim: carriers`, - then the lookup will be converted to a boolean array with the dimensions ['techs', 'carriers']. - """ _group: ClassVar[COMPONENTS_T] = "lookups" @@ -397,15 +389,15 @@ def slices(self) -> _SubExpressions: _group: ClassVar[COMPONENTS_T] = "piecewise_constraints" -class LinearExpressionDef(_MathIndexedComponent, _MathEquationComponent): +class ExpressionDef(_MathIndexedComponent, _MathEquationComponent): """ - Schema for named global expressions. + 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 global expressions). + and other expressions). - NOTE: If expecting to use global expression `A` in global expression `B`, `A` must + NOTE: If expecting to use expression `A` in expression `B`, `A` must be defined above `B`. """ @@ -414,15 +406,15 @@ class LinearExpressionDef(_MathIndexedComponent, _MathEquationComponent): default: NumericVal = float("nan") """If set, will be the default value for the expression.""" equations: _Equations = _Equations() - """Global expression math equations.""" + """Expression math equations.""" sub_expressions: _SubExpressions = _SubExpressions() - """Global expression named sub-expressions.""" + """Expression named sub-expressions.""" slices: _SubExpressions = _SubExpressions() - """Global expression named index slices.""" + """Expression named index slices.""" order: int = 0 - """Order in which to apply this global expression relative to all others, if different to its definition order.""" + """Order in which to apply this expression relative to all others, if different to its definition order.""" - _group: ClassVar[COMPONENTS_T] = "global_expressions" + _group: ClassVar[COMPONENTS_T] = "expressions" class _Bounds(LinopyBaseModel): @@ -499,11 +491,11 @@ def mask(self) -> str: _group: ClassVar[COMPONENTS_T] = "objectives" -class PostprocessedExpressionDef(LinearExpressionDef): +class PostprocessedExpressionDef(ExpressionDef): """ Schema for postprocessed expressions. - Can be used to combine parameters, variables, and global expressions into a single expression solving the model. + 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`. @@ -549,10 +541,10 @@ class VariableDefs(LinopyDictModel): root: dict[AttrStr, VariableDef] = Field(default_factory=dict) -class LinearExpressionDefs(LinopyDictModel): - """Linopy model global_expressions dictionary.""" +class ExpressionDefs(LinopyDictModel): + """Linopy model expressions dictionary.""" - root: dict[AttrStr, LinearExpressionDef] = Field(default_factory=dict) + root: dict[AttrStr, ExpressionDef] = Field(default_factory=dict) class ConstraintDefs(LinopyDictModel): @@ -604,8 +596,8 @@ class MathModel(LinopyBaseModel): """All lookups to include in the optimisation problem.""" variables: VariableDefs = VariableDefs() """All decision variables to include in the optimisation problem.""" - global_expressions: LinearExpressionDefs = LinearExpressionDefs() - """All global expressions that can be applied to 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() @@ -618,7 +610,7 @@ class MathModel(LinopyBaseModel): """Checks to apply before building the optimisation problem.""" @model_validator(mode="after") - def unique_component_names(self): + def unique_component_names(self) -> Self: """Ensure all component names are unique.""" groups = sorted( ( @@ -651,7 +643,7 @@ def parsing_components(self) -> dict[str, dict[str, set[str]]]: parsing_components = { "dimensions": ["dimensions"], "inputs": ["lookups", "parameters"], - "results": ["variables", "global_expressions"], + "results": ["variables", "expressions"], } def _names(): @@ -690,7 +682,7 @@ def find( MATH_DEFS_T = ( ConstraintDef | VariableDef - | LinearExpressionDef + | ExpressionDef | ObjectiveDef | PiecewiseConstraintDef ) diff --git a/linopy/expressions.py b/linopy/expressions.py index 318682341..84fa17319 100644 --- a/linopy/expressions.py +++ b/linopy/expressions.py @@ -14,7 +14,7 @@ from collections.abc import Callable, Hashable, Iterator, Mapping, Sequence from dataclasses import dataclass, field from itertools import product, zip_longest -from typing import TYPE_CHECKING, Any, TypeAlias, TypeVar, cast, overload +from typing import TYPE_CHECKING, Any, TypeAlias, TypeVar, cast, overload, ItemsView from warnings import warn import numpy as np @@ -56,6 +56,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, @@ -65,6 +66,7 @@ is_constant, iterate_slices, maybe_group_terms_polars, + save_join, to_dataframe, to_polars, ) @@ -387,6 +389,7 @@ def __init__(self, data: Dataset | Any | None, model: Model) -> None: if not isinstance(model, Model): raise ValueError("model must be an instance of linopy.Model") + data = data.assign_attrs(name=None) self._model = model self._data = cast(Dataset, data) @@ -897,6 +900,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 @@ -2454,6 +2464,125 @@ 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/model.py b/linopy/model.py index 48a8200b0..9ff943668 100644 --- a/linopy/model.py +++ b/linopy/model.py @@ -56,6 +56,7 @@ CSRConstraint, ) from linopy.expressions import ( + Expressions, LinearExpression, QuadraticExpression, ScalarLinearExpression, @@ -197,6 +198,7 @@ class Model: _solver: solvers.Solver | None _variables: Variables + _expressions: Expressions _constraints: Constraints _objective: Objective _parameters: Dataset @@ -207,6 +209,7 @@ class Model: _xCounter: int _cCounter: int _varnameCounter: int + _exprnameCounter: int _connameCounter: int _pwlCounter: int _blocks: DataArray | None @@ -218,6 +221,7 @@ class Model: __slots__ = ( # containers "_variables", + "_expressions", "_constraints", "_objective", "_parameters", @@ -230,6 +234,7 @@ class Model: "_xCounter", "_cCounter", "_varnameCounter", + "_exprnameCounter", "_connameCounter", "_pwlCounter", "_blocks", @@ -290,6 +295,7 @@ def __init__( linopy.Model """ 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() @@ -299,6 +305,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 @@ -358,6 +365,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: """ @@ -563,6 +577,7 @@ def scalar_attrs(self) -> list[str]: "_xCounter", "_cCounter", "_varnameCounter", + "_exprnameCounter", "_connameCounter", "_pwlCounter", "force_dim_names", @@ -581,11 +596,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}" @@ -820,6 +837,81 @@ 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, From 9d23afdf3f31911196247aa4cf08615e340dda9a Mon Sep 17 00:00:00 2001 From: Bryn Pickering <17178478+brynpickering@users.noreply.github.com> Date: Fri, 17 Jul 2026 17:19:16 +0100 Subject: [PATCH 04/12] [WIP] parser route separation (raw/expr/math_string) + declarative parsing tests Snapshot of in-progress work before a structural rewrite of linopy/declarative. Co-Authored-By: Claude Fable 5 --- linopy/declarative/build.py | 66 ++--- linopy/declarative/expression_parser.py | 334 +++++++++++------------- linopy/declarative/helper_functions.py | 55 ++-- linopy/declarative/mask_parser.py | 59 ++--- linopy/declarative/parsing.py | 163 ++++++++++-- test/test_declarative_parsing.py | 333 +++++++++++++++++++++++ 6 files changed, 711 insertions(+), 299 deletions(-) create mode 100644 test/test_declarative_parsing.py diff --git a/linopy/declarative/build.py b/linopy/declarative/build.py index cc44e4e01..c37860bc6 100644 --- a/linopy/declarative/build.py +++ b/linopy/declarative/build.py @@ -2,26 +2,24 @@ import time import typing -from requests import exceptions -from tqdm.asyncio import tqdm +import numpy as np import xarray as xr +from tqdm.asyncio import tqdm -from linopy.expressions import merge from linopy.declarative import eval_attrs, helper_functions, parsing from linopy.declarative.schema import ( LOGGER, ConfigModel, ConstraintDef, + ExpressionDef, MathModel, ObjectiveDef, VariableDef, - ExpressionDef, ) -from linopy.expressions import LinearExpression -from linopy.io import TQDM_COLOR +from linopy.expressions import LinearExpression, merge from linopy.io import TQDM_COLOR from linopy.model import Model -import numpy as np + ORDERED_COMPONENTS_T = typing.Literal[ "variables", "expressions", @@ -59,7 +57,8 @@ def __init__(self, math_def: dict, input_data: xr.Dataset, config: dict): self._check_inputs() def _update_dtypes(self, ds: xr.Dataset, id_: str = "") -> xr.Dataset: - """Update data types of coordinates or data variables in the dataset. + """ + Update data types of coordinates or data variables in the dataset. Args: ds (xr.Dataset): Dataset to update. @@ -92,7 +91,11 @@ def _update_dtypes(self, ds: xr.Dataset, id_: str = "") -> xr.Dataset: ) match dtype_str: case "string": - updated_var = var_data.astype(dtype).where(var_data.notnull()).where(var_data != "") + updated_var = ( + var_data.astype(dtype) + .where(var_data.notnull()) + .where(var_data != "") + ) case "datetime": updated_var = time._datetime_index( var_data.to_series(), self.config.datetime_format @@ -122,7 +125,7 @@ def _sorted_by_order( def _check_inputs(self) -> None: data_checks = self.math.checks - check_results = {"raise": [], "warn": []} + check_results: dict[str, list[str]] = {"raise": [], "warn": []} parser_ = parsing.mask_parser.generate_mask_string_parser( **self.math.parsing_components["mask"] ) @@ -131,18 +134,21 @@ def _check_inputs(self) -> None: "math": self.math, "input_data": self.input_data, "config": self.config, - "helper_functions": helper_functions._registry["mask"], + "helper_functions": helper_functions._registry["mask"], } + active = self.input_data.get("active", xr.DataArray(True)) for name, check in data_checks.root.items(): if check.active: parsed_ = parser_.parse_string(check.mask, parse_all=True) eval_attrs_ = eval_attrs.EvalAttrs(equation_name=name, **eval_kwargs) - evaluated = parsed_[0].eval("array", eval_attrs_) - if evaluated.any() and (evaluated & self.input_data.active).any(): + evaluated = parsed_[0].eval("raw", eval_attrs_) + if (evaluated & active).any(): check_results[check.errors].append(check.message) print_warnings_and_raise_errors( - check_results["warn"], check_results["raise"], during="model input data checks" + check_results["warn"], + check_results["raise"], + during="model input data checks", ) def add_variable(self, name: str, definition: VariableDef) -> None: @@ -215,10 +221,7 @@ def add_expression(self, name: str, definition: ExpressionDef) -> None: mask=sub_mask, references=references, ) - if isinstance(expr_to_fill, xr.DataArray): - expr = expr.fillna(expr_to_fill) - else: - expr = merge([expr, expr_to_fill.where(sub_mask)]) + expr = merge([expr, expr_to_fill.where(sub_mask)]) if not expr.isnull().all(): self.model.add_expressions(name=name, data=expr, mask=all_mask) self.model.expressions[name].attrs["references"] = references @@ -250,7 +253,9 @@ def add_constraint(self, name: str, definition: ConstraintDef) -> None: rhs = LinearExpression(float("nan"), self.model).where(mask) all_mask = mask.copy() if not mask.any(): - LOGGER.info(f"constraints:{name} | No valid data points after applying mask. Constraint not added to model.") + LOGGER.info( + f"constraints:{name} | No valid data points after applying mask. Constraint not added to model." + ) return None equations = parsed_component.parse_equations() @@ -264,7 +269,9 @@ def add_constraint(self, name: str, definition: ConstraintDef) -> None: references=references, ) if not sub_mask.any(): - LOGGER.info(f"constraints:{equation.name} | No valid data points after applying mask. Constraint not added to model.") + LOGGER.info( + f"constraints:{equation.name} | No valid data points after applying mask. Constraint not added to model." + ) continue sub_mask = parsed_component.drop_dims_not_in_foreach(sub_mask) if (sign.notnull() & sub_mask).any(): @@ -273,25 +280,21 @@ def add_constraint(self, name: str, definition: ConstraintDef) -> None: "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 = equation.evaluate_expression( + lhs_to_fill, sign_to_fill, rhs_to_fill = equation.evaluate_equation( self.input_data, self.model, self.math, mask=sub_mask, references=references, ) - if isinstance(lhs_to_fill, xr.DataArray): - lhs_to_fill = LinearExpression(lhs_to_fill, self.model) lhs = merge([lhs, lhs_to_fill]) - - if isinstance(rhs_to_fill, xr.DataArray): - rhs_to_fill = LinearExpression(rhs_to_fill, self.model) rhs = merge([rhs, rhs_to_fill]) - sign = sign.fillna(sign_to_fill) if sign.isnull().all(): - LOGGER.info(f"constraints:{name} | No valid data points after applying mask. Constraint not added to model.") + LOGGER.info( + f"constraints:{name} | No valid data points after applying mask. Constraint not added to model." + ) return None self.model.add_constraints( @@ -364,9 +367,7 @@ def build(self) -> Model: start = time.time() getattr(self, f"add_{component}")(name, definition) end = time.time() - start - LOGGER.debug( - f"{components}:{name} | Built in {end:.4f}s" - ) + LOGGER.debug(f"{components}:{name} | Built in {end:.4f}s") LOGGER.info(f"{components} | Generated.") return self.model @@ -377,7 +378,8 @@ def print_warnings_and_raise_errors( during: str = "model processing", bullet: str = " * ", ) -> None: - """Process collections of warnings/errors. + """ + Process collections of warnings/errors. Prints warnings / raises errors with a bullet point list of the concatenated collections. diff --git a/linopy/declarative/expression_parser.py b/linopy/declarative/expression_parser.py index 7891d0641..eb0885af9 100644 --- a/linopy/declarative/expression_parser.py +++ b/linopy/declarative/expression_parser.py @@ -35,7 +35,7 @@ from abc import ABC, abstractmethod from collections.abc import Callable, Iterable, Iterator from dataclasses import replace -from typing import Any, Literal, overload, TypeVar +from typing import Any, Literal, TypeVar, overload import numpy as np import pandas as pd @@ -54,7 +54,7 @@ EXPR_T = LinearExpression | QuadraticExpression EXPRVAR_T = TypeVar("EXPRVAR_T", LinearExpression, QuadraticExpression) ARRAY_T = TypeVar("ARRAY_T", LinearExpression, QuadraticExpression, xr.DataArray) -RETURN_T = Literal["expr", "array", "math_string"] +RETURN_T = Literal["expr", "raw", "math_string"] class EvalString(ABC): @@ -79,90 +79,57 @@ def error_msg(self, message: str) -> ValueError: ) -class EvalArrayOrMath(EvalString): - """Abstract class to evaluate expressions as either arrays or math strings.""" - - @abstractmethod - def as_math_string(self) -> str: - """Evaluate and return expression as LaTeX.""" - - @abstractmethod - def as_array(self) -> xr.DataArray | list[xr.DataArray]: - """ - Evaluate and return expression as a DataArray or list. - - If the evaluated expression returns a simple string or number, - this value will be assigned as both the `name` and the data of the returned DataArray. - The purpose of this is to be able to access the string/number value whether we query the array name or its data. - """ - - def as_expr(self) -> xr.DataArray | list[xr.DataArray]: - """ - Evaluate and return expression as a LinearExpression or QuadraticExpression. - """ - return self.as_array() - - # Math strings evaluate to strings. - @overload - def eval( - self, return_type: Literal["math_string"], eval_attrs: EvalAttrs - ) -> str: ... - - # Arrays evaluate to arrays - @overload - def eval( - self, return_type: Literal["array"], eval_attrs: EvalAttrs - ) -> xr.DataArray | list[xr.DataArray]: ... +def _to_linexpr(obj: Any) -> Any: + """ + Normalise a model object to a linopy expression. - def eval( - self, return_type: RETURN_T, eval_attrs: EvalAttrs - ) -> str | xr.DataArray | list[xr.DataArray] | EXPR_T: - """ - Evaluate math string expression. + ``Variable`` objects are converted to ``LinearExpression`` via ``to_linexpr``; + ``LinearExpression``/``QuadraticExpression``/``xr.DataArray`` objects are + returned unchanged. This is the single place where the ``Variable`` -> + ``LinearExpression`` coercion is performed on the expression route. + """ + if isinstance(obj, Variable): + return obj.to_linexpr() + return obj - Args: - return_type (Literal[math_string, input, array]): - Dictates how the expression should be evaluated (see `Returns` section). - eval_attrs (EvalAttrs): Evaluation attributes. - Returns: - str | list[str | float] | xr.DataArray: - If `math_string` is desired, returns a valid LaTex math string. - If `array` is desired, returns xarray DataArray or a list of strings/numbers (if the expression represents a list). - """ - self.eval_attrs = eval_attrs - evaluated: str | list[str | float] | xr.DataArray | list[xr.DataArray] | EXPR_T - if return_type == "array": - evaluated = self.as_array() - elif return_type == "math_string": - evaluated = self.as_math_string() - elif return_type == "expr": - evaluated = self.as_expr() +class EvalNode(EvalString): + """ + Base class for nodes evaluated as math strings, raw data, or expressions. - return evaluated + Three evaluation modes are supported (see :meth:`eval`): -class EvalArrayOrMathExpr(EvalString): - """Abstract class to evaluate expressions as either arrays or math strings.""" + - ``math_string``: a LaTeX string (:meth:`as_math_string`). + - ``raw``: the underlying data without any route-specific transformation + (:meth:`as_raw`) - an ``xr.DataArray`` for parameters/lookups/dimensions and + the raw model object (``Variable``/``LinearExpression``) for model entries. + - ``expr``: a linopy expression suitable for arithmetic composition + (:meth:`as_expr`). The default implementation returns the raw data; nodes that + compose arithmetic override it to guarantee an expression is returned. + """ @abstractmethod def as_math_string(self) -> str: """Evaluate and return expression as LaTeX.""" @abstractmethod - def as_array(self) -> xr.DataArray | list[xr.DataArray]: + def as_raw(self) -> xr.DataArray | list[xr.DataArray]: """ - Evaluate and return expression as a DataArray or list. + Evaluate and return the underlying data without route-specific transformation. If the evaluated expression returns a simple string or number, this value will be assigned as both the `name` and the data of the returned DataArray. The purpose of this is to be able to access the string/number value whether we query the array name or its data. """ - @abstractmethod - def as_expr(self) -> EXPR_T: + def as_expr(self) -> Any: """ Evaluate and return expression as a LinearExpression or QuadraticExpression. + + The default implementation returns the raw data (:meth:`as_raw`); nodes that + compose arithmetic (operands, signs, functions, (sub-)components) override it. """ + return self.as_raw() # Math strings evaluate to strings. @overload @@ -170,41 +137,38 @@ def eval( self, return_type: Literal["math_string"], eval_attrs: EvalAttrs ) -> str: ... - # Arrays evaluate to arrays + # Raw evaluation returns the underlying data. @overload def eval( - self, return_type: Literal["array"], eval_attrs: EvalAttrs + self, return_type: Literal["raw"], eval_attrs: EvalAttrs ) -> xr.DataArray | list[xr.DataArray]: ... - def eval( - self, return_type: RETURN_T, eval_attrs: EvalAttrs - ) -> str | xr.DataArray | list[xr.DataArray] | EXPR_T: + def eval(self, return_type: RETURN_T, eval_attrs: EvalAttrs) -> Any: """ - Evaluate math string expression. + Evaluate a parsed expression node. Args: - return_type (Literal[math_string, input, array]): + return_type (Literal["math_string", "raw", "expr"]): Dictates how the expression should be evaluated (see `Returns` section). eval_attrs (EvalAttrs): Evaluation attributes. Returns: - str | list[str | float] | xr.DataArray: - If `math_string` is desired, returns a valid LaTex math string. - If `array` is desired, returns xarray DataArray or a list of strings/numbers (if the expression represents a list). + If `math_string`, a valid LaTeX math string. + If `raw`, the underlying data (`xr.DataArray`, list, or raw model object). + If `expr`, a linopy `LinearExpression`/`QuadraticExpression`. """ self.eval_attrs = eval_attrs - evaluated: str | list[str | float] | xr.DataArray | list[xr.DataArray] | EXPR_T - if return_type == "array": - evaluated = self.as_array() + evaluated: Any + if return_type == "raw": + evaluated = self.as_raw() elif return_type == "math_string": evaluated = self.as_math_string() elif return_type == "expr": evaluated = self.as_expr() - return evaluated -class EvalComparisonOp(EvalString): +class EvalComparisonOp(EvalNode): """Class for processing comparison operations.""" OP_TRANSLATOR = {"<=": r" \leq ", ">=": r" \geq ", "==": " = ", "=": " = "} @@ -234,15 +198,17 @@ def __repr__(self) -> str: @overload def _eval(self, return_type: Literal["math_string"]) -> tuple[str, str]: ... - # array return + # raw return @overload def _eval( - self, return_type: Literal["array"] + self, return_type: Literal["raw"] ) -> tuple[xr.DataArray, xr.DataArray]: ... - def _eval( - self, return_type: RETURN_T - ) -> tuple[str, str] | tuple[xr.DataArray, xr.DataArray]: + # expression return + @overload + def _eval(self, return_type: Literal["expr"]) -> tuple[Any, Any]: ... + + def _eval(self, return_type: RETURN_T) -> tuple[Any, Any]: """Evaluate the LHS and RHS of the comparison.""" lhs = self.lhs.eval(return_type, self.eval_attrs) rhs = self.rhs.eval(return_type, self.eval_attrs) @@ -252,9 +218,9 @@ def as_math_string(self) -> str: # noqa: D102, override lhs, rhs = self._eval("math_string") return lhs + self.OP_TRANSLATOR[self.op] + rhs - def as_array(self) -> xr.DataArray: # noqa: D102, override + def as_raw(self) -> xr.DataArray: # noqa: D102, override self.eval_attrs = replace(self.eval_attrs, apply_mask=False) - lhs, rhs = self._eval("array") + lhs, rhs = self._eval("raw") match self.op: case "<=": comparison = lhs <= rhs @@ -271,6 +237,7 @@ def as_array(self) -> xr.DataArray: # noqa: D102, override def as_expr( self, ) -> tuple[LinearExpression, xr.DataArray, LinearExpression | xr.DataArray]: + """Evaluate the comparison as a ``(lhs, sign, rhs)`` tuple for constraint assembly.""" lhs, rhs = self._eval("expr") mask = self.eval_attrs.mask for side, arr in {"left": lhs, "right": rhs}.items(): @@ -279,41 +246,11 @@ def as_expr( raise self.error_msg( f"The {side}-hand side of the equation is indexed over dimensions not present in `foreach`: {extra_dims}" ) - lhs_masked = lhs.where(mask) - rhs_masked = rhs.where(mask) - if isinstance(lhs_masked, Variable): - lhs_masked = lhs_masked.to_linexpr() - if isinstance(rhs_masked, Variable): - rhs_masked = rhs_masked.to_linexpr() + lhs_masked = _to_linexpr(lhs.where(mask)) + rhs_masked = _to_linexpr(rhs.where(mask)) sign_masked = xr.DataArray(self.op).where(mask) return lhs_masked, sign_masked, rhs_masked - def eval( - self, return_type: RETURN_T, eval_attrs: EvalAttrs - ) -> tuple[LinearExpression, xr.DataArray, LinearExpression | xr.DataArray] | str: - """ - Evaluate math string expression. - - Args: - return_type (Literal[math_string, input, array]): - Dictates how the expression should be evaluated (see `Returns` section). - eval_attrs (EvalAttrs): Evaluation attributes. - - Returns: - str | list[str | float] | xr.DataArray: - If `math_string` is desired, returns a valid LaTex math string. - If `array` is desired, returns xarray DataArray or a list of strings/numbers (if the expression represents a list). - """ - self.eval_attrs = eval_attrs - evaluated: str | tuple[LinearExpression, xr.DataArray, LinearExpression | xr.DataArray] - if return_type == "array": - evaluated = self.as_array() - elif return_type == "math_string": - evaluated = self.as_math_string() - elif return_type == "expr": - evaluated = self.as_expr() - - return evaluated class EvalToCallable(EvalString): """Parent class for callable functionality.""" @@ -339,7 +276,7 @@ def eval(self, return_type: RETURN_T, eval_attrs: EvalAttrs) -> Callable: return evaluated -class EvalOperatorOperand(EvalArrayOrMathExpr): +class EvalOperatorOperand(EvalNode): """Evaluation of math operands.""" LATEX_OPERATOR_LOOKUP: dict[str, str] = { @@ -429,7 +366,9 @@ def _operate( @staticmethod def _operate( - val: xr.DataArray | EXPRVAR_T, evaluated_operand: xr.DataArray | EXPRVAR_T, operator_: str + val: xr.DataArray | EXPRVAR_T, + evaluated_operand: xr.DataArray | EXPRVAR_T, + operator_: str, ) -> xr.DataArray | EXPRVAR_T: """Apply evaluated operation on two DataArrays.""" match operator_: @@ -465,11 +404,11 @@ def as_math_string(self) -> str: # noqa: D102, override ) return val - def as_array(self) -> xr.DataArray: # noqa: D102, override - val = self._apply_mask(self.value[0].eval("array", self.eval_attrs)) + def as_raw(self) -> xr.DataArray: # noqa: D102, override + val = self._apply_mask(self.value[0].eval("raw", self.eval_attrs)) for operator_, operand in self._operator_operands(self.value[1:]): - evaluated_operand = self._apply_mask(operand.eval("array", self.eval_attrs)) + evaluated_operand = self._apply_mask(operand.eval("raw", self.eval_attrs)) val = self._operate(val, evaluated_operand, operator_) return val @@ -482,8 +421,7 @@ def as_expr(self) -> EXPR_T: # noqa: D102, override return val - -class EvalSignOp(EvalArrayOrMathExpr): +class EvalSignOp(EvalNode): """Class for processing expressions with + or -.""" def __init__(self, instring: str, loc: int, tokens: pp.ParseResults) -> None: @@ -512,7 +450,7 @@ def _eval(self, return_type: Literal["math_string"]) -> str: ... # array return @overload - def _eval(self, return_type: Literal["array"]) -> xr.DataArray: ... + def _eval(self, return_type: Literal["raw"]) -> xr.DataArray: ... # expression return @overload @@ -525,8 +463,8 @@ def _eval(self, return_type: RETURN_T) -> xr.DataArray | EXPR_T | str: def as_math_string(self) -> str: # noqa: D102 return self.sign + self._eval("math_string") - def as_array(self) -> xr.DataArray: # noqa: D102, override - evaluated = self._eval("array") + def as_raw(self) -> xr.DataArray: # noqa: D102, override + evaluated = self._eval("raw") if self.sign == "-": evaluated = -1 * evaluated return evaluated @@ -538,7 +476,7 @@ def as_expr(self) -> EXPR_T: # noqa: D102, override return evaluated -class EvalFunction(EvalArrayOrMathExpr): +class EvalFunction(EvalNode): """Class to process parsed functions.""" def __init__(self, instring: str, loc: int, tokens: pp.ParseResults) -> None: @@ -573,7 +511,7 @@ def _arg_eval(self, return_type: Literal["math_string"], arg: Any) -> str: ... @overload def _arg_eval( - self, return_type: Literal["array"], arg: Any + self, return_type: Literal["raw"], arg: Any ) -> xr.DataArray | list[str | float]: ... def _arg_eval( @@ -585,7 +523,7 @@ def _arg_eval( elif isinstance(arg, list): evaluated = [self._arg_eval(return_type, arg_) for arg_ in arg] elif isinstance(arg, ListParser): - evaluated = arg.eval("array", self.eval_attrs) + evaluated = arg.eval("raw", self.eval_attrs) else: evaluated = arg.eval(return_type, self.eval_attrs) if isinstance(evaluated, xr.DataArray) and isinstance(arg, EvalGenericString): @@ -596,21 +534,33 @@ def _arg_eval( def _eval(self, return_type: Literal["math_string"]) -> str: ... @overload - def _eval(self, return_type: Literal["array"]) -> xr.DataArray: ... + def _eval(self, return_type: Literal["raw"]) -> xr.DataArray: ... def _eval(self, return_type: RETURN_T) -> str | xr.DataArray | EXPR_T: - """Pass evaluated arguments to evaluated helper function.""" + """ + Pass evaluated arguments to evaluated helper function. + + The helper function itself is created with the enclosing ``return_type`` so + that expression-route helpers can dispatch to their ``as_expr`` implementation. + Its arguments, however, are always evaluated in ``raw`` mode (never ``expr``): + helper functions must receive un-normalised inputs (``xr.DataArray`` for + parameters/lookups/dimensions and the raw model object for variables/expressions) + rather than values that have been coerced to boolean masks or ``LinearExpression``. + """ helper_function = self.func_name.eval(return_type, self.eval_attrs) if helper_function.ignore_mask: self.eval_attrs = replace(self.eval_attrs, mask=xr.DataArray(True)) + arg_return_type: RETURN_T = ( + "math_string" if return_type == "math_string" else "raw" + ) args_ = [] for arg in self.args: - args_.append(self._arg_eval(return_type, arg)) + args_.append(self._arg_eval(arg_return_type, arg)) kwargs_ = {} for kwarg_name, kwarg_val in self.kwargs.items(): - kwargs_[kwarg_name] = self._arg_eval(return_type, kwarg_val) + kwargs_[kwarg_name] = self._arg_eval(arg_return_type, kwarg_val) evaluated = helper_function(*args_, **kwargs_) return evaluated @@ -618,8 +568,8 @@ def _eval(self, return_type: RETURN_T) -> str | xr.DataArray | EXPR_T: def as_math_string(self) -> str: # noqa: D102, override return self._eval("math_string") - def as_array(self) -> xr.DataArray: # noqa: D102, override - return self._eval("array") + def as_raw(self) -> xr.DataArray: # noqa: D102, override + return self._eval("raw") def as_expr(self) -> EXPR_T: # noqa: D102, override return self._eval("expr") @@ -657,15 +607,15 @@ def as_callable(self, return_type: RETURN_T) -> Callable: helper_functions = self.eval_attrs.helper_functions if self.name not in helper_functions.keys(): raise self.error_msg(f"Invalid helper function defined: {self.name}") - elif not isinstance(helper_functions[self.name], type(ParsingHelperFunction)): + elif not issubclass(helper_functions[self.name], ParsingHelperFunction): raise self.error_msg( - f"Helper function must be subclassed from calliope.backend.helper_functions.ParsingHelperFunction: {self.name}" + f"Helper function must be subclassed from linopy.declarative.helper_functions.ParsingHelperFunction: {self.name}" ) else: return helper_functions[self.name](return_type, self.eval_attrs) -class EvalSlicedComponent(EvalArrayOrMathExpr): +class EvalSlicedComponent(EvalNode): """For processing of sliced parameters / decision variables.""" def __init__(self, instring: str, loc: int, tokens: pp.ParseResults) -> None: @@ -723,7 +673,7 @@ def __replace(term: pp.ParseResults) -> str: def _eval(self, return_type: Literal["math_string"]) -> tuple[str, dict]: ... @overload - def _eval(self, return_type: Literal["array"]) -> tuple[xr.DataArray, dict]: ... + def _eval(self, return_type: Literal["raw"]) -> tuple[xr.DataArray, dict]: ... @overload def _eval(self, return_type: Literal["expr"]) -> tuple[EXPR_T, dict]: ... @@ -757,8 +707,8 @@ def as_math_string(self) -> str: # noqa: D102, override obj_parser.set_parse_action(self._replace_rule(singular_slice_refs)) return obj_parser.parse_string(evaluated, parse_all=True)[0] - def as_array(self) -> xr.DataArray: # noqa: D102, override - evaluated, slices = self._eval("array") + def as_raw(self) -> xr.DataArray: # noqa: D102, override + evaluated, slices = self._eval("raw") return evaluated.sel(**slices) def as_expr(self) -> EXPR_T: @@ -766,7 +716,7 @@ def as_expr(self) -> EXPR_T: return evaluated.sel(**slices) -class EvalIndexSlice(EvalArrayOrMath): +class EvalIndexSlice(EvalNode): """For processing `$slice` expressions.""" def __init__(self, instring: str, loc: int, tokens: pp.ParseResults) -> None: @@ -794,7 +744,7 @@ def _eval(self, return_type: Literal["math_string"], as_values: bool) -> str: .. @overload def _eval( - self, return_type: Literal["array"], as_values: bool + self, return_type: Literal["raw"], as_values: bool ) -> xr.DataArray | list[xr.DataArray]: ... def _eval( @@ -809,13 +759,12 @@ def _eval( def as_math_string(self) -> str: # noqa: D102, override return self._eval("math_string", False) - def as_array(self) -> xr.DataArray | list[xr.DataArray]: # noqa: D102, override - evaluated = self._eval("array", True) + def as_raw(self) -> xr.DataArray | list[xr.DataArray]: # noqa: D102, override + evaluated = self._eval("raw", True) return evaluated - -class EvalSubExpressions(EvalArrayOrMathExpr): +class EvalSubExpressions(EvalNode): """For processing sub-expressions.""" def __init__(self, instring: str, loc: int, tokens: pp.ParseResults) -> None: @@ -842,9 +791,12 @@ def __repr__(self) -> str: def _eval(self, return_type: Literal["math_string"]) -> str: ... @overload - def _eval(self, return_type: Literal["array"]) -> xr.DataArray: ... + def _eval(self, return_type: Literal["raw"]) -> xr.DataArray: ... + + @overload + def _eval(self, return_type: Literal["expr"]) -> EXPR_T: ... - def _eval(self, return_type: RETURN_T) -> str | xr.DataArray: + def _eval(self, return_type: RETURN_T) -> str | xr.DataArray | EXPR_T: """Evaluate the referenced sub_expression.""" return self.eval_attrs.sub_expression_dict[self.name][0].eval( return_type, self.eval_attrs @@ -853,15 +805,14 @@ def _eval(self, return_type: RETURN_T) -> str | xr.DataArray: def as_math_string(self) -> str: # noqa: D102, override return self._eval("math_string") - def as_array(self) -> xr.DataArray: # noqa: D102, override - return self._eval("array") + def as_raw(self) -> xr.DataArray: # noqa: D102, override + return self._eval("raw") - def as_expr(self) -> EXPR_T: + def as_expr(self) -> EXPR_T: # noqa: D102, override return self._eval("expr") - -class EvalNumber(EvalArrayOrMath): +class EvalNumber(EvalNode): """For processing numbers.""" def __init__(self, instring: str, loc: int, tokens: pp.ParseResults) -> None: @@ -894,11 +845,11 @@ def as_math_string(self) -> str: # noqa: D102, override f"{float(self.value):.6g}", ) - def as_array(self) -> xr.DataArray: # noqa: D102, override + def as_raw(self) -> xr.DataArray: # noqa: D102, override return xr.DataArray(float(self.value), name=float(self.value)) -class ListParser(EvalArrayOrMath): +class ListParser(EvalNode): """For parsing lists.""" def __init__(self, instring: str, loc: int, tokens: pp.ParseResults) -> None: @@ -923,17 +874,17 @@ def __repr__(self) -> str: return f"{self.val}" def as_math_string(self) -> str: # noqa: D102, override - input_list = self.as_array() + input_list = self.as_raw() return "[" + ",".join(str(i.name) for i in input_list) + "]" - def as_array(self) -> list[xr.DataArray]: # noqa: D102, override - values = [val.eval("array", self.eval_attrs) for val in self.val] + def as_raw(self) -> list[xr.DataArray]: # noqa: D102, override + values = [val.eval("raw", self.eval_attrs) for val in self.val] # strings and numbers are returned as xarray arrays of size 1, # so we extract those values. return values -class EvalUnslicedComponent(EvalArrayOrMathExpr): +class EvalUnslicedComponent(EvalNode): """Evaluation of unsliced components.""" def __init__(self, instring: str, loc: int, tokens: pp.ParseResults) -> None: @@ -961,7 +912,7 @@ def __repr__(self) -> str: def as_math_string(self) -> str: # noqa: D102, override self.eval_attrs = replace(self.eval_attrs, as_values=False) - evaluated = self.as_array() + evaluated = self.as_raw() self.eval_attrs.references.add(self.name) if "math_repr" in evaluated.attrs: @@ -971,50 +922,57 @@ def as_math_string(self) -> str: # noqa: D102, override return data_var_string - def as_array(self) -> xr.DataArray: # noqa: D102, override - group = self.eval_attrs.math.find(self.name)._group - if group in ["parameters", "lookups"]: - evaluated = self.eval_attrs.input_data[self.name] - elif group == "dimensions": + def as_raw(self) -> xr.DataArray: # noqa: D102, override + math_def = self.eval_attrs.math.find(self.name) + group = math_def._group + self.eval_attrs.references.add(self.name) + if group in ["parameters", "lookups", "dimensions"]: + # A parameter/lookup/dimension defined in the math but absent from the + # input data resolves to its default (NaN if none is set). try: evaluated = self.eval_attrs.input_data[self.name] except KeyError: evaluated = xr.DataArray(np.nan) else: - evaluated = getattr(self.eval_attrs.model, group)[self.name] - if evaluated.isnull().any() and pd.notna( - default := self.eval_attrs.math.find(self.name)["default"] - ): + # Model entries (variables / expressions): return the raw model object + # unchanged. Defaults only apply to parameters/lookups/dimensions. A model + # entry that was never built (e.g. skipped because its mask was empty) + # resolves to a NaN expression rather than raising. + try: + return getattr(self.eval_attrs.model, group)[self.name] + except KeyError: + return LinearExpression(xr.DataArray(np.nan), self.eval_attrs.model) + if evaluated.isnull().any() and pd.notna(default := math_def["default"]): evaluated = evaluated.fillna(default) - - self.eval_attrs.references.add(self.name) return evaluated def as_expr(self) -> xr.DataArray | EXPR_T: # noqa: D102, override - group = self.eval_attrs.math.find(self.name)._group + math_def = self.eval_attrs.math.find(self.name) + group = math_def._group + self.eval_attrs.references.add(self.name) if group in ["parameters", "lookups", "dimensions"]: + # Parameters / lookups / dimensions are scalar coefficients: keep as DataArray. try: evaluated = self.eval_attrs.input_data[self.name] except KeyError: evaluated = xr.DataArray(np.nan) else: + # Model entries: normalise Variable -> LinearExpression immediately so that + # arithmetic composition on the expression route always sees an expression. try: - evaluated = getattr(self.eval_attrs.model, group)[self.name] + evaluated = _to_linexpr( + getattr(self.eval_attrs.model, group)[self.name] + ) except KeyError: - evaluated = LinearExpression(xr.DataArray(np.nan), self.eval_attrs.model) - if evaluated.isnull().any() and pd.notna( - default := self.eval_attrs.math.find(self.name)["default"] - ): - if isinstance(evaluated, Variable): - evaluated = evaluated.to_linexpr() + evaluated = LinearExpression( + xr.DataArray(np.nan), self.eval_attrs.model + ) + if evaluated.isnull().any() and pd.notna(default := math_def["default"]): evaluated = evaluated.fillna(default) - - self.eval_attrs.references.add(self.name) return evaluated - -class EvalGenericString(EvalArrayOrMath): +class EvalGenericString(EvalNode): """For generic string parsing.""" def __init__(self, instring: str, loc: int, tokens: pp.ParseResults) -> None: @@ -1041,7 +999,7 @@ def __repr__(self) -> str: def as_math_string(self): # noqa: D102, override return str(self.val) - def as_array(self) -> xr.DataArray: # noqa: D102, override + def as_raw(self) -> xr.DataArray: # noqa: D102, override return xr.DataArray(str(self.val), name=str(self.val)) diff --git a/linopy/declarative/helper_functions.py b/linopy/declarative/helper_functions.py index b876ffb68..7604ad58f 100644 --- a/linopy/declarative/helper_functions.py +++ b/linopy/declarative/helper_functions.py @@ -14,7 +14,6 @@ from typing import Any, Literal, overload import numpy as np -import pandas as pd import xarray as xr from linopy.declarative.eval_attrs import EvalAttrs @@ -37,7 +36,7 @@ class ParsingHelperFunction(ABC): """Abstract base class for helper function parsing.""" def __init__( - self, return_type: Literal["array", "math_string"], attrs: "EvalAttrs" + self, return_type: Literal["raw", "expr", "math_string"], attrs: "EvalAttrs" ) -> None: """ Abstract helper function class, which all helper functions must subclass. @@ -71,28 +70,44 @@ def as_math_string(self, *args: Any, **kwargs: Any) -> str: """ @abstractmethod - def as_array(self, *args: Any, **kwargs: Any) -> LinearExpression | xr.DataArray: + def as_raw(self, *args: Any, **kwargs: Any) -> LinearExpression | xr.DataArray: """ Method to apply the helper function to provide an n-dimensional array output. - This method is called when the class is initialised with ``return_type=array``. + This method is called when the class is initialised with ``return_type=raw`` and, + by default, ``return_type=expr`` (see :meth:`as_expr`). """ + def as_expr(self, *args: Any, **kwargs: Any) -> LinearExpression | xr.DataArray: + """ + Method to apply the helper function on the expression route. + + This method is called when the class is initialised with ``return_type=expr``. + By default it 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 __call__(self, *args: Any, **kwargs: Any) -> Any: """ When a helper function is accessed by evaluating a parsing string, this method is called. The value of `return_type` on initialisation of the class defines whether this method returns either: - - a string (``return_type=math_string``) - - :meth:xr.DataArray (``return_type=array``) + - a string (``return_type=math_string``, via :meth:`as_math_string`) + - the raw data (``return_type=raw``, via :meth:`as_raw`) + - a linopy expression (``return_type=expr``, via :meth:`as_expr`) """ if self._return_type == "math_string": return self.as_math_string(*args, **kwargs) - elif self._return_type == "array": - return self.as_array(*args, **kwargs) + elif self._return_type == "raw": + return self.as_raw(*args, **kwargs) elif self._return_type == "expr": - return self.as_array(*args, **kwargs) + return self.as_expr(*args, **kwargs) + else: + raise ValueError( + f"Unknown helper function return type: {self._return_type!r}" + ) def __init_subclass__(cls) -> None: """ @@ -229,7 +244,7 @@ def as_math_string( # noqa: D102, override # Using bigvee for "collective-or" return rf"\bigvee\limits_{{{substack_overstring}}} ({array})" - def as_array( + def as_raw( self, input_component: xr.DataArray, *, over: xr.DataArray | list[xr.DataArray] ) -> xr.DataArray: """ @@ -275,7 +290,7 @@ def as_math_string( # noqa: D102, override else: return rf"\bigwedge({', '.join(substrings)})" - def as_array( + def as_raw( self, *, within: xr.DataArray, how: Literal["all", "any"], **dims: str ) -> xr.DataArray: """ @@ -399,7 +414,7 @@ def as_math_string( # noqa: D102, override substack_overstring = rf"\substack{{{overstring}}}" return rf"\sum\limits_{{{substack_overstring}}} ({array})" - def as_array( + def as_raw( self, array: xr.DataArray, *, over: xr.DataArray | list[xr.DataArray] ) -> xr.DataArray: """ @@ -437,7 +452,7 @@ def as_math_string(self, array: str, **lookup_arrays: str) -> str: # noqa: D102 array = self._update_iterator(array, new_strings, "add") return array - def as_array( + def as_raw( self, array: xr.DataArray, **lookup_arrays: xr.DataArray ) -> xr.DataArray: """ @@ -539,7 +554,7 @@ 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_array(self, **dim_idx_mapping: int) -> xr.DataArray: + def as_raw(self, **dim_idx_mapping: int) -> xr.DataArray: """ Get value of a model dimension at a given integer index. @@ -577,7 +592,7 @@ def as_array(self, **dim_idx_mapping: int) -> xr.DataArray: dim, idx = self._mapping_to_dim_idx(**dim_idx_mapping) return self._attrs.input_data.coords[dim][int(idx)] - # For as_array + # For as_raw @overload @staticmethod def _mapping_to_dim_idx(**dim_idx_mapping: int) -> tuple[str, int]: ... @@ -614,7 +629,7 @@ def as_math_string(self, array: str, **roll_kwargs: str) -> str: # noqa: D102, component = self._update_iterator(array, new_strings, "add") return component - def as_array(self, array: xr.DataArray, **roll_kwargs: int) -> xr.DataArray: + def as_raw(self, array: xr.DataArray, **roll_kwargs: int) -> xr.DataArray: """ Roll (a.k.a., shift) the array along the given dimension(s) by the given number of places. @@ -654,7 +669,7 @@ class Mask(ParsingHelperFunction): def as_math_string(self, array: str, condition: str) -> str: # noqa: D102, override return rf"({array} \text{{if }} {condition} == True)" - def as_array(self, array: xr.DataArray, condition: xr.DataArray) -> xr.DataArray: + def as_raw(self, array: xr.DataArray, condition: xr.DataArray) -> xr.DataArray: """ Apply a `mask` condition to a math array within an expression string. @@ -716,7 +731,7 @@ def as_math_string(self, array: str, groupby: str, group_dim: str) -> str: # no overstring = rf"\substack{{{foreach_string}}}" return rf"\sum\limits_{{{overstring}}} ({array})" - def as_array( + def as_raw( self, array: xr.DataArray, groupby: xr.DataArray, group_dim: xr.DataArray ) -> xr.DataArray: """ @@ -798,7 +813,7 @@ def as_math_string(self, array: str, over: str, group: str) -> str: # noqa: D10 return rf"\sum\limits_{{{overstring}}} ({array})" - def as_array( + def as_raw( self, array: xr.DataArray, over: xr.DataArray, group: xr.DataArray ) -> xr.DataArray: """ @@ -894,7 +909,7 @@ def as_math_string(self, array: str, over: str, N: int) -> str: # noqa: D102, o return rf"\sum\limits_{{\text{{{new_iterator}}}={over_singular}}}^{{{over_singular}+{N}}} ({updated_iterator_array})" - def as_array(self, array: xr.DataArray, over: xr.DataArray, N: int) -> xr.DataArray: + def as_raw(self, array: xr.DataArray, over: xr.DataArray, N: int) -> xr.DataArray: """ Sum values from current up to N from current on the dimension `over`. diff --git a/linopy/declarative/mask_parser.py b/linopy/declarative/mask_parser.py index 25042c2ad..8f4589b98 100644 --- a/linopy/declarative/mask_parser.py +++ b/linopy/declarative/mask_parser.py @@ -45,21 +45,19 @@ def get_dot_attr(var: Any, attr: str) -> Any: return value -class EvalNot(expression_parser.EvalSignOp, expression_parser.EvalArrayOrMath): +class EvalNot(expression_parser.EvalSignOp): """Parse action to process successfully parsed expressions with a leading `not`.""" def as_math_string(self) -> str: # noqa: D102, override evaluated = self.value.eval("math_string", self.eval_attrs) return rf"\neg ({evaluated})" - def as_array(self) -> xr.DataArray: # noqa: D102, override - evaluated = self.value.eval("array", self.eval_attrs) + def as_raw(self) -> xr.DataArray: # noqa: D102, override + evaluated = self.value.eval("raw", self.eval_attrs) return ~evaluated -class EvalAndOr( - expression_parser.EvalOperatorOperand, expression_parser.EvalArrayOrMath -): +class EvalAndOr(expression_parser.EvalOperatorOperand): """ Processing of successfully parsed expressions with and/or operators. @@ -94,11 +92,11 @@ def _apply_mask(self, evaluated: xr.DataArray) -> xr.DataArray: def as_math_string(self) -> str: # noqa: D102, override return super().as_math_string() - def as_array(self) -> xr.DataArray: # noqa: D102, override - return super().as_array() + def as_raw(self) -> xr.DataArray: # noqa: D102, override + return super().as_raw() -class ConfigOptionParser(expression_parser.EvalArrayOrMath): +class ConfigOptionParser(expression_parser.EvalNode): """Parsing of configuration options.""" def __init__(self, instring: str, loc: int, tokens: pp.ParseResults) -> None: @@ -124,7 +122,7 @@ def __repr__(self): def as_math_string(self) -> str: # noqa: D102, override return rf"\text{{config.{self.config_option}}}" - def as_array(self) -> xr.DataArray: # noqa: D102, override + def as_raw(self) -> xr.DataArray: # noqa: D102, override config_val = get_dot_attr(self.eval_attrs.config, self.config_option) if not isinstance(config_val, int | float | str | bool | np.bool_): @@ -136,7 +134,7 @@ def as_array(self) -> xr.DataArray: # noqa: D102, override return xr.DataArray(config_val) -class ResultArrayParser(expression_parser.EvalArrayOrMath): +class ResultArrayParser(expression_parser.EvalNode): """Variable/Expression array processing.""" def __init__(self, instring: str, loc: int, tokens: pp.ParseResults) -> None: @@ -167,7 +165,7 @@ def as_math_string(self) -> str: # noqa: D102, override return math_repr - def as_array(self) -> xr.DataArray: # noqa: D102, override + def as_raw(self) -> xr.DataArray: # noqa: D102, override self.eval_attrs.references.add(self.array_name) da = self.eval_attrs.model[self.array_name] if self.eval_attrs.apply_mask: @@ -175,7 +173,7 @@ def as_array(self) -> xr.DataArray: # noqa: D102, override return da -class InputArrayParser(expression_parser.EvalArrayOrMath): +class InputArrayParser(expression_parser.EvalNode): """Input array processing.""" def __init__(self, instring: str, loc: int, tokens: pp.ParseResults) -> None: @@ -208,7 +206,7 @@ def as_math_string(self) -> str: # noqa: D102, override math_repr = rf"\exists ({math_repr})" return math_repr - def as_array(self) -> xr.DataArray: # noqa: D102, override + def as_raw(self) -> xr.DataArray: # noqa: D102, override self.eval_attrs.references.add(self.array_name) da = self.eval_attrs.input_data.get(self.array_name, xr.DataArray(False)) if self.eval_attrs.apply_mask and da.dtype.kind != "b": @@ -220,7 +218,7 @@ def as_array(self) -> xr.DataArray: # noqa: D102, override return da -class DimensionArrayParser(expression_parser.EvalArrayOrMath): +class DimensionArrayParser(expression_parser.EvalNode): """Dimension array processing.""" def __init__(self, instring: str, loc: int, tokens: pp.ParseResults) -> None: @@ -246,15 +244,13 @@ def __repr__(self): def as_math_string(self) -> str: # noqa: D102, override return self.array_name - def as_array(self) -> xr.DataArray: # noqa: D102, override + def as_raw(self) -> xr.DataArray: # noqa: D102, override # We want the mask string to evaluate successfully even if a dimension hasn't been defined. da = self.eval_attrs.input_data.get(self.array_name, xr.DataArray()) return da -class ComparisonParser( - expression_parser.EvalComparisonOp, expression_parser.EvalArrayOrMath -): +class ComparisonParser(expression_parser.EvalComparisonOp): """Parse action to process successfully parsed strings of the form x=y.""" OP_TRANSLATOR = { @@ -276,9 +272,9 @@ def as_math_string(self) -> str: # noqa: D102, override rhs = rf"\text{{{rhs}}}" return lhs + self.OP_TRANSLATOR[self.op] + rhs - def as_array(self) -> xr.DataArray: # noqa: D102, override + def as_raw(self) -> xr.DataArray: # noqa: D102, override self.eval_attrs = replace(self.eval_attrs, apply_mask=False) - lhs, rhs = self._eval("array") + lhs, rhs = self._eval("raw") match self.op: case "<=": comparison = lhs <= rhs @@ -293,7 +289,7 @@ def as_array(self) -> xr.DataArray: # noqa: D102, override return xr.DataArray(comparison) -class SubsetParser(expression_parser.EvalArrayOrMath): +class SubsetParser(expression_parser.EvalNode): """Dimension subset parsing.""" def __init__(self, instring: str, loc: int, tokens: pp.ParseResults) -> None: @@ -318,7 +314,7 @@ def __repr__(self): def _eval(self) -> list[str | float]: """Evaluate each element of the subset list.""" - values = [val.eval("array", self.eval_attrs) for val in self.val] + values = [val.eval("raw", self.eval_attrs) for val in self.val] return [val.item() if isinstance(val, xr.DataArray) else val for val in values] def as_math_string(self) -> str: # noqa: D102, override @@ -328,15 +324,15 @@ def as_math_string(self) -> str: # noqa: D102, override subset_string = "[" + ",".join(str(i) for i in subset) + "]" return rf"\text{{{iterator}}} \in \text{{{subset_string}}}" - def as_array(self) -> xr.DataArray: # noqa: D102, override + def as_raw(self) -> xr.DataArray: # noqa: D102, override subset = self._eval() self.eval_attrs = replace(self.eval_attrs, apply_mask=False) - da = self.set_name.eval("array", replace(self.eval_attrs, apply_mask=False)) + da = self.set_name.eval("raw", replace(self.eval_attrs, apply_mask=False)) set_item_in_subset = da.isin(subset) return set_item_in_subset -class BoolOperandParser(expression_parser.EvalArrayOrMath): +class BoolOperandParser(expression_parser.EvalNode): """Boolean operand parsing.""" def __init__(self, instring: str, loc: int, tokens: pp.ParseResults) -> None: @@ -360,7 +356,7 @@ def __repr__(self): def as_math_string(self): # noqa: D102, override return self.val - def as_array(self) -> xr.DataArray: # noqa: D102, override + def as_raw(self) -> xr.DataArray: # noqa: D102, override if self.val == "true": bool_val = xr.DataArray(np.True_) elif self.val == "false": @@ -398,14 +394,14 @@ def eval(self, *args, **kwargs) -> str: def data_var_parser( - names: Iterable, parse_action: type[expression_parser.EvalArrayOrMath] + names: Iterable, parse_action: type[expression_parser.EvalNode] ) -> pp.ParserElement: """ Process model data variables which can be any valid python identifier (string + "_"). Args: names (Iterable): List of valid component names. - parse_action (type[expression_parser.EvalArrayOrMath]): Parse action to evaluate the parsed string. + parse_action (type[expression_parser.EvalNode]): Parse action to evaluate the parsed string. Returns: pp.ParserElement: parser for model data variables which will access the data @@ -570,7 +566,10 @@ def generate_mask_string_parser( number, unique_evaluatable_string, dimensions_parser ) subset = subset_parser( - [dimensions_parser, inputs_parser], config_option, number, general_evaluatable_string + [dimensions_parser, inputs_parser], + config_option, + number, + general_evaluatable_string, ) arithmetic = pp.Forward() diff --git a/linopy/declarative/parsing.py b/linopy/declarative/parsing.py index 00246fdc4..226ee0ad9 100644 --- a/linopy/declarative/parsing.py +++ b/linopy/declarative/parsing.py @@ -9,7 +9,7 @@ import logging import operator from collections.abc import Callable, Iterable -from typing import TYPE_CHECKING, Literal, overload +from typing import TYPE_CHECKING, Any, Literal, overload import pyparsing as pp import xarray as xr @@ -22,6 +22,7 @@ ) from linopy.declarative.schema import MATH_DEFS_T, ConfigModel, MathModel, _Equations from linopy.expressions import LinearExpression + if TYPE_CHECKING: from linopy.model import Model TRUE_ARRAY = xr.DataArray(True) @@ -195,7 +196,7 @@ def evaluate_mask( math: MathModel, config: ConfigModel, *, - return_type: Literal["array"] = "array", + return_type: Literal["raw"] = "raw", references: set | None = None, initial_mask: xr.DataArray = TRUE_ARRAY, ) -> xr.DataArray: ... @@ -220,7 +221,7 @@ def evaluate_mask( math: MathModel, config: ConfigModel, *, - return_type: str = "array", + return_type: str = "raw", references: set | None = None, initial_mask: xr.DataArray = TRUE_ARRAY, ) -> xr.DataArray | str: @@ -232,9 +233,9 @@ def evaluate_mask( model (Model): Linopy model. math (MathModel): Calliope math definitions. config (ConfigModel): Build configuration options. - return_type (str, optional): If "array", return xarray.DataArray. + return_type (str, optional): If "raw", return xarray.DataArray. If "math_string", return LaTex math string. - Defaults to "array". + Defaults to "raw". references (set | None, optional): List of references to use in evaluation. Defaults to None. initial_mask (xr.DataArray, optional): If given, the mask array resulting @@ -286,7 +287,37 @@ def drop_dims_not_in_foreach(self, mask: xr.DataArray) -> xr.DataArray: unwanted_dims = set(mask.dims).difference(self.sets) return (mask.sum(unwanted_dims) > 0).astype(bool).transpose(*self.sets) - # Expecting anything (most likely an array) if not requesting latex string. + def _evaluate( + self, + input_data: xr.Dataset, + model: Model, + math: MathModel, + *, + return_type: Literal["expr", "math_string"], + references: set | None, + mask: xr.DataArray, + ) -> Any: + """ + Evaluate the parsed expression tree. + + Shared by :meth:`evaluate_expression` (arithmetic roots) and + :meth:`evaluate_equation` (comparison roots). + """ + eval_attrs_ = { + "equation_name": self.name, + "slice_dict": self.slices, + "sub_expression_dict": self.sub_expressions, + "input_data": input_data, + "model": model, + "math": math, + "mask": mask, + "helper_functions": helper_functions._registry["expression"], + } + if references is not None: + eval_attrs_["references"] = references + return self.expression[0].eval(return_type, eval_attrs.EvalAttrs(**eval_attrs_)) + + # Expecting a linopy expression if not requesting latex string. @overload def evaluate_expression( self, @@ -322,17 +353,17 @@ def evaluate_expression( mask: xr.DataArray = TRUE_ARRAY, ) -> LinearExpression | str: """ - Evaluate a math string to produce an array backend objects or a LaTex math string. + Evaluate an arithmetic math string (expressions/objectives). Args: input_data (xr.Dataset): Model input data. - model (xr.Dataset): Backend interface component dataset. - math (MathModel): Calliope math definitions. + model (Model): Linopy model. + math (MathModel): Linopy math definitions. Keyword Args: return_type (str, optional): - If "array", return xarray.DataArray. If "math_string", return LaTex math string. - Defaults to "array". + If "expr", return a linopy expression. If "math_string", return a LaTeX + math string. Defaults to "expr". references (set | None, optional): If given, any references in the math string to other model components will be logged here. Defaults to None. @@ -341,28 +372,102 @@ def evaluate_expression( Defaults to xr.DataArray(True). Returns: - xr.DataArray | str: - If return_type == `array`: array of backend expression objects. - If return_type == `math_string`: Valid LaTeX math string defining the - "mask" conditions using logic notation. + LinearExpression | str: + If return_type == `expr`: a linopy expression. A pure-parameter + expression (evaluated to an ``xr.DataArray``) is coerced to a + ``LinearExpression``. + If return_type == `math_string`: a valid LaTeX math string. """ - eval_attrs_ = { - "equation_name": self.name, - "slice_dict": self.slices, - "sub_expression_dict": self.sub_expressions, - "input_data": input_data, - "model": model, - "math": math, - "mask": mask, - "helper_functions": helper_functions._registry["expression"], - } - if references is not None: - eval_attrs_["references"] = references - evaluated = self.expression[0].eval( - return_type, eval_attrs.EvalAttrs(**eval_attrs_) + evaluated = self._evaluate( + input_data, + model, + math, + return_type=return_type, + references=references, + mask=mask, ) + if return_type == "expr" and isinstance(evaluated, xr.DataArray): + evaluated = LinearExpression(evaluated, model) return evaluated + # Expecting a (lhs, sign, rhs) tuple if not requesting latex string. + @overload + def evaluate_equation( + self, + input_data: xr.Dataset, + model: Model, + math: MathModel, + *, + return_type: Literal["expr"] = "expr", + references: set | None = None, + mask: xr.DataArray = TRUE_ARRAY, + ) -> tuple[LinearExpression, xr.DataArray, LinearExpression]: ... + + # Expecting string if requesting latex string. + @overload + def evaluate_equation( + self, + input_data: xr.Dataset, + model: Model, + math: MathModel, + *, + return_type: Literal["math_string"], + references: set | None = None, + ) -> str: ... + + def evaluate_equation( + self, + input_data: xr.Dataset, + model: Model, + math: MathModel, + *, + return_type: Literal["expr", "math_string"] = "expr", + references: set | None = None, + mask: xr.DataArray = TRUE_ARRAY, + ) -> tuple[LinearExpression, xr.DataArray, LinearExpression] | str: + """ + Evaluate a comparison math string (constraints) of the form ``LHS OP RHS``. + + Args: + input_data (xr.Dataset): Model input data. + model (Model): Linopy model. + math (MathModel): Linopy math definitions. + + Keyword Args: + return_type (str, optional): + If "expr", return a ``(lhs, sign, rhs)`` tuple for constraint assembly. + If "math_string", return a LaTeX math string. Defaults to "expr". + references (set | None, optional): + If given, any references in the math string to other model components + will be logged here. Defaults to None. + mask (xr.DataArray, optional): + If given, should be a boolean array with which to mask any produced arrays. + Defaults to xr.DataArray(True). + + Returns: + tuple[LinearExpression, xr.DataArray, LinearExpression] | str: + If return_type == `expr`: a ``(lhs, sign, rhs)`` tuple, where ``lhs``/``rhs`` + are linopy expressions (a pure-parameter side is coerced to a + ``LinearExpression``) and ``sign`` is a DataArray of the comparison operator. + If return_type == `math_string`: a valid LaTeX math string. + """ + evaluated = self._evaluate( + input_data, + model, + math, + return_type=return_type, + references=references, + mask=mask, + ) + if return_type == "math_string": + return evaluated + lhs, sign, rhs = evaluated + if isinstance(lhs, xr.DataArray): + lhs = LinearExpression(lhs, model) + if isinstance(rhs, xr.DataArray): + rhs = LinearExpression(rhs, model) + return lhs, sign, rhs + def raise_error_on_mask_expr_mismatch( self, expression: xr.DataArray, mask: xr.DataArray ) -> None: diff --git a/test/test_declarative_parsing.py b/test/test_declarative_parsing.py new file mode 100644 index 000000000..2e0335dde --- /dev/null +++ b/test/test_declarative_parsing.py @@ -0,0 +1,333 @@ +# Copyright (C) since 2013 Calliope contributors listed in AUTHORS. +# Licensed under the Apache 2.0 License (see LICENSE file). +""" +Tests for the declarative math parser route separation (mask / expr / raw). + +These tests guard the contracts established by the parser-route refactor: + +- mask evaluation always returns a boolean ``xr.DataArray``; +- expression evaluation always returns a linopy expression; +- equation (comparison) evaluation returns a ``(lhs, sign, rhs)`` tuple; +- helper-function arguments are always evaluated in ``raw`` mode. +""" + +from __future__ import annotations + +import numpy as np +import pytest +import xarray as xr + +from linopy.declarative import helper_functions +from linopy.declarative.build import DeclarativeModelBuilder, declarative_model +from linopy.declarative.parsing import ParsedBackendComponent +from linopy.expressions import LinearExpression +from linopy.variables import Variable + +NODES = ["a", "b", "c"] + + +def _math() -> dict: + """Minimal but representative math definition exercising every route.""" + return { + "dimensions": {"node": {"dtype": "string", "iterator": "n"}}, + "parameters": { + "cost": {"default": 0}, + "cap_max": {"default": float("inf")}, + }, + "variables": { + "flow": { + "foreach": ["node"], + "bounds": {"lower": 0, "upper": float("inf")}, + }, + }, + "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", + }, + }, + } + + +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}, + ) + + +@pytest.fixture +def builder_with_flow() -> 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 + + +def _first_equation(builder: DeclarativeModelBuilder, group: str, name: str): + """Parse a component and return (component, first_equation, foreach_sub_mask).""" + definition = getattr(builder.math, group)[name] + component = ParsedBackendComponent( + group, name, definition, builder.math.parsing_components + ) + mask = component.generate_top_level_mask( + builder.input_data, + builder.model, + builder.math, + builder.config, + references=set(), + ) + equation = component.parse_equations()[0] + sub_mask = equation.evaluate_mask( + builder.input_data, + builder.model, + builder.math, + builder.config, + initial_mask=mask, + ) + sub_mask = component.drop_dims_not_in_foreach(sub_mask) + return component, equation, sub_mask + + +# --------------------------------------------------------------------------- # +# Mask route +# --------------------------------------------------------------------------- # + + +def test_top_level_mask_returns_boolean_dataarray(builder_with_flow): + component, _, 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(): + math = _math() + 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"]) + component = ParsedBackendComponent( + "constraints", + "cap", + builder.math.constraints["cap"], + builder.math.parsing_components, + ) + mask = component.generate_top_level_mask( + builder.input_data, + builder.model, + builder.math, + builder.config, + references=set(), + ) + equation = component.parse_equations()[0] + result = equation.evaluate_mask( + builder.input_data, + builder.model, + builder.math, + builder.config, + initial_mask=mask, + ) + assert isinstance(result, xr.DataArray) + assert result.dtype == bool + + +# --------------------------------------------------------------------------- # +# Expression route +# --------------------------------------------------------------------------- # + + +def test_expression_with_variable_returns_linexpr(builder_with_flow): + _, equation, sub_mask = _first_equation( + builder_with_flow, "expressions", "total_cost" + ) + result = equation.evaluate_expression( + builder_with_flow.input_data, + builder_with_flow.model, + builder_with_flow.math, + mask=sub_mask, + ) + assert isinstance(result, LinearExpression) + + +def test_pure_parameter_expression_coerced_to_linexpr(builder_with_flow): + _, equation, sub_mask = _first_equation( + builder_with_flow, "expressions", "cost_plus_one" + ) + result = equation.evaluate_expression( + builder_with_flow.input_data, + builder_with_flow.model, + builder_with_flow.math, + 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(builder_with_flow): + _, equation, sub_mask = _first_equation( + builder_with_flow, "expressions", "sub_expr_test" + ) + result = equation.evaluate_expression( + builder_with_flow.input_data, + builder_with_flow.model, + builder_with_flow.math, + mask=sub_mask, + ) + assert isinstance(result, LinearExpression) + + +# --------------------------------------------------------------------------- # +# Equation (constraint) route +# --------------------------------------------------------------------------- # + + +def test_equation_returns_lhs_sign_rhs_tuple(builder_with_flow): + _, equation, sub_mask = _first_equation(builder_with_flow, "constraints", "cap") + lhs, sign, rhs = equation.evaluate_equation( + builder_with_flow.input_data, + builder_with_flow.model, + builder_with_flow.math, + 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)) <= {"<="} + + +# --------------------------------------------------------------------------- # +# Helper-function argument evaluation (raw mode) +# --------------------------------------------------------------------------- # + + +class _RecordArgs(helper_functions.ParsingHelperFunction): + """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, **kwargs): # noqa: D102 + return "record_args" + + def as_raw(self, *args, **kwargs): # 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] + + +def test_helper_arguments_are_evaluated_raw(builder_with_flow): + math = _math() + # flow is a variable, cost is a parameter -> raw mode must preserve both types. + math["expressions"]["total_cost"]["equations"][0]["expression"] = ( + "record_args(flow, cost)" + ) + builder = DeclarativeModelBuilder(math, _inputs(), {}) + builder.add_variable("flow", builder.math.variables["flow"]) + _, equation, sub_mask = _first_equation(builder, "expressions", "total_cost") + + _RecordArgs.received = [] + equation.evaluate_expression( + builder.input_data, builder.model, builder.math, mask=sub_mask + ) + assert _RecordArgs.received, "helper was not called" + # The variable arrives un-normalised (raw Variable, not LinearExpression); + # the parameter arrives as a raw DataArray (not coerced/masked to booleans). + assert Variable in _RecordArgs.received + assert xr.DataArray in _RecordArgs.received + assert LinearExpression not in _RecordArgs.received + + +def test_invalid_helper_function_rejected(builder_with_flow): + """A registry entry that is not a ParsingHelperFunction subclass is rejected.""" + registry = helper_functions._registry["expression"] + registry["not_a_helper"] = str # type: ignore[assignment] + try: + math = _math() + math["expressions"]["total_cost"]["equations"][0]["expression"] = ( + "not_a_helper(flow)" + ) + builder = DeclarativeModelBuilder(math, _inputs(), {}) + builder.add_variable("flow", builder.math.variables["flow"]) + _, equation, sub_mask = _first_equation(builder, "expressions", "total_cost") + with pytest.raises(ValueError, match="must be subclassed"): + equation.evaluate_expression( + builder.input_data, builder.model, builder.math, mask=sub_mask + ) + finally: + registry.pop("not_a_helper", None) + + +# --------------------------------------------------------------------------- # +# Input-data checks +# --------------------------------------------------------------------------- # + + +def test_checks_run_without_active_variable(): + """`_check_inputs` must not require an `active` variable in the input data.""" + math = _math() + 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(): + math = _math() + 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(), {}) + + +# --------------------------------------------------------------------------- # +# End-to-end build +# --------------------------------------------------------------------------- # + + +def test_declarative_model_end_to_end(): + 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"} From f33031a51b2cc83325f512d142114e06fa306296 Mon Sep 17 00:00:00 2001 From: Bryn Pickering <17178478+brynpickering@users.noreply.github.com> Date: Thu, 23 Jul 2026 16:24:06 +0100 Subject: [PATCH 05/12] Refactor initial Calliope DecMath port Co-Authored-By: Claude --- linopy/declarative/__init__.py | 22 +- linopy/declarative/build.py | 593 ++++----- linopy/declarative/eval_attrs.py | 62 - linopy/declarative/evaluate.py | 571 +++++++++ linopy/declarative/expression_parser.py | 1461 ----------------------- linopy/declarative/grammar.py | 635 ++++++++++ linopy/declarative/helper_functions.py | 964 --------------- linopy/declarative/helpers.py | 673 +++++++++++ linopy/declarative/latex.py | 446 +++++++ linopy/declarative/mask_parser.py | 618 ---------- linopy/declarative/parsing.py | 1413 +++++++++------------- linopy/declarative/schema.py | 234 +--- test/test_declarative_parsing.py | 760 +++++++++--- 13 files changed, 3749 insertions(+), 4703 deletions(-) delete mode 100644 linopy/declarative/eval_attrs.py create mode 100644 linopy/declarative/evaluate.py delete mode 100644 linopy/declarative/expression_parser.py create mode 100644 linopy/declarative/grammar.py delete mode 100644 linopy/declarative/helper_functions.py create mode 100644 linopy/declarative/helpers.py create mode 100644 linopy/declarative/latex.py delete mode 100644 linopy/declarative/mask_parser.py diff --git a/linopy/declarative/__init__.py b/linopy/declarative/__init__.py index 80bfc3f97..bbcf650d2 100644 --- a/linopy/declarative/__init__.py +++ b/linopy/declarative/__init__.py @@ -1 +1,21 @@ -"""Linopy declarative text math interface.""" \ No newline at end of file +""" +Linopy declarative math interface. + +Build a linopy model from a declarative math definition (typically loaded from +YAML) and an xarray dataset of input data, via :func:`declarative_model`. +""" + +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 index c37860bc6..15c649191 100644 --- a/linopy/declarative/build.py +++ b/linopy/declarative/build.py @@ -1,14 +1,28 @@ -import textwrap +""" +Linopy declarative model-build module. + +This module contains the entry point to build a linopy optimisation model from a +declarative math definition (a dictionary, typically loaded from YAML) and an +xarray dataset of input data. +""" + +from __future__ import annotations + +import logging import time -import typing +from collections.abc import Iterable, Iterator, Mapping +from dataclasses import replace +from typing import Any, Literal, get_args -import numpy as np import xarray as xr -from tqdm.asyncio import tqdm +from tqdm.auto import tqdm -from linopy.declarative import eval_attrs, helper_functions, parsing +from linopy.declarative import parsing +from linopy.declarative.evaluate import Context, evaluate +from linopy.declarative.grammar import Component, find_refs +from linopy.declarative.helpers import HelperFunction, build_registry from linopy.declarative.schema import ( - LOGGER, + DTYPE_OPTIONS, ConfigModel, ConstraintDef, ExpressionDef, @@ -20,62 +34,97 @@ from linopy.io import TQDM_COLOR from linopy.model import Model -ORDERED_COMPONENTS_T = typing.Literal[ +LOGGER = logging.getLogger(__name__) + +ORDERED_COMPONENTS_T = Literal[ "variables", "expressions", "constraints", - # "piecewise_constraints", "objectives", ] -DTYPE_OPTIONS = { - "string": str, - "float": float, - "bool": bool, - "datetime": np.datetime64, - "date": np.datetime64, - "integer": int, -} - -DATETIME_DTYPE = "M" -"""Numpy type kind for datetime arrays""" +_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) -> Model: - """Build a Linopy Model from declarative math definitions and input data.""" - builder = DeclarativeModelBuilder(math_def, input_data, config) - return builder.build() +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 DeclarativeModelBuilder: - def __init__(self, math_def: dict, input_data: xr.Dataset, config: dict): + """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. + """ self.model = Model() self.math = MathModel.model_validate(math_def) self.input_data = self._update_dtypes(input_data) self.config = ConfigModel.model_validate(config) - + self._ctx = Context( + model=self.model, + input_data=self.input_data, + math=self.math, + config=self.config, + helpers=build_registry(helpers), + ) self._check_inputs() def _update_dtypes(self, ds: xr.Dataset, id_: str = "") -> xr.Dataset: """ - Update data types of coordinates or data variables in the dataset. + Coerce dataset variables to the dtypes given by their math definitions. - Args: - ds (xr.Dataset): Dataset to update. - math (math_schema.CalliopeBuildMath): Model math definition. - id_ (str, optional): ID of the dataset being updated, for logging purposes. Defaults to an empty string. - - Raises: - ValueError: If there is a mismatch between the provided variable and its definition in the model math. - - Returns: - xr.Dataset: `ds` with data types updated. + 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( - var_name, subset=["lookups", "parameters", "dimensions"] + str(var_name), subset=["lookups", "parameters", "dimensions"] ) except KeyError: LOGGER.info( @@ -84,7 +133,9 @@ def _update_dtypes(self, ds: xr.Dataset, id_: str = "") -> xr.Dataset: ) continue - dtype_str = math_def.dtype # type: ignore + 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" @@ -96,18 +147,6 @@ def _update_dtypes(self, ds: xr.Dataset, id_: str = "") -> xr.Dataset: .where(var_data.notnull()) .where(var_data != "") ) - case "datetime": - updated_var = time._datetime_index( - var_data.to_series(), self.config.datetime_format - ).to_xarray() - case "date": - updated_var = ( - time._datetime_index( - var_data.to_series(), self.config.date_format - ) - .to_xarray() - .assign_attrs(var_data.attrs) - ) case "bool": updated_var = var_data.fillna(False).astype(dtype) case _: @@ -116,327 +155,203 @@ def _update_dtypes(self, ds: xr.Dataset, id_: str = "") -> xr.Dataset: 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, check in self.math.checks.root.items(): + if not check.active: + continue + mask_node = parsing.parse_mask(check.mask, self.math, name) + check_ctx = replace(self._ctx, route="mask", equation_name=name) + evaluated = evaluate(mask_node, check_ctx) + if (evaluated & active).any(): + messages = error_msgs if check.errors == "raise" else warn_msgs + messages.append(check.message) + + 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: typing.Mapping[str, typing.Any], - ) -> list[tuple[str, typing.Any]]: - """Return (name, obj) pairs from a root mapping, sorted by obj.order.""" + 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 _check_inputs(self) -> None: - data_checks = self.math.checks - check_results: dict[str, list[str]] = {"raise": [], "warn": []} - parser_ = parsing.mask_parser.generate_mask_string_parser( - **self.math.parsing_components["mask"] - ) - eval_kwargs = { - "model": self.model, - "math": self.math, - "input_data": self.input_data, - "config": self.config, - "helper_functions": helper_functions._registry["mask"], - } - active = self.input_data.get("active", xr.DataArray(True)) - for name, check in data_checks.root.items(): - if check.active: - parsed_ = parser_.parse_string(check.mask, parse_all=True) - eval_attrs_ = eval_attrs.EvalAttrs(equation_name=name, **eval_kwargs) - evaluated = parsed_[0].eval("raw", eval_attrs_) - if (evaluated & active).any(): - check_results[check.errors].append(check.message) - - print_warnings_and_raise_errors( - check_results["warn"], - check_results["raise"], - during="model input data checks", - ) + def _references( + self, definition: Any, equations: Iterable[parsing.Equation] = () + ) -> list[str]: + """Return the sorted names of all math components a component references.""" + mask_node = parsing.parse_mask(getattr(definition, "mask", "True"), self.math) + refs = find_refs(mask_node, Component) + for equation in equations: + refs |= equation.references() + return sorted(refs) + + def _iter_equations( + self, + equations: list[parsing.Equation], + group: parsing.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: - references: set[str] = set() - parsed_component = parsing.ParsedBackendComponent( - "variables", name, definition, self.math.parsing_components - ) - mask = parsed_component.generate_top_level_mask( - self.input_data, - self.model, - self.math, - self.config, - align_to_foreach_sets=True, - break_early=True, - references=references, + """Add a decision variable to the model, masked by its math definition.""" + mask = parsing.component_mask("variables", name, definition, 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", ) - kwargs = { - "upper": definition.bounds.upper, - "lower": definition.bounds.lower, - "integer": definition.domain == "integer", - "binary": definition.domain == "binary", - } - if mask.any(): - self.model.add_variables(coords=mask.coords, name=name, mask=mask, **kwargs) - self.model.variables[name].attrs["references"] = references - else: - LOGGER.info( - f"variables:{name} | No valid data points after applying mask. Variable not added to model." - ) + # Variable.attrs values are typed Hashable, but a sorted list serializes best. + self.model.variables[name].attrs["references"] = self._references(definition) # type: ignore[assignment] def add_expression(self, name: str, definition: ExpressionDef) -> None: - references: set[str] = set() - parsed_component = parsing.ParsedBackendComponent( - "expressions", name, definition, self.math.parsing_components - ) - mask = parsed_component.generate_top_level_mask( - self.input_data, - self.model, - self.math, - self.config, - align_to_foreach_sets=True, - break_early=True, - references=references, - ) - expr = LinearExpression(float("nan"), self.model).where(mask) - all_mask = mask.copy() - if mask.any(): - equations = parsed_component.parse_equations() - for equation in equations: - sub_mask = equation.evaluate_mask( - self.input_data, - self.model, - self.math, - self.config, - initial_mask=mask, - references=references, - ) - if not sub_mask.any(): - continue - sub_mask = parsed_component.drop_dims_not_in_foreach(sub_mask) - if (~expr.isnull() & 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." - ) - expr_to_fill = equation.evaluate_expression( - self.input_data, - self.model, - self.math, - mask=sub_mask, - references=references, - ) - expr = merge([expr, expr_to_fill.where(sub_mask)]) - if not expr.isnull().all(): - self.model.add_expressions(name=name, data=expr, mask=all_mask) - self.model.expressions[name].attrs["references"] = references - else: - LOGGER.info( - f"expressions:{name} | No valid data points after applying mask. Expression not added to model." + """Add a named expression to the model, merging its equation variants.""" + mask = parsing.component_mask("expressions", name, definition, 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) + equations = parsing.parse_component("expressions", name, definition, self.math) + for equation, sub_mask in self._iter_equations(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." ) - else: - LOGGER.info( - f"expressions:{name} | No valid data points after applying mask. Expression not added to model." - ) + 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( + definition, equations + ) def add_constraint(self, name: str, definition: ConstraintDef) -> None: - references: set[str] = set() - parsed_component = parsing.ParsedBackendComponent( - "constraints", name, definition, self.math.parsing_components - ) - mask = parsed_component.generate_top_level_mask( - self.input_data, - self.model, - self.math, - self.config, - align_to_foreach_sets=True, - break_early=True, - references=references, - ) - lhs = LinearExpression(float("nan"), self.model).where(mask) - sign = xr.DataArray().where(parsed_component.drop_dims_not_in_foreach(mask)) - rhs = LinearExpression(float("nan"), self.model).where(mask) - all_mask = mask.copy() + """Add a constraint to the model, merging its equation variants.""" + mask = parsing.component_mask("constraints", name, definition, self._ctx) if not mask.any(): - LOGGER.info( - f"constraints:{name} | No valid data points after applying mask. Constraint not added to model." - ) - return None - - equations = parsed_component.parse_equations() - for equation in equations: - sub_mask = equation.evaluate_mask( - self.input_data, - self.model, - self.math, - self.config, - initial_mask=mask, - references=references, - ) - if not sub_mask.any(): - LOGGER.info( - f"constraints:{equation.name} | No valid data points after applying mask. Constraint not added to model." - ) - continue - sub_mask = parsed_component.drop_dims_not_in_foreach(sub_mask) + 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) + equations = parsing.parse_component("constraints", name, definition, self.math) + for equation, sub_mask in self._iter_equations(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." + 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 = equation.evaluate_equation( - self.input_data, - self.model, - self.math, - mask=sub_mask, - references=references, + 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} | No valid data points after applying mask. Constraint not added to model." - ) - return None - + LOGGER.info(f"constraints:{name} | {_SKIP_MESSAGE}") + return self.model.add_constraints( - coords=all_mask.coords, + coords=mask.coords, name=name, lhs=lhs, - sign=sign.fillna( - "==" - ), # Default to equality to avoid errors; will be masked. + # Default to equality to avoid errors on masked-out points. + sign=sign.fillna("=="), rhs=rhs, - mask=all_mask, + mask=mask, + ) + self.model.constraints[name].attrs["references"] = self._references( + definition, equations ) - self.model.constraints[name].attrs["references"] = references def add_objective(self, name: str, definition: ObjectiveDef) -> None: - references: set[str] = set() - parsed_component = parsing.ParsedBackendComponent( - "objectives", name, definition, self.math.parsing_components - ) - mask = parsed_component.generate_top_level_mask( - self.input_data, - self.model, - self.math, - self.config, - align_to_foreach_sets=True, - break_early=True, - references=references, - ) - expr = LinearExpression(float("nan"), self.model).where(mask) - if mask.any(): - equations = parsed_component.parse_equations() - for equation in equations: - sub_mask = equation.evaluate_mask( - self.input_data, - self.model, - self.math, - self.config, - initial_mask=mask, - references=references, - ) - if not sub_mask.any(): - continue - sub_mask = parsed_component.drop_dims_not_in_foreach(sub_mask) - if (~expr.isnull() & 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." - ) - expr_to_fill = equation.evaluate_expression( - self.input_data, - self.model, - self.math, - mask=sub_mask, - references=references, + """Set the model objective, merging its equation variants.""" + mask = parsing.component_mask("objectives", name, definition, self._ctx) + if not mask.any(): + LOGGER.info(f"objectives:{name} | {_SKIP_MESSAGE}") + return + pieces: list[tuple[LinearExpression, xr.DataArray]] = [] + filled = xr.DataArray(False) + equations = parsing.parse_component("objectives", name, definition, self.math) + for equation, sub_mask in self._iter_equations(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." ) - expr = expr_to_fill - self.model.add_objective(expr=expr, sense=definition.sense) - self.model.objective.attrs["references"] = references + 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( + definition, equations + ) def build(self) -> Model: - for components in typing.get_args(ORDERED_COMPONENTS_T): - component = components.removesuffix("s") - ordered_items = self._sorted_by_order(self.math[components].root) - ordered_items_tqdm = tqdm( - ordered_items, - desc=f"Building {components}.", - colour=TQDM_COLOR, + """ + 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 name, definition in ordered_items_tqdm: + for group in get_args(ORDERED_COMPONENTS_T): + component = group.removesuffix("s") + ordered_items = self._sorted_by_order(self.math[group].root) + for name, definition in tqdm( + ordered_items, desc=f"Building {group}.", colour=TQDM_COLOR + ): start = time.time() getattr(self, f"add_{component}")(name, definition) - end = time.time() - start - LOGGER.debug(f"{components}:{name} | Built in {end:.4f}s") - LOGGER.info(f"{components} | Generated.") + LOGGER.debug(f"{group}:{name} | Built in {time.time() - start:.4f}s") + LOGGER.info(f"{group} | Generated.") return self.model - - -def print_warnings_and_raise_errors( - warnings: list[str] | dict[str, list[str]] | None = None, - errors: list[str] | dict[str, list[str]] | None = None, - during: str = "model processing", - bullet: str = " * ", -) -> None: - """ - Process collections of warnings/errors. - - Prints warnings / raises errors with a bullet point list of the concatenated - collections. - - Lists will return simple bullet lists: - E.g. warnings=["foo", "bar"] becomes: - - Possible issues found during model processing: - * foo - * bar - - Dicts of lists will return nested bullet lists: - E.g. errors={"foo": ["foobar", "foobaz"]} becomes: - - Errors during model processing: - * foo - * foobar - * foobaz - - Args: - warnings (list[str] | dict[str, list[str]] | None, optional): - List of warning strings or dictionary of warning strings. - If None or an empty list, no warnings will be printed. - Defaults to None. - errors (list[str] | dict[str, list[str]] | None, optional): - List of error strings or dictionary of error strings. - If None or an empty list, no errors will be raised. - Defaults to None. - during (str, optional): - Substring that will be placed at the top of the concatenated list of warnings/errors to point to during which phase of data processing they occurred. - Defaults to "model processing". - bullet (str, optional): Type of bullet points to use. Defaults to " * ". - - Raises: - ModelError: If errors is not None or is a non-empty list/dict - - """ - spacer = " " * len(bullet) - - def _sort_strings(stringlist: list[str]) -> list[str]: - return sorted(list(set(stringlist))) - - def _predicate(string_: str) -> bool: - return not string_.startswith((bullet, spacer)) - - def _indenter(strings: list[str] | dict[str, list[str]]) -> str: - if isinstance(strings, dict): - sorted_strings = [] - for k, v in strings.items(): - sorted_strings.append(str(k) + ":") - sorted_strings.extend(_sort_strings([spacer + bullet + i for i in v])) - else: - sorted_strings = _sort_strings(strings) - return textwrap.indent("\n".join(sorted_strings), bullet, predicate=_predicate) - - if warnings: - LOGGER.info(f"Possible issues found during {during}:\n" + _indenter(warnings)) - - if errors: - raise ValueError(f"Errors during {during}:\n" + _indenter(errors)) diff --git a/linopy/declarative/eval_attrs.py b/linopy/declarative/eval_attrs.py deleted file mode 100644 index e525d5294..000000000 --- a/linopy/declarative/eval_attrs.py +++ /dev/null @@ -1,62 +0,0 @@ -# Copyright (C) since 2013 Calliope contributors listed in AUTHORS. -# Licensed under the Apache 2.0 License (see LICENSE file). -"""Parsing evaluation attributes.""" - -from __future__ import annotations - -import logging -from collections.abc import Callable -from dataclasses import dataclass, field -from typing import TYPE_CHECKING - -import xarray as xr - -from linopy.declarative.schema import ConfigModel, MathModel - -if TYPE_CHECKING: - from linopy.model import Model -TRUE_ARRAY = xr.DataArray(True) - -LOGGER = logging.getLogger(__name__) - - -@dataclass -class EvalAttrs: - """Attributes required for evaluating parsed expressions.""" - - model: Model - """Backend interface component dataset.""" - - equation_name: str = "" - """Name of the equation being evaluated.""" - - helper_functions: dict[str, Callable] = field(default_factory=dict) - """Helper functions available for evaluations.""" - - input_data: xr.Dataset = field(default_factory=xr.Dataset) - """Model input data.""" - - math: MathModel = field(default_factory=MathModel) - """Linopy math definitions.""" - - apply_mask: bool = True - """Whether to apply the 'where' condition.""" - - as_values: bool = False - """If True, return with the array contents evaluated to base Python objects. - If False, return with the array contents as they are in the backend dataset.""" - - config: ConfigModel = field(default_factory=ConfigModel) - """Build configuration options.""" - - references: set[str] = field(default_factory=set) - """References to dimensions/lookups/parameters/variables/global expressions used in the expression.""" - - slice_dict: dict = field(default_factory=dict) - """Dictionary to look up array slice expressions if referenced in the evaluated string.""" - - sub_expression_dict: dict = field(default_factory=dict) - """Dictionary to look up sub-expressions if referenced in the evaluated string.""" - - mask: xr.DataArray = field(default_factory=lambda: xr.DataArray(True)) - """Boolean array defining where the expression should be applied.""" diff --git a/linopy/declarative/evaluate.py b/linopy/declarative/evaluate.py new file mode 100644 index 000000000..6ad7bbc83 --- /dev/null +++ b/linopy/declarative/evaluate.py @@ -0,0 +1,571 @@ +""" +Linopy declarative math evaluation module. + +This module contains the evaluation context and the two tree walkers that turn a +parsed math AST (see :mod:`linopy.declarative.grammar`) into either a LaTeX math +string (:func:`to_math_string`) or data — an `xr.DataArray` or a linopy +expression (:func:`evaluate`). +""" + +from __future__ import annotations + +import operator +import re +from dataclasses import dataclass, field, 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.grammar import ( + Arith, + Call, + Compare, + Component, + ConfigRef, + Constant, + ListNode, + Node, + Sliced, + SliceRef, + SubExprRef, + Subset, + Unary, +) +from linopy.declarative.helpers import KIND_T, HelperFunction, 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) + +ROUTE_T = Literal["expression", "mask"] + +_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).""" + + route: ROUTE_T = "expression" + """Whether the AST being evaluated came from an expression or a mask string.""" + + apply_mask: bool = True + """On the mask route, 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 error(ctx: Context, node: Node, message: str) -> ValueError: + """Return a ValueError contextualised with the equation name and source string.""" + return ValueError(f"({ctx.equation_name}, {node.instring}) | {message}") + + +def get_dot_attr(var: Any, attr: str) -> Any: + """ + Get a nested attribute in dot notation (e.g. "foo.bar"). + + Works for nested objects: dictionaries, pydantic models, etc. + """ + levels = attr.split(".", 1) + value = var[levels[0]] if isinstance(var, dict) else getattr(var, levels[0]) + if len(levels) > 1: + value = get_dot_attr(value, levels[1]) + return value + + +def _to_linexpr(obj: Any) -> Any: + """ + Normalise a model object to a linopy expression. + + `Variable` objects are converted to `LinearExpression`; everything else is + returned unchanged. This is the single place where the `Variable` -> + `LinearExpression` coercion is performed on the expression route. + """ + if isinstance(obj, Variable): + return obj.to_linexpr() + return obj + + +def _apply_mask(evaluated: Any, mask: xr.DataArray) -> Any: + """Mask an evaluated operand, broadcasting first if it cannot be masked directly.""" + try: + return evaluated.where(mask) + except AttributeError: + return evaluated.broadcast_like(mask).where(mask) + + +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 + ) + + +# --------------------------------------------------------------------------- +# Data evaluation +# --------------------------------------------------------------------------- + + +def evaluate(node: Node, ctx: Context, expr: bool = False) -> Any: + """ + Evaluate a math AST node to data. + + Parameters + ---------- + node : Node + AST node to evaluate. + ctx : Context + Evaluation context. + expr : bool, default: False + If False ("raw" mode), return the underlying data without route-specific + transformation: an `xr.DataArray` for parameters/lookups/dimensions, the + raw model object (`Variable`/`LinearExpression`) for model entries, and a + boolean array on the mask route. If True ("expr" mode), guarantee a + linopy-expression-compatible result: `Variable` objects are coerced to + `LinearExpression` and a top-level :class:`Compare` returns a masked + `(lhs, sign, rhs)` tuple for constraint assembly. + """ + match node: + case Constant(value=bool() as val): + return xr.DataArray(np.bool_(val)) + case Constant(value=str() as val): + return val + case Constant(value=val): + return xr.DataArray(float(val), name=float(val)) + case ListNode(items=items): + return [evaluate(item, ctx) for item in items] + case Component(): + return _evaluate_component(node, ctx, expr) + case ConfigRef(): + return _evaluate_config(node, ctx) + case SubExprRef(name=name): + return evaluate(ctx.sub_expressions[name], ctx, expr) + case SliceRef(name=name): + return evaluate(ctx.slices[name], ctx) + case Sliced(obj=obj, slices=slices): + evaluated_slices = { + dim: [_unwrap(i) for i in vals] + if isinstance(vals := evaluate(slicer, ctx), list) + else vals + for dim, slicer in slices.items() + } + return evaluate(obj, ctx, expr).sel(**evaluated_slices) + case Call(): + return _evaluate_call(node, ctx, expr) + case Unary(op=op, operand=operand): + if op == "not": + return ~evaluate(operand, ctx) + evaluated = evaluate(operand, ctx, expr) + return -1 * evaluated if op == "-" else evaluated + case Arith(first=first, rest=rest): + boolean = rest[0][0] in ("and", "or") + val = evaluate(first, ctx, expr) + if not boolean: + val = _apply_mask(val, ctx.mask) + for op, operand in rest: + evaluated = evaluate(operand, ctx, expr) + if not boolean: + evaluated = _apply_mask(evaluated, ctx.mask) + val = _OPERATIONS[op](val, evaluated) + return val + case Compare() if expr: + return _evaluate_equation(node, ctx) + case Compare(lhs=lhs, op=op, rhs=rhs): + unmasked_ctx = replace(ctx, apply_mask=False) + comparison = _OPERATIONS[op]( + evaluate(lhs, unmasked_ctx), evaluate(rhs, unmasked_ctx) + ) + return xr.DataArray(comparison) + case Subset(items=items, dim=dim): + subset = [_unwrap(evaluate(item, ctx)) for item in items] + dim_array = evaluate(dim, replace(ctx, apply_mask=False)) + return dim_array.isin(subset) + case _: + raise error( + ctx, node, f"Cannot evaluate node of type {type(node).__name__}" + ) + + +def _evaluate_component(node: Component, ctx: Context, expr: bool) -> Any: + """Evaluate a component reference according to its category and the evaluation mode.""" + name = node.name + if node.category == "dimension": + # The mask string should evaluate successfully even if a dimension isn't defined. + return ctx.input_data.get(name, xr.DataArray()) + if node.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 node.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)) + 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. On the expression route 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 not expr: + return evaluated + evaluated = _to_linexpr(evaluated) + if evaluated.isnull().any() and pd.notna(default := math_def["default"]): + evaluated = evaluated.fillna(default) + return evaluated + + +def _evaluate_config(node: ConfigRef, ctx: Context) -> xr.DataArray: + """Evaluate a config option reference to a dimensionless array.""" + config_val = get_dot_attr(ctx.config, node.option) + if not isinstance(config_val, int | float | str | bool | np.bool_): + raise error( + ctx, + node, + f"Configuration option resolves to invalid type " + f"`{type(config_val).__name__}`, expected a number, string, or boolean.", + ) + return xr.DataArray(config_val) + + +def _lookup_helper(node: Call, ctx: Context) -> type[HelperFunction]: + """Return the helper class for a function call, validating it exists in the registry.""" + kind: KIND_T = "mask" if ctx.route == "mask" else "expression" + helpers = ctx.helpers.get(kind, {}) + if node.func not in helpers: + raise error(ctx, node, f"Invalid helper function defined: {node.func}") + helper_cls = helpers[node.func] + if not (isinstance(helper_cls, type) and issubclass(helper_cls, HelperFunction)): + raise error( + ctx, + node, + "Helper function must be subclassed from " + f"linopy.declarative.helpers.HelperFunction: {node.func}", + ) + return helper_cls + + +def _evaluate_call(node: Call, ctx: Context, expr: bool) -> Any: + """ + Evaluate a helper-function call. + + The helper itself is instantiated with the enclosing mode so that + expression-route helpers can dispatch to their `as_expr` implementation. + Its arguments, however, are always evaluated in 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 that have been coerced to boolean masks or `LinearExpression`. + """ + helper_cls = _lookup_helper(node, ctx) + helper = helper_cls("expr" if expr else "raw", ctx) + if helper_cls.ignore_mask: + ctx = replace(ctx, mask=TRUE_ARRAY) + args = [evaluate(arg, ctx) for arg in node.args] + kwargs = {name: evaluate(val, ctx) for name, val in node.kwargs.items()} + return helper(*args, **kwargs) + + +def _evaluate_equation(node: Compare, ctx: Context) -> tuple[Any, xr.DataArray, Any]: + """Evaluate an equation to a masked `(lhs, sign, rhs)` tuple for constraint assembly.""" + lhs = evaluate(node.lhs, ctx, expr=True) + rhs = evaluate(node.rhs, ctx, expr=True) + 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, + node, + 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)) + rhs_masked = _to_linexpr(rhs.where(ctx.mask)) + sign_masked = xr.DataArray(node.op).where(ctx.mask) + return lhs_masked, sign_masked, rhs_masked + + +# --------------------------------------------------------------------------- +# LaTeX math-string evaluation +# --------------------------------------------------------------------------- + + +def to_math_string(node: Node, ctx: Context) -> str: + """ + Evaluate a math AST node to a LaTeX math string. + + Parameters + ---------- + node : Node + AST node to evaluate. + ctx : Context + Evaluation context. + """ + match node: + case Constant(value=bool() as val): + return str(val).lower() + case Constant(value=str() as val): + return val + case Constant(value=val): + return re.sub( + r"([\d]+?)e([+-])([\d]+)", + r"\1\\mathord{\\times}10^{\2\3}", + f"{float(val):.6g}", + ) + case ListNode(items=items): + return "[" + ",".join(_plain_string(item, ctx) for item in items) + "]" + case Component(): + return _component_math_string(node, ctx) + case ConfigRef(option=option): + return rf"\text{{config.{option}}}" + case SubExprRef(name=name): + return to_math_string(ctx.sub_expressions[name], ctx) + case SliceRef(name=name): + return to_math_string(ctx.slices[name], ctx) + case Sliced(): + return _sliced_math_string(node, ctx) + case Call(): + helper = _lookup_helper(node, ctx)("math_string", ctx) + args = [_call_arg_math_string(arg, ctx) for arg in node.args] + kwargs = { + name: _call_arg_math_string(val, ctx) + for name, val in node.kwargs.items() + } + return helper(*args, **kwargs) + case Unary(op="not", operand=operand): + return rf"\neg ({to_math_string(operand, ctx)})" + case Unary(op=op, operand=operand): + return op + to_math_string(operand, ctx) + case Arith(first=first, rest=rest): + val = to_math_string(first, ctx) + for op, operand in rest: + evaluated = to_math_string(operand, ctx) + # We ignore identity elements that do nothing (e.g. `0 + flow` is `flow`) + if evaluated == _LATEX_IDENTITIES.get(op): + continue + if isinstance(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 + case Compare(lhs=lhs, op=op, rhs=rhs) if ctx.route == "expression": + lhs_str = to_math_string(lhs, ctx) + rhs_str = to_math_string(rhs, ctx) + return lhs_str + _LATEX_EQUATION_OPERATORS[op] + rhs_str + case Compare(lhs=lhs, op=op, rhs=rhs): + unmasked_ctx = replace(ctx, apply_mask=False) + lhs_str = to_math_string(lhs, unmasked_ctx) + rhs_str = to_math_string(rhs, unmasked_ctx) + if r"\text" not in rhs_str: + rhs_str = rf"\text{{{rhs_str}}}" + return lhs_str + _LATEX_MASK_OPERATORS[op] + rhs_str + case Subset(items=items, dim=dim): + subset = [_unwrap(evaluate(item, ctx)) for item in 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 = ( + dim.name if isinstance(dim, Component) else to_math_string(dim, 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}}}" + case _: + raise error( + ctx, node, f"Cannot render node of type {type(node).__name__} as LaTeX" + ) + + +def _plain_string(item: Node, ctx: Context) -> str: + """Return a plain-text representation of a list item for LaTeX rendering.""" + evaluated = evaluate(item, ctx) + 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 evaluate(arg, ctx) + return to_math_string(arg, ctx) + + +def _component_math_string(node: Component, ctx: Context) -> str: + """ + Render a 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 = node.name + custom = ctx.math_reprs.get(name) + if node.category == "dimension": + return name + if node.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 node.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 = evaluate(node, ctx) + attrs = getattr(evaluated, "attrs", {}) + return str(attrs.get("math_repr", name)) + + +def _sliced_math_string(node: Sliced, ctx: Context) -> str: + r""" + Render a sliced component as LaTeX. + + If the component's LaTeX representation carries an iterator substring (from a + `math_repr` data attribute, e.g. `\textbf{flow}_\text{n}`), the slices are + injected into it by re-parsing (e.g. `\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): to_math_string(slicer, ctx) + for dim, slicer in node.slices.items() + } + obj_string = to_math_string(node.obj, ctx) + + def _replace(term: pp.ParseResults) -> Any: + if len(term) == 1: + return term + replacers = {k: f"{k}={v}" for k, v in slice_strings.items()} + return ( + term[0] + term[1] + ",".join(replacers.get(k, k) for k in term[2]) + term[3] + ) + + id_ = pp.Combine( + pp.Word(pp.alphas, pp.alphanums) + + pp.ZeroOrMore("_" + pp.Word(pp.alphanums)) + + pp.Opt("_") + ) + id_formatted = pp.Combine("\\" + pp.Word(pp.alphas) + "{" + id_ + "}") + obj_parser = id_formatted + pp.Opt( + r"_\text{" + pp.Group(pp.DelimitedList(id_)) + "}" + ) + obj_parser.set_parse_action(_replace) + try: + return obj_parser.parse_string(obj_string, parse_all=True)[0] + except pp.ParseException: + subscript = ",".join(f"{k}={v}" for k, v in slice_strings.items()) + return rf"{obj_string}_\text{{{subscript}}}" diff --git a/linopy/declarative/expression_parser.py b/linopy/declarative/expression_parser.py deleted file mode 100644 index eb0885af9..000000000 --- a/linopy/declarative/expression_parser.py +++ /dev/null @@ -1,1461 +0,0 @@ -# Copyright (C) since 2013 Calliope contributors listed in AUTHORS. -# Licensed under the Apache 2.0 License (see LICENSE file). - -## -# Part of the code in this file is adapted from -# https://github.com/pyparsing/pyparsing/blob/master/examples/eval_arith.py -# available under the MIT license -## -# Copyright 2009, 2011 Paul McGuire -# -# Permission is hereby granted, free of charge, to any person obtaining -# a copy of this software and associated documentation files (the -# "Software"), to deal in the Software without restriction, including -# without limitation the rights to use, copy, modify, merge, publish, -# distribute, sublicense, and/or sell copies of the Software, and to -# permit persons to whom the Software is furnished to do so, subject to -# the following conditions: - -# The above copyright notice and this permission notice shall be -# included in all copies or substantial portions of the Software. - -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, -# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF -# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. -# IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY -# CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, -# TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE -# SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. -## -"""Expression parsing functionality.""" - -from __future__ import annotations - -import re -from abc import ABC, abstractmethod -from collections.abc import Callable, Iterable, Iterator -from dataclasses import replace -from typing import Any, Literal, TypeVar, overload - -import numpy as np -import pandas as pd -import pyparsing as pp -import xarray as xr - -from linopy.declarative.eval_attrs import EvalAttrs -from linopy.declarative.helper_functions import ParsingHelperFunction -from linopy.expressions import LinearExpression, QuadraticExpression -from linopy.variables import Variable - -pp.ParserElement.enable_packrat() - -SUB_EXPRESSION_CLASSIFIER = "$" - -EXPR_T = LinearExpression | QuadraticExpression -EXPRVAR_T = TypeVar("EXPRVAR_T", LinearExpression, QuadraticExpression) -ARRAY_T = TypeVar("ARRAY_T", LinearExpression, QuadraticExpression, xr.DataArray) -RETURN_T = Literal["expr", "raw", "math_string"] - - -class EvalString(ABC): - """Parent class for all string evaluation classes - used in type hinting.""" - - name: str - eval_attrs: EvalAttrs - instring: str - - def __eq__(self, other: Any) -> bool: - """Functionality for '==' operations.""" - return self.__repr__() == other - - @abstractmethod - def __repr__(self) -> str: - """Return string representation of the parsed grammar.""" - - def error_msg(self, message: str) -> ValueError: - """Raise an error message with context.""" - return ValueError( - f"({self.eval_attrs.equation_name}, {self.instring}) | {message}" - ) - - -def _to_linexpr(obj: Any) -> Any: - """ - Normalise a model object to a linopy expression. - - ``Variable`` objects are converted to ``LinearExpression`` via ``to_linexpr``; - ``LinearExpression``/``QuadraticExpression``/``xr.DataArray`` objects are - returned unchanged. This is the single place where the ``Variable`` -> - ``LinearExpression`` coercion is performed on the expression route. - """ - if isinstance(obj, Variable): - return obj.to_linexpr() - return obj - - -class EvalNode(EvalString): - """ - Base class for nodes evaluated as math strings, raw data, or expressions. - - Three evaluation modes are supported (see :meth:`eval`): - - - ``math_string``: a LaTeX string (:meth:`as_math_string`). - - ``raw``: the underlying data without any route-specific transformation - (:meth:`as_raw`) - an ``xr.DataArray`` for parameters/lookups/dimensions and - the raw model object (``Variable``/``LinearExpression``) for model entries. - - ``expr``: a linopy expression suitable for arithmetic composition - (:meth:`as_expr`). The default implementation returns the raw data; nodes that - compose arithmetic override it to guarantee an expression is returned. - """ - - @abstractmethod - def as_math_string(self) -> str: - """Evaluate and return expression as LaTeX.""" - - @abstractmethod - def as_raw(self) -> xr.DataArray | list[xr.DataArray]: - """ - Evaluate and return the underlying data without route-specific transformation. - - If the evaluated expression returns a simple string or number, - this value will be assigned as both the `name` and the data of the returned DataArray. - The purpose of this is to be able to access the string/number value whether we query the array name or its data. - """ - - def as_expr(self) -> Any: - """ - Evaluate and return expression as a LinearExpression or QuadraticExpression. - - The default implementation returns the raw data (:meth:`as_raw`); nodes that - compose arithmetic (operands, signs, functions, (sub-)components) override it. - """ - return self.as_raw() - - # Math strings evaluate to strings. - @overload - def eval( - self, return_type: Literal["math_string"], eval_attrs: EvalAttrs - ) -> str: ... - - # Raw evaluation returns the underlying data. - @overload - def eval( - self, return_type: Literal["raw"], eval_attrs: EvalAttrs - ) -> xr.DataArray | list[xr.DataArray]: ... - - def eval(self, return_type: RETURN_T, eval_attrs: EvalAttrs) -> Any: - """ - Evaluate a parsed expression node. - - Args: - return_type (Literal["math_string", "raw", "expr"]): - Dictates how the expression should be evaluated (see `Returns` section). - eval_attrs (EvalAttrs): Evaluation attributes. - - Returns: - If `math_string`, a valid LaTeX math string. - If `raw`, the underlying data (`xr.DataArray`, list, or raw model object). - If `expr`, a linopy `LinearExpression`/`QuadraticExpression`. - """ - self.eval_attrs = eval_attrs - evaluated: Any - if return_type == "raw": - evaluated = self.as_raw() - elif return_type == "math_string": - evaluated = self.as_math_string() - elif return_type == "expr": - evaluated = self.as_expr() - return evaluated - - -class EvalComparisonOp(EvalNode): - """Class for processing comparison operations.""" - - OP_TRANSLATOR = {"<=": r" \leq ", ">=": r" \geq ", "==": " = ", "=": " = "} - - def __init__(self, instring: str, loc: int, tokens: pp.ParseResults) -> None: - """ - Parse action to process successfully parsed equations of the form LHS OPERATOR RHS. - - Args: - instring (str): String that was parsed (used in error message). - loc (int): - Location in parsed string where parsing error was logged. - This is not used; we include it as pyparsing injects it alongside `instring` when setting the parse action. - tokens (pp.ParseResults): - Contains a list with an RHS (pp.ParseResults), operator (str), and LHS (pp.ParseResults). - """ - self.lhs, self.op, self.rhs = tokens - self.instring = instring - self.loc = loc - self.values = tokens - - def __repr__(self) -> str: - """Programming / official string representation.""" - return f"{self.lhs.__repr__()} {self.op} {self.rhs.__repr__()}" - - # string return - @overload - def _eval(self, return_type: Literal["math_string"]) -> tuple[str, str]: ... - - # raw return - @overload - def _eval( - self, return_type: Literal["raw"] - ) -> tuple[xr.DataArray, xr.DataArray]: ... - - # expression return - @overload - def _eval(self, return_type: Literal["expr"]) -> tuple[Any, Any]: ... - - def _eval(self, return_type: RETURN_T) -> tuple[Any, Any]: - """Evaluate the LHS and RHS of the comparison.""" - lhs = self.lhs.eval(return_type, self.eval_attrs) - rhs = self.rhs.eval(return_type, self.eval_attrs) - return lhs, rhs - - def as_math_string(self) -> str: # noqa: D102, override - lhs, rhs = self._eval("math_string") - return lhs + self.OP_TRANSLATOR[self.op] + rhs - - def as_raw(self) -> xr.DataArray: # noqa: D102, override - self.eval_attrs = replace(self.eval_attrs, apply_mask=False) - lhs, rhs = self._eval("raw") - match self.op: - case "<=": - comparison = lhs <= rhs - case ">=": - comparison = lhs >= rhs - case "<": - comparison = lhs < rhs - case ">": - comparison = lhs > rhs - case "==": - comparison = lhs == rhs - return xr.DataArray(comparison) - - def as_expr( - self, - ) -> tuple[LinearExpression, xr.DataArray, LinearExpression | xr.DataArray]: - """Evaluate the comparison as a ``(lhs, sign, rhs)`` tuple for constraint assembly.""" - lhs, rhs = self._eval("expr") - mask = self.eval_attrs.mask - for side, arr in {"left": lhs, "right": rhs}.items(): - extra_dims = set(arr.dims).difference(set(mask.dims) | {"_term"}) - if extra_dims: - raise self.error_msg( - f"The {side}-hand side of the equation is indexed over dimensions not present in `foreach`: {extra_dims}" - ) - lhs_masked = _to_linexpr(lhs.where(mask)) - rhs_masked = _to_linexpr(rhs.where(mask)) - sign_masked = xr.DataArray(self.op).where(mask) - return lhs_masked, sign_masked, rhs_masked - - -class EvalToCallable(EvalString): - """Parent class for callable functionality.""" - - @abstractmethod - def as_callable(self, return_type: RETURN_T) -> Callable: - """Callable processing.""" - ... - - def eval(self, return_type: RETURN_T, eval_attrs: EvalAttrs) -> Callable: - """ - Evaluate math string expression. - - Args: - return_type (str): Whether to return a math string or xarray DataArray. - eval_attrs (EvalAttrs): Evaluation attributes. - - Returns: - Callable: returns helper function. - """ - self.eval_attrs = eval_attrs - evaluated = self.as_callable(return_type) - return evaluated - - -class EvalOperatorOperand(EvalNode): - """Evaluation of math operands.""" - - LATEX_OPERATOR_LOOKUP: dict[str, str] = { - "**": "{val}^{{{operand}}}", - "*": r"{val} \times {operand}", - "/": r"\frac{{ {val} }}{{ {operand} }}", - "+": "{val} + {operand}", - "-": "{val} - {operand}", - } - SKIP_IF: list[str] = ["+", "-"] - - def __init__(self, instring: str, loc: int, tokens: pp.ParseResults) -> None: - """ - Process successfully parsed expressions with operands separated by an operator. - - I.e.: OPERAND OPERATOR OPERAND OPERATOR OPERAND ... - - Args: - instring (str): String that was parsed (used in error message). - loc (int): - Location in parsed string where parsing error was logged. - This is not used; we include it as pyparsing injects it alongside `instring` when setting the parse action. - tokens (pp.ParseResults): - Contains a list of the form [operand (pp.ParseResults), operator (str), - operand (pp.ParseResults), operator (str), ...]. - """ - self.value: pp.ParseResults = tokens[0] - self.values = tokens - self.instring = instring - - def __repr__(self) -> str: - """String representation.""" - first_operand = self.value[0].__repr__() - operand_operator_pairs = " ".join( - op + " " + val.__repr__() - for op, val in self._operator_operands(self.value[1:]) - ) - arithmetic_string = f"({first_operand} {operand_operator_pairs})" - return arithmetic_string - - def _operator_operands( - self, token_list: list - ) -> Iterator[tuple[str, pp.ParseResults]]: - """Generator to extract operators and operands in pairs.""" - it = iter(token_list) - while 1: - try: - yield (next(it), next(it)) - except StopIteration: - break - - def _apply_mask(self, evaluated: ARRAY_T) -> ARRAY_T: - """Util function to apply mask arrays to non-latex strings.""" - mask = self.eval_attrs.mask - try: - evaluated = evaluated.where(mask) - except AttributeError: - evaluated = evaluated.broadcast_like(mask).where(mask) - - return evaluated - - def _skip_component_on_conditional(self, component: str, operator_: str) -> bool: - """ - Conditional to skip adding to math string if element evaluates to zero. - - E.g., "0 + flow_cap" is better evaluated as simply "flow_cap". - """ - return component == "0" and operator_ in self.SKIP_IF - - @overload - @staticmethod - def _operate( - val: xr.DataArray, evaluated_operand: xr.DataArray, operator_: str - ) -> xr.DataArray: ... - - @overload - @staticmethod - def _operate( - val: xr.DataArray, evaluated_operand: EXPRVAR_T, operator_: str - ) -> EXPRVAR_T: ... - - @overload - @staticmethod - def _operate( - val: EXPRVAR_T, evaluated_operand: xr.DataArray, operator_: str - ) -> EXPRVAR_T: ... - - @staticmethod - def _operate( - val: xr.DataArray | EXPRVAR_T, - evaluated_operand: xr.DataArray | EXPRVAR_T, - operator_: str, - ) -> xr.DataArray | EXPRVAR_T: - """Apply evaluated operation on two DataArrays.""" - match operator_: - case "**": - val = val**evaluated_operand - case "*": - val = val * evaluated_operand - case "/": - val = val / evaluated_operand - case "+": - val = val + evaluated_operand - case "-": - val = val - evaluated_operand - return val - - def as_math_string(self) -> str: # noqa: D102, override - val = self.value[0].eval("math_string", self.eval_attrs) - - for operator_, operand in self._operator_operands(self.value[1:]): - evaluated_operand = operand.eval("math_string", self.eval_attrs) - # We ignore zeros that do nothing - if self._skip_component_on_conditional(evaluated_operand, operator_): - continue - if isinstance(self.value[0], type(self)): - val = "(" + val + ")" - if isinstance(operand, type(self)): - evaluated_operand = "(" + evaluated_operand + ")" - if self._skip_component_on_conditional(val, operator_): - val = evaluated_operand - else: - val = self.LATEX_OPERATOR_LOOKUP[operator_].format( - val=val, operand=evaluated_operand - ) - return val - - def as_raw(self) -> xr.DataArray: # noqa: D102, override - val = self._apply_mask(self.value[0].eval("raw", self.eval_attrs)) - - for operator_, operand in self._operator_operands(self.value[1:]): - evaluated_operand = self._apply_mask(operand.eval("raw", self.eval_attrs)) - val = self._operate(val, evaluated_operand, operator_) - return val - - def as_expr(self) -> EXPR_T: # noqa: D102, override - val = self._apply_mask(self.value[0].eval("expr", self.eval_attrs)) - - for operator_, operand in self._operator_operands(self.value[1:]): - evaluated_operand = self._apply_mask(operand.eval("expr", self.eval_attrs)) - val = self._operate(val, evaluated_operand, operator_) - return val - - -class EvalSignOp(EvalNode): - """Class for processing expressions with + or -.""" - - def __init__(self, instring: str, loc: int, tokens: pp.ParseResults) -> None: - """ - Parse action to process successfully parsed expressions with a leading + or - sign. - - Args: - instring (str): String that was parsed (used in error message). - loc (int): - Location in parsed string where parsing error was logged. - This is not used; we include it as pyparsing injects it alongside `instring` when setting the parse action. - tokens (pp.ParseResults): - Contains a list of the form [sign (str), operand (pp.ParseResults)]. - """ - self.sign, self.value = tokens[0] - self.values = tokens - self.instring = instring - - def __repr__(self) -> str: - """Programming / official string representation.""" - return str(f"({self.sign}){self.value.__repr__()}") - - # string return - @overload - def _eval(self, return_type: Literal["math_string"]) -> str: ... - - # array return - @overload - def _eval(self, return_type: Literal["raw"]) -> xr.DataArray: ... - - # expression return - @overload - def _eval(self, return_type: Literal["expr"]) -> EXPR_T: ... - - def _eval(self, return_type: RETURN_T) -> xr.DataArray | EXPR_T | str: - """Evaluate the element that will have the sign attached to it.""" - return self.value.eval(return_type, self.eval_attrs) - - def as_math_string(self) -> str: # noqa: D102 - return self.sign + self._eval("math_string") - - def as_raw(self) -> xr.DataArray: # noqa: D102, override - evaluated = self._eval("raw") - if self.sign == "-": - evaluated = -1 * evaluated - return evaluated - - def as_expr(self) -> EXPR_T: # noqa: D102, override - evaluated = self._eval("expr") - if self.sign == "-": - evaluated = -1 * evaluated - return evaluated - - -class EvalFunction(EvalNode): - """Class to process parsed functions.""" - - def __init__(self, instring: str, loc: int, tokens: pp.ParseResults) -> None: - """ - Parse action to process successfully parsed helper function strings. - - Strings must be in the following form: helper_function_name(*args, **kwargs). - - Args: - instring (str): String that was parsed (used in error message). - loc (int): - Location in parsed string where parsing error was logged. - This is not used; we include it as pyparsing injects it alongside `instring` when setting the parse action. - tokens (pp.ParseResults): - Has a dictionary component with the parsed elements: - helper_function_name (pp.ParseResults), args (list), kwargs (dict). - """ - token_dict = tokens.as_dict() - self.func_name: pp.ParseResults = token_dict["helper_function_name"] - self.args: list = token_dict["args"] - self.kwargs: dict = token_dict["kwargs"] - self.values = tokens - self.instring = instring - - def __repr__(self) -> str: - """Programming / official string representation.""" - _kwargs = ", ".join(f"{k}={v}" for k, v in self.kwargs.items()) - return f"{str(self.func_name)}(args={self.args}, kwargs={{{_kwargs}}})" - - @overload - def _arg_eval(self, return_type: Literal["math_string"], arg: Any) -> str: ... - - @overload - def _arg_eval( - self, return_type: Literal["raw"], arg: Any - ) -> xr.DataArray | list[str | float]: ... - - def _arg_eval( - self, return_type: RETURN_T, arg: Any - ) -> str | xr.DataArray | list[str | float]: - """Evaluate the arguments of the helper function.""" - if isinstance(arg, pp.ParseResults): - evaluated = arg[0].eval(return_type, self.eval_attrs) - elif isinstance(arg, list): - evaluated = [self._arg_eval(return_type, arg_) for arg_ in arg] - elif isinstance(arg, ListParser): - evaluated = arg.eval("raw", self.eval_attrs) - else: - evaluated = arg.eval(return_type, self.eval_attrs) - if isinstance(evaluated, xr.DataArray) and isinstance(arg, EvalGenericString): - evaluated = evaluated.item() - return evaluated - - @overload - def _eval(self, return_type: Literal["math_string"]) -> str: ... - - @overload - def _eval(self, return_type: Literal["raw"]) -> xr.DataArray: ... - - def _eval(self, return_type: RETURN_T) -> str | xr.DataArray | EXPR_T: - """ - Pass evaluated arguments to evaluated helper function. - - The helper function itself is created with the enclosing ``return_type`` so - that expression-route helpers can dispatch to their ``as_expr`` implementation. - Its arguments, however, are always evaluated in ``raw`` mode (never ``expr``): - helper functions must receive un-normalised inputs (``xr.DataArray`` for - parameters/lookups/dimensions and the raw model object for variables/expressions) - rather than values that have been coerced to boolean masks or ``LinearExpression``. - """ - helper_function = self.func_name.eval(return_type, self.eval_attrs) - if helper_function.ignore_mask: - self.eval_attrs = replace(self.eval_attrs, mask=xr.DataArray(True)) - - arg_return_type: RETURN_T = ( - "math_string" if return_type == "math_string" else "raw" - ) - args_ = [] - for arg in self.args: - args_.append(self._arg_eval(arg_return_type, arg)) - - kwargs_ = {} - for kwarg_name, kwarg_val in self.kwargs.items(): - kwargs_[kwarg_name] = self._arg_eval(arg_return_type, kwarg_val) - - evaluated = helper_function(*args_, **kwargs_) - return evaluated - - def as_math_string(self) -> str: # noqa: D102, override - return self._eval("math_string") - - def as_raw(self) -> xr.DataArray: # noqa: D102, override - return self._eval("raw") - - def as_expr(self) -> EXPR_T: # noqa: D102, override - return self._eval("expr") - - -class EvalHelperFuncName(EvalToCallable): - """For processing parsed helper function names.""" - - def __init__(self, instring: str, loc: int, tokens: pp.ParseResults) -> None: - """ - Parse action to process successfully parsed helper function names. - - This is a unique parse action so that we can catch invalid helper functions - most safely. - - Args: - instring (str): String that was parsed (used in error message). - loc (int): - Location in parsed string where parsing error was logged. - This is not used; we include it as pyparsing injects it alongside `instring` when setting the parse action. - tokens (pp.ParseResults): - Has one parsed element: helper_function_name (str). - """ - self.name = self.value = tokens[0] - self.instring = instring - self.loc = loc - self.values = tokens - - def __repr__(self) -> str: - """Programming / official string representation.""" - return str(self.name) - - def as_callable(self, return_type: RETURN_T) -> Callable: - """Evalluate and return the callable action of the helper function.""" - helper_functions = self.eval_attrs.helper_functions - if self.name not in helper_functions.keys(): - raise self.error_msg(f"Invalid helper function defined: {self.name}") - elif not issubclass(helper_functions[self.name], ParsingHelperFunction): - raise self.error_msg( - f"Helper function must be subclassed from linopy.declarative.helper_functions.ParsingHelperFunction: {self.name}" - ) - else: - return helper_functions[self.name](return_type, self.eval_attrs) - - -class EvalSlicedComponent(EvalNode): - """For processing of sliced parameters / decision variables.""" - - def __init__(self, instring: str, loc: int, tokens: pp.ParseResults) -> None: - """ - Process successfully parsed sliced parameters or decision variables. - - In the form of param_or_var[*slices]. - - Args: - instring (str): String that was parsed (used in error message). - loc (int): - Location in parsed string where parsing error was logged. - This is not used; we include it as pyparsing injects it alongside `instring` when setting the parse action. - tokens (pp.ParseResults): - Has a dictionary component with the parsed elements: - param_or_var_name (str), slices (list of strings). - """ - token_dict = tokens.as_dict() - self.obj_name: pp.ParseResults = token_dict["param_or_var_name"] - - self.slices: dict[str, pp.ParseResults] = { - idx["set_name"][0]: idx["slicer"][0] for idx in token_dict["slices"] - } - self.values = tokens - self.instring = instring - - def __repr__(self) -> str: - """Programming / official string representation.""" - slices = ", ".join(f"{k}={v.__repr__()}" for k, v in self.slices.items()) - return f"SLICED_{self.obj_name}[{slices}]" - - @staticmethod - def _replace_rule(index_slices: dict) -> Callable: - """ - String parsing rule to catch and replace dimension names with the names + their slices. - - E.g., `techs` -> `techs=pv`. - """ - - def __replace(term: pp.ParseResults) -> str: - if len(term) == 1: - return term - else: - replacers = {k: f"{k}={v}" for k, v in index_slices.items()} - return ( - term[0] - + term[1] - + ",".join(replacers.get(k, k) for k in term[2]) - + term[3] - ) - - return __replace - - @overload - def _eval(self, return_type: Literal["math_string"]) -> tuple[str, dict]: ... - - @overload - def _eval(self, return_type: Literal["raw"]) -> tuple[xr.DataArray, dict]: ... - - @overload - def _eval(self, return_type: Literal["expr"]) -> tuple[EXPR_T, dict]: ... - - def _eval(self, return_type: RETURN_T) -> tuple[str | xr.DataArray | EXPR_T, dict]: - """Evaluate the slice dim and vals of each slice element.""" - slices: dict[str, Any] = { - k: xr.concat(slice_, dim=k) - if isinstance(slice_ := v.eval(return_type, self.eval_attrs), list) - else slice_ - for k, v in self.slices.items() - } - - evaluated = self.obj_name.eval(return_type, self.eval_attrs) - return evaluated, slices - - def as_math_string(self) -> str: # noqa: D102, override - evaluated, slices = self._eval("math_string") - singular_slice_refs = { - self.eval_attrs.math.dimensions[k].iterator: v for k, v in slices.items() - } - id_ = pp.Combine( - pp.Word(pp.alphas, pp.alphanums) - + pp.ZeroOrMore("_" + pp.Word(pp.alphanums)) - + pp.Opt("_") - ) - id_formatted = pp.Combine("\\" + pp.Word(pp.alphas) + "{" + id_ + "}") - obj_parser = id_formatted + pp.Opt( - r"_\text{" + pp.Group(pp.DelimitedList(id_)) + "}" - ) - obj_parser.set_parse_action(self._replace_rule(singular_slice_refs)) - return obj_parser.parse_string(evaluated, parse_all=True)[0] - - def as_raw(self) -> xr.DataArray: # noqa: D102, override - evaluated, slices = self._eval("raw") - return evaluated.sel(**slices) - - def as_expr(self) -> EXPR_T: - evaluated, slices = self._eval("expr") - return evaluated.sel(**slices) - - -class EvalIndexSlice(EvalNode): - """For processing `$slice` expressions.""" - - def __init__(self, instring: str, loc: int, tokens: pp.ParseResults) -> None: - """ - Process successfully parsed expression index `$slice` references. - - Args: - instring (str): String that was parsed (used in error message). - loc (int): - Location in parsed string where parsing error was logged. - This is not used; we include it as pyparsing injects it alongside `instring` when setting the parse action. - tokens (pp.ParseResults): - Has one parsed element containing the index slice name (str). - """ - self.name: str = tokens[0] - self.values = tokens - self.instring = instring - - def __repr__(self) -> str: - """Programming / official string representation.""" - return "REFERENCE:" + str(self.name) - - @overload - def _eval(self, return_type: Literal["math_string"], as_values: bool) -> str: ... - - @overload - def _eval( - self, return_type: Literal["raw"], as_values: bool - ) -> xr.DataArray | list[xr.DataArray]: ... - - def _eval( - self, return_type: RETURN_T, as_values: bool - ) -> str | xr.DataArray | list[xr.DataArray]: - """Evaluate the referenced `slice`.""" - self.eval_attrs = replace(self.eval_attrs, as_values=as_values) - return self.eval_attrs.slice_dict[self.name][0].eval( - return_type, self.eval_attrs - ) - - def as_math_string(self) -> str: # noqa: D102, override - return self._eval("math_string", False) - - def as_raw(self) -> xr.DataArray | list[xr.DataArray]: # noqa: D102, override - evaluated = self._eval("raw", True) - return evaluated - - -class EvalSubExpressions(EvalNode): - """For processing sub-expressions.""" - - def __init__(self, instring: str, loc: int, tokens: pp.ParseResults) -> None: - """ - Process successfully parsed `$sub_expressions`. - - Args: - instring (str): String that was parsed (used in error message). - loc (int): - Location in parsed string where parsing error was logged. - This is not used; we include it as pyparsing injects it alongside `instring` when setting the parse action. - tokens (pp.ParseResults): - Has one parsed element containing the sub_expression name (str). - """ - self.name: str = tokens[0] - self.values = tokens - self.instring = instring - - def __repr__(self) -> str: - """Programming / official string representation.""" - return "SUB_EXPRESSION:" + str(self.name) - - @overload - def _eval(self, return_type: Literal["math_string"]) -> str: ... - - @overload - def _eval(self, return_type: Literal["raw"]) -> xr.DataArray: ... - - @overload - def _eval(self, return_type: Literal["expr"]) -> EXPR_T: ... - - def _eval(self, return_type: RETURN_T) -> str | xr.DataArray | EXPR_T: - """Evaluate the referenced sub_expression.""" - return self.eval_attrs.sub_expression_dict[self.name][0].eval( - return_type, self.eval_attrs - ) - - def as_math_string(self) -> str: # noqa: D102, override - return self._eval("math_string") - - def as_raw(self) -> xr.DataArray: # noqa: D102, override - return self._eval("raw") - - def as_expr(self) -> EXPR_T: # noqa: D102, override - return self._eval("expr") - - -class EvalNumber(EvalNode): - """For processing numbers.""" - - def __init__(self, instring: str, loc: int, tokens: pp.ParseResults) -> None: - """ - Process successfully parsed numbers. - - Catches integers (1), floats (1.), and in scientific notation (1e1). - Also capture infinity (inf/.inf). - - Args: - instring (str): String that was parsed (used in error message). - loc (int): - Location in parsed string where parsing error was logged. - This is not used; we include it as pyparsing injects it alongside `instring` when setting the parse action. - tokens (pp.ParseResults): - Has one parsed element containing the number (str). - """ - self.value = tokens[0] - self.values = tokens - self.instring = instring - - def __repr__(self) -> str: - """Programming / official string representation.""" - return "NUM:" + str(self.value) - - def as_math_string(self) -> str: # noqa: D102, override - return re.sub( - r"([\d]+?)e([+-])([\d]+)", - r"\1\\mathord{\\times}10^{\2\3}", - f"{float(self.value):.6g}", - ) - - def as_raw(self) -> xr.DataArray: # noqa: D102, override - return xr.DataArray(float(self.value), name=float(self.value)) - - -class ListParser(EvalNode): - """For parsing lists.""" - - def __init__(self, instring: str, loc: int, tokens: pp.ParseResults) -> None: - """ - Process successfully parsed lists of generic strings. - - This is required since we call "eval()" on all elements of the where string, - so lists of strings need to be evaluatable as a whole "package". - - Args: - instring (str): String that was parsed (used in error message). - loc (int): - Location in parsed string where parsing error was logged. - This is not used; we include it as pyparsing injects it alongside `instring` when setting the parse action. - tokens (pp.ParseResults): a list of parsed string elements. - """ - self.val = tokens - self.instring = instring - - def __repr__(self) -> str: - """Programming / official string representation.""" - return f"{self.val}" - - def as_math_string(self) -> str: # noqa: D102, override - input_list = self.as_raw() - return "[" + ",".join(str(i.name) for i in input_list) + "]" - - def as_raw(self) -> list[xr.DataArray]: # noqa: D102, override - values = [val.eval("raw", self.eval_attrs) for val in self.val] - # strings and numbers are returned as xarray arrays of size 1, - # so we extract those values. - return values - - -class EvalUnslicedComponent(EvalNode): - """Evaluation of unsliced components.""" - - def __init__(self, instring: str, loc: int, tokens: pp.ParseResults) -> None: - """ - Parse action to process successfully parsed generic strings. - - This is required since we call "eval()" on all elements of the where string, - so even arbitrary strings (used in comparison operations) need to be evaluatable. - - Args: - instring (str): String that was parsed (used in error message). - loc (int): - Location in parsed string where parsing error was logged. - This is not used; we include it as pyparsing injects it alongside `instring` when setting the parse action. - tokens (pp.ParseResults): Has one parsed element: string name (str). - """ - self.val = tokens[0] - self.name = str(self.val) - self.values = tokens - self.instring = instring - - def __repr__(self) -> str: - """Programming / official string representation.""" - return f"COMPONENT:{self.name}" - - def as_math_string(self) -> str: # noqa: D102, override - self.eval_attrs = replace(self.eval_attrs, as_values=False) - evaluated = self.as_raw() - self.eval_attrs.references.add(self.name) - - if "math_repr" in evaluated.attrs: - data_var_string = evaluated.attrs["math_repr"] - else: - data_var_string = self.name - - return data_var_string - - def as_raw(self) -> xr.DataArray: # noqa: D102, override - math_def = self.eval_attrs.math.find(self.name) - group = math_def._group - self.eval_attrs.references.add(self.name) - if group in ["parameters", "lookups", "dimensions"]: - # A parameter/lookup/dimension defined in the math but absent from the - # input data resolves to its default (NaN if none is set). - try: - evaluated = self.eval_attrs.input_data[self.name] - except KeyError: - evaluated = xr.DataArray(np.nan) - else: - # Model entries (variables / expressions): return the raw model object - # unchanged. Defaults only apply to parameters/lookups/dimensions. A model - # entry that was never built (e.g. skipped because its mask was empty) - # resolves to a NaN expression rather than raising. - try: - return getattr(self.eval_attrs.model, group)[self.name] - except KeyError: - return LinearExpression(xr.DataArray(np.nan), self.eval_attrs.model) - if evaluated.isnull().any() and pd.notna(default := math_def["default"]): - evaluated = evaluated.fillna(default) - return evaluated - - def as_expr(self) -> xr.DataArray | EXPR_T: # noqa: D102, override - math_def = self.eval_attrs.math.find(self.name) - group = math_def._group - self.eval_attrs.references.add(self.name) - if group in ["parameters", "lookups", "dimensions"]: - # Parameters / lookups / dimensions are scalar coefficients: keep as DataArray. - try: - evaluated = self.eval_attrs.input_data[self.name] - except KeyError: - evaluated = xr.DataArray(np.nan) - else: - # Model entries: normalise Variable -> LinearExpression immediately so that - # arithmetic composition on the expression route always sees an expression. - try: - evaluated = _to_linexpr( - getattr(self.eval_attrs.model, group)[self.name] - ) - except KeyError: - evaluated = LinearExpression( - xr.DataArray(np.nan), self.eval_attrs.model - ) - if evaluated.isnull().any() and pd.notna(default := math_def["default"]): - evaluated = evaluated.fillna(default) - return evaluated - - -class EvalGenericString(EvalNode): - """For generic string parsing.""" - - def __init__(self, instring: str, loc: int, tokens: pp.ParseResults) -> None: - """ - Process successfully parsed generic strings. - - This is required since we call "eval()" on all elements of the where string, - so even arbitrary strings (used in comparison operations) need to be evaluatable. - - Args: - instring (str): String that was parsed (used in error message). - loc (int): - Location in parsed string where parsing error was logged. - This is not used; we include it as pyparsing injects it alongside `instring` when setting the parse action. - tokens (pp.ParseResults): Has one parsed element: string name (str). - """ - self.val = tokens[0] - self.instring = instring - - def __repr__(self) -> str: - """Programming / official string representation.""" - return f"STRING:{self.val}" - - def as_math_string(self): # noqa: D102, override - return str(self.val) - - def as_raw(self) -> xr.DataArray: # noqa: D102, override - return xr.DataArray(str(self.val), name=str(self.val)) - - -def helper_function_parser( - *args: pp.ParserElement, - generic_identifier: pp.ParserElement, - allow_function_in_function: bool = False, -) -> pp.ParserElement: - """ - Process helper functions of the form `helper_function(*args, **kwargs)`. - - Helper functions can accept other parser elements as arguments, - i.e., components, parameters or variables, numbers, and other functions. - - Available helper functions are predefined in calliope.backend.helper_functions. - - Calling an unavailable helper will lead to a raised exception on evaluating the - parsed element. - - Based partially on: # https://stackoverflow.com/questions/61807705/pyparsing-generic-python-function-args-and-kwargs - - Args: - *args (pp.ParserElement): - Parser elements that can be arguments in the function (e.g., "number", "sliced_param_or_var"). - NOTE: the order of inclusion in the args list matters. The parser will parse based on first matches. - generic_identifier (pp.ParserElement): - Parser for valid python variables without leading underscore and not called "inf". - This parser has no parse action. - allow_function_in_function (bool, optional): - If True, allows functions to be defined inside functions. - Nested functions are evaluated from the greatest level of nesting up to the main helper function. - Defaults to True. - - Returns: - pp.ParserElement: - Parser for functions which will call the function with the specified - arguments on evaluation. - """ - helper_function = pp.Forward() - allowed_parser_elements_in_args = list(args) - lpar = pp.Suppress("(") - rpar = pp.Suppress(")") - - helper_function_name = generic_identifier.set_results_name("helper_function_name") - helper_function_name.set_parse_action(EvalHelperFuncName) - - if allow_function_in_function: - allowed_parser_elements_in_args.insert(0, helper_function) - - arg_values = pp.MatchFirst(allowed_parser_elements_in_args) + pp.NotAny("=") - - # define function arguments - arglist = pp.DelimitedList(arg_values.copy()) - args_ = pp.Group(arglist).set_results_name("args") - - # define function keyword arguments - key = generic_identifier + pp.Suppress("=") - kwarg_list = pp.DelimitedList(pp.dict_of(key, arg_values)) - kwargs_ = pp.Group(kwarg_list).set_results_name("kwargs") - - # build generic function - helper_func_args = args_ + pp.Suppress(",") + kwargs_ | pp.Opt( - args_, default=[] - ) + pp.Opt(kwargs_, default={}) - helper_function << ( - pp.Combine(helper_function_name + lpar) + helper_func_args + rpar - ) - - helper_function.set_parse_action(EvalFunction) - - return helper_function - - -def sliced_param_or_var_parser( - slicer: Iterable[pp.ParserElement], - generic_identifier: pp.ParserElement, - unsliced_object: pp.ParserElement, - allow_slice_references: bool = True, -) -> pp.ParserElement: - """ - Process strings representing sliced model parameters or variables. - - E.g. "source_use_max[node, tech]". - - If a parameter, must be a data variable in the Model.inputs xarray dataset. - - If a variable, must be an optimisation problem decision variable. - - The parser will not verify whether it has parsed a valid parameter or variable until - evaluation. - - Args: - slicer (Iterable[pp.ParserElement]): - List of parsers that can be used to define the slice of a parameter or variable. - E.g., "number", "evaluatable_identifier", "unsliced_param_or_var". - These elements will be parsed in the order they are given. - generic_identifier (pp.ParserElement): - Parser that evaluates to a string. - unsliced_object (pp.ParserElement): - Parser for valid backend objects. - On evaluation, this parser will access the backend object from the backend dataset. - allow_slice_references (bool): - If True, allow reference to `slice` expressions (e.g. `$bar` in `foo[bars=$bar]`). - Defaults to True. - - Returns: - pp.ParserElement: - Parser which returns a dictionary with name of parameter/variable and list - of index items as separate entries on. - """ - lspar = pp.Suppress("[") - rspar = pp.Suppress("]") - - direct_slicer = pp.MatchFirst(slicer) - if allow_slice_references: - slicer_ref = pp.Suppress(SUB_EXPRESSION_CLASSIFIER) + generic_identifier - slicer_ref.set_parse_action(EvalIndexSlice) - slicer = (slicer_ref | direct_slicer)("slicer") - else: - slicer = direct_slicer("slicer") - - slice = pp.Group(generic_identifier("set_name") + pp.Suppress("=") + slicer) - - slices = pp.Group(pp.DelimitedList(slice))("slices") - sliced_object_name = unsliced_object("param_or_var_name") - - sliced_param_or_var = pp.Combine(sliced_object_name + lspar) + slices + rspar - sliced_param_or_var.set_parse_action(EvalSlicedComponent) - - return sliced_param_or_var - - -def sub_expression_parser(generic_identifier: pp.ParserElement) -> pp.ParserElement: - """ - Parse strings prepended with the YAML constraint sub-expression classifier `$`. - - E.g. "$my_sub_expr" - - Args: - generic_identifier (pp.ParserElement): - Parser for valid python variables without leading underscore and not called "inf". - This parser has no parse action. - - Returns: - pp.ParserElement: - Parser which produces a dictionary of the form {"sub_expression": "my_sub_expression"} on evaluation. - """ - sub_expression = pp.Combine( - pp.Suppress(SUB_EXPRESSION_CLASSIFIER) + generic_identifier - ) - sub_expression.set_parse_action(EvalSubExpressions) - - return sub_expression - - -def unsliced_object_parser(valid_component_names: Iterable[str]) -> pp.ParserElement: - """ - Parse unsliced objects and identify their corresponding parse actions. - - Creates a copy of the generic identifier and sets a parse action to find the string in - the list of input parameters or optimisation decision variables. - - Args: - valid_component_names (Iterable[str]): A - All backend object names, to ensure they are captured by this parser function. - - Returns: - pp.ParserElement: - Copy of input parser with added parse action to lookup an unsliced - parameter/variable value - """ - unsliced_param_or_var = pp.one_of(valid_component_names, as_keyword=True) - unsliced_param_or_var.set_parse_action(EvalUnslicedComponent) - - return unsliced_param_or_var - - -def evaluatable_identifier_parser( - identifier: pp.ParserElement, valid_component_names: Iterable -) -> pp.ParserElement: - """ - Create an evaluatable copy of the generic identifier that will return a string or a model component as an array. - - Args: - identifier (pp.ParserElement): - Parser for valid python variables without leading underscore and not called "inf". - This parser has no parse action. - valid_component_names (Iterable[str]): A - All backend object names, to ensure they are *not* captured by this parser function. - - Returns: - pp.ParserElement: - Parser for valid python variables without leading underscore and not called "inf". - Evaluates to a string or an array (if it is a model component). - """ - evaluatable_identifier = ( - ~pp.one_of(valid_component_names, as_keyword=True) + identifier - ).set_parse_action(EvalGenericString) - - return evaluatable_identifier - - -def list_parser(*args: pp.ParserElement) -> pp.ParserElement: - """ - Parse strings which define a list of other strings or numbers. - - Lists are defined as anything wrapped in square brackets (`[]`). - - Args: - *args (pp.ParserElement): - Parser elements that can be list elements (e.g., "number", "evaluatable_identifier", "unsliced_param_or_var"). - These elements will be parsed in the order they are given. - - Returns: - pp.ParserElement: Parser for valid lists of strings and/or numbers. - """ - list_elements = pp.MatchFirst(args) - id_list = pp.Suppress("[") + pp.DelimitedList(list_elements) + pp.Suppress("]") - id_list.set_parse_action(ListParser) - return id_list - - -def setup_base_parser_elements() -> tuple[pp.ParserElement, pp.ParserElement]: - """ - Setup parser elements that will be components of other parsers. - - Returns: - tuple[pp.ParserElement, pp.ParserElement]: (number, generic_identifier) - number: parser for numbers (integer, float, scentific notation, "inf"/".inf"). - generic_identifier: parser for valid python variables without leading - underscore and not called "inf". This parser has no parse action. - """ - inf_kw = pp.Combine(pp.Opt(pp.Suppress(".")) + pp.Keyword("inf", caseless=True)) - number = pp.pyparsing_common.number | inf_kw - generic_identifier = ~inf_kw + pp.Word(pp.alphas, pp.alphanums + "_") - - number.set_parse_action(EvalNumber) - - return number, generic_identifier - - -def arithmetic_parser(*args, arithmetic: pp.Forward | None = None) -> pp.Forward: - """ - Parsing grammar to combine equation elements using basic arithmetic (+, -, *, /, **). - - Can handle the difference between a sign (e.g., -1,+1) and a addition/subtraction (0 - 1, 0 + 1). - Whitespace is ignored on parsing (i.e., "1+1+foo" is equivalent to "1 + 1 + foo"). - - Args: - *args: arguments in the form of a list. These can be: - helper_function (pp.ParserElement): parsing grammar to process helper functions - of the form `helper_function(*args, **kwargs)`. - sliced_param_or_var (pp.ParserElement): parser for sliced parameters or variables, e.g. "foo[bar]" - sub_expression (pp.ParserElement): parser for constraint sub expressions, e.g. "$foo" - unsliced_param_or_var (pp.ParserElement): parser for unsliced parameters or variables, e.g. "foo" - number (pp.ParserElement): parser for numbers (integer, float, scientific notation, "inf"/".inf"). - arithmetic (pp.Forward | None, optional): If given, add arithmetic rules to this - existing parsing rule (otherwise, arithmetic rules will be a newly generated rule). - Defaults to None. - - Returns: - pp.Forward: parser for strings which use arithmetic operations to combine other parser elements. - """ - signop = pp.one_of(["+", "-"]) - multop = pp.one_of(["*", "/"]) - expop = pp.Literal("**") - if arithmetic is None: - arithmetic = pp.Forward() - - arithmetic <<= pp.infixNotation( - # the order matters if two could capture the same string, e.g. "inf". - pp.MatchFirst(args), - [ - (signop, 1, pp.opAssoc.RIGHT, EvalSignOp), - (expop, 2, pp.opAssoc.LEFT, EvalOperatorOperand), - (multop, 2, pp.opAssoc.LEFT, EvalOperatorOperand), - (signop, 2, pp.opAssoc.LEFT, EvalOperatorOperand), - ], - ) - - return arithmetic - - -def equation_comparison_parser(arithmetic: pp.ParserElement) -> pp.ParserElement: - """ - Parsing grammar to combine equation elements either side of a comparison operator (<= >= ==). - - Whitespace is ignored on parsing (i.e., "1+foo==$bar" is equivalent to "1 + 1 == $bar"). - - Args: - arithmetic (pp.ParserElement): - Parser for arithmetic operations to combine other parser elements. - - Returns: - pp.ParserElement: - Parser for strings of the form "LHS OPERATOR RHS". - """ - comparison_operators = pp.one_of(["<=", ">=", "="]) - equation_comparison = arithmetic + comparison_operators + arithmetic - equation_comparison.set_parse_action(EvalComparisonOp) - - return equation_comparison - - -def generate_slice_parser(valid_component_names: Iterable) -> pp.ParserElement: - """ - Create parser for index slice reference expressions. - - These expressions are linked to the equation expression by e.g. `$bar` in `foo[bars=$bar]`. - Unlike sub-expressions and equation expressions, these strings cannot contain arithmetic - nor references to sub expressions. - - Args: - valid_component_names (Iterable): - Allowed names for optimisation problem components (parameters, decision variables, expressions), - to allow the parser to separate these from generic strings. - - Returns: - pp.ParserElement: Parser for expression strings under the constraint key "slices". - """ - number, identifier = setup_base_parser_elements() - evaluatable_identifier = evaluatable_identifier_parser( - identifier, valid_component_names - ) - unsliced_param = unsliced_object_parser(valid_component_names) - helper_func_list = list_parser(number, evaluatable_identifier, unsliced_param) - slice_list = list_parser(number, evaluatable_identifier) - sliced_param = sliced_param_or_var_parser( - [number, evaluatable_identifier, slice_list], - identifier, - unsliced_param, - allow_slice_references=False, - ) - - helper_function = helper_function_parser( - sliced_param, - unsliced_param, - number, - helper_func_list, - evaluatable_identifier, - generic_identifier=identifier, - allow_function_in_function=True, - ) - - return ( - helper_function - | sliced_param - | unsliced_param - | number - | slice_list - | evaluatable_identifier - ) - - -def generate_sub_expression_parser(valid_component_names: Iterable) -> pp.Forward: - """ - Create parser for sub expressions. - - These expressions are linked to the equation expression by e.g. `$bar`. - This parser allows arbitrarily nested arithmetic and function calls (and arithmetic inside function calls) - and reference to index slice expressions. - - Args: - valid_component_names (Iterable): - Allowed names for optimisation problem components (parameters, decision variables, expressions), - to allow the parser to separate these from generic strings. - - Returns: - pp.ParserElement: Parser for expression strings under the constraint key "sub_expressions". - """ - number, identifier = setup_base_parser_elements() - evaluatable_identifier = evaluatable_identifier_parser( - identifier, valid_component_names - ) - unsliced_param = unsliced_object_parser(valid_component_names) - helper_func_list = list_parser(number, evaluatable_identifier, unsliced_param) - slice_list = list_parser(number, evaluatable_identifier) - sliced_param = sliced_param_or_var_parser( - [number, evaluatable_identifier, slice_list], identifier, unsliced_param - ) - - arithmetic = pp.Forward() - helper_function = helper_function_parser( - arithmetic, - helper_func_list, - evaluatable_identifier, - generic_identifier=identifier, - ) - arithmetic = arithmetic_parser( - helper_function, sliced_param, number, unsliced_param, arithmetic=arithmetic - ) - return arithmetic - - -def generate_arithmetic_parser(valid_component_names: Iterable) -> pp.ParserElement: - """ - Create parser for arithmetic expressions (+, -, /, *, **). - - This parser allows arbitrarily nested arithmetic and function calls (and arithmetic inside function calls) - and reference to sub-expressions and index slice expressions. - - Args: - valid_component_names (Iterable): - Allowed names for optimisation problem components (parameters, decision variables, expressions), - to allow the parser to separate these from generic strings. - - Returns: - pp.ParserElement: Partial parser for expression strings under the constraint key "equation/equations". - """ - number, identifier = setup_base_parser_elements() - evaluatable_identifier = evaluatable_identifier_parser( - identifier, valid_component_names - ) - unsliced_param = unsliced_object_parser(valid_component_names) - helper_func_list = list_parser(number, evaluatable_identifier, unsliced_param) - slice_list = list_parser(number, evaluatable_identifier) - sliced_param = sliced_param_or_var_parser( - [number, evaluatable_identifier, slice_list], identifier, unsliced_param - ) - sub_expression = sub_expression_parser(identifier) - - arithmetic = pp.Forward() - helper_function = helper_function_parser( - arithmetic, - helper_func_list, - evaluatable_identifier, - generic_identifier=identifier, - ) - arithmetic = arithmetic_parser( - helper_function, - sub_expression, - sliced_param, - number, - unsliced_param, - arithmetic=arithmetic, - ) - - return arithmetic - - -def generate_equation_parser(valid_component_names: Iterable) -> pp.ParserElement: - """ - Create parser for equation expressions of the form LHS OPERATOR RHS (e.g. `foo == 1 + bar`). - - This parser allows arbitrarily nested arithmetic and function calls (and arithmetic inside function calls) - and reference to sub-expressions and index slice expressions. - - Args: - valid_component_names (Iterable): - Allowed names for optimisation problem components (parameters, decision variables, expressions), - to allow the parser to separate these from generic strings. - - Returns: - pp.ParserElement: Parser for expression strings under the constraint key "equation/equations". - """ - arithmetic = generate_arithmetic_parser(valid_component_names) - equation_comparison = equation_comparison_parser(arithmetic) - - return equation_comparison diff --git a/linopy/declarative/grammar.py b/linopy/declarative/grammar.py new file mode 100644 index 000000000..f936ea42e --- /dev/null +++ b/linopy/declarative/grammar.py @@ -0,0 +1,635 @@ +""" +Linopy declarative math grammar module. + +This module contains the AST node types produced when parsing declarative math +strings, and the pyparsing grammars that produce them. Nodes are pure data; +all evaluation logic lives in :mod:`linopy.declarative.evaluate`. + +The infix-notation grammar structure is adapted from the pyparsing `eval_arith.py` +example (https://github.com/pyparsing/pyparsing/blob/master/examples/eval_arith.py, +MIT licensed). +""" + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass, field, fields +from functools import cache +from typing import Literal + +import pyparsing as pp + +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.""" + +COMPONENT_CATEGORY_T = Literal["any", "dimension", "input", "result"] + + +# --------------------------------------------------------------------------- +# AST nodes +# --------------------------------------------------------------------------- + + +@dataclass(frozen=True, kw_only=True) +class Node: + """ + Base class of all declarative math AST nodes. + + Nodes are immutable data produced by the grammars in this module and consumed + by the walkers in :mod:`linopy.declarative.evaluate`. + """ + + instring: str = field(repr=False, compare=False) + """The full source string this node was parsed from (used in error messages).""" + + +@dataclass(frozen=True, kw_only=True) +class Constant(Node): + """A literal number (including `inf`), boolean, or generic string.""" + + value: float | bool | str + + +@dataclass(frozen=True, kw_only=True) +class ListNode(Node): + """A literal list of items, e.g. `[a, b, 1]`.""" + + items: tuple[Node, ...] + + +@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" + + +@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] + + +@dataclass(frozen=True, kw_only=True) +class SliceRef(Node): + """A `$name` reference to a named slicer, valid only inside slice brackets.""" + + name: str + + +@dataclass(frozen=True, kw_only=True) +class SubExprRef(Node): + """A `$name` reference to a named sub-expression.""" + + name: str + + +@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] + + +@dataclass(frozen=True, kw_only=True) +class Unary(Node): + """A unary operation: leading `+`/`-` sign or boolean `not`.""" + + op: str + operand: Node + + +@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], ...] + + +@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 + + +@dataclass(frozen=True, kw_only=True) +class Subset(Node): + """A dimension subset condition, e.g. `[a, b] in node`.""" + + items: tuple[Node, ...] + dim: Node + + +@dataclass(frozen=True, kw_only=True) +class ConfigRef(Node): + """A reference to a build-configuration option, e.g. `config.foo`.""" + + option: str + + +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] + + +# --------------------------------------------------------------------------- +# Parse actions +# --------------------------------------------------------------------------- + + +def _number_action(instring: str, loc: int, tokens: pp.ParseResults) -> Constant: + return Constant(value=float(tokens[0]), instring=instring) + + +def _string_action(instring: str, loc: int, tokens: pp.ParseResults) -> Constant: + return Constant(value=str(tokens[0]), instring=instring) + + +def _bool_action(instring: str, loc: int, tokens: pp.ParseResults) -> Constant: + return Constant(value=str(tokens[0]).lower() == "true", instring=instring) + + +def _list_action(instring: str, loc: int, tokens: pp.ParseResults) -> ListNode: + return ListNode(items=tuple(tokens), instring=instring) + + +def _component_action( + category: COMPONENT_CATEGORY_T, +) -> Callable[[str, int, pp.ParseResults], Component]: + def _action(instring: str, loc: int, tokens: pp.ParseResults) -> Component: + return Component(name=str(tokens[0]), category=category, instring=instring) + + return _action + + +def _sliced_action(instring: str, loc: int, tokens: pp.ParseResults) -> Sliced: + slices = {str(grp["set_name"][0]): grp["slicer"][0] for grp in tokens["slices"]} + return Sliced(obj=tokens["obj"], slices=slices, instring=instring) + + +def _slice_ref_action(instring: str, loc: int, tokens: pp.ParseResults) -> SliceRef: + return SliceRef(name=str(tokens[0]), instring=instring) + + +def _sub_expr_ref_action( + instring: str, loc: int, tokens: pp.ParseResults +) -> SubExprRef: + return SubExprRef(name=str(tokens[0]), instring=instring) + + +def _call_action(instring: str, loc: int, tokens: pp.ParseResults) -> Call: + 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 Call(func=token_dict["func"], args=args, kwargs=kwargs, instring=instring) + + +def _unary_action(instring: str, loc: int, tokens: pp.ParseResults) -> Unary: + op, operand = tokens[0] + return Unary(op=str(op), operand=operand, instring=instring) + + +def _arith_action(instring: str, loc: int, tokens: pp.ParseResults) -> Arith: + items = tokens[0] + rest = tuple( + (str(op), operand) for op, operand in zip(items[1::2], items[2::2], strict=True) + ) + return Arith(first=items[0], rest=rest, instring=instring) + + +def _compare_action(instring: str, loc: int, tokens: pp.ParseResults) -> Compare: + lhs, op, rhs = tokens + return Compare(lhs=lhs, op=str(op), rhs=rhs, instring=instring) + + +def _subset_action(instring: str, loc: int, tokens: pp.ParseResults) -> Subset: + items, dim = tokens + return Subset(items=tuple(items), dim=dim, instring=instring) + + +def _config_action(instring: str, loc: int, tokens: pp.ParseResults) -> ConfigRef: + return ConfigRef(option=str(tokens[0]), instring=instring) + + +# --------------------------------------------------------------------------- +# 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(_number_action) + 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_action(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( + _string_action + ) + + +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(_list_action) + + +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_action) + + +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(_slice_ref_action) + 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_action) + + +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(_sub_expr_ref_action) + + +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_action), + (expop, 2, pp.opAssoc.LEFT, _arith_action), + (multop, 2, pp.opAssoc.LEFT, _arith_action), + (signop, 2, pp.opAssoc.LEFT, _arith_action), + ], + ) + return arithmetic + + +# --------------------------------------------------------------------------- +# Grammar entry points +# --------------------------------------------------------------------------- + + +@cache +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. + """ + 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=False, + ) + call = _call_parser( + sliced, + component, + number, + call_list, + string, + identifier=identifier, + allow_nested_calls=True, + ) + return call | sliced | component | number | slicer_list | string + + +@cache +def sub_expression_grammar(component_names: frozenset[str]) -> pp.Forward: + """ + 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. + """ + 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 + ) + arithmetic = pp.Forward() + call = _call_parser(arithmetic, call_list, string, identifier=identifier) + return _arithmetic_rules(call, sliced, number, component, arithmetic=arithmetic) + + +@cache +def arithmetic_grammar(component_names: frozenset[str]) -> pp.Forward: + """ + 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. + """ + 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 + ) + sub_expression = _sub_expression_ref_parser(identifier) + arithmetic = pp.Forward() + call = _call_parser(arithmetic, call_list, string, identifier=identifier) + return _arithmetic_rules( + call, sub_expression, sliced, number, component, arithmetic=arithmetic + ) + + +@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_action) + + +@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( + _config_action + ) + bool_operand = ( + pp.Keyword("True", caseless=True) | pp.Keyword("False", caseless=True) + ).set_parse_action(_bool_action) + 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_action) + + 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_action) + + 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_action), + (andorop, 2, pp.opAssoc.LEFT, _arith_action), + ], + ) diff --git a/linopy/declarative/helper_functions.py b/linopy/declarative/helper_functions.py deleted file mode 100644 index 7604ad58f..000000000 --- a/linopy/declarative/helper_functions.py +++ /dev/null @@ -1,964 +0,0 @@ -# Copyright (C) since 2013 Calliope contributors listed in AUTHORS. -# Licensed under the Apache 2.0 License (see LICENSE file). - -""" -Functions that can be used to process data in math `mask` and `expression` strings. - -`NAME` is the function name to use in the math strings. -""" - -import functools -import re -from abc import ABC, abstractmethod -from collections.abc import Mapping -from typing import Any, Literal, overload - -import numpy as np -import xarray as xr - -from linopy.declarative.eval_attrs import EvalAttrs -from linopy.expressions import LinearExpression - -DTYPE_OPTIONS = { - "string": str, - "float": float, - "bool": bool, - "datetime": np.datetime64, - "date": np.datetime64, - "integer": int, -} -_registry: dict[ - Literal["mask", "expression"], dict[str, type["ParsingHelperFunction"]] -] = {"mask": {}, "expression": {}} - - -class ParsingHelperFunction(ABC): - """Abstract base class for helper function parsing.""" - - def __init__( - self, return_type: Literal["raw", "expr", "math_string"], attrs: "EvalAttrs" - ) -> None: - """ - Abstract helper function class, which all helper functions must subclass. - - The abstract properties and methods defined here must be defined by all helper functions. - """ - self._return_type = return_type - self._attrs = attrs - - @property - @abstractmethod - def ALLOWED_IN(self) -> list[Literal["mask", "expression"]]: - """List of parseable math strings that this function can be accessed from.""" - - @property - @abstractmethod - def NAME(self) -> str: - """Helper function name that is used in the math expression/mask string.""" - - @property - def ignore_mask(self) -> bool: - """If True, `mask` arrays will not be applied to the incoming data variables (valid for expression helpers).""" - return False - - @abstractmethod - def as_math_string(self, *args: Any, **kwargs: Any) -> str: - """ - Method to update LaTeX math strings to include the action applied by the helper function. - - This method is called when the class is initialised with ``return_type=math_string``. - """ - - @abstractmethod - def as_raw(self, *args: Any, **kwargs: Any) -> LinearExpression | xr.DataArray: - """ - Method to apply the helper function to provide an n-dimensional array output. - - This method is called when the class is initialised with ``return_type=raw`` and, - by default, ``return_type=expr`` (see :meth:`as_expr`). - """ - - def as_expr(self, *args: Any, **kwargs: Any) -> LinearExpression | xr.DataArray: - """ - Method to apply the helper function on the expression route. - - This method is called when the class is initialised with ``return_type=expr``. - By default it 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 __call__(self, *args: Any, **kwargs: Any) -> Any: - """ - When a helper function is accessed by evaluating a parsing string, this method is called. - - The value of `return_type` on initialisation of the class defines whether this - method returns either: - - a string (``return_type=math_string``, via :meth:`as_math_string`) - - the raw data (``return_type=raw``, via :meth:`as_raw`) - - a linopy expression (``return_type=expr``, via :meth:`as_expr`) - """ - if self._return_type == "math_string": - return self.as_math_string(*args, **kwargs) - elif self._return_type == "raw": - return self.as_raw(*args, **kwargs) - elif self._return_type == "expr": - return self.as_expr(*args, **kwargs) - else: - raise ValueError( - f"Unknown helper function return type: {self._return_type!r}" - ) - - def __init_subclass__(cls) -> None: - """ - Override subclass definition. - - 1. Do not allow new helper functions to have a name that is already defined (be it a built-in function or a custom function). - 2. Wrap helper function __call__ in a check for the function being allowed in specific parsing string types. - """ - super().__init_subclass__() - for allowed in cls.ALLOWED_IN: - if cls.NAME in _registry[allowed].keys(): - raise ValueError( - f"`{allowed}` string helper function `{cls.NAME}` already exists" - ) - for allowed in cls.ALLOWED_IN: - _registry[allowed][cls.NAME] = cls - - @staticmethod - def _update_iterator( - instring: str, - iterator_converter: dict[str, str], - method: Literal["add", "replace"], - ) -> str: - r""" - Utility function for generating latex strings in multiple helper functions. - - Find an iterator in the iterator substring of the component string - (anything wrapped in `_text{}`). Other parts of the iterator substring can be anything - except curly braces, e.g. the standalone `foo` will be found here and acted upon: - `\\textit{my_param}_\text{bar,foo,foo=bar,foo+1}` - - Args: - instring (str): String in which the iterator substring can be found. - iterator_converter (dict[str, str]): - key: the iterator to search for. - val: The new string to **append** to the iterator name (if method = add) or **replace**. - method (Literal[add, replace]): Whether to add to the iterator or replace it entirely - Returns: - str: `instring`, but with `iterator` replaced with `iterator + new_string` - """ - - def __replace_in_iterator(matched): - iterator_list = matched.group(2).split(",") - new_iterator_list = [] - for it in iterator_list: - if it in iterator_converter: - it = ( - it + iterator_converter[it] - if method == "add" - else iterator_converter[it] - ) - new_iterator_list.append(it) - - return matched.group(1) + ",".join(new_iterator_list) + matched.group(3) - - return re.sub(r"(_\\text{)([^{}]*?)(})", __replace_in_iterator, instring) - - def _get_dims_from_iterators(self, instring: str) -> list[str]: - """ - For a given math string describing a math component, extract the iterators and return the dimensions (a.k.a., sets) that they are members of. - - Args: - instring (str): string describing a math component. - - Returns: - list[str]: List of dimensions over which the math component is iterating. - - """ - - def __extract_dims(matched) -> str: - iterators = matched.group(2) - # Split on `,`, add 's' back in to singular iterators to refer to dimension names, - # then rejoin `,` as we must return a string. - return ",".join( - [ - dim_name - for i in iterators.split(",") - for dim_name, dim_math in self._attrs.math.dimensions.root.items() - if dim_math.iterator == i - ] - ) - - dims = re.sub(r"^.*(_\\text{)([^{}]*?)(})", __extract_dims, instring) - return dims.split(",") - - def _instr(self, dim: str) -> str: - """ - Utility function for generating latex strings in multiple helper functions. - - Args: - dim (str): Dimension suffixed with a "s" (e.g., "techs") - - Returns: - str: LaTeX string for iterator in a set (e.g., "tech in techs") - """ - iterator = self._dim_iterator(dim) - return rf"\text{{{iterator}}} \in \text{{{dim}}}" - - def _to_str_list( - self, vals: list[str | xr.DataArray] | str | xr.DataArray | list[xr.DataArray] - ) -> list[str]: - """ - Force a string to a list of length one if not already provided as a list. - - Args: - vals (list[str] | str): Values (or single value) to force to a list. - - Returns: - list[str]: Input forced to a list. - """ - if not isinstance(vals, list): - vals = [vals] - return [str(i.name) if isinstance(i, xr.DataArray) else i for i in vals] - - def _dim_iterator(self, dim: str) -> str: - return self._attrs.math.dimensions[dim].iterator - - -class MaskAny(ParsingHelperFunction): - """Apply `any` over a dimension in `mask` string.""" - - # Class name doesn't match NAME to avoid a clash with typing.Any - #: - NAME = "any" - #: - ALLOWED_IN = ["mask"] - - def as_math_string( # noqa: D102, override - self, array: str, *, over: str | list[str | xr.DataArray] - ) -> str: - over_list = self._to_str_list(over) - overstring = r" \\ ".join(self._instr(i) for i in over_list) - substack_overstring = rf"\substack{{{overstring}}}" - # 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 the boolean mask array of a model input by applying `any` over some dimension(s). - - If the component exists in the model, returns a boolean array with dimensions reduced - by applying a boolean OR operation along the dimensions given in `over`. - If the component does not exist, returns a dimensionless False array. - - Args: - input_component (str): Reference to a model input. - over (str | list[str]): dimension(s) over which to apply `any`. - - Returns: - xr.DataArray: resulting array. - """ - 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(self._to_str_list(over)) - - return input_component.any(dim=available_dims, keep_attrs=True) - - -class Defined(ParsingHelperFunction): - """Find all items of one dimension that are defined in an item of another dimension.""" - - #: - NAME = "defined" - #: - ALLOWED_IN = ["mask"] - - def as_math_string( # noqa: D102, override - self, *, within: xr.DataArray, how: Literal["all", "any"], **dims - ) -> str: - substrings = [] - for name, vals in dims.items(): - substrings.append(self._latex_substring(how, name, vals, str(within.name))) - if len(substrings) == 1: - return substrings[0] - else: - return rf"\bigwedge({', '.join(substrings)})" - - def as_raw( - self, *, within: xr.DataArray, how: Literal["all", "any"], **dims: str - ) -> xr.DataArray: - """ - Find whether members of a model dimension are defined inside another. - - For instance, whether a node defines a specific tech (or group of techs). - Or, whether a tech defines a specific carrier. - - Args: - within (str): the model dimension to check. - how (Literal[all, any]): Whether to return True for `any` match of nested members or for `all` nested members. - **dims (str): - **key**: dimension whose members will be searched for as being defined under the primary dimension (`within`). - Must be one of the core model dimensions: [nodes, techs, carriers] - **value**: subset of the dimension members to find. - Transmission techs can be called using the base tech name (e.g., `ac_transmission`) and all link techs will be collected (e.g., [`ac_transmission:region1`, `ac_transmission:region2`]). - - - Returns: - xr.DataArray: - For each member of `within`, True if any/all member(s) in `dims` is nested within that member. - - Examples: - Check for any of a list of techs being defined at nodes. - Assuming a YAML definition of: - - ```yaml - nodes: - node1: - techs: - tech1: - tech3: - node2: - techs: - tech2: - tech3: - ``` - Then: - ``` - >>> defined(techs=[tech1, tech2], within=nodes, how=any) - [out] - array([ True, False]) - Coordinates: - * nodes (nodes) >> defined(techs=[tech1, tech2], within=nodes, how=all) - [out] - array([ False, False]) - Coordinates: - * nodes (nodes) set: - """ - From the definition matrix, get the dimensions that have not been defined. - - This includes dimensions not defined as keys of `dims` or as the value of `within`. - - Args: - dim_names (list[str]): Keys of `dims`. - within (str): dimension whose members are being checked. - - Raises: - ValueError: Can only define dimensions that exist in model.definition_matrix. - - Returns: - set: Undefined dimensions to remove from the definition matrix. - """ - definition_matrix = self._attrs.input_data.definition_matrix - missing_dims = set([*dim_names, within]).difference(definition_matrix.dims) - if missing_dims: - raise ValueError( - f"Unexpected model dimension referenced in `{self.NAME}` helper function. " - "Only dimensions given by `model.inputs.definition_matrix` can be used. " - f"Received: {missing_dims}" - ) - return set(definition_matrix.dims).difference([*dim_names, within]) - - def _latex_substring( - self, how: Literal["all", "any"], dim: str, vals: str | list[str], within: str - ) -> str: - if how == "all": - # Using wedge for "collective-and" - tex_how = "wedge" - elif how == "any": - # Using vee for "collective-or" - tex_how = "vee" - vals = self._to_str_list(vals) - within_iterator = self._dim_iterator(within) - dim_iterator = self._dim_iterator(dim) - selection = rf"\text{{{dim_iterator}}} \in \text{{[{','.join(vals)}]}}" - - return rf"\big{tex_how}\limits_{{\substack{{{selection}}}}}\text{{{dim_iterator} defined in {within_iterator}}}" - - -class Sum(ParsingHelperFunction): - """Apply a summation over dimension(s) in math expressions.""" - - NAME = "sum" - #: - ALLOWED_IN = ["expression", "mask"] - - def as_math_string( # noqa: D102, override - self, array: str, *, over: str | list[str | xr.DataArray] - ) -> str: - over_list = self._to_str_list(over) - overstring = r" \\ ".join(self._instr(i) for i in over_list) - substack_overstring = rf"\substack{{{overstring}}}" - 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). - - Args: - array (xr.DataArray): expression array - over (xr.DataArray | list[xr.DataArray]): - Dimension(s) over which to apply `sum`. - Array names will be extracted from the DataArray objects. - - Returns: - xr.DataArray: - Array with dimensions reduced by applying a summation over the dimensions given in `over`. - NaNs are ignored (xarray.DataArray.sum arg: `skipna: True`) and if all values along the dimension(s) are NaN, - the summation will lead to a NaN (xarray.DataArray.sum arg: `min_count=1`). - """ - filtered_over = set(self._to_str_list(over)).intersection(array.dims) - return array.sum(filtered_over) - - -class SelectFromLookupArrays(ParsingHelperFunction): - """N-dimensional indexing functionality.""" - - #: - NAME = "select_from_lookup_arrays" - #: - ALLOWED_IN = ["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() - } - array = self._update_iterator(array, new_strings, "add") - return array - - 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. - - Args: - array (xr.DataArray): Array on which to apply vectorised indexing. - **lookup_arrays (xr.DataArray): - key: dimension on which to apply vectorised indexing - value: array whose values are either NaN or values from the dimension given in the key. - - Raises: - ValueError: `array` must be indexed over the dimensions given in the `lookup_arrays` dict keys. - ValueError: All `lookup_arrays` must be indexed over all the dimensions given in the `lookup_arrays` dict keys. - - 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. - - Examples: - >>> coords = {"foo": ["A", "B", "C"]} - >>> array = xr.DataArray([1, 2, 3], coords=coords) - >>> lookup_array = xr.DataArray( - ... np.array(["B", "A", np.nan], dtype="O"), coords=coords, name="bar" - ... ) - >>> model_data = xr.Dataset({"bar": lookup_array}) - >>> select_from_lookup_arrays = SelectFromLookupArrays( - ... model_data=model_data - ... ) - >>> select_from_lookup_arrays(array, foo=lookup_array) - - array([ 2., 1., nan]) - Coordinates: - * foo (foo) object 'A' 'B' 'C' - - The lookup array assigns the value at "B" to "A" and vice versa. - "C" is masked since the lookup array value is NaN. - """ - # 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} 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}` 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._attrs.input_data[index.name].stack({dim: dims}) - ix = array.indexes[index_dim].get_indexer(stacked_lookup) - if (ix == -1).all(): - received_lookup = ( - self._attrs.input_data[index.name].to_series().dropna() - ) - raise IndexError( - f"Trying to select items on the dimension {index_dim} from the {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) - - # Create a mask to nullify any lookup values that are not given (i.e., are np.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 - unstacked_result = result.drop_vars(dims).unstack(dim) - return unstacked_result - - -class GetValAtIndex(ParsingHelperFunction): - """Getter functionality for obtaining values at specific integer indices.""" - - #: - NAME = "get_val_at_index" - #: - ALLOWED_IN = ["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 value of a model dimension at a given integer index. - - This function is primarily useful for timeseries data. - - Args: - **dim_idx_mapping (int): kwargs with - key (str): Model dimension in which to extract value. - value (int): Integer index of the value to extract (assuming zero-indexing). - - Returns: - xr.DataArray: Dimensionless array containing one value. - - Examples: - >>> coords = { - ... "timesteps": [ - ... "2000-01-01 00:00", - ... "2000-01-01 01:00", - ... "2000-01-01 02:00", - ... ] - ... } - >>> model_data = xr.Dataset(coords=coords) - >>> get_val_at_index = GetValAtIndex(model_data=model_data) - >>> get_val_at_index(model_data)(timesteps=0) - - array('2000-01-01 00:00', dtype='>> get_val_at_index(model_data)(timesteps=-1) - - array('2000-01-01 00:00', dtype=' tuple[str, int]: ... - - # For as_math_string - @overload - @staticmethod - def _mapping_to_dim_idx(**dim_idx_mapping: str) -> tuple[str, str]: ... - - @staticmethod - def _mapping_to_dim_idx(**dim_idx_mapping) -> tuple[str, str | int]: - 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(ParsingHelperFunction): - """Roll (a.k.a. shift) items along ordered dimensions.""" - - #: - NAME = "roll" - #: - ALLOWED_IN = ["expression"] - - @property - def ignore_mask(self) -> bool: - """Whether or not to ignore `mask` functionality.""" - return 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() - } - component = self._update_iterator(array, new_strings, "add") - return component - - def as_raw(self, array: xr.DataArray, **roll_kwargs: int) -> xr.DataArray: - """ - Roll (a.k.a., shift) 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. - - Args: - array (xr.DataArray): Array on which to roll data. - **roll_kwargs (int): kwargs with the following - key (str): name of dimension on which to roll. - value (int): number of places to roll data. - - Returns: - xr.DataArray: `array` with rolled data. - - Examples: - >>> array = xr.DataArray([1, 2, 3], coords={"foo": ["A", "B", "C"]}) - >>> model_data = xr.Dataset({"bar": array}) - >>> roll = Roll() - >>> roll("bar", foo=1) - - array([3, 1, 2]) - Coordinates: - * foo (foo) 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. - - Args: - array (xr.DataArray): Math component array. - condition (xr.DataArray): - Boolean mask array. - If not `bool` type, NaNs and 0 will be assumed as False and all other values will be assumed as True. - - Returns: - xr.DataArray: - Returns the input array with the condition applied, - including having been broadcast across any new dimensions provided by the condition. - - Examples: - One common use-case is to introduce a new dimension to the variable which represents subsets of one of the main model dimensions. - In this case, each member of `cap_node_groups` is a subset of `nodes` and we want to sum `flow_cap` over each of those subsets and set a maximum value. - - input: - ```yaml - data_definitions: - node_grouping: - data: True - index: [[group_1, region1], [group_1, region1_1], [group_2, region1_2], [group_2, region1_3], [group_3, region2]] - dims: [cap_node_groups, nodes] - node_group_max: - data: [1, 2, 3] - index: [group_1, group_2, group_3] - dims: cap_node_groups - ``` - - math: - ```yaml - constraints: - my_new_constraint: - foreach: [techs, cap_node_groups] - equations: - - expression: sum(mask(flow_cap, node_grouping), over=nodes) <= node_group_max - ``` - """ - return array.where(condition.fillna(False).astype(bool)) - - -class GroupSum(ParsingHelperFunction): - """Apply a summation over an array grouping.""" - - #: - NAME = "group_sum" - #: - ALLOWED_IN = ["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]) - overstring = rf"\substack{{{foreach_string}}}" - return rf"\sum\limits_{{{overstring}}} ({array})" - - def as_raw( - self, array: xr.DataArray, groupby: xr.DataArray, group_dim: xr.DataArray - ) -> xr.DataArray: - """ - Sum an array over the given groupings. - - Args: - array (xr.DataArray): expression array - groupby (xr.DataArray): Array with which to group the array. - group_dim (str): Name of dimension that the `groupby` values are members of. - This will become a new dimension over which the array is indexed once grouping is complete. - - Returns: - xr.DataArray: - Array with dimension(s) aggregated over the `groupby`. - - Note: - - The array is returned with all dimensions over which `groupby` is indexed replaced by a new dimension named by `group_dim`. - - To groupby datetime periods (weeks, months, dates, etc.), consider using `group_datetime` for convenience, as you do not need to define a separate `groupby` array. - - Examples: - To get the sum over an ad-hoc combination of techs at nodes, e.g. to limit their overall outflow in any given timestep, you would do the following: - - 1. Define an array linking node-tech combinations with a group: - ```yaml - data_definitions: - # You may prefer to define this in a CSV file or when referring to the techs within the `nodes` model definition. - power_plant_groups: - data: [low_emission_plant, low_emission_plant, high_emission_plant, high_emission_plant] - index: [ - [tech_1, node_1], - [tech_2, node_1], - [tech_1, node_2], - [tech_2, node_2], - ] - dims: [techs, nodes] - ``` - 2. Define a set of outflow limits: - ```yaml - data_definitions: - emission_limits: - data: [20, 10] - index: [low_emission_plant, high_emission_plant] - dims: [emission_groups] - ``` - 3. Define the math to link the two, using `group_sum`: - ```yaml - constraints: - node_tech_emission_group_max: - foreach: [emission_groups, carriers, timesteps] - mask: emission_limits - equations: - - expression: group_sum(flow_out, power_plant_groups, emission_groups) <= emission_limits - ``` - """ - # We can't apply typical xarray rolling window functionality - - grouping_dims = groupby.dims - groups = array.stack(_stacked=grouping_dims).groupby( - groupby.rename(group_dim.name).stack(_stacked=grouping_dims) - ) - - grouped = groups.sum("_stacked") - return grouped - - -class GroupDatetime(ParsingHelperFunction): - """Apply a summation over a datetime group on a datetime dimension in math expressions.""" - - NAME = "group_datetime" - #: - ALLOWED_IN = ["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}}}(\text{{{self._dim_iterator(over)}}}) = \text{{{self._dim_iterator(group)}}}" - overstring = rf"\substack{{{foreach_string}}}" - - return rf"\sum\limits_{{{overstring}}} ({array})" - - def as_raw( - self, array: xr.DataArray, over: xr.DataArray, group: xr.DataArray - ) -> xr.DataArray: - """ - Sum an expression array over the given dimension(s). - - Args: - array (xr.DataArray): expression array - over (xr.DataArray): dimension name over which to group - group (xr.DataArray): datetime grouper. - Any xarray/pandas datetime grouper options - datetime grouper options include 'date', 'dayofweek', 'month', etc. - - - Returns: - xr.DataArray: - Array with datetime dimension aggregated over the grouper. - - Note: - - The array is returned with the `over` dimension replaced by the name of the grouper. - So, if you select to resample to monthly, the returned array will include the `month` dimension. - - the `date`/`time` groupers will return the date/time as a string in ISO8601 format (e.g. "2025-01-01"/"01:00:00"). - All other groupers will return integer values (e.g. month 1, 2, 3, etc.). - - Examples: - One common use-case is to allow demand to be met at any point on a given date. - For such a demand tech, the daily demand should be indexed over `date`, e.g.: - - sink_use_equals_daily.csv - ``` - date,sink_use_equals_daily - 2000-01-01,10 - 2000-01-02,15 - ... - ``` - - Then, to set the daily flow into the demand tech to those values: - ```yaml - constraints: - daily_demand: - foreach: [nodes, techs, carriers, date] - mask: sink_use_equals_daily - equations: - - expression: "group_datetime(flow_in, timesteps, date) == sink_use_equals_daily" - ``` - - Similarly, a monthly maximum resource to a supply technology might be used, to simulate e.g. biofuel feedstock availability: - - source_use_max_monthly.csv - ``` - month,source_use_max_monthly - 1,10 - 2,15 - ... - ``` - - Then, to set the daily flow into the demand tech to those values: - ```yaml - constraints: - daily_demand: - foreach: [nodes, techs, carriers, month] - mask: source_use_max_monthly - equations: - - expression: "group_datetime(flow_in, timesteps, month) <= source_use_max_monthly" - ``` - """ - dtype = DTYPE_OPTIONS[self._attrs.math.dimensions[group.name].dtype] - group_sum_helper = GroupSum(self._return_type, self._attrs) - array = group_sum_helper( - array, getattr(array[over.name].dt, group.name).astype(dtype), group - ) - - return array - - -class SumNextN(ParsingHelperFunction): - """ - Sum the 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 = ["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 = self._update_iterator( - array, {self._dim_iterator(over): new_iterator}, "replace" - ) - - return rf"\sum\limits_{{\text{{{new_iterator}}}={over_singular}}}^{{{over_singular}+{N}}} ({updated_iterator_array})" - - def as_raw(self, array: xr.DataArray, over: xr.DataArray, N: int) -> xr.DataArray: - """ - Sum values from current up to N from current on the dimension `over`. - - Works best for ordered arrays (datetime, integer). - - - Args: - array (xr.DataArray): Math component array. - over (str): Dimension over which to sum - N (int): number of items beyond the current value to sum from - - Returns: - xr.DataArray: - Returns the input array with the condition applied, - including having been broadcast across any new dimensions provided by the condition. - - Note: - - The rolling window does not wrap around to the start of the set when reaching the end. - That is, if you have N = 4 then for a dimension of length T, at T - 1 it will sum over dimension positions (T - 1, T), not (T - 1, T, 0, 1). - - You will find that this over-constrains the model unless you limit the constraint (using the `mask` string) to only apply over `len(over) - N`. - This is linked to the abovementioned lack of wrapping. - E.g. `mask: timesteps<=get_val_at_index(timesteps=-24)` if N == 24. - - This function is based on an integer number of steps from the current step. - For datetime dimensions like `timesteps`, you will (a) need to be using a regular time frequency (e.g. hourly) and (b) update `N` to reflect the resolution of your time dimension - (N = 4 in if resample.timesteps=`1h` -> N = 2 if resample.timesteps=`2h`). - - Examples: - One common use-case is to collate N timesteps beyond a given timestep to apply a constraint to it - (e.g., demand must be less than X in the next 24 hours): - - For such a demand tech, the portion of its demand that is flexible should be separated from `sink_use_equals` to e.g., - a `sink_use_flexible` timeseries parameter which we will use in the DSR constraint: - - ```yaml - constraints: - 4hr_demand_side_response: - foreach: ["nodes", "techs", "carriers", "timesteps"] - mask: "carrier_in AND sink_use_flexible AND timesteps<=get_val_at_index(timesteps=-24)" - equations: - - expression: sum_next_n(flow_in, timesteps, 4) == sum_next_n(sink_use_flexible, timesteps, 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 - ) - ) - final_array = xr.concat(results, dim=over).broadcast_like(array) - return final_array diff --git a/linopy/declarative/helpers.py b/linopy/declarative/helpers.py new file mode 100644 index 000000000..7695ccee4 --- /dev/null +++ b/linopy/declarative/helpers.py @@ -0,0 +1,673 @@ +""" +Linopy declarative math helper-functions module. + +This module contains the helper functions that can be called in declarative math +`mask` and `expression` strings (by their `NAME`), the abstract base class from +which users can define their own, and the registry builder that makes them +available to a model build. +""" + +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.evaluate import Context + +MODE_T = Literal["raw", "expr", "math_string"] +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, mode: MODE_T, context: Context) -> None: + """ + Initialise the helper for one evaluation mode. + + Parameters + ---------- + mode : Literal["raw", "expr", "math_string"] + The evaluation mode this instance will dispatch to when called. + context : Context + The evaluation context (input data, math definition, config, ...). + """ + self._mode = mode + 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 __call__(self, *args: Any, **kwargs: Any) -> Any: + """Dispatch to the `as_*` method matching the mode set at initialisation.""" + if self._mode == "math_string": + return self.as_math_string(*args, **kwargs) + elif self._mode == "raw": + return self.as_raw(*args, **kwargs) + elif self._mode == "expr": + return self.as_expr(*args, **kwargs) + else: + raise ValueError(f"Unknown helper function mode: {self._mode!r}") + + 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._mode, self._context) + return group_sum_helper( + 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 000000000..8df71c73e --- /dev/null +++ b/linopy/declarative/latex.py @@ -0,0 +1,446 @@ +""" +Linopy declarative LaTeX math documentation module. + +This module builds a human-readable mathematical formulation document from a +declarative math definition, without building an optimisation problem: every +active math component is rendered to LaTeX (equations, mask conditions, foreach +sets, bounds) together with its metadata and cross-references, and the result +can be generated as Markdown, reStructuredText, or LaTeX source. +""" + +from __future__ import annotations + +import math as pymath +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.evaluate import Context, to_math_string +from linopy.declarative.grammar import Component, find_refs +from linopy.declarative.helpers import HelperFunction, build_registry, dim_iterator +from linopy.declarative.schema import ConfigModel, MathModel +from linopy.model import Model + +FORMAT_T = Literal["md", "rst", "tex"] + +_DOCUMENTED_GROUPS: dict[str, str] = { + "parameters": "Parameters", + "lookups": "Lookups", + "variables": "Variables", + "expressions": "Expressions", + "constraints": "Constraints", + "objectives": "Objectives", +} + +_EQUATION_GROUPS = ("expressions", "constraints", "objectives") + +_REPR_STYLES = { + "parameters": "textit", + "lookups": "textit", + "variables": "textbf", + "expressions": "textbf", +} + + +@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, ...).""" + + +def _number_string(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 f"{value:.6g}" + + +class LatexModelBuilder: + """ + 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. + """ + self.math = MathModel.model_validate(math_def) + self.input_data = input_data if input_data is not None else xr.Dataset() + self.config = ConfigModel.model_validate(config or {}) + self.components: dict[str, dict[str, RenderedComponent]] = {} + self._ctx = Context( + model=Model(), + input_data=self.input_data, + math=self.math, + config=self.config, + helpers=build_registry(helpers), + math_reprs=self._build_math_reprs(), + ) + + def _iterator(self, dim: str) -> str: + """Return the LaTeX iterator of a dimension (its own name if not declared).""" + return dim_iterator(self.math, dim) + + 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(): + dims = getattr(definition, "foreach", None) + if dims is None: + dims = ( + list(self.input_data[name].dims) + if name in self.input_data + else [] + ) + iterators = ",".join(self._iterator(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: object) -> str: + r"""Return the LaTeX `\forall` line body for a component's `foreach` sets.""" + sets = getattr(definition, "foreach", []) + if not sets: + return "" + instrs = ", ".join( + rf"\text{{{self._iterator(dim)}}} \in \text{{{dim}}}" for dim in sets + ) + return rf"\forall{{}} {instrs}" + + def _mask_string(self, definition: object, name: str) -> str: + """Return the LaTeX rendering of a component's top-level mask ("" if true).""" + mask_node = parsing.parse_mask( + getattr(definition, "mask", "True"), self.math, name + ) + rendered = to_math_string( + mask_node, replace(self._ctx, route="mask", equation_name=name) + ) + return "" if rendered == "true" else rendered + + def _render_metadata(self, definition: object) -> dict[str, str]: + """Return the documentable metadata (unit, default, ...) of a definition.""" + extras: dict[str, str] = {} + unit = getattr(definition, "unit", "") + if unit: + extras["Unit"] = unit + default = getattr(definition, "default", None) + if default is not None and pd.notna(default): + if isinstance(default, int | float) and pymath.isinf(default): + extras["Default"] = "inf" if default > 0 else "-inf" + else: + extras["Default"] = str(default) + return extras + + def add_component(self, group: str, name: str, definition: object) -> None: + """Render one math component and store it under `self.components`.""" + rendered = RenderedComponent( + group=group, + name=name, + title=getattr(definition, "title", ""), + description=getattr(definition, "description", ""), + foreach=self._foreach_string(definition), + mask=self._mask_string(definition, f"{group}:{name}"), + extras=self._render_metadata(definition), + ) + mask_node = parsing.parse_mask(getattr(definition, "mask", "True"), self.math) + uses = find_refs(mask_node, Component) + + if group in _EQUATION_GROUPS: + equations = parsing.parse_component(group, name, definition, self.math) # type: ignore[arg-type] + for equation in equations: + equation_ctx = replace(self._ctx, equation_name=equation.name) + rendered.equations.append( + { + "mask": parsing.as_latex(equation, equation_ctx, what="mask"), + "expression": parsing.as_latex(equation, equation_ctx), + } + ) + uses |= equation.references() + elif group == "variables": + rendered.extras["Domain"] = definition.domain # type: ignore[attr-defined] + rendered.equations.append( + {"mask": "", "expression": self._bounds_string(name, definition)} + ) + uses |= { + bound + for bound in (definition.bounds.lower, definition.bounds.upper) # type: ignore[attr-defined] + if isinstance(bound, str) + } + if group == "objectives": + rendered.extras["Sense"] = ( + "minimise" if definition.sense == "min" else "maximise" # type: ignore[attr-defined] + ) + + # 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. + rendered.foreach = _escape_text_mode(rendered.foreach) + rendered.mask = _escape_text_mode(rendered.mask) + for rendered_eq in rendered.equations: + rendered_eq["mask"] = _escape_text_mode(rendered_eq["mask"]) + rendered_eq["expression"] = _escape_text_mode(rendered_eq["expression"]) + + self.components.setdefault(group, {})[name] = rendered + + def _bounds_string(self, name: str, definition: object) -> str: + """Return the LaTeX bounds equation of a decision variable.""" + bounds = definition.bounds # type: ignore[attr-defined] + reprs = self._ctx.math_reprs + lower, upper = ( + reprs.get(bound, rf"\textit{{{bound}}}") + if isinstance(bound, str) + else _number_string(bound) + for bound in (bounds.lower, bounds.upper) + ) + return rf"{lower} \leq {reprs[name]} \leq {upper}" + + 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 _math_block(format: FORMAT_T, lines: list[str]) -> list[str]: + """Return a display-math block wrapping an `array` of the given LaTeX lines.""" + joined = " \\\\\n ".join(lines) + array = f"\\begin{{array}}{{l}}\n {joined}\n\\end{{array}}" + if format == "md": + return ["$$", array, "$$", ""] + if format == "rst": + indented = "\n".join(f" {line}" for line in array.split("\n")) + return [".. math::", "", indented, ""] + return [r"\begin{equation}", array, 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("") + for equation in component.equations: + lines = [] + if component.foreach: + lines.append(component.foreach) + for mask in (component.mask, equation["mask"]): + if mask: + lines.append(rf"\text{{if }} {mask}") + lines.append(equation["expression"]) + blocks.extend(_math_block(format, lines)) + return blocks diff --git a/linopy/declarative/mask_parser.py b/linopy/declarative/mask_parser.py deleted file mode 100644 index 8f4589b98..000000000 --- a/linopy/declarative/mask_parser.py +++ /dev/null @@ -1,618 +0,0 @@ -# Copyright (C) since 2013 Calliope contributors listed in AUTHORS. -# Licensed under the Apache 2.0 License (see LICENSE file). -"""Parsing for 'mask' statements.""" - -from __future__ import annotations - -import operator -from collections.abc import Iterable -from dataclasses import replace -from typing import Any - -import numpy as np -import pandas as pd -import pyparsing as pp -import xarray as xr - -from linopy.declarative import expression_parser - -pp.ParserElement.enable_packrat() -BOOLEANTYPE = np.bool_ | np.typing.NDArray[np.bool_] - - -def get_dot_attr(var: Any, attr: str) -> Any: - """ - Get nested attributes in dot notation. - - Works for nested objects (e.g., dictionaries, pydantic models). - - Args: - var (Any): Object to extract nested attributes from. - attr (str): Name of the attribute (e.g., "foo.bar"). - - Returns: - Any: Value at the given location. - """ - levels = attr.split(".", 1) - - if isinstance(var, dict): - value = var[levels[0]] - else: - value = getattr(var, levels[0]) - - if len(levels) > 1: - value = get_dot_attr(value, levels[1]) - return value - - -class EvalNot(expression_parser.EvalSignOp): - """Parse action to process successfully parsed expressions with a leading `not`.""" - - def as_math_string(self) -> str: # noqa: D102, override - evaluated = self.value.eval("math_string", self.eval_attrs) - return rf"\neg ({evaluated})" - - def as_raw(self) -> xr.DataArray: # noqa: D102, override - evaluated = self.value.eval("raw", self.eval_attrs) - return ~evaluated - - -class EvalAndOr(expression_parser.EvalOperatorOperand): - """ - Processing of successfully parsed expressions with and/or operators. - - E.g., "OPERAND OPERATOR OPERAND OPERATOR OPERAND ..." - """ - - LATEX_OPERATOR_LOOKUP: dict[str, str] = { - "and": r"{val} \land {operand}", - "or": r"{val} \lor {operand}", - } - SKIP_IF = ["and", "or"] - - def _skip_component_on_conditional(self, component: str, operator_: str) -> bool: - return component == "true" and operator_ in self.SKIP_IF - - @staticmethod - def _operate( - val: xr.DataArray, evaluated_operand: xr.DataArray, operator_: str - ) -> xr.DataArray: - """Apply bitwise comparison between boolean xarray dataarrays.""" - match operator_: - case "and": - val = operator.and_(val, evaluated_operand) - case "or": - val = operator.or_(val, evaluated_operand) - return val - - def _apply_mask(self, evaluated: xr.DataArray) -> xr.DataArray: - """Override func from parent class to effectively do nothing.""" - return evaluated - - def as_math_string(self) -> str: # noqa: D102, override - return super().as_math_string() - - def as_raw(self) -> xr.DataArray: # noqa: D102, override - return super().as_raw() - - -class ConfigOptionParser(expression_parser.EvalNode): - """Parsing of configuration options.""" - - def __init__(self, instring: str, loc: int, tokens: pp.ParseResults) -> None: - """ - Parse action to process successfully parsed configuration option names. - - Args: - instring (str): String that was parsed (used in error message). - loc (int): - Location in parsed string mask parsing error was logged. - This is not used; we include it as pyparsing injects it alongside `instring` when setting the parse action. - tokens (pp.ParseResults): - Has two parsed elements: config group name (str) and config option (str). - """ - self.config_option = tokens[0] - self.instring = instring - self.loc = loc - - def __repr__(self): - """Programming / official string representation.""" - return f"CONFIG:{self.config_option}" - - def as_math_string(self) -> str: # noqa: D102, override - return rf"\text{{config.{self.config_option}}}" - - def as_raw(self) -> xr.DataArray: # noqa: D102, override - config_val = get_dot_attr(self.eval_attrs.config, self.config_option) - - if not isinstance(config_val, int | float | str | bool | np.bool_): - raise self.error_msg( - f"mask string | Configuration option resolves to invalid " - f"type `{type(config_val).__name__}`, expected a number, string, or boolean." - ) - else: - return xr.DataArray(config_val) - - -class ResultArrayParser(expression_parser.EvalNode): - """Variable/Expression array processing.""" - - def __init__(self, instring: str, loc: int, tokens: pp.ParseResults) -> None: - """ - Parse action to process successfully parsed model variable/global expression names. - - Args: - instring (str): String that was parsed (used in error message). - loc (int): - Location in parsed string mask parsing error was logged. - This is not used; we include it as pyparsing injects it alongside `instring` when setting the parse action. - tokens (pp.ParseResults): - Has one parsed element: model data variable name (str). - """ - self.array_name = tokens[0] - self.instring = instring - self.loc = loc - - def __repr__(self): - """Programming / official string representation.""" - return f"RESULT:{self.array_name}" - - def as_math_string(self) -> str: # noqa: D102, override - self.eval_attrs.references.add(self.array_name) - math_repr = self.eval_attrs.model[self.array_name].attrs.get( - "math_repr", rf"\exists (\textbf{{{self.array_name}}})" - ) - - return math_repr - - def as_raw(self) -> xr.DataArray: # noqa: D102, override - self.eval_attrs.references.add(self.array_name) - da = self.eval_attrs.model[self.array_name] - if self.eval_attrs.apply_mask: - da = ~da.isnull() - return da - - -class InputArrayParser(expression_parser.EvalNode): - """Input array processing.""" - - def __init__(self, instring: str, loc: int, tokens: pp.ParseResults) -> None: - """ - Parse action to process successfully parsed model input array names. - - Args: - instring (str): String that was parsed (used in error message). - loc (int): - Location in parsed string mask parsing error was logged. - This is not used; we include it as pyparsing injects it alongside `instring` when setting the parse action. - tokens (pp.ParseResults): - Has one parsed element: model data variable name (str). - """ - self.array_name = tokens[0] - self.instring = instring - self.loc = loc - - def __repr__(self): - """Programming / official string representation.""" - return f"INPUT:{self.array_name}" - - def as_math_string(self) -> str: # noqa: D102, override - self.eval_attrs.references.add(self.array_name) - - math_repr = self.eval_attrs.input_data[self.array_name].attrs.get( - "math_repr", rf"\textit{{{self.array_name}}}" - ) - if self.eval_attrs.apply_mask: - math_repr = rf"\exists ({math_repr})" - return math_repr - - def as_raw(self) -> xr.DataArray: # noqa: D102, override - self.eval_attrs.references.add(self.array_name) - da = self.eval_attrs.input_data.get(self.array_name, xr.DataArray(False)) - if self.eval_attrs.apply_mask and da.dtype.kind != "b": - da = da.notnull() & (da != np.inf) & (da != -np.inf) - elif da.isnull().any() and pd.notnull( - default := self.eval_attrs.math.find(self.array_name).default - ): - da = da.fillna(default) - return da - - -class DimensionArrayParser(expression_parser.EvalNode): - """Dimension array processing.""" - - def __init__(self, instring: str, loc: int, tokens: pp.ParseResults) -> None: - """ - Parse action to process successfully parsed model dimension names. - - Args: - instring (str): String that was parsed (used in error message). - loc (int): - Location in parsed string mask parsing error was logged. - This is not used; we include it as pyparsing injects it alongside `instring` when setting the parse action. - tokens (pp.ParseResults): - Has one parsed element: model data variable name (str). - """ - self.array_name = tokens[0] - self.instring = instring - self.loc = loc - - def __repr__(self): - """Programming / official string representation.""" - return f"DIM:{self.array_name}" - - def as_math_string(self) -> str: # noqa: D102, override - return self.array_name - - def as_raw(self) -> xr.DataArray: # noqa: D102, override - # We want the mask string to evaluate successfully even if a dimension hasn't been defined. - da = self.eval_attrs.input_data.get(self.array_name, xr.DataArray()) - return da - - -class ComparisonParser(expression_parser.EvalComparisonOp): - """Parse action to process successfully parsed strings of the form x=y.""" - - OP_TRANSLATOR = { - "<=": r"\mathord{\leq}", - ">=": r"\mathord{\geq}", - "==": r"\mathord{==}", - "<": r"\mathord{<}", - ">": r"\mathord{>}", - } - - def __repr__(self): - """Return string representation of the parsed grammar.""" - return f"{self.lhs}{self.op}{self.rhs}" - - def as_math_string(self) -> str: # noqa: D102, override - self.eval_attrs = replace(self.eval_attrs, apply_mask=False) - lhs, rhs = self._eval("math_string") - if r"\text" not in rhs: - rhs = rf"\text{{{rhs}}}" - return lhs + self.OP_TRANSLATOR[self.op] + rhs - - def as_raw(self) -> xr.DataArray: # noqa: D102, override - self.eval_attrs = replace(self.eval_attrs, apply_mask=False) - lhs, rhs = self._eval("raw") - match self.op: - case "<=": - comparison = lhs <= rhs - case ">=": - comparison = lhs >= rhs - case "<": - comparison = lhs < rhs - case ">": - comparison = lhs > rhs - case "==": - comparison = lhs == rhs - return xr.DataArray(comparison) - - -class SubsetParser(expression_parser.EvalNode): - """Dimension subset parsing.""" - - def __init__(self, instring: str, loc: int, tokens: pp.ParseResults) -> None: - """ - Parse action to process successfully parsed dimension subsetting. - - Args: - instring (str): String that was parsed (used in error message). - loc (int): - Location in parsed string mask parsing error was logged. - This is not used; we include it as pyparsing injects it alongside `instring` when setting the parse action. - tokens (pp.ParseResults): - Has two parsed elements: model set name (str), set items (Any). - """ - self.val, self.set_name = tokens - self.instring = instring - self.loc = loc - - def __repr__(self): - """Return string representation of the parsed grammar.""" - return f"SUBSET:{self.set_name}{self.val}" - - def _eval(self) -> list[str | float]: - """Evaluate each element of the subset list.""" - values = [val.eval("raw", self.eval_attrs) for val in self.val] - return [val.item() if isinstance(val, xr.DataArray) else val for val in values] - - def as_math_string(self) -> str: # noqa: D102, override - subset = self._eval() - dim = self.set_name.eval("math_string", self.eval_attrs) - iterator = self.eval_attrs.math.dimensions[dim].iterator - subset_string = "[" + ",".join(str(i) for i in subset) + "]" - return rf"\text{{{iterator}}} \in \text{{{subset_string}}}" - - def as_raw(self) -> xr.DataArray: # noqa: D102, override - subset = self._eval() - self.eval_attrs = replace(self.eval_attrs, apply_mask=False) - da = self.set_name.eval("raw", replace(self.eval_attrs, apply_mask=False)) - set_item_in_subset = da.isin(subset) - return set_item_in_subset - - -class BoolOperandParser(expression_parser.EvalNode): - """Boolean operand parsing.""" - - def __init__(self, instring: str, loc: int, tokens: pp.ParseResults) -> None: - """ - Parse action to process successfully parsed boolean strings. - - Args: - instring (str): String that was parsed (used in error message). - loc (int): - Location in parsed string mask parsing error was logged. - This is not used; we include it as pyparsing injects it alongside `instring` when setting the parse action. - tokens (pp.ParseResults): Has one parsed element: boolean (str). - """ - self.val = tokens[0].lower() - self.instring = instring - - def __repr__(self): - """Programming / official string representation.""" - return f"BOOL:{self.val}" - - def as_math_string(self): # noqa: D102, override - return self.val - - def as_raw(self) -> xr.DataArray: # noqa: D102, override - if self.val == "true": - bool_val = xr.DataArray(np.True_) - elif self.val == "false": - bool_val = xr.DataArray(np.False_) - return bool_val - - -class GenericStringParser(expression_parser.EvalString): - """Parsing of generic strings.""" - - def __init__(self, instring: str, loc: int, tokens: pp.ParseResults) -> None: - """ - Parse action to process successfully parsed generic strings. - - This is required since we call "eval()" on all elements of the mask string, - so even arbitrary strings (used in comparison operations) need to be evaluatable. - - Args: - instring (str): String that was parsed (used in error message). - loc (int): - Location in parsed string mask parsing error was logged. - This is not used; we include it as pyparsing injects it alongside `instring` when setting the parse action. - tokens (pp.ParseResults): Has one parsed element: string name (str). - """ - self.val = tokens[0] - self.instring = instring - - def __repr__(self) -> str: - """Return string representation of the parsed grammar.""" - return f"STRING:{self.val}" - - def eval(self, *args, **kwargs) -> str: - """Evaluation just returns the string of values.""" - return str(self.val) - - -def data_var_parser( - names: Iterable, parse_action: type[expression_parser.EvalNode] -) -> pp.ParserElement: - """ - Process model data variables which can be any valid python identifier (string + "_"). - - Args: - names (Iterable): List of valid component names. - parse_action (type[expression_parser.EvalNode]): Parse action to evaluate the parsed string. - - Returns: - pp.ParserElement: parser for model data variables which will access the data - variable from the Calliope model dataset. - """ - data_var = pp.one_of(names, as_keyword=True) - data_var.set_parse_action(parse_action) - - return data_var - - -def config_option_parser(generic_identifier: pp.ParserElement) -> pp.ParserElement: - """ - Parsing grammar to process model configuration option key names of the form "x.y.z". - - Args: - generic_identifier (pp.ParserElement): - Parser for valid python variables without leading underscore and not called "inf". - This parser has no parse action. - - Returns: - pp.ParserElement: - Parser for configuration options which will be accessed from the configuration - dictionary attached to the attributes of the Calliope model dataset. - """ - data_var = pp.Suppress("config.") + generic_identifier - data_var.set_parse_action(ConfigOptionParser) - - return data_var - - -def bool_parser() -> pp.ParserElement: - """Parsing grammar for True/False (any case), which will evaluate to np.bool_.""" - TRUE = pp.Keyword("True", caseless=True) - FALSE = pp.Keyword("False", caseless=True) - bool_operand = TRUE | FALSE - bool_operand.set_parse_action(BoolOperandParser) - - return bool_operand - - -def evaluatable_string_parser( - generic_identifier: pp.ParserElement, valid_components: Iterable -) -> pp.ParserElement: - """Parsing grammar to make generic strings used in comparison operations evaluatable.""" - evaluatable_identifier = ( - ~pp.one_of(valid_components, as_keyword=True) + generic_identifier - ) - evaluatable_identifier.set_parse_action(GenericStringParser) - - return evaluatable_identifier - - -def comparison_parser( - lhs: list[pp.ParserElement], rhs: list[pp.ParserElement] -) -> pp.ParserElement: - """ - Parsing grammar to process comparisons of the form `variable_or_config=comparator`. - - Args: - lhs (list[pp.ParserElement]): - Parsers that can be included on the left-hand side of the comparison; will be matched in the order provided. - rhs (list[pp.ParserElement]): - Parsers that can be included on the right-hand side of the comparison; will be matched in the order provided. - - Returns: - pp.ParserElement: - Parser which will return a bool/boolean array as a result of the comparison. - """ - comparison_operators = pp.oneOf(["<", ">", "==", ">=", "<="]) - comparison_expression = ( - pp.MatchFirst(lhs) + comparison_operators + pp.MatchFirst(rhs) - ) - comparison_expression.set_parse_action(ComparisonParser) - - return comparison_expression - - -def subset_parser( - data_vars: list[pp.ParserElement], *subset_items: pp.ParserElement -) -> pp.ParserElement: - """ - Parsing grammar to process subsets. - - Args: - data_var (pp.ParserElement): data variable parser - *subset_items (pp.ParserElement): parsers that can be included in the subset list; will be matched in the order provided. - - Returns: - pp.ParserElement: subset parser. - """ - subset = pp.Group(pp.delimited_list(pp.MatchFirst(subset_items))) - subset_expression = ( - pp.Suppress("[") - + subset - + pp.Suppress("]") - + pp.Suppress(pp.White(" ", min=1)) - + pp.Suppress("in") - + pp.Suppress(pp.White(" ", min=1)) - + pp.MatchFirst(data_vars) - ) - subset_expression.set_parse_action(SubsetParser) - - return subset_expression - - -def mask_parser(*args: pp.ParserElement) -> pp.ParserElement: - """ - Parser for strings which use AND/OR/NOT operators to combine other parser elements. - - Args: - *args (pp.ParserElement): - parsers that can be included in the mask string; will be matched in the order provided. - - Returns: - pp.ParserElement: mask parser. - """ - notop = pp.Keyword("not", caseless=True) - andorop = pp.Keyword("and", caseless=True) | pp.Keyword("or", caseless=True) - - mask_rules = pp.infixNotation( - pp.MatchFirst(args), - [ - (notop, 1, pp.opAssoc.RIGHT, EvalNot), - (andorop, 2, pp.opAssoc.LEFT, EvalAndOr), - ], - ) - - return mask_rules - - -def generate_mask_string_parser( - dimensions: Iterable, - inputs: Iterable, - results: Iterable, - postprocessed: Iterable | None = None, -) -> pp.ParserElement: - """ - Creates and executes the mask parser. - - Args: - dimensions (Iterable): List of valid dimension names. - inputs (Iterable): List of valid input names. - results (Iterable): List of valid variable/global expression names. - postprocessed (Iterable | None): List of valid postprocessed expression names. Defaults to None. - - Returns: - pp.ParseResults: evaluatable to a bool/boolean array. - """ - postprocessed = postprocessed if postprocessed is not None else set() - number, generic_identifier = expression_parser.setup_base_parser_elements() - dimensions_parser = data_var_parser(dimensions, DimensionArrayParser) - inputs_parser = data_var_parser(inputs, InputArrayParser) - results_parser = data_var_parser(results | postprocessed, ResultArrayParser) - config_option = config_option_parser(generic_identifier) - bool_operand = bool_parser() - unique_evaluatable_string = evaluatable_string_parser( - generic_identifier, set().union(dimensions, inputs, results, postprocessed) - ) - general_evaluatable_string = evaluatable_string_parser(generic_identifier, []) - id_list = expression_parser.list_parser( - number, unique_evaluatable_string, dimensions_parser - ) - subset = subset_parser( - [dimensions_parser, inputs_parser], - config_option, - number, - general_evaluatable_string, - ) - - arithmetic = pp.Forward() - comparison_helper_function = expression_parser.helper_function_parser( - unique_evaluatable_string, - number, - id_list, - arithmetic, - generic_identifier=generic_identifier, - ) - arithmetic_elements = [ - comparison_helper_function, - number, - dimensions_parser, - inputs_parser, - config_option, - ] - if postprocessed: - arithmetic_elements.insert(-2, results_parser) - - comparison_arithmetic = expression_parser.arithmetic_parser( - *arithmetic_elements, arithmetic=arithmetic - ) - comparison = comparison_parser( - lhs=[comparison_arithmetic], - rhs=[ - comparison_helper_function, - bool_operand, - number, - general_evaluatable_string, - ], - ) - - helper_function = expression_parser.helper_function_parser( - unique_evaluatable_string, - number, - id_list, - dimensions_parser, - inputs_parser, - results_parser, - config_option, - generic_identifier=generic_identifier, - ) - return mask_parser( - bool_operand, comparison, helper_function, subset, inputs_parser, results_parser - ) diff --git a/linopy/declarative/parsing.py b/linopy/declarative/parsing.py index 226ee0ad9..eaeac359a 100644 --- a/linopy/declarative/parsing.py +++ b/linopy/declarative/parsing.py @@ -1,6 +1,12 @@ -# Copyright (C) since 2013 Calliope contributors listed in AUTHORS. -# Licensed under the Apache 2.0 License (see LICENSE file). -"""Methods for math syntax parsing.""" +""" +Linopy declarative math parsing module. + +This module turns a validated math component definition into a list of +:class:`Equation` objects — pure data holding the parsed expression/mask ASTs +with all `$name` sub-expression and slicer references resolved — and provides +the typed entry points that evaluate an equation to a boolean mask array, a +linopy expression, a constraint tuple, or a LaTeX math string. +""" from __future__ import annotations @@ -8,928 +14,583 @@ import itertools import logging import operator -from collections.abc import Callable, Iterable -from typing import TYPE_CHECKING, Any, Literal, overload +from dataclasses import dataclass, field, replace +from typing import Any, Literal import pyparsing as pp import xarray as xr -from linopy.declarative import ( - eval_attrs, - expression_parser, - helper_functions, - mask_parser, +from linopy.declarative import grammar +from linopy.declarative.evaluate import ( + TRUE_ARRAY, + Context, + evaluate, + to_math_string, +) +from linopy.declarative.grammar import Node +from linopy.declarative.schema import ( + MATH_DEFS_T, + ConstraintDef, + ExpressionDef, + MathModel, + ObjectiveDef, + _Equations, ) -from linopy.declarative.schema import MATH_DEFS_T, ConfigModel, MathModel, _Equations from linopy.expressions import LinearExpression -if TYPE_CHECKING: - from linopy.model import Model -TRUE_ARRAY = xr.DataArray(True) - LOGGER = logging.getLogger(__name__) +GROUP_T = Literal[ + "variables", + "expressions", + "constraints", + "piecewise_constraints", + "objectives", + "postprocessed", +] -class ParsedBackendEquation: - """Backend equation parser.""" +EQUATION_DEFS_T = ConstraintDef | ExpressionDef | ObjectiveDef +"""Math component definitions that carry `equations`/`sub_expressions`/`slices` keys.""" - def __init__( - self, - equation_name: str, - sets: list[str], - expression: pp.ParseResults, - mask_list: list[pp.ParseResults], - sub_expressions: dict[str, pp.ParseResults] | None = None, - slices: dict[str, pp.ParseResults] | None = None, - ) -> None: - """ - For parsing equation expressions and corresponding "mask" strings. - - Args: - equation_name (str): Name of equation. - sets (list[str]): - Model data sets with which to create the initial multi-dimensional masking array - of the evaluated "mask" string. - expression (pp.ParseResults): - Parsed arithmetic/equation expression. - mask_list (list[pp.ParseResults]): - List of parsed mask strings. - sub_expressions (dict[str, pp.ParseResults] | None, optional): - Dictionary of parsed sub-expressions with which to replace sub-expression references - on evaluation of the parsed expression. Defaults to None. - slices (dict[str, pp.ParseResults] | None, optional): - Dictionary of parsed array slices with which to replace slice references - on evaluation of the parsed expression / sub-expression. Defaults to None. - """ - self.name = equation_name - self.mask = mask_list - self.expression = expression - self.sub_expressions = ( - sub_expressions if sub_expressions is not None else dict() - ) - self.slices = slices if slices is not None else dict() - self.sets = sets +_ERR_BULLET = " * " - def find_sub_expressions(self) -> set[str]: - """ - Identify all the references to sub_expressions in the parsed expression. - Returns: - set[str]: Unique sub-expression references. - """ - valid_eval_classes: tuple = ( - expression_parser.EvalOperatorOperand, - expression_parser.EvalFunction, - ) - to_find = expression_parser.EvalSubExpressions - elements: list - if isinstance(self.expression[0], to_find): - elements = [self.expression[0]] - else: - elements = [self.expression[0].values] - - return self._find_items_in_expression(elements, to_find, valid_eval_classes) - - def find_slices(self) -> set[str]: - """ - Finds all references to array slices in the expression and sub-expressions. +@dataclass(frozen=True) +class Equation: + """ + One fully-resolved equation of a math component. - Returns: - set[str]: Unique slice references. - """ - valid_eval_classes = tuple( - [ - expression_parser.EvalOperatorOperand, - expression_parser.EvalFunction, - expression_parser.EvalSlicedComponent, - ] - ) - to_find = expression_parser.EvalIndexSlice - elements: list = [ - self.expression[0].values, - *list(self.sub_expressions.values()), - ] + Produced by :func:`parse_component`: each combination of sub-expression and + slicer variants referenced by a user-defined equation yields one `Equation`. + """ - return self._find_items_in_expression(elements, to_find, valid_eval_classes) + name: str + """Unique equation name, including the chosen sub-expression/slicer variants.""" - @staticmethod - def _find_items_in_expression( - parser_elements: list | pp.ParseResults, - to_find: type[expression_parser.EvalString], - valid_eval_classes: tuple[type[expression_parser.EvalString], ...], - ) -> set[str]: - """ - Recursively find sub-expressions / index items defined in an equation expression. + sets: tuple[str, ...] + """The component's `foreach` dimensions.""" - Args: - parser_elements (list | pp.ParseResults): list of parser elements to check. - to_find (type[expression_parser.EvalString]): type of equation element to search for. - valid_eval_classes (tuple[type[expression_parser.EvalString], ...]): Other expression - elements that can be recursively searched + expression: Node + """Parsed expression AST.""" - Returns: - set[str]: All unique component / index item names. - """ - items: list = [] - recursive_func = functools.partial( - ParsedBackendEquation._find_items_in_expression, - to_find=to_find, - valid_eval_classes=valid_eval_classes, - ) - for parser_element in parser_elements: - if isinstance(parser_element, to_find): - items.append(parser_element.name) - - elif isinstance(parser_element, pp.ParseResults | list): - items.extend(recursive_func(parser_elements=parser_element)) - - elif isinstance(parser_element, valid_eval_classes): - items.extend(recursive_func(parser_elements=parser_element.values)) - return set(items) - - def add_expression_group_combination( - self, - expression_group_name: Literal["sub_expressions", "slices"], - expression_group_combination: Iterable[ParsedBackendEquation], - ) -> ParsedBackendEquation: - """ - Add parsed sub-expressions/index slices to a copy of self with updated names and mask lists. + masks: tuple[Node, ...] + """Parsed mask ASTs: the equation's own mask plus those of the chosen variants.""" - Args: - expression_group_name (Literal[sub_expressions, slices]): - Which of `sub-expressions`/`index slices` is being added. - expression_group_combination (Iterable[ParsedBackendEquation]): - All items of expression_group_name to be added. + sub_expressions: dict[str, Node] = field(default_factory=dict) + """Resolved `$name` sub-expression AST per name.""" - Returns: - ParsedBackendEquation: Copy of self with added sub-expressions/index slice dictionary and updated name - and mask list to include those corresponding to the dictionary entries. - """ - new_mask_list = [*self.mask] - for expr in expression_group_combination: - new_mask_list.extend(expr.mask) - new_name = f"{self.name}-{'-'.join([expr.name for expr in expression_group_combination])}" - expression_group_dict = { - expression_group_name: { - expr.name.split(":")[0]: expr.expression - for expr in expression_group_combination - } - } - return ParsedBackendEquation( - equation_name=new_name, - sets=self.sets, - expression=self.expression, - mask_list=new_mask_list, - **{ - "sub_expressions": self.sub_expressions, - "slices": self.slices, - **expression_group_dict, # type: ignore - }, - ) + slices: dict[str, Node] = field(default_factory=dict) + """Resolved `$name` slicer AST per name.""" - # Expecting array if not requesting latex string - @overload - def evaluate_mask( - self, - input_data: xr.Dataset, - model: Model, - math: MathModel, - config: ConfigModel, - *, - return_type: Literal["raw"] = "raw", - references: set | None = None, - initial_mask: xr.DataArray = TRUE_ARRAY, - ) -> xr.DataArray: ... - - # Expecting string if requesting latex string. - @overload - def evaluate_mask( - self, - input_data: xr.Dataset, - model: Model, - math: MathModel, - config: ConfigModel, - *, - return_type: Literal["math_string"], - references: set | None = None, - ) -> str: ... - - def evaluate_mask( - self, - input_data: xr.Dataset, - model: Model, - math: MathModel, - config: ConfigModel, - *, - return_type: str = "raw", - references: set | None = None, - initial_mask: xr.DataArray = TRUE_ARRAY, - ) -> xr.DataArray | str: - """ - Evaluate parsed backend object dictionary `mask` string. - - Args: - input_data (xr.Dataset): Model input data. - model (Model): Linopy model. - math (MathModel): Calliope math definitions. - config (ConfigModel): Build configuration options. - return_type (str, optional): If "raw", return xarray.DataArray. - If "math_string", return LaTex math string. - Defaults to "raw". - references (set | None, optional): List of references to use in evaluation. - Defaults to None. - initial_mask (xr.DataArray, optional): If given, the mask array resulting - from evaluation will be further masked by this array. - Defaults to xr.DataArray(True) (i.e., no effect). - - Returns: - xr.DataArray | str: - If return_type == `array`: Boolean array defining on which index items a parsed component should be built. - If return_type == `math_string`: Valid LaTeX math string defining the "mask" conditions using logic notation. - """ - eval_attrs_ = { - "equation_name": self.name, - "helper_functions": helper_functions._registry["mask"], - "input_data": input_data, - "model": model, - "math": math, - "config": config, - } - if references is not None: - eval_attrs_["references"] = references - - evaluated_masks = [ - mask[0].eval(return_type, eval_attrs.EvalAttrs(**eval_attrs_)) - for mask in self.mask + 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(), ] - if return_type == "math_string": - return r"\land{}".join(f"({i})" for i in evaluated_masks if i != "true") - else: - mask = xr.DataArray( - functools.reduce(operator.and_, [initial_mask, *evaluated_masks]) - ) - if not mask.any(): - self.log_not_added("'mask' does not apply anywhere.") - return mask - - def drop_dims_not_in_foreach(self, mask: xr.DataArray) -> xr.DataArray: - """ - Remove all dimensions not included in "foreach" from the input array. - - Args: - mask (xr.DataArray): Array with potentially unwanted dimensions - - Returns: - xr.DataArray: - Array with same dimensions as the user-defined foreach sets. - Dimensions are ordered to match the order given by the sets. - """ - unwanted_dims = set(mask.dims).difference(self.sets) - return (mask.sum(unwanted_dims) > 0).astype(bool).transpose(*self.sets) - - def _evaluate( - self, - input_data: xr.Dataset, - model: Model, - math: MathModel, - *, - return_type: Literal["expr", "math_string"], - references: set | None, - mask: xr.DataArray, - ) -> Any: - """ - Evaluate the parsed expression tree. - - Shared by :meth:`evaluate_expression` (arithmetic roots) and - :meth:`evaluate_equation` (comparison roots). - """ - eval_attrs_ = { - "equation_name": self.name, - "slice_dict": self.slices, - "sub_expression_dict": self.sub_expressions, - "input_data": input_data, - "model": model, - "math": math, - "mask": mask, - "helper_functions": helper_functions._registry["expression"], - } - if references is not None: - eval_attrs_["references"] = references - return self.expression[0].eval(return_type, eval_attrs.EvalAttrs(**eval_attrs_)) - - # Expecting a linopy expression if not requesting latex string. - @overload - def evaluate_expression( - self, - input_data: xr.Dataset, - model: Model, - math: MathModel, - *, - return_type: Literal["expr"] = "expr", - references: set | None = None, - mask: xr.DataArray = TRUE_ARRAY, - ) -> LinearExpression: ... - - # Expecting string if requesting latex string. - @overload - def evaluate_expression( - self, - input_data: xr.Dataset, - model: Model, - math: MathModel, - *, - return_type: Literal["math_string"], - references: set | None = None, - ) -> str: ... - - def evaluate_expression( - self, - input_data: xr.Dataset, - model: Model, - math: MathModel, - *, - return_type: Literal["expr", "math_string"] = "expr", - references: set | None = None, - mask: xr.DataArray = TRUE_ARRAY, - ) -> LinearExpression | str: - """ - Evaluate an arithmetic math string (expressions/objectives). - - Args: - input_data (xr.Dataset): Model input data. - model (Model): Linopy model. - math (MathModel): Linopy math definitions. - - Keyword Args: - return_type (str, optional): - If "expr", return a linopy expression. If "math_string", return a LaTeX - math string. Defaults to "expr". - references (set | None, optional): - If given, any references in the math string to other model components - will be logged here. Defaults to None. - mask (xr.DataArray, optional): - If given, should be a boolean array with which to mask any produced arrays. - Defaults to xr.DataArray(True). - - Returns: - LinearExpression | str: - If return_type == `expr`: a linopy expression. A pure-parameter - expression (evaluated to an ``xr.DataArray``) is coerced to a - ``LinearExpression``. - If return_type == `math_string`: a valid LaTeX math string. - """ - evaluated = self._evaluate( - input_data, - model, - math, - return_type=return_type, - references=references, - mask=mask, - ) - if return_type == "expr" and isinstance(evaluated, xr.DataArray): - evaluated = LinearExpression(evaluated, model) - return evaluated - - # Expecting a (lhs, sign, rhs) tuple if not requesting latex string. - @overload - def evaluate_equation( - self, - input_data: xr.Dataset, - model: Model, - math: MathModel, - *, - return_type: Literal["expr"] = "expr", - references: set | None = None, - mask: xr.DataArray = TRUE_ARRAY, - ) -> tuple[LinearExpression, xr.DataArray, LinearExpression]: ... - - # Expecting string if requesting latex string. - @overload - def evaluate_equation( - self, - input_data: xr.Dataset, - model: Model, - math: MathModel, - *, - return_type: Literal["math_string"], - references: set | None = None, - ) -> str: ... - - def evaluate_equation( - self, - input_data: xr.Dataset, - model: Model, - math: MathModel, - *, - return_type: Literal["expr", "math_string"] = "expr", - references: set | None = None, - mask: xr.DataArray = TRUE_ARRAY, - ) -> tuple[LinearExpression, xr.DataArray, LinearExpression] | str: - """ - Evaluate a comparison math string (constraints) of the form ``LHS OP RHS``. - - Args: - input_data (xr.Dataset): Model input data. - model (Model): Linopy model. - math (MathModel): Linopy math definitions. - - Keyword Args: - return_type (str, optional): - If "expr", return a ``(lhs, sign, rhs)`` tuple for constraint assembly. - If "math_string", return a LaTeX math string. Defaults to "expr". - references (set | None, optional): - If given, any references in the math string to other model components - will be logged here. Defaults to None. - mask (xr.DataArray, optional): - If given, should be a boolean array with which to mask any produced arrays. - Defaults to xr.DataArray(True). - - Returns: - tuple[LinearExpression, xr.DataArray, LinearExpression] | str: - If return_type == `expr`: a ``(lhs, sign, rhs)`` tuple, where ``lhs``/``rhs`` - are linopy expressions (a pure-parameter side is coerced to a - ``LinearExpression``) and ``sign`` is a DataArray of the comparison operator. - If return_type == `math_string`: a valid LaTeX math string. - """ - evaluated = self._evaluate( - input_data, - model, - math, - return_type=return_type, - references=references, - mask=mask, - ) - if return_type == "math_string": - return evaluated - lhs, sign, rhs = evaluated - if isinstance(lhs, xr.DataArray): - lhs = LinearExpression(lhs, model) - if isinstance(rhs, xr.DataArray): - rhs = LinearExpression(rhs, model) - return lhs, sign, rhs - - def raise_error_on_mask_expr_mismatch( - self, expression: xr.DataArray, mask: xr.DataArray - ) -> None: - """ - Checks if an evaluated expression is consistent with the `mask` array. - - Args: - expression (xr.DataArray): array of linear expressions or one side of a constraint equation. - mask (xr.DataArray): mask array; there should be a valid expression value for all True elements. - - Raises: - BackendError: - Raised if there is a dimension in the expression that is not in the mask. - BackendError: - Raised if the expression has any NaN mask the mask applies. - """ - broadcast_dims_mask = set(expression.dims).difference(set(mask.dims)) - if broadcast_dims_mask: - raise ValueError( - f"{self.name} | The linear expression array is indexed over dimensions not present in `foreach`: {broadcast_dims_mask}" - ) - # Check whether expression has NaN values in elements mask the expression should be valid. - incomplete_constraints = expression.isnull() & mask - if incomplete_constraints.any(): - raise ValueError( - f"{self.name} | Missing a linear expression for some coordinates selected by 'mask'. Adapting 'mask' might help." - ) - - def log_not_added( - self, - message: str, - level: Literal["info", "warning", "debug", "error", "critical"] = "debug", - ): - """ - Log to module-level logger with some prettification of the message. - - Args: - message (str): Message to log. - level (Literal["info", "warning", "debug", "error", "critical"], optional): - Log level. Defaults to "debug". - """ - getattr(LOGGER, level)( - f"Math parsing | {self.name} | Component not added; {message}" + return set().union( + *(grammar.find_refs(tree, grammar.Component) for tree in trees) ) -class ParsedBackendComponent(ParsedBackendEquation): - """Backend component parser.""" +# --------------------------------------------------------------------------- +# Parsing +# --------------------------------------------------------------------------- - _ERR_BULLET: str = " * " - _ERR_STRING_ORDER: list[str] = ["expression_group", "id", "expr_or_mask"] - PARSERS: dict[str, Callable] = { - "constraints": expression_parser.generate_equation_parser, - "expressions": expression_parser.generate_arithmetic_parser, - "postprocessed": expression_parser.generate_arithmetic_parser, - "objectives": expression_parser.generate_arithmetic_parser, - "piecewise_constraints": expression_parser.generate_arithmetic_parser, - } - def __init__( - self, - group: Literal[ - "variables", - "expressions", - "constraints", - "piecewise_constraints", - "objectives", - "postprocessed", - ], - name: str, - unparsed_data: MATH_DEFS_T, - parsing_components: dict[str, dict[str, set[str]]], - ) -> None: - """ - Parse an optimisation problem configuration. - - Defined in a dictionary of strings loaded from YAML into a series of Python - objects that can be passed onto a solver interface like Pyomo or Gurobipy. - - Args: - group (Literal["variables", "expressions", "constraints", "objectives"]): - Optimisation problem component group to which the unparsed data belongs. - name (str): Name of the optimisation problem component - unparsed_data (T): Unparsed math formulation. Expected structure depends on - the group to which the optimisation problem component belongs. - parsing_components (dict[str, dict[str, Iterable[str]]]): - Dictionary of valid component names for different categories of model data to use in parsing `mask` and `expression` strings. - """ - self.name = f"{group}:{name}" - self.group = group - self._unparsed = unparsed_data - self._mask_components = parsing_components["mask"] - self._expression_components = set().union( - *parsing_components["expression"].values() - ) - self.mask: list[pp.ParseResults] = [] - self.equations: list[ParsedBackendEquation] = [] - self.equation_expression_parser: Callable = self.PARSERS.get( - group, lambda x: None - ) - - # capture errors to dump after processing, - # to make it easier for a user to fix the constraint YAML. - self._errors: list = [] - self._tracker = self._init_tracker() +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())) - # Initialise switches - self._is_valid: bool = True - # Add objects that are used by shared functions - self.sets: set[str] = set(unparsed_data.foreach) +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"]), + ) - def get_parsing_position(self): - """Create "." separated list from tracked strings.""" - return ".".join( - filter(None, [self._tracker[i] for i in self._ERR_STRING_ORDER]) - ) - def reset_tracker(self): - """Re-initialise error string tracking.""" - self._tracker = self._init_tracker() +class _ErrorCollector: + """Collect parse errors with their positions, to raise a single error at the end.""" - def _init_tracker(self): - """Initialise error string tracking as dictionary of `key: None`.""" - return {i: None for i in self._ERR_STRING_ORDER} + def __init__(self, component_name: str) -> None: + self.component_name = component_name + self.errors: list[str] = [] - def parse_top_level_mask( - self, errors: Literal["raise", "ignore"] = "raise" - ) -> None: + def parse( + self, parser: pp.ParserElement, string: str, position: str + ) -> Node | None: """ - Parse the "mask" string that is (optionally) given as a top-level key of the math component dictionary. - - Args: - errors (Literal["raise", "ignore"], optional): - Collected parsing errors can be raised directly or ignored. - If errors exist and are ignored, the parsed component cannot be successfully evaluated. Defaults to "raise". - """ - top_level_mask = self.parse_mask_string(self._unparsed.mask) - - if errors == "raise": - self.raise_caught_errors() - - if self._is_valid: - self.mask = [top_level_mask] + Parse `string`, returning its AST root or None if parsing fails. - def parse_equations( - self, errors: Literal["raise", "ignore"] = "raise" - ) -> list[ParsedBackendEquation]: + Failures are stored with a caret marker pointing at the parse position, + for raising later via :meth:`raise_errors`. """ - Parse `expression` and `mask` strings of math component dictionary. - - Args: - errors (Literal["raise", "ignore"], optional): - Collected parsing errors can be raised directly or ignored. - If errors exist and are ignored, the parsed component cannot be successfully evaluated. Defaults to "raise". - - Returns: - list[ParsedBackendEquation]: - List of parsed equations ready to be evaluated. - The length of the list depends on the product of provided equations and sub-expression/slice references. - """ - equations = self.generate_expression_list( - expression_parser=self.equation_expression_parser( - self._expression_components - ), - expression_list=self._unparsed.equations, - expression_group="equations", - id_prefix=self.name, + try: + return parser.parse_string(string, parse_all=True)[0] + except pp.ParseException as excinfo: + pointer = f"{position} (line {excinfo.lineno}, char {excinfo.col}): " + marker_pos = " " * (len(pointer) + 2 * len(_ERR_BULLET) + excinfo.col - 1) + self.errors.append(f"{pointer}{excinfo.line}\n{marker_pos}^") + return None + + def raise_errors(self) -> None: + """Raise all collected parse errors as a single bullet-point ValueError.""" + if self.errors: + raise ValueError(f"- {self.component_name}: {self.errors}") + + +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 = _ErrorCollector(name) + parsed = collector.parse(_mask_grammar(math), mask_string, "mask") + collector.raise_errors() + assert parsed is not None + return parsed + + +def _parse_variants( + collector: _ErrorCollector, + 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, f"{position_id}.mask") + expression = collector.parse( + parser, item.expression, f"{position_id}.expression" ) - - sub_expression_dict = { - c_name: self.generate_expression_list( - expression_parser=expression_parser.generate_sub_expression_parser( - self._expression_components - ), - expression_list=c_list, - expression_group="sub_expressions", - id_prefix=c_name, - ) - for c_name, c_list in self._unparsed.sub_expressions.root.items() - } - slice_dict = { - idx_name: self.generate_expression_list( - expression_parser=expression_parser.generate_slice_parser( - self._expression_components - ), - expression_list=idx_list, - expression_group="slices", - id_prefix=idx_name, - ) - for idx_name, idx_list in self._unparsed.slices.root.items() - } - - if errors == "raise": - self.raise_caught_errors() - - equations_with_sub_expressions = [] - for equation in equations: - equations_with_sub_expressions.extend( - self.extend_equation_list_with_expression_group( - equation, sub_expression_dict, "sub_expressions" + if expression is not None and mask is not None: + equations.append( + Equation( + name=f"{name_prefix}:{idx}", + sets=sets, + expression=expression, + masks=(mask,), ) ) - equations_with_sub_expressions_and_slices: list[ParsedBackendEquation] = [] - for equation in equations_with_sub_expressions: - equations_with_sub_expressions_and_slices.extend( - self.extend_equation_list_with_expression_group( - equation, slice_dict, "slices" - ) + return equations + + +def _expand( + component_name: str, + equations: list[Equation], + candidates: dict[str, list[Equation]], + kind: Literal["sub_expressions", "slices"], +) -> 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, with the chosen variants' masks and + ASTs merged in. + """ + expanded = [] + for equation in equations: + ref_type = grammar.SubExprRef if kind == "sub_expressions" else grammar.SliceRef + trees = [equation.expression, *equation.sub_expressions.values()] + refs = set().union(*(grammar.find_refs(tree, ref_type) for tree in trees)) + if not refs: + expanded.append(equation) + continue + undefined = refs.difference(candidates.keys()) + if undefined: + raise KeyError( + f"{component_name}: Undefined {kind} found in equation: {undefined}" ) - - return equations_with_sub_expressions_and_slices - - def _parse_string( - self, parser: pp.ParserElement, parse_string: str - ) -> pp.ParseResults: - """ - Parse equation string according to predefined parsing grammar. - - Args: - parser (pp.ParserElement): Parsing grammar. - parse_string (str): String to parse according to parser grammar. - - Returns: - Optional[pp.ParseResults]: - Parsed string. If any parsing errors are caught, - they will be logged to `self._errors` to raise later. - """ - try: - parsed = parser.parse_string(parse_string, parse_all=True) - except pp.ParseException as excinfo: - parsed = pp.ParseResults([]) - self._is_valid = False - pointer = f"{self.get_parsing_position()} (line {excinfo.lineno}, char {excinfo.col}): " - marker_pos = " " * ( - len(pointer) + 2 * len(self._ERR_BULLET) + excinfo.col - 1 + 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), ) - self._errors.append(f"{pointer}{excinfo.line}\n{marker_pos}^") - - return parsed - - def parse_mask_string(self, mask_string: str = "True") -> pp.ParseResults: - """ - Parse a "mask" string of the form "CONDITION OPERATOR CONDITION". - - The operator can be "and"/"or"/"not and"/"not or". - - Args: - mask_string (str): - string value from a math dictionary "mask" key. - Defaults to "True", to have no effect on the subsequent subsetting. - - Returns: - pp.ParseResults: Parsed string. If any parsing errors are caught, - they will be logged to `self._errors` to raise later. - """ - parser = mask_parser.generate_mask_string_parser(**self._mask_components) - self._tracker["expr_or_mask"] = "mask" - return self._parse_string(parser, mask_string) - - def generate_expression_list( - self, - expression_parser: pp.ParserElement, - expression_list: _Equations, - expression_group: Literal["equations", "sub_expressions", "slices"], - id_prefix: str = "", - ) -> list[ParsedBackendEquation]: - """ - Align user-defined constraint equations/sub-expressions. - - Achieved by parsing expressions, specifying a default "mask" string if not - defined, and providing an ID to enable returning to the initial dictionary. - - Args: - expression_parser (pp.ParserElement): parser to use. - expression_list (list[UnparsedEquation]): list of constraint equations - or sub-expressions with arithmetic expression string and optional - mask string. - expression_group (Literal["equations", "sub_expressions", "slices"]): - For error reporting, the constraint dict key corresponding to the parse_string. - id_prefix (str, optional): Extends the ID from a number corresponding to the - expression_list position `idx` to a tuple of the form (id_prefix, idx). - Defaults to "". - - Returns: - list[ParsedBackendEquation]: Aligned expression dictionaries with parsed - expression strings. - """ - parsed_equation_list = [] - - if expression_group == "equations": - to_track = {"expression_group": f"{expression_group}[{{id}}]"} - else: - to_track = { - "expression_group": expression_group, - "id": f"{id_prefix}[{{id}}]", + resolved = { + variant.name.split(":")[0]: variant.expression + for variant in combination } - - for idx, expression_data in enumerate(expression_list): - self._tracker.update({k: v.format(id=idx) for k, v in to_track.items()}) - - parsed_mask = self.parse_mask_string(expression_data.mask) - - self._tracker["expr_or_mask"] = "expression" - parsed_expression = self._parse_string( - expression_parser, expression_data.expression - ) - if len(parsed_expression) > 0: - parsed_equation_list.append( - ParsedBackendEquation( - equation_name=":".join(filter(None, [id_prefix, str(idx)])), - sets=self.sets, - mask_list=[parsed_mask], - expression=parsed_expression, - ) + if kind == "sub_expressions": + new_equation = replace( + equation, name=new_name, masks=new_masks, sub_expressions=resolved ) - self.reset_tracker() - - return parsed_equation_list - - def extend_equation_list_with_expression_group( - self, - parsed_equation: ParsedBackendEquation, - parsed_items: dict[str, list[ParsedBackendEquation]], - expression_group: Literal["sub_expressions", "slices"], - ) -> list[ParsedBackendEquation]: - """ - Extend equation expressions with sub-expression data. - - Finds all sub-expressions referenced in an equation expression and returns a - product of the sub-expression data. - - Args: - parsed_equation (ParsedBackendEquation): Equation data dictionary. - parsed_items (dict[str, list[ParsedBackendEquation]]): - Dictionary of expressions to replace within the equation data dictionary. - expression_group (Literal["sub_expressions", "slices"]): - Name of expression group that the parsed_items dict is referencing. - - Returns: - list[ParsedBackendEquation]: Expanded list of parsed equations with the - product of all references to items from the `expression_group` - producing a new equation object. E.g., if the input equation object has - a reference to an slice which itself has two expression options, two - equation objects will be added to the return list. - """ - if expression_group == "sub_expressions": - equation_items = parsed_equation.find_sub_expressions() - elif expression_group == "slices": - equation_items = parsed_equation.find_slices() - if not equation_items: - return [parsed_equation] - - invalid_items = equation_items.difference(parsed_items.keys()) - if invalid_items: - raise KeyError( - f"{self.name}: Undefined {expression_group} found in equation: {invalid_items}" - ) - - parsed_item_product = itertools.product( - *[parsed_items[k] for k in equation_items] + else: + new_equation = replace( + equation, name=new_name, masks=new_masks, slices=resolved + ) + expanded.append(new_equation) + return expanded + + +def parse_component( + group: GROUP_T, name: str, definition: EQUATION_DEFS_T, math: MathModel +) -> 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 (syntax errors across all of them are + collected and raised together), then every equation is expanded with the + cartesian product of the sub-expression and slicer variants it references. + + Parameters + ---------- + group : 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. + + Returns + ------- + list[Equation] + One equation per user-defined equation and referenced variant combination. + """ + component_name = f"{group}:{name}" + names = _expression_names(math) + equation_parser = ( + grammar.equation_grammar(names) + if group == "constraints" + else grammar.arithmetic_grammar(names) + ) + mask_parser = _mask_grammar(math) + # Objectives are adimensional: they carry no `foreach` key. + sets = tuple(getattr(definition, "foreach", ())) + collector = _ErrorCollector(component_name) + + equations = _parse_variants( + collector, + equation_parser, + mask_parser, + definition.equations, + sets, + "equations", + component_name, + ) + sub_expressions = { + sub_name: _parse_variants( + collector, + 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, + 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() + } + collector.raise_errors() - return [ - parsed_equation.add_expression_group_combination( - expression_group, parsed_item_combination - ) - for parsed_item_combination in parsed_item_product - ] + equations = _expand(component_name, equations, sub_expressions, "sub_expressions") + return _expand(component_name, equations, slices, "slices") - def foreach_matrix(self, input_data: xr.Dataset) -> xr.DataArray: - """ - Generate a multi-dimensional array mask a constraint will be built. - The multi-dimensional boolean array is based on the sets over which the - constraint is to be built (`foreach`) and the model `exists` array. +# --------------------------------------------------------------------------- +# Component-level masking +# --------------------------------------------------------------------------- - Args: - input_data (xr.Dataset): Calliope model dataset. - Returns: - xr.DataArray: boolean array indexed over ["nodes", "techs", "carriers"] - + any additional dimensions provided by `foreach`. - """ - if self.sets.difference(input_data.dims): - self.log_not_added( - f"indexed over unidentified set names: `{self.sets.difference(input_data.dims)}`." - ) - return xr.DataArray(False) - if not self.sets: - return xr.DataArray(True) - else: - exists_and_foreach = [input_data[i].notnull() for i in self.sets] - return functools.reduce(operator.and_, exists_and_foreach) - - def generate_top_level_mask( - self, - input_data: xr.Dataset, - model: Model, - math: MathModel, - config: ConfigModel, - *, - align_to_foreach_sets: bool = True, - break_early: bool = True, - references: set | None = None, - ) -> xr.DataArray: - """ - Generate a multi-dimentional "mask" array. - - The multi-dimensional array is created using model inputs and component sets - defined in foreach. The component top-level "mask" is then applied to the - array. - - Args: - input_data (xr.Dataset): Model input data. - model (xr.Dataset): Backend interface component dataset. - math (MathModel): Calliope math definitions. - config (ConfigModel): Build configuration options. - align_to_foreach_sets (bool, optional): - By default, all foreach arrays have the dimensions ("nodes", "techs", "carriers") - as well as any additional dimensions provided by the component's "foreach" key. - If this argument is True, the dimensions not included in "foreach" are removed from the array. - Defaults to True. - break_early (bool, optional): - If any intermediate array has no valid elements (i.e. all are False), - the function will return that array rather than continuing - saving - time and memory on large models. Defaults to True. - references (set | None, optional): references to use during evaluation. Defaults to None. - - Returns: - xr.DataArray: Boolean array defining on which index items a parsed component should be built. - """ - foreach_mask = self.foreach_matrix(input_data) - - if not foreach_mask.any(): - self.log_not_added("'foreach' does not apply anywhere.") - - if break_early and not foreach_mask.any(): - return foreach_mask - - self.parse_top_level_mask() - mask = self.evaluate_mask( - input_data, - model, - math, - config, - initial_mask=foreach_mask, - references=references if references is not None else set(), - ) - if break_early and not mask.any(): - return mask +def foreach_mask(sets: tuple[str, ...], input_data: xr.Dataset) -> xr.DataArray: + """ + Return the initial boolean array spanning a component's `foreach` dimensions. - if align_to_foreach_sets: - mask = self.drop_dims_not_in_foreach(mask) + 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 component_mask( + group: GROUP_T, + name: str, + definition: MATH_DEFS_T, + 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 (optional) + top-level `mask` string, breaking early if no valid element remains. + + Parameters + ---------- + group : GROUP_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. + 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}" + # Objectives are adimensional: they carry no `foreach` or `mask` keys. + sets = tuple(getattr(definition, "foreach", ())) + mask_string = getattr(definition, "mask", "True") + initial_mask = foreach_mask(sets, ctx.input_data) + if not initial_mask.any(): + LOGGER.debug( + f"Math parsing | {component_name} | Component not added; " + "'foreach' does not apply anywhere." + ) + return initial_mask + + mask_node = parse_mask(mask_string, ctx.math, component_name) + mask_ctx = replace(ctx, route="mask", equation_name=component_name) + mask = xr.DataArray(initial_mask & evaluate(mask_node, mask_ctx)) + if not mask.any(): + LOGGER.debug( + f"Math parsing | {component_name} | Component not added; " + "'mask' does not apply anywhere." + ) return mask - def raise_caught_errors(self): - """Pipe parsing errors to the ModelError bullet point list generator.""" - errors = [] - if not self._is_valid: - errors.append({f"{self.name}": self._errors}) - if errors: - raise ValueError( - "\n".join(f"- {k}: {v}" for err in errors for k, v in err.items()) - ) + if align_to_foreach_sets: + mask = drop_dims_not_in_foreach(mask, sets) + return mask + + +# --------------------------------------------------------------------------- +# Typed evaluation entry points +# --------------------------------------------------------------------------- + + +def _equation_ctx( + equation: Equation, + ctx: Context, + route: Literal["expression", "mask"], + **kwargs: Any, +) -> Context: + """Return a context copy carrying the equation's name and resolved references.""" + return replace( + ctx, + equation_name=equation.name, + route=route, + 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 = [evaluate(mask, mask_ctx) for mask in equation.masks] + mask = xr.DataArray(functools.reduce(operator.and_, [initial_mask, *evaluated])) + if not mask.any(): + LOGGER.debug( + f"Math parsing | {equation.name} | Component not added; " + "'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, "expression", mask=mask) + evaluated = evaluate(equation.expression, expr_ctx, expr=True) + if isinstance(evaluated, xr.DataArray): + evaluated = LinearExpression(evaluated, ctx.model) + return evaluated + + +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; pure-parameter sides are + coerced to `LinearExpression` and `sign` is an array of the comparison + operator. + """ + expr_ctx = _equation_ctx(equation, ctx, "expression", mask=mask) + lhs, sign, rhs = evaluate(equation.expression, expr_ctx, expr=True) + if isinstance(lhs, xr.DataArray): + lhs = LinearExpression(lhs, ctx.model) + if isinstance(rhs, xr.DataArray): + rhs = LinearExpression(rhs, ctx.model) + return lhs, sign, rhs + + +def as_latex( + equation: Equation, + ctx: Context, + *, + what: Literal["expression", "mask"] = "expression", +) -> str: + """ + Render an equation's expression or 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. + """ + if what == "mask": + mask_ctx = _equation_ctx(equation, ctx, "mask") + strings = [to_math_string(mask, mask_ctx) for mask in equation.masks] + return r"\land{}".join(f"({s})" for s in strings if s != "true") + expr_ctx = _equation_ctx(equation, ctx, "expression") + return to_math_string(equation.expression, expr_ctx) + + +def check_mask_expr_consistency( + name: str, expression: xr.DataArray, mask: xr.DataArray +) -> None: + """ + Check that an evaluated expression is consistent with its mask array. + + Parameters + ---------- + name : str + Name to identify the equation by in error messages. + expression : xr.DataArray + Array of linear expressions or one side of a constraint equation. + mask : xr.DataArray + Boolean mask; there should be a valid expression value wherever it is True. + + Raises + ------ + ValueError + If the expression is indexed over dimensions not present in the mask, or + has missing (NaN) entries where the mask applies. + """ + broadcast_dims_mask = set(expression.dims).difference(set(mask.dims)) + if broadcast_dims_mask: + raise ValueError( + f"{name} | The linear expression array is indexed over dimensions " + f"not present in `foreach`: {broadcast_dims_mask}" + ) + incomplete_constraints = expression.isnull() & mask + if incomplete_constraints.any(): + raise ValueError( + f"{name} | Missing a linear expression for some coordinates selected " + "by 'mask'. Adapting 'mask' might help." + ) diff --git a/linopy/declarative/schema.py b/linopy/declarative/schema.py index 05c8a44a5..a17afaef0 100644 --- a/linopy/declarative/schema.py +++ b/linopy/declarative/schema.py @@ -1,22 +1,37 @@ -# Copyright (C) since 2013 Calliope contributors listed in AUTHORS. -# Licensed under the Apache 2.0 License (see LICENSE file). -"""Schema for Calliope mathematical definition.""" +""" +Linopy declarative math schema module. + +This module contains the pydantic models that validate declarative math and +build-configuration definitions. +""" + +from __future__ import annotations import logging from collections.abc import Hashable, Iterable from functools import cached_property -from typing import Annotated, ClassVar, Literal, Self, TypeVar +from typing import Annotated, Any, ClassVar, Literal, Self, TypeVar +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__) -LOGGER.setLevel(logging.INFO) -# == + # 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", @@ -58,24 +73,24 @@ def _validate_unique_list(v: list) -> list: class LinopyDictModel(RootModel): - """Pydantic Model that is used to store dictionaries with user-defined keys and Calliope pydantic model values.""" + """Pydantic model storing a dictionary of user-named component definitions.""" - def __setitem__(self, *args, **kwargs) -> None: + 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__} directly. Use the `update` method instead, which will return a copy.", + f"Cannot set a {self.__class__.__name__} item directly. Re-validate a new definition dictionary instead.", ) - def __getitem__(self, key): + def __getitem__(self, key: str) -> Any: """Expose the root attribute when getting an item by key.""" return self.root[key] - def __repr__(self, *args, **kwargs): + 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): + def __rich_repr__(self) -> Iterable: """Prettyprint the __repr__ of the root attribute when requesting the prettyprint of the class.""" yield from self.root.items() @@ -84,82 +99,29 @@ def _active(self) -> dict[str, BaseModel]: """Return only active components.""" return {k: v for k, v in self.root.items() if v.active} - def update( - self, update_def: dict | BaseModel, deep: bool = False, overwrite: bool = True - ) -> Self: - """ - Return a new iteration of the model with updated fields. - - Args: - update_def (dict | BaseModel): Dictionary or pydantic model with which to update the base model. - deep (bool, optional): Set to True to make a deep copy of the model. Defaults to False. - overwrite (bool, optional): Set to False to only update fields that are not already set in the base model. Defaults to True. - - Returns: - BaseModel: New model instance. - """ - update_dict: dict = ( - update_def.model_dump(exclude_unset=True) - if isinstance(update_def, BaseModel) - else update_def - ) - new_dict = dict() - # Iterate through dict to be updated and convert any sub-dicts into their respective pydantic model objects. - for key, val in update_dict.items(): - key_class = self.root.get(key, None) - if isinstance(key_class, LinopyBaseModel): - new_dict[key] = key_class.update(val, deep=deep, overwrite=overwrite) - elif isinstance(key_class, LinopyListModel): - if overwrite: - new_dict[key] = key_class.update(val) - else: - continue - elif key_class == val: - continue - else: - if key not in self.root or overwrite: - LOGGER.debug(f"Adding {self.__class__.__name__} entry: `{key}`") - new_dict[key] = self.model_validate({key: val})[key] - - return self.model_validate(self.root | new_dict) - class LinopyListModel(RootModel): - """Pydantic Model that is used to store lists of Linopy pydantic models.""" + """Pydantic model storing a list of definitions.""" - def __iter__(self): + def __iter__(self) -> Any: """Iterate over root attribute contents when iterating over class.""" return iter(self.root) - def __getitem__(self, item: int): + 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, **kwargs): + 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): + def __rich_repr__(self) -> Iterable: """Prettyprint the __repr__ of the root attribute when requesting the prettyprint of the class.""" yield from self.root - def update(self, update_list: list) -> Self: - """ - Return a new iteration of the model fields entirely replaced. - - We do not allow updating individual items in the list as it's hard to guarantee the order of items in the list. - - Args: - update_list (list): List with which to update the base model. - - Returns: - BaseModel: New model instance. - """ - return self.model_validate(update_list) - class LinopyBaseModel(BaseModel): - """A base class for creating pydantic models for Linopy models.""" + """Base class for declarative math pydantic models.""" model_config = { "extra": "forbid", @@ -168,69 +130,10 @@ class LinopyBaseModel(BaseModel): "use_attribute_docstrings": True, } - def __getitem__(self, item): + def __getitem__(self, item: str) -> Any: """Allow attribute access via item lookup.""" return getattr(self, item) - def update( - self, - update_def: dict | BaseModel, - deep: bool = False, - overwrite: bool = True, - _suppress_log: bool = False, - ) -> Self: - """ - Return a new iteration of the model with updated fields. - - Args: - update_def (dict | BaseModel): Dictionary or pydantic model with which to update the base model. - deep (bool, optional): Set to True to make a deep copy of the model. Defaults to False. - overwrite (bool, optional): Set to False to only update fields that are not already set in the base model. Defaults to True. - _suppress_log (bool, optional): - Set to True to suppress logging of updated fields. - This is an internal method argument used to avoid logging updates when the update method is called recursively. - Defaults to False. - - Returns: - BaseModel: New model instance. - """ - new_dict = dict() - # Iterate through dict to be updated and convert any sub-dicts into their respective pydantic model objects. - # Wrapped in `AttrDict` to allow users to define dot notation nested configuration. - # We revert to dict format to avoid issues with the `model_copy` method later. - update_dict = ( - update_def.model_dump(exclude_unset=True) - if isinstance(update_def, BaseModel) - else update_def - ) - for key, val in update_dict.items(): - key_class = getattr(self, key, None) - if isinstance(key_class, LinopyBaseModel | LinopyDictModel): - new_dict[key] = key_class.update(val, deep=deep, overwrite=overwrite) - elif isinstance(key_class, LinopyListModel): - if overwrite: - new_dict[key] = key_class.update(val) - else: - continue - elif key_class == val: - continue - else: - if not _suppress_log and ( - key not in self.model_fields_set - or (key in self.model_fields_set and overwrite) - ): - LOGGER.debug( - f"Updating {self.__class__.__name__} `{key}`: {key_class} -> {val}" - ) - new_dict[key] = val - updated = super().model_copy(update=new_dict, deep=deep) - if not overwrite: - extra_update = super().model_dump(exclude_unset=True, serialize_as_any=True) - updated = updated.update(extra_update, deep=deep, _suppress_log=True) - return updated.model_validate( - updated.model_dump(exclude_unset=True, serialize_as_any=True) - ) - class _ExpressionItem(LinopyBaseModel): """Schema for equations, _subexpressions and slices.""" @@ -371,21 +274,6 @@ class PiecewiseConstraintDef(_MathIndexedComponent): y_values: str """Y parameter name containing data, indexed over the `breakpoints` dimension.""" - @property - def equations(self) -> _Equations: - """Dummy property to satisfy type hinting.""" - return _Equations() - - @property - def sub_expressions(self) -> _SubExpressions: - """Dummy property to satisfy type hinting.""" - return _SubExpressions() - - @property - def slices(self) -> _SubExpressions: - """Dummy property to satisfy type hinting.""" - return _SubExpressions() - _group: ClassVar[COMPONENTS_T] = "piecewise_constraints" @@ -448,21 +336,6 @@ class VariableDef(_MathIndexedComponent): Either real (a.k.a. continuous) or integer.""" bounds: _Bounds = _Bounds() - @property - def equations(self) -> _Equations: - """Dummy property to satisfy type hinting.""" - return _Equations() - - @property - def sub_expressions(self) -> _SubExpressions: - """Dummy property to satisfy type hinting.""" - return _SubExpressions() - - @property - def slices(self) -> _SubExpressions: - """Dummy property to satisfy type hinting.""" - return _SubExpressions() - _group: ClassVar[COMPONENTS_T] = "variables" @@ -478,16 +351,6 @@ class ObjectiveDef(_MathEquationComponent): """Whether the objective function should be minimised or maximised in the optimisation.""" - @property - def foreach(self) -> UniqueList[AttrStr]: - """Objectives are always adimensional.""" - return [] - - @property - def mask(self) -> str: - """Dummy property to satisfy type hinting.""" - return "True" - _group: ClassVar[COMPONENTS_T] = "objectives" @@ -579,11 +442,10 @@ class Checks(LinopyDictModel): class MathModel(LinopyBaseModel): """ - Mathematical definition of Calliope math. + Declarative definition of a linopy optimisation problem. - Contains mathematical programming components available for optimising with Calliope. - Can contain partial definitions if they are meant to be layered on top of another. - E.g.: layering 'base' and 'operate' math. + Contains all mathematical programming components from which a linopy model + can be built. """ model_config = {"title": "Model math schema"} @@ -619,8 +481,8 @@ def unique_component_names(self) -> Self: ), key=len, ) - seen = set() - duplicates = set() + seen: set[str] = set() + duplicates: set[str] = set() for field_names in groups: duplicates |= field_names & seen seen |= field_names @@ -634,11 +496,13 @@ def unique_component_names(self) -> Self: @cached_property def parsing_components(self) -> dict[str, dict[str, set[str]]]: """ - Return a set of valid component names in the model to use in `mask` string parsing. + Return the valid component names available to each parser. - Returns: - dict[Literal["dimension_names", "input_names", "result_names"], set[str]]: - Set of valid names grouped by location in the math in which they are defined. + 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"], @@ -646,7 +510,7 @@ def parsing_components(self) -> dict[str, dict[str, set[str]]]: "results": ["variables", "expressions"], } - def _names(): + 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() @@ -680,11 +544,7 @@ def find( MATH_DEFS_T = ( - ConstraintDef - | VariableDef - | ExpressionDef - | ObjectiveDef - | PiecewiseConstraintDef + ConstraintDef | VariableDef | ExpressionDef | ObjectiveDef | PiecewiseConstraintDef ) diff --git a/test/test_declarative_parsing.py b/test/test_declarative_parsing.py index 2e0335dde..8eeb0c937 100644 --- a/test/test_declarative_parsing.py +++ b/test/test_declarative_parsing.py @@ -1,25 +1,30 @@ -# Copyright (C) since 2013 Calliope contributors listed in AUTHORS. -# Licensed under the Apache 2.0 License (see LICENSE file). """ -Tests for the declarative math parser route separation (mask / expr / raw). +Tests for the declarative math interface. -These tests guard the contracts established by the parser-route refactor: - -- mask evaluation always returns a boolean ``xr.DataArray``; -- expression evaluation always returns a linopy expression; -- equation (comparison) evaluation returns a ``(lhs, sign, rhs)`` tuple; -- helper-function arguments are always evaluated in ``raw`` mode. +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 + import numpy as np import pytest import xarray as xr +import yaml -from linopy.declarative import helper_functions +from linopy.declarative import evaluate, grammar, parsing from linopy.declarative.build import DeclarativeModelBuilder, declarative_model -from linopy.declarative.parsing import ParsedBackendComponent +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 MathModel from linopy.expressions import LinearExpression from linopy.variables import Variable @@ -82,153 +87,304 @@ def _inputs() -> xr.Dataset: ) +def _ctx(builder: DeclarativeModelBuilder, **kwargs) -> evaluate.Context: + """Build a fresh evaluation context from a builder's validated components.""" + return evaluate.Context( + model=builder.model, + input_data=builder.input_data, + math=builder.math, + config=builder.config, + helpers=kwargs.pop("helpers", build_registry()), + **kwargs, + ) + + +def _first_equation( + builder: DeclarativeModelBuilder, group: str, name: str +) -> tuple[parsing.Equation, xr.DataArray, evaluate.Context]: + """Parse a component and return its first equation, sub-mask, and context.""" + ctx = _ctx(builder) + definition = getattr(builder.math, group)[name] + mask = parsing.component_mask(group, name, definition, ctx) + equation = parsing.parse_component(group, name, definition, builder.math)[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() -> DeclarativeModelBuilder: - """A builder with the ``flow`` variable already added to the model.""" + """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 -def _first_equation(builder: DeclarativeModelBuilder, group: str, name: str): - """Parse a component and return (component, first_equation, foreach_sub_mask).""" - definition = getattr(builder.math, group)[name] - component = ParsedBackendComponent( - group, name, definition, builder.math.parsing_components - ) - mask = component.generate_top_level_mask( - builder.input_data, - builder.model, - builder.math, - builder.config, - references=set(), - ) - equation = component.parse_equations()[0] - sub_mask = equation.evaluate_mask( - builder.input_data, - builder.model, - builder.math, - builder.config, - initial_mask=mask, - ) - sub_mask = component.drop_dims_not_in_foreach(sub_mask) - return component, equation, sub_mask +class TestGrammar: + """String -> AST parsing.""" + NAMES = frozenset({"flow", "cost", "cap_max", "node"}) -# --------------------------------------------------------------------------- # -# Mask route -# --------------------------------------------------------------------------- # + def test_equation_tree_shape(self): + 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): + import pyparsing as pp -def test_top_level_mask_returns_boolean_dataarray(builder_with_flow): - component, _, 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()) + with pytest.raises(pp.ParseException): + grammar.equation_grammar(self.NAMES).parse_string( + "flow < cap_max", parse_all=True + ) + def test_arithmetic_tree_shape(self): + 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): + 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): + import pyparsing as pp + + with pytest.raises(pp.ParseException): + grammar.sub_expression_grammar(self.NAMES).parse_string( + "$foo + 1", parse_all=True + ) -def test_mask_comparison_and_subset_and_helper_return_bool(): - math = _math() - 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"]) - component = ParsedBackendComponent( - "constraints", - "cap", - builder.math.constraints["cap"], - builder.math.parsing_components, - ) - mask = component.generate_top_level_mask( - builder.input_data, - builder.model, - builder.math, - builder.config, - references=set(), - ) - equation = component.parse_equations()[0] - result = equation.evaluate_mask( - builder.input_data, - builder.model, - builder.math, - builder.config, - initial_mask=mask, - ) - assert isinstance(result, xr.DataArray) - assert result.dtype == bool + def test_find_refs_in_call_kwargs(self): + 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_parse_error_carries_position_marker(self): + math = _math() + math["constraints"]["cap"]["equations"][0]["expression"] = "flow <= <=" + builder = DeclarativeModelBuilder(math, _inputs(), {}) + with pytest.raises(ValueError, match="equations\\[0\\].expression"): + parsing.parse_component( + "constraints", "cap", builder.math.constraints["cap"], builder.math + ) -# --------------------------------------------------------------------------- # -# Expression route -# --------------------------------------------------------------------------- # +class TestMaskRoute: + """Mask strings evaluate to boolean arrays.""" -def test_expression_with_variable_returns_linexpr(builder_with_flow): - _, equation, sub_mask = _first_equation( - builder_with_flow, "expressions", "total_cost" - ) - result = equation.evaluate_expression( - builder_with_flow.input_data, - builder_with_flow.model, - builder_with_flow.math, - mask=sub_mask, + def test_top_level_mask_returns_boolean_dataarray(self, builder_with_flow): + _, 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 = _math() + 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"], ) - assert isinstance(result, LinearExpression) + def test_mask_atoms_return_bool(self, builder_with_flow, mask_string): + node = parsing.parse_mask(mask_string, builder_with_flow.math) + result = evaluate.evaluate(node, _ctx(builder_with_flow, route="mask")) + assert isinstance(result, xr.DataArray) + assert result.dtype == bool + def test_existence_coercion(self, builder_with_flow): + """Bare input references coerce to existence booleans on the mask route.""" + node = parsing.parse_mask("cap_max", builder_with_flow.math) + result = evaluate.evaluate(node, _ctx(builder_with_flow, route="mask")) + assert result.values.tolist() == [True, True, True] -def test_pure_parameter_expression_coerced_to_linexpr(builder_with_flow): - _, equation, sub_mask = _first_equation( - builder_with_flow, "expressions", "cost_plus_one" - ) - result = equation.evaluate_expression( - builder_with_flow.input_data, - builder_with_flow.model, - builder_with_flow.math, - mask=sub_mask, - ) - # No decision variable is involved, but the contract is still LinearExpression. - assert isinstance(result, LinearExpression) +class TestExpressionRoute: + """Expression strings evaluate to linopy expressions.""" -def test_sub_expression_reference_returns_linexpr(builder_with_flow): - _, equation, sub_mask = _first_equation( - builder_with_flow, "expressions", "sub_expr_test" - ) - result = equation.evaluate_expression( - builder_with_flow.input_data, - builder_with_flow.model, - builder_with_flow.math, - mask=sub_mask, - ) - assert isinstance(result, LinearExpression) + def test_expression_with_variable_returns_linexpr(self, builder_with_flow): + 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): + 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) -# --------------------------------------------------------------------------- # -# Equation (constraint) route -# --------------------------------------------------------------------------- # + def test_sub_expression_reference_returns_linexpr(self, builder_with_flow): + 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): + """Two variants of one sub-expression yield two equations with merged masks.""" + math = _math() + 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 = _math() + math["expressions"]["sub_expr_test"]["sub_expressions"] = { + "bar": [{"expression": "flow"}] + } + builder = DeclarativeModelBuilder(math, _inputs(), {}) + with pytest.raises(KeyError, match="Undefined sub_expressions"): + parsing.parse_component( + "expressions", + "sub_expr_test", + builder.math.expressions["sub_expr_test"], + builder.math, + ) -def test_equation_returns_lhs_sign_rhs_tuple(builder_with_flow): - _, equation, sub_mask = _first_equation(builder_with_flow, "constraints", "cap") - lhs, sign, rhs = equation.evaluate_equation( - builder_with_flow.input_data, - builder_with_flow.model, - builder_with_flow.math, - 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_plain_and_list_slices(self, builder_with_flow): + ctx = _ctx(builder_with_flow, 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 = evaluate.evaluate(scalar_sliced, ctx, expr=True) + assert isinstance(result, LinearExpression) + + list_sliced = arith.parse_string("flow[node=[a, b]]", parse_all=True)[0] + result = evaluate.evaluate(list_sliced, ctx, expr=True) + assert result.data.sizes["node"] == 2 + + def test_slicer_reference(self, builder_with_flow): + """`$name` slicer references resolve like sub-expressions (feature parity).""" + math = _math() + 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 -# --------------------------------------------------------------------------- # -# Helper-function argument evaluation (raw mode) -# --------------------------------------------------------------------------- # +class TestConstraintRoute: + """Constraint equations evaluate to (lhs, sign, rhs) tuples.""" + + def test_equation_returns_lhs_sign_rhs_tuple(self, builder_with_flow): + 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 = _math() + # `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.""" -class _RecordArgs(helper_functions.ParsingHelperFunction): + def test_equation_latex(self, builder_with_flow): + equation, _, ctx = _first_equation(builder_with_flow, "constraints", "cap") + assert parsing.as_latex(equation, ctx) == r"flow \leq cap_max" + + def test_sum_latex(self, builder_with_flow): + equation, _, ctx = _first_equation(builder_with_flow, "objectives", "obj") + assert ( + parsing.as_latex(equation, ctx) + == r"\sum\limits_{\substack{\text{n} \in \text{node}}} (total_cost)" + ) + + def test_mask_latex(self): + math = _math() + 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(equation, ctx, what="mask") == ( + r"(\textit{cost}\mathord{>}\text{1} \land \text{n} \in \text{[a,b]})" + ) + + def test_sliced_component_latex(self, builder_with_flow): + ctx = _ctx(builder_with_flow) + arith = grammar.arithmetic_grammar(frozenset({"flow", "node"})) + tree = arith.parse_string("flow[node=a]", parse_all=True)[0] + assert evaluate.to_math_string(tree, ctx) == r"flow_\text{n=a}" + + def test_identity_operands_are_skipped(self, builder_with_flow): + ctx = _ctx(builder_with_flow) + arith = grammar.arithmetic_grammar(frozenset({"flow"})) + tree = arith.parse_string("0 + flow", parse_all=True)[0] + assert evaluate.to_math_string(tree, ctx) == "flow" + + +class _RecordArgs(HelperFunction): """Test-only helper that records the types of the arguments it receives.""" NAME = "record_args" @@ -244,90 +400,304 @@ def as_raw(self, *args, **kwargs): # noqa: D102 return args[0] -def test_helper_arguments_are_evaluated_raw(builder_with_flow): - math = _math() - # flow is a variable, cost is a parameter -> raw mode must preserve both types. - math["expressions"]["total_cost"]["equations"][0]["expression"] = ( - "record_args(flow, cost)" - ) - builder = DeclarativeModelBuilder(math, _inputs(), {}) - builder.add_variable("flow", builder.math.variables["flow"]) - _, equation, sub_mask = _first_equation(builder, "expressions", "total_cost") +class _Double(HelperFunction): + """Test-only helper doubling its argument.""" - _RecordArgs.received = [] - equation.evaluate_expression( - builder.input_data, builder.model, builder.math, mask=sub_mask - ) - assert _RecordArgs.received, "helper was not called" - # The variable arrives un-normalised (raw Variable, not LinearExpression); - # the parameter arrives as a raw DataArray (not coerced/masked to booleans). - assert Variable in _RecordArgs.received - assert xr.DataArray in _RecordArgs.received - assert LinearExpression not in _RecordArgs.received - - -def test_invalid_helper_function_rejected(builder_with_flow): - """A registry entry that is not a ParsingHelperFunction subclass is rejected.""" - registry = helper_functions._registry["expression"] - registry["not_a_helper"] = str # type: ignore[assignment] - try: + NAME = "double" + ALLOWED_IN = ["expression"] + + def as_math_string(self, array): # noqa: D102 + return rf"2 \times {array}" + + def as_raw(self, array): # noqa: D102 + return 2 * array + + +class TestHelpers: + """Helper-function registration and argument evaluation.""" + + def test_helper_arguments_are_evaluated_raw(self): + """Helper args arrive un-normalised: raw Variable/DataArray, not LinearExpression.""" + math = _math() + 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): + with pytest.raises(ValueError, match="must be subclassed"): + build_registry([str]) # type: ignore[list-item] + + def test_non_subclass_rejected_at_evaluation(self): + """A hand-built registry with an invalid entry is rejected at call time.""" math = _math() math["expressions"]["total_cost"]["equations"][0]["expression"] = ( "not_a_helper(flow)" ) builder = DeclarativeModelBuilder(math, _inputs(), {}) builder.add_variable("flow", builder.math.variables["flow"]) - _, equation, sub_mask = _first_equation(builder, "expressions", "total_cost") + registry = build_registry() + registry["expression"]["not_a_helper"] = str # type: ignore[assignment] + ctx = _ctx(builder, helpers=registry) + definition = builder.math.expressions["total_cost"] + equation = parsing.parse_component( + "expressions", "total_cost", definition, builder.math + )[0] with pytest.raises(ValueError, match="must be subclassed"): - equation.evaluate_expression( - builder.input_data, builder.model, builder.math, mask=sub_mask - ) - finally: - registry.pop("not_a_helper", None) + parsing.as_expression(equation, ctx) + def test_unknown_helper_rejected(self): + math = _math() + 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)) -# --------------------------------------------------------------------------- # -# Input-data checks -# --------------------------------------------------------------------------- # + def test_duplicate_name_rejected(self): + class _ClashingSum(HelperFunction): + NAME = "sum" + ALLOWED_IN = ["expression"] + def as_math_string(self, *args, **kwargs): # noqa: D102 + return "" -def test_checks_run_without_active_variable(): - """`_check_inputs` must not require an `active` variable in the input data.""" - math = _math() - 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 as_raw(self, *args, **kwargs): # noqa: D102 + return xr.DataArray() + with pytest.raises(ValueError, match="already exists"): + build_registry([_ClashingSum]) -def test_check_raises_when_triggered_without_active(): - math = _math() - math["checks"] = { - "too_expensive": { - "mask": "cost > 0", - "message": "cost too high", - "errors": "raise", + def test_custom_helper_end_to_end(self): + math = _math() + 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): + 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 evaluate.evaluate(tree, ctx).item() == "a" + + +class TestBuilder: + """Model assembly from parsed math.""" + + def test_overlapping_equation_masks_rejected(self): + math = _math() + 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 = _math() + 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): + 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 = _math() + math["lookups"] = { + "flag": {"dtype": "bool", "default": False}, + "label": {"dtype": "string"}, } - } - with pytest.raises(ValueError, match="cost too high"): + inputs = _inputs() + 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): + """Input checks must not require an `active` variable in the input data.""" + math = _math() + 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 = _math() + 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(), {}) -# --------------------------------------------------------------------------- # -# End-to-end build -# --------------------------------------------------------------------------- # + def test_check_warns(self, caplog): + math = _math() + 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 + + +class TestLatexDoc: + """LaTeX math documentation building.""" + + def test_components_render_with_decorated_reprs(self): + 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 = _math() + 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): + # `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): + # 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): + # `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): + 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 = _math() + 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_sub_expression_variants_render_as_separate_blocks(self): + math = _math() + 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_declarative_model_end_to_end(): - 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"} + def test_markdown_document_structure(self): + 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): + doc = latex_math_doc(_math(), _inputs(), format="rst") + assert ".. math::" in doc + assert "Math formulation\n================" in doc + + def test_tex_document_structure(self): + 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_declarative_model_end_to_end(self): + 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"} + + def test_repo_math_yaml_validates_and_parses(self): + """The demo math.yaml at the repo root validates and every component parses.""" + math_path = Path(__file__).parent.parent / "math.yaml" + if not math_path.exists(): + pytest.skip("repo-root math.yaml not present") + math = MathModel.model_validate(yaml.safe_load(math_path.read_text())) + for group in ("expressions", "constraints", "objectives"): + for name, definition in getattr(math, group).root.items(): + equations = parsing.parse_component(group, name, definition, math) + assert equations, f"{group}:{name} produced no equations" + # And the full LaTeX math documentation generates in every format. + for fmt in ("md", "rst", "tex"): + doc = latex_math_doc(yaml.safe_load(math_path.read_text()), format=fmt) + name = r"storage\_balance" if fmt == "tex" else "storage_balance" + assert name in doc From 822cd81bbcf8eff4086207d3121a9dce0c2da1b1 Mon Sep 17 00:00:00 2001 From: Bryn Pickering <17178478+brynpickering@users.noreply.github.com> Date: Thu, 23 Jul 2026 17:22:37 +0100 Subject: [PATCH 06/12] Refactor for maintainability. Co-Authored-By: Claude --- linopy/declarative/build.py | 94 ++-- linopy/declarative/evaluate.py | 571 --------------------- linopy/declarative/grammar.py | 448 +++++------------ linopy/declarative/helpers.py | 25 +- linopy/declarative/latex.py | 78 +-- linopy/declarative/nodes.py | 828 +++++++++++++++++++++++++++++++ linopy/declarative/parsing.py | 133 ++--- linopy/declarative/schema.py | 23 +- test/test_declarative_parsing.py | 65 ++- 9 files changed, 1142 insertions(+), 1123 deletions(-) delete mode 100644 linopy/declarative/evaluate.py create mode 100644 linopy/declarative/nodes.py diff --git a/linopy/declarative/build.py b/linopy/declarative/build.py index 15c649191..53ac77be5 100644 --- a/linopy/declarative/build.py +++ b/linopy/declarative/build.py @@ -12,17 +12,18 @@ import time from collections.abc import Iterable, Iterator, Mapping from dataclasses import replace -from typing import Any, Literal, get_args +from typing import Any import xarray as xr from tqdm.auto import tqdm from linopy.declarative import parsing -from linopy.declarative.evaluate import Context, evaluate -from linopy.declarative.grammar import Component, find_refs 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, @@ -36,13 +37,6 @@ LOGGER = logging.getLogger(__name__) -ORDERED_COMPONENTS_T = Literal[ - "variables", - "expressions", - "constraints", - "objectives", -] - _SKIP_MESSAGE = "No valid data points after applying mask. Not added to model." @@ -76,41 +70,85 @@ def declarative_model( return DeclarativeModelBuilder(math_def, input_data, config, helpers).build() -class DeclarativeModelBuilder: - """Builder turning a declarative math definition into a linopy Model.""" +class _DeclarativeBase: + """Shared validation and context setup of the model and LaTeX builders.""" def __init__( self, math_def: dict, - input_data: xr.Dataset, - config: 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, ready to `build()`. + Validate the math definition, input data, and config. Parameters ---------- math_def : dict Declarative math definition. - input_data : xr.Dataset + input_data : xr.Dataset, optional Model input data. - config : dict + 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.input_data = self._update_dtypes(input_data) - self.config = ConfigModel.model_validate(config) + 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, definition: Any, equations: Iterable[parsing.Equation] = () + ) -> list[str]: + """Return the sorted names of all math components a component references.""" + mask_node = parsing.parse_mask(getattr(definition, "mask", "True"), self.math) + refs = find_refs(mask_node, Component) + for equation in 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: @@ -164,8 +202,8 @@ def _check_inputs(self) -> None: if not check.active: continue mask_node = parsing.parse_mask(check.mask, self.math, name) - check_ctx = replace(self._ctx, route="mask", equation_name=name) - evaluated = evaluate(mask_node, check_ctx) + 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) @@ -184,20 +222,10 @@ 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 _references( - self, definition: Any, equations: Iterable[parsing.Equation] = () - ) -> list[str]: - """Return the sorted names of all math components a component references.""" - mask_node = parsing.parse_mask(getattr(definition, "mask", "True"), self.math) - refs = find_refs(mask_node, Component) - for equation in equations: - refs |= equation.references() - return sorted(refs) - def _iter_equations( self, equations: list[parsing.Equation], - group: parsing.GROUP_T, + group: EQUATION_GROUP_T, mask: xr.DataArray, ) -> Iterator[tuple[parsing.Equation, xr.DataArray]]: """ @@ -344,7 +372,7 @@ def build(self) -> Model: raise ValueError( f"Only one active objective is supported, found: {active_objectives}" ) - for group in get_args(ORDERED_COMPONENTS_T): + for group in BUILD_ORDER: component = group.removesuffix("s") ordered_items = self._sorted_by_order(self.math[group].root) for name, definition in tqdm( diff --git a/linopy/declarative/evaluate.py b/linopy/declarative/evaluate.py deleted file mode 100644 index 6ad7bbc83..000000000 --- a/linopy/declarative/evaluate.py +++ /dev/null @@ -1,571 +0,0 @@ -""" -Linopy declarative math evaluation module. - -This module contains the evaluation context and the two tree walkers that turn a -parsed math AST (see :mod:`linopy.declarative.grammar`) into either a LaTeX math -string (:func:`to_math_string`) or data — an `xr.DataArray` or a linopy -expression (:func:`evaluate`). -""" - -from __future__ import annotations - -import operator -import re -from dataclasses import dataclass, field, 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.grammar import ( - Arith, - Call, - Compare, - Component, - ConfigRef, - Constant, - ListNode, - Node, - Sliced, - SliceRef, - SubExprRef, - Subset, - Unary, -) -from linopy.declarative.helpers import KIND_T, HelperFunction, 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) - -ROUTE_T = Literal["expression", "mask"] - -_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).""" - - route: ROUTE_T = "expression" - """Whether the AST being evaluated came from an expression or a mask string.""" - - apply_mask: bool = True - """On the mask route, 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 error(ctx: Context, node: Node, message: str) -> ValueError: - """Return a ValueError contextualised with the equation name and source string.""" - return ValueError(f"({ctx.equation_name}, {node.instring}) | {message}") - - -def get_dot_attr(var: Any, attr: str) -> Any: - """ - Get a nested attribute in dot notation (e.g. "foo.bar"). - - Works for nested objects: dictionaries, pydantic models, etc. - """ - levels = attr.split(".", 1) - value = var[levels[0]] if isinstance(var, dict) else getattr(var, levels[0]) - if len(levels) > 1: - value = get_dot_attr(value, levels[1]) - return value - - -def _to_linexpr(obj: Any) -> Any: - """ - Normalise a model object to a linopy expression. - - `Variable` objects are converted to `LinearExpression`; everything else is - returned unchanged. This is the single place where the `Variable` -> - `LinearExpression` coercion is performed on the expression route. - """ - if isinstance(obj, Variable): - return obj.to_linexpr() - return obj - - -def _apply_mask(evaluated: Any, mask: xr.DataArray) -> Any: - """Mask an evaluated operand, broadcasting first if it cannot be masked directly.""" - try: - return evaluated.where(mask) - except AttributeError: - return evaluated.broadcast_like(mask).where(mask) - - -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 - ) - - -# --------------------------------------------------------------------------- -# Data evaluation -# --------------------------------------------------------------------------- - - -def evaluate(node: Node, ctx: Context, expr: bool = False) -> Any: - """ - Evaluate a math AST node to data. - - Parameters - ---------- - node : Node - AST node to evaluate. - ctx : Context - Evaluation context. - expr : bool, default: False - If False ("raw" mode), return the underlying data without route-specific - transformation: an `xr.DataArray` for parameters/lookups/dimensions, the - raw model object (`Variable`/`LinearExpression`) for model entries, and a - boolean array on the mask route. If True ("expr" mode), guarantee a - linopy-expression-compatible result: `Variable` objects are coerced to - `LinearExpression` and a top-level :class:`Compare` returns a masked - `(lhs, sign, rhs)` tuple for constraint assembly. - """ - match node: - case Constant(value=bool() as val): - return xr.DataArray(np.bool_(val)) - case Constant(value=str() as val): - return val - case Constant(value=val): - return xr.DataArray(float(val), name=float(val)) - case ListNode(items=items): - return [evaluate(item, ctx) for item in items] - case Component(): - return _evaluate_component(node, ctx, expr) - case ConfigRef(): - return _evaluate_config(node, ctx) - case SubExprRef(name=name): - return evaluate(ctx.sub_expressions[name], ctx, expr) - case SliceRef(name=name): - return evaluate(ctx.slices[name], ctx) - case Sliced(obj=obj, slices=slices): - evaluated_slices = { - dim: [_unwrap(i) for i in vals] - if isinstance(vals := evaluate(slicer, ctx), list) - else vals - for dim, slicer in slices.items() - } - return evaluate(obj, ctx, expr).sel(**evaluated_slices) - case Call(): - return _evaluate_call(node, ctx, expr) - case Unary(op=op, operand=operand): - if op == "not": - return ~evaluate(operand, ctx) - evaluated = evaluate(operand, ctx, expr) - return -1 * evaluated if op == "-" else evaluated - case Arith(first=first, rest=rest): - boolean = rest[0][0] in ("and", "or") - val = evaluate(first, ctx, expr) - if not boolean: - val = _apply_mask(val, ctx.mask) - for op, operand in rest: - evaluated = evaluate(operand, ctx, expr) - if not boolean: - evaluated = _apply_mask(evaluated, ctx.mask) - val = _OPERATIONS[op](val, evaluated) - return val - case Compare() if expr: - return _evaluate_equation(node, ctx) - case Compare(lhs=lhs, op=op, rhs=rhs): - unmasked_ctx = replace(ctx, apply_mask=False) - comparison = _OPERATIONS[op]( - evaluate(lhs, unmasked_ctx), evaluate(rhs, unmasked_ctx) - ) - return xr.DataArray(comparison) - case Subset(items=items, dim=dim): - subset = [_unwrap(evaluate(item, ctx)) for item in items] - dim_array = evaluate(dim, replace(ctx, apply_mask=False)) - return dim_array.isin(subset) - case _: - raise error( - ctx, node, f"Cannot evaluate node of type {type(node).__name__}" - ) - - -def _evaluate_component(node: Component, ctx: Context, expr: bool) -> Any: - """Evaluate a component reference according to its category and the evaluation mode.""" - name = node.name - if node.category == "dimension": - # The mask string should evaluate successfully even if a dimension isn't defined. - return ctx.input_data.get(name, xr.DataArray()) - if node.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 node.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)) - 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. On the expression route 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 not expr: - return evaluated - evaluated = _to_linexpr(evaluated) - if evaluated.isnull().any() and pd.notna(default := math_def["default"]): - evaluated = evaluated.fillna(default) - return evaluated - - -def _evaluate_config(node: ConfigRef, ctx: Context) -> xr.DataArray: - """Evaluate a config option reference to a dimensionless array.""" - config_val = get_dot_attr(ctx.config, node.option) - if not isinstance(config_val, int | float | str | bool | np.bool_): - raise error( - ctx, - node, - f"Configuration option resolves to invalid type " - f"`{type(config_val).__name__}`, expected a number, string, or boolean.", - ) - return xr.DataArray(config_val) - - -def _lookup_helper(node: Call, ctx: Context) -> type[HelperFunction]: - """Return the helper class for a function call, validating it exists in the registry.""" - kind: KIND_T = "mask" if ctx.route == "mask" else "expression" - helpers = ctx.helpers.get(kind, {}) - if node.func not in helpers: - raise error(ctx, node, f"Invalid helper function defined: {node.func}") - helper_cls = helpers[node.func] - if not (isinstance(helper_cls, type) and issubclass(helper_cls, HelperFunction)): - raise error( - ctx, - node, - "Helper function must be subclassed from " - f"linopy.declarative.helpers.HelperFunction: {node.func}", - ) - return helper_cls - - -def _evaluate_call(node: Call, ctx: Context, expr: bool) -> Any: - """ - Evaluate a helper-function call. - - The helper itself is instantiated with the enclosing mode so that - expression-route helpers can dispatch to their `as_expr` implementation. - Its arguments, however, are always evaluated in 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 that have been coerced to boolean masks or `LinearExpression`. - """ - helper_cls = _lookup_helper(node, ctx) - helper = helper_cls("expr" if expr else "raw", ctx) - if helper_cls.ignore_mask: - ctx = replace(ctx, mask=TRUE_ARRAY) - args = [evaluate(arg, ctx) for arg in node.args] - kwargs = {name: evaluate(val, ctx) for name, val in node.kwargs.items()} - return helper(*args, **kwargs) - - -def _evaluate_equation(node: Compare, ctx: Context) -> tuple[Any, xr.DataArray, Any]: - """Evaluate an equation to a masked `(lhs, sign, rhs)` tuple for constraint assembly.""" - lhs = evaluate(node.lhs, ctx, expr=True) - rhs = evaluate(node.rhs, ctx, expr=True) - 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, - node, - 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)) - rhs_masked = _to_linexpr(rhs.where(ctx.mask)) - sign_masked = xr.DataArray(node.op).where(ctx.mask) - return lhs_masked, sign_masked, rhs_masked - - -# --------------------------------------------------------------------------- -# LaTeX math-string evaluation -# --------------------------------------------------------------------------- - - -def to_math_string(node: Node, ctx: Context) -> str: - """ - Evaluate a math AST node to a LaTeX math string. - - Parameters - ---------- - node : Node - AST node to evaluate. - ctx : Context - Evaluation context. - """ - match node: - case Constant(value=bool() as val): - return str(val).lower() - case Constant(value=str() as val): - return val - case Constant(value=val): - return re.sub( - r"([\d]+?)e([+-])([\d]+)", - r"\1\\mathord{\\times}10^{\2\3}", - f"{float(val):.6g}", - ) - case ListNode(items=items): - return "[" + ",".join(_plain_string(item, ctx) for item in items) + "]" - case Component(): - return _component_math_string(node, ctx) - case ConfigRef(option=option): - return rf"\text{{config.{option}}}" - case SubExprRef(name=name): - return to_math_string(ctx.sub_expressions[name], ctx) - case SliceRef(name=name): - return to_math_string(ctx.slices[name], ctx) - case Sliced(): - return _sliced_math_string(node, ctx) - case Call(): - helper = _lookup_helper(node, ctx)("math_string", ctx) - args = [_call_arg_math_string(arg, ctx) for arg in node.args] - kwargs = { - name: _call_arg_math_string(val, ctx) - for name, val in node.kwargs.items() - } - return helper(*args, **kwargs) - case Unary(op="not", operand=operand): - return rf"\neg ({to_math_string(operand, ctx)})" - case Unary(op=op, operand=operand): - return op + to_math_string(operand, ctx) - case Arith(first=first, rest=rest): - val = to_math_string(first, ctx) - for op, operand in rest: - evaluated = to_math_string(operand, ctx) - # We ignore identity elements that do nothing (e.g. `0 + flow` is `flow`) - if evaluated == _LATEX_IDENTITIES.get(op): - continue - if isinstance(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 - case Compare(lhs=lhs, op=op, rhs=rhs) if ctx.route == "expression": - lhs_str = to_math_string(lhs, ctx) - rhs_str = to_math_string(rhs, ctx) - return lhs_str + _LATEX_EQUATION_OPERATORS[op] + rhs_str - case Compare(lhs=lhs, op=op, rhs=rhs): - unmasked_ctx = replace(ctx, apply_mask=False) - lhs_str = to_math_string(lhs, unmasked_ctx) - rhs_str = to_math_string(rhs, unmasked_ctx) - if r"\text" not in rhs_str: - rhs_str = rf"\text{{{rhs_str}}}" - return lhs_str + _LATEX_MASK_OPERATORS[op] + rhs_str - case Subset(items=items, dim=dim): - subset = [_unwrap(evaluate(item, ctx)) for item in 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 = ( - dim.name if isinstance(dim, Component) else to_math_string(dim, 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}}}" - case _: - raise error( - ctx, node, f"Cannot render node of type {type(node).__name__} as LaTeX" - ) - - -def _plain_string(item: Node, ctx: Context) -> str: - """Return a plain-text representation of a list item for LaTeX rendering.""" - evaluated = evaluate(item, ctx) - 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 evaluate(arg, ctx) - return to_math_string(arg, ctx) - - -def _component_math_string(node: Component, ctx: Context) -> str: - """ - Render a 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 = node.name - custom = ctx.math_reprs.get(name) - if node.category == "dimension": - return name - if node.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 node.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 = evaluate(node, ctx) - attrs = getattr(evaluated, "attrs", {}) - return str(attrs.get("math_repr", name)) - - -def _sliced_math_string(node: Sliced, ctx: Context) -> str: - r""" - Render a sliced component as LaTeX. - - If the component's LaTeX representation carries an iterator substring (from a - `math_repr` data attribute, e.g. `\textbf{flow}_\text{n}`), the slices are - injected into it by re-parsing (e.g. `\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): to_math_string(slicer, ctx) - for dim, slicer in node.slices.items() - } - obj_string = to_math_string(node.obj, ctx) - - def _replace(term: pp.ParseResults) -> Any: - if len(term) == 1: - return term - replacers = {k: f"{k}={v}" for k, v in slice_strings.items()} - return ( - term[0] + term[1] + ",".join(replacers.get(k, k) for k in term[2]) + term[3] - ) - - id_ = pp.Combine( - pp.Word(pp.alphas, pp.alphanums) - + pp.ZeroOrMore("_" + pp.Word(pp.alphanums)) - + pp.Opt("_") - ) - id_formatted = pp.Combine("\\" + pp.Word(pp.alphas) + "{" + id_ + "}") - obj_parser = id_formatted + pp.Opt( - r"_\text{" + pp.Group(pp.DelimitedList(id_)) + "}" - ) - obj_parser.set_parse_action(_replace) - try: - return obj_parser.parse_string(obj_string, parse_all=True)[0] - except pp.ParseException: - subscript = ",".join(f"{k}={v}" for k, v in slice_strings.items()) - return rf"{obj_string}_\text{{{subscript}}}" diff --git a/linopy/declarative/grammar.py b/linopy/declarative/grammar.py index f936ea42e..216aa9a25 100644 --- a/linopy/declarative/grammar.py +++ b/linopy/declarative/grammar.py @@ -1,9 +1,10 @@ """ Linopy declarative math grammar module. -This module contains the AST node types produced when parsing declarative math -strings, and the pyparsing grammars that produce them. Nodes are pure data; -all evaluation logic lives in :mod:`linopy.declarative.evaluate`. +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 `eval_arith.py` example (https://github.com/pyparsing/pyparsing/blob/master/examples/eval_arith.py, @@ -12,13 +13,56 @@ from __future__ import annotations -from collections.abc import Callable -from dataclasses import dataclass, field, fields from functools import cache -from typing import Literal 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 = "$" @@ -30,259 +74,6 @@ MASK_OPERATORS = ("<", ">", "==", ">=", "<=") """Comparison operators allowed in mask strings.""" -COMPONENT_CATEGORY_T = Literal["any", "dimension", "input", "result"] - - -# --------------------------------------------------------------------------- -# AST nodes -# --------------------------------------------------------------------------- - - -@dataclass(frozen=True, kw_only=True) -class Node: - """ - Base class of all declarative math AST nodes. - - Nodes are immutable data produced by the grammars in this module and consumed - by the walkers in :mod:`linopy.declarative.evaluate`. - """ - - instring: str = field(repr=False, compare=False) - """The full source string this node was parsed from (used in error messages).""" - - -@dataclass(frozen=True, kw_only=True) -class Constant(Node): - """A literal number (including `inf`), boolean, or generic string.""" - - value: float | bool | str - - -@dataclass(frozen=True, kw_only=True) -class ListNode(Node): - """A literal list of items, e.g. `[a, b, 1]`.""" - - items: tuple[Node, ...] - - -@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" - - -@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] - - -@dataclass(frozen=True, kw_only=True) -class SliceRef(Node): - """A `$name` reference to a named slicer, valid only inside slice brackets.""" - - name: str - - -@dataclass(frozen=True, kw_only=True) -class SubExprRef(Node): - """A `$name` reference to a named sub-expression.""" - - name: str - - -@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] - - -@dataclass(frozen=True, kw_only=True) -class Unary(Node): - """A unary operation: leading `+`/`-` sign or boolean `not`.""" - - op: str - operand: Node - - -@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], ...] - - -@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 - - -@dataclass(frozen=True, kw_only=True) -class Subset(Node): - """A dimension subset condition, e.g. `[a, b] in node`.""" - - items: tuple[Node, ...] - dim: Node - - -@dataclass(frozen=True, kw_only=True) -class ConfigRef(Node): - """A reference to a build-configuration option, e.g. `config.foo`.""" - - option: str - - -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] - - -# --------------------------------------------------------------------------- -# Parse actions -# --------------------------------------------------------------------------- - - -def _number_action(instring: str, loc: int, tokens: pp.ParseResults) -> Constant: - return Constant(value=float(tokens[0]), instring=instring) - - -def _string_action(instring: str, loc: int, tokens: pp.ParseResults) -> Constant: - return Constant(value=str(tokens[0]), instring=instring) - - -def _bool_action(instring: str, loc: int, tokens: pp.ParseResults) -> Constant: - return Constant(value=str(tokens[0]).lower() == "true", instring=instring) - - -def _list_action(instring: str, loc: int, tokens: pp.ParseResults) -> ListNode: - return ListNode(items=tuple(tokens), instring=instring) - - -def _component_action( - category: COMPONENT_CATEGORY_T, -) -> Callable[[str, int, pp.ParseResults], Component]: - def _action(instring: str, loc: int, tokens: pp.ParseResults) -> Component: - return Component(name=str(tokens[0]), category=category, instring=instring) - - return _action - - -def _sliced_action(instring: str, loc: int, tokens: pp.ParseResults) -> Sliced: - slices = {str(grp["set_name"][0]): grp["slicer"][0] for grp in tokens["slices"]} - return Sliced(obj=tokens["obj"], slices=slices, instring=instring) - - -def _slice_ref_action(instring: str, loc: int, tokens: pp.ParseResults) -> SliceRef: - return SliceRef(name=str(tokens[0]), instring=instring) - - -def _sub_expr_ref_action( - instring: str, loc: int, tokens: pp.ParseResults -) -> SubExprRef: - return SubExprRef(name=str(tokens[0]), instring=instring) - - -def _call_action(instring: str, loc: int, tokens: pp.ParseResults) -> Call: - 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 Call(func=token_dict["func"], args=args, kwargs=kwargs, instring=instring) - - -def _unary_action(instring: str, loc: int, tokens: pp.ParseResults) -> Unary: - op, operand = tokens[0] - return Unary(op=str(op), operand=operand, instring=instring) - - -def _arith_action(instring: str, loc: int, tokens: pp.ParseResults) -> Arith: - items = tokens[0] - rest = tuple( - (str(op), operand) for op, operand in zip(items[1::2], items[2::2], strict=True) - ) - return Arith(first=items[0], rest=rest, instring=instring) - - -def _compare_action(instring: str, loc: int, tokens: pp.ParseResults) -> Compare: - lhs, op, rhs = tokens - return Compare(lhs=lhs, op=str(op), rhs=rhs, instring=instring) - - -def _subset_action(instring: str, loc: int, tokens: pp.ParseResults) -> Subset: - items, dim = tokens - return Subset(items=tuple(items), dim=dim, instring=instring) - - -def _config_action(instring: str, loc: int, tokens: pp.ParseResults) -> ConfigRef: - return ConfigRef(option=str(tokens[0]), instring=instring) - # --------------------------------------------------------------------------- # Grammar primitives @@ -292,7 +83,7 @@ def _config_action(instring: str, loc: int, tokens: pp.ParseResults) -> ConfigRe 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(_number_action) + number = (pp.pyparsing_common.number | inf_kw).set_parse_action(Constant.number) identifier = ~inf_kw + pp.Word(pp.alphas, pp.alphanums + "_") return number, identifier @@ -308,7 +99,7 @@ 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_action(category)) + return _names_parser(names).set_parse_action(Component.from_tokens_as(category)) def _string_parser( @@ -316,7 +107,7 @@ def _string_parser( ) -> pp.ParserElement: """Return a parser for generic strings that are not in `excluded_names`.""" return (~_names_parser(excluded_names) + identifier).set_parse_action( - _string_action + Constant.string ) @@ -324,7 +115,7 @@ 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(_list_action) + return id_list.set_parse_action(ListNode.from_tokens) def _call_parser( @@ -361,7 +152,7 @@ def _call_parser( ) + pp.Opt(kwarg_list, default={}) call <<= func_name + call_args + pp.Suppress(")") - return call.set_parse_action(_call_action) + return call.set_parse_action(Call.from_tokens) def _sliced_component_parser( @@ -387,7 +178,7 @@ def _sliced_component_parser( slicer: pp.ParserElement = pp.MatchFirst(slicers) if allow_slice_references: slice_ref = pp.Suppress(REFERENCE_CLASSIFIER) + identifier - slice_ref.set_parse_action(_slice_ref_action) + slice_ref.set_parse_action(SliceRef.from_tokens) slicer = slice_ref | slicer one_slice = pp.Group( @@ -395,13 +186,13 @@ def _sliced_component_parser( ) slices = pp.Group(pp.DelimitedList(one_slice))("slices") sliced = pp.Combine(component("obj") + pp.Suppress("[")) + slices + pp.Suppress("]") - return sliced.set_parse_action(_sliced_action) + 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(_sub_expr_ref_action) + return ref.set_parse_action(SubExprRef.from_tokens) def _arithmetic_rules( @@ -427,10 +218,10 @@ def _arithmetic_rules( # the order matters if two could capture the same string, e.g. "inf". pp.MatchFirst(operands), [ - (signop, 1, pp.opAssoc.RIGHT, _unary_action), - (expop, 2, pp.opAssoc.LEFT, _arith_action), - (multop, 2, pp.opAssoc.LEFT, _arith_action), - (signop, 2, pp.opAssoc.LEFT, _arith_action), + (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 @@ -442,18 +233,28 @@ def _arithmetic_rules( @cache -def slice_grammar(component_names: frozenset[str]) -> pp.ParserElement: +def _expression_grammar( + component_names: frozenset[str], + *, + arithmetic: bool = True, + slice_refs: bool = True, + sub_expr_refs: bool = False, +) -> 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. + 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) @@ -464,22 +265,45 @@ def slice_grammar(component_names: frozenset[str]) -> pp.ParserElement: [number, string, slicer_list], identifier, component, - allow_slice_references=False, + allow_slice_references=slice_refs, ) - call = _call_parser( - sliced, - component, - number, - call_list, - string, - identifier=identifier, - allow_nested_calls=True, - ) - return call | sliced | component | number | slicer_list | string + 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) -@cache -def sub_expression_grammar(component_names: frozenset[str]) -> pp.Forward: +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. @@ -492,21 +316,10 @@ def sub_expression_grammar(component_names: frozenset[str]) -> pp.Forward: component_names : frozenset[str] Valid math component names, to separate them from generic strings. """ - 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 - ) - arithmetic = pp.Forward() - call = _call_parser(arithmetic, call_list, string, identifier=identifier) - return _arithmetic_rules(call, sliced, number, component, arithmetic=arithmetic) + return _expression_grammar(component_names) -@cache -def arithmetic_grammar(component_names: frozenset[str]) -> pp.Forward: +def arithmetic_grammar(component_names: frozenset[str]) -> pp.ParserElement: """ Return the grammar for arithmetic expressions (`+ - * / **`). @@ -518,20 +331,7 @@ def arithmetic_grammar(component_names: frozenset[str]) -> pp.Forward: component_names : frozenset[str] Valid math component names, to separate them from generic strings. """ - 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 - ) - sub_expression = _sub_expression_ref_parser(identifier) - arithmetic = pp.Forward() - call = _call_parser(arithmetic, call_list, string, identifier=identifier) - return _arithmetic_rules( - call, sub_expression, sliced, number, component, arithmetic=arithmetic - ) + return _expression_grammar(component_names, sub_expr_refs=True) @cache @@ -549,7 +349,7 @@ def equation_grammar(component_names: frozenset[str]) -> pp.ParserElement: """ arithmetic = arithmetic_grammar(component_names) equation = arithmetic + pp.one_of(list(EQUATION_OPERATORS)) + arithmetic - return equation.set_parse_action(_compare_action) + return equation.set_parse_action(Compare.from_tokens) @cache @@ -578,11 +378,11 @@ def mask_grammar( input_ = _component_parser(inputs, "input") result = _component_parser(results, "result") config_option = (pp.Suppress("config.") + identifier).set_parse_action( - _config_action + ConfigRef.from_tokens ) bool_operand = ( pp.Keyword("True", caseless=True) | pp.Keyword("False", caseless=True) - ).set_parse_action(_bool_action) + ).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) @@ -598,7 +398,7 @@ def mask_grammar( + pp.Suppress("in") + pp.Suppress(pp.White(" ", min=1)) + pp.MatchFirst([dimension, input_]) - ).set_parse_action(_subset_action) + ).set_parse_action(Subset.from_tokens) arithmetic = pp.Forward() comparison_call = _call_parser( @@ -611,7 +411,7 @@ def mask_grammar( arithmetic + pp.one_of(list(MASK_OPERATORS)) + pp.MatchFirst([comparison_call, bool_operand, number, general_string]) - ).set_parse_action(_compare_action) + ).set_parse_action(Compare.from_tokens) call = _call_parser( unique_string, @@ -629,7 +429,7 @@ def mask_grammar( return pp.infix_notation( pp.MatchFirst([bool_operand, comparison, call, subset, input_, result]), [ - (notop, 1, pp.opAssoc.RIGHT, _unary_action), - (andorop, 2, pp.opAssoc.LEFT, _arith_action), + (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 index 7695ccee4..a6f9244ab 100644 --- a/linopy/declarative/helpers.py +++ b/linopy/declarative/helpers.py @@ -22,9 +22,8 @@ from linopy.expressions import LinearExpression if TYPE_CHECKING: - from linopy.declarative.evaluate import Context + from linopy.declarative.nodes import Context -MODE_T = Literal["raw", "expr", "math_string"] KIND_T = Literal["mask", "expression"] @@ -107,18 +106,15 @@ class HelperFunction(ABC): ignore_mask: ClassVar[bool] = False """If True, `mask` arrays are not applied to the function's incoming arguments.""" - def __init__(self, mode: MODE_T, context: Context) -> None: + def __init__(self, context: Context) -> None: """ - Initialise the helper for one evaluation mode. + Initialise the helper. Parameters ---------- - mode : Literal["raw", "expr", "math_string"] - The evaluation mode this instance will dispatch to when called. context : Context The evaluation context (input data, math definition, config, ...). """ - self._mode = mode self._context = context @abstractmethod @@ -139,17 +135,6 @@ def as_expr(self, *args: Any, **kwargs: Any) -> LinearExpression | xr.DataArray: """ return self.as_raw(*args, **kwargs) - def __call__(self, *args: Any, **kwargs: Any) -> Any: - """Dispatch to the `as_*` method matching the mode set at initialisation.""" - if self._mode == "math_string": - return self.as_math_string(*args, **kwargs) - elif self._mode == "raw": - return self.as_raw(*args, **kwargs) - elif self._mode == "expr": - return self.as_expr(*args, **kwargs) - else: - raise ValueError(f"Unknown helper function mode: {self._mode!r}") - def _dim_iterator(self, dim: str) -> str: """Return the LaTeX iterator name of dimension `dim`.""" return dim_iterator(self._context.math, dim) @@ -543,8 +528,8 @@ def as_raw( """ group_name = str(group.name) dtype = DTYPE_OPTIONS[self._context.math.dimensions[group_name].dtype] - group_sum_helper = GroupSum(self._mode, self._context) - return group_sum_helper( + group_sum_helper = GroupSum(self._context) + return group_sum_helper.as_raw( array, getattr(array[str(over.name)].dt, group_name).astype(dtype), group ) diff --git a/linopy/declarative/latex.py b/linopy/declarative/latex.py index 8df71c73e..8358cc42c 100644 --- a/linopy/declarative/latex.py +++ b/linopy/declarative/latex.py @@ -20,25 +20,13 @@ import xarray as xr from linopy.declarative import parsing -from linopy.declarative.evaluate import Context, to_math_string -from linopy.declarative.grammar import Component, find_refs -from linopy.declarative.helpers import HelperFunction, build_registry, dim_iterator -from linopy.declarative.schema import ConfigModel, MathModel -from linopy.model import Model +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 DOCUMENTED_GROUPS, EQUATION_GROUPS FORMAT_T = Literal["md", "rst", "tex"] -_DOCUMENTED_GROUPS: dict[str, str] = { - "parameters": "Parameters", - "lookups": "Lookups", - "variables": "Variables", - "expressions": "Expressions", - "constraints": "Constraints", - "objectives": "Objectives", -} - -_EQUATION_GROUPS = ("expressions", "constraints", "objectives") - _REPR_STYLES = { "parameters": "textit", "lookups": "textit", @@ -82,16 +70,7 @@ class RenderedComponent: """Additional metadata to document (unit, default, sense, ...).""" -def _number_string(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 f"{value:.6g}" - - -class LatexModelBuilder: +class LatexModelBuilder(_DeclarativeBase): """ Builder turning a declarative math definition into LaTeX math documentation. @@ -128,22 +107,9 @@ def __init__( helpers : Iterable[type[HelperFunction]], optional User-defined helper functions, in addition to the built-in ones. """ - self.math = MathModel.model_validate(math_def) - self.input_data = input_data if input_data is not None else xr.Dataset() - self.config = ConfigModel.model_validate(config or {}) + super().__init__(math_def, input_data, config, helpers) self.components: dict[str, dict[str, RenderedComponent]] = {} - self._ctx = Context( - model=Model(), - input_data=self.input_data, - math=self.math, - config=self.config, - helpers=build_registry(helpers), - math_reprs=self._build_math_reprs(), - ) - - def _iterator(self, dim: str) -> str: - """Return the LaTeX iterator of a dimension (its own name if not declared).""" - return dim_iterator(self.math, dim) + self._ctx = replace(self._ctx, math_reprs=self._build_math_reprs()) def _build_math_reprs(self) -> dict[str, str]: r""" @@ -163,7 +129,7 @@ def _build_math_reprs(self) -> dict[str, str]: if name in self.input_data else [] ) - iterators = ",".join(self._iterator(str(dim)) for dim in dims) + 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 @@ -174,17 +140,15 @@ def _foreach_string(self, definition: object) -> str: if not sets: return "" instrs = ", ".join( - rf"\text{{{self._iterator(dim)}}} \in \text{{{dim}}}" for dim in sets + rf"\text{{{dim_iterator(self.math, dim)}}} \in \text{{{dim}}}" + for dim in sets ) return rf"\forall{{}} {instrs}" - def _mask_string(self, definition: object, name: str) -> str: - """Return the LaTeX rendering of a component's top-level mask ("" if true).""" - mask_node = parsing.parse_mask( - getattr(definition, "mask", "True"), self.math, name - ) - rendered = to_math_string( - mask_node, replace(self._ctx, route="mask", equation_name=name) + 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 @@ -204,19 +168,21 @@ def _render_metadata(self, definition: object) -> dict[str, str]: def add_component(self, group: str, name: str, definition: object) -> None: """Render one math component and store it under `self.components`.""" + mask_node = parsing.parse_mask( + getattr(definition, "mask", "True"), self.math, f"{group}:{name}" + ) rendered = RenderedComponent( group=group, name=name, title=getattr(definition, "title", ""), description=getattr(definition, "description", ""), foreach=self._foreach_string(definition), - mask=self._mask_string(definition, f"{group}:{name}"), + mask=self._mask_string(mask_node, f"{group}:{name}"), extras=self._render_metadata(definition), ) - mask_node = parsing.parse_mask(getattr(definition, "mask", "True"), self.math) uses = find_refs(mask_node, Component) - if group in _EQUATION_GROUPS: + if group in EQUATION_GROUPS: equations = parsing.parse_component(group, name, definition, self.math) # type: ignore[arg-type] for equation in equations: equation_ctx = replace(self._ctx, equation_name=equation.name) @@ -263,7 +229,7 @@ def _bounds_string(self, name: str, definition: object) -> str: lower, upper = ( reprs.get(bound, rf"\textit{{{bound}}}") if isinstance(bound, str) - else _number_string(bound) + else latex_number(bound) for bound in (bounds.lower, bounds.upper) ) return rf"{lower} \leq {reprs[name]} \leq {upper}" @@ -278,7 +244,7 @@ def build(self) -> LatexModelBuilder: Itself, with `self.components` filled, so that document generation can be chained (`builder.build().generate_math_doc()`). """ - for group in _DOCUMENTED_GROUPS: + for group in DOCUMENTED_GROUPS: for name, definition in getattr(self.math, group)._active.items(): self.add_component(group, name, definition) @@ -309,7 +275,7 @@ def generate_math_doc(self, format: FORMAT_T = "md") -> str: if not self.components: self.build() blocks = [_heading(format, 1, "Math formulation"), ""] - for group, group_title in _DOCUMENTED_GROUPS.items(): + for group, group_title in DOCUMENTED_GROUPS.items(): group_components = self.components.get(group) if not group_components: continue diff --git a/linopy/declarative/nodes.py b/linopy/declarative/nodes.py new file mode 100644 index 000000000..b60666190 --- /dev/null +++ b/linopy/declarative/nodes.py @@ -0,0 +1,828 @@ +""" +Linopy declarative math AST module. + +This module contains the AST node types produced when parsing declarative math +strings, together with — per node — the pyparsing parse action(s) that build it +(`from_tokens`), the data evaluator (`evaluate`, returning an `xr.DataArray` or +a linopy expression), and the LaTeX renderer (`to_latex`). The evaluation +context (:class:`Context`) and shared utilities live here too; the pyparsing +grammars that produce the nodes live in :mod:`linopy.declarative.grammar`. +""" + +from __future__ import annotations + +import operator +import re +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: + """ + 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).""" + + def evaluate(self, ctx: Context) -> Any: + """Evaluate this node to data (an `xr.DataArray` or a linopy expression).""" + raise NotImplementedError + + def to_latex(self, ctx: Context) -> str: + """Render this node as a LaTeX math string.""" + raise NotImplementedError + + +@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)) + 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) + if r"\text" 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 index eaeac359a..f51a5d674 100644 --- a/linopy/declarative/parsing.py +++ b/linopy/declarative/parsing.py @@ -21,14 +21,20 @@ import xarray as xr from linopy.declarative import grammar -from linopy.declarative.evaluate import ( +from linopy.declarative.nodes import ( + MODE_T, TRUE_ARRAY, + Component, Context, - evaluate, - to_math_string, + Node, + SliceRef, + SubExprRef, + find_refs, + to_linexpr, ) -from linopy.declarative.grammar import Node from linopy.declarative.schema import ( + COMPONENTS_T, + EQUATION_GROUP_T, MATH_DEFS_T, ConstraintDef, ExpressionDef, @@ -40,15 +46,6 @@ LOGGER = logging.getLogger(__name__) -GROUP_T = Literal[ - "variables", - "expressions", - "constraints", - "piecewise_constraints", - "objectives", - "postprocessed", -] - EQUATION_DEFS_T = ConstraintDef | ExpressionDef | ObjectiveDef """Math component definitions that carry `equations`/`sub_expressions`/`slices` keys.""" @@ -90,9 +87,7 @@ def references(self) -> set[str]: *self.sub_expressions.values(), *self.slices.values(), ] - return set().union( - *(grammar.find_refs(tree, grammar.Component) for tree in trees) - ) + return set().union(*(find_refs(tree, Component) for tree in trees)) # --------------------------------------------------------------------------- @@ -210,9 +205,9 @@ def _expand( """ expanded = [] for equation in equations: - ref_type = grammar.SubExprRef if kind == "sub_expressions" else grammar.SliceRef + ref_type = SubExprRef if kind == "sub_expressions" else SliceRef trees = [equation.expression, *equation.sub_expressions.values()] - refs = set().union(*(grammar.find_refs(tree, ref_type) for tree in trees)) + refs = set().union(*(find_refs(tree, ref_type) for tree in trees)) if not refs: expanded.append(equation) continue @@ -244,7 +239,7 @@ def _expand( def parse_component( - group: GROUP_T, name: str, definition: EQUATION_DEFS_T, math: MathModel + group: EQUATION_GROUP_T, name: str, definition: EQUATION_DEFS_T, math: MathModel ) -> list[Equation]: """ Parse a math component's equations into fully-resolved :class:`Equation` objects. @@ -256,7 +251,7 @@ def parse_component( Parameters ---------- - group : GROUP_T + group : EQUATION_GROUP_T Component group the definition belongs to (defines the equation grammar: comparisons for constraints, arithmetic otherwise). name : str @@ -361,8 +356,16 @@ def drop_dims_not_in_foreach(mask: xr.DataArray, sets: tuple[str, ...]) -> xr.Da 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: GROUP_T, + group: COMPONENTS_T, name: str, definition: MATH_DEFS_T, ctx: Context, @@ -377,7 +380,7 @@ def component_mask( Parameters ---------- - group : GROUP_T + group : COMPONENTS_T Component group the definition belongs to. name : str Name of the math component. @@ -394,21 +397,15 @@ def component_mask( sets = tuple(getattr(definition, "foreach", ())) mask_string = getattr(definition, "mask", "True") initial_mask = foreach_mask(sets, ctx.input_data) - if not initial_mask.any(): - LOGGER.debug( - f"Math parsing | {component_name} | Component not added; " - "'foreach' does not apply anywhere." - ) + if _mask_is_empty( + initial_mask, component_name, "'foreach' does not apply anywhere" + ): return initial_mask mask_node = parse_mask(mask_string, ctx.math, component_name) - mask_ctx = replace(ctx, route="mask", equation_name=component_name) - mask = xr.DataArray(initial_mask & evaluate(mask_node, mask_ctx)) - if not mask.any(): - LOGGER.debug( - f"Math parsing | {component_name} | Component not added; " - "'mask' does not apply anywhere." - ) + 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: @@ -424,14 +421,14 @@ def component_mask( def _equation_ctx( equation: Equation, ctx: Context, - route: Literal["expression", "mask"], + 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, - route=route, + mode=mode, sub_expressions=equation.sub_expressions, slices=equation.slices, **kwargs, @@ -460,13 +457,9 @@ def as_mask( Boolean array defining on which index items the equation applies. """ mask_ctx = _equation_ctx(equation, ctx, "mask") - evaluated = [evaluate(mask, mask_ctx) for mask in equation.masks] + evaluated = [mask.evaluate(mask_ctx) for mask in equation.masks] mask = xr.DataArray(functools.reduce(operator.and_, [initial_mask, *evaluated])) - if not mask.any(): - LOGGER.debug( - f"Math parsing | {equation.name} | Component not added; " - "'mask' does not apply anywhere." - ) + _mask_is_empty(mask, equation.name, "'mask' does not apply anywhere") return mask @@ -491,11 +484,8 @@ def as_expression( The evaluated expression; a pure-parameter expression (evaluated to an `xr.DataArray`) is coerced to a `LinearExpression`. """ - expr_ctx = _equation_ctx(equation, ctx, "expression", mask=mask) - evaluated = evaluate(equation.expression, expr_ctx, expr=True) - if isinstance(evaluated, xr.DataArray): - evaluated = LinearExpression(evaluated, ctx.model) - return evaluated + expr_ctx = _equation_ctx(equation, ctx, "expr", mask=mask) + return to_linexpr(equation.expression.evaluate(expr_ctx), ctx.model) def as_constraint( @@ -520,12 +510,8 @@ def as_constraint( coerced to `LinearExpression` and `sign` is an array of the comparison operator. """ - expr_ctx = _equation_ctx(equation, ctx, "expression", mask=mask) - lhs, sign, rhs = evaluate(equation.expression, expr_ctx, expr=True) - if isinstance(lhs, xr.DataArray): - lhs = LinearExpression(lhs, ctx.model) - if isinstance(rhs, xr.DataArray): - rhs = LinearExpression(rhs, ctx.model) + expr_ctx = _equation_ctx(equation, ctx, "expr", mask=mask) + lhs, sign, rhs = equation.expression.evaluate(expr_ctx) return lhs, sign, rhs @@ -555,42 +541,7 @@ def as_latex( """ if what == "mask": mask_ctx = _equation_ctx(equation, ctx, "mask") - strings = [to_math_string(mask, mask_ctx) for mask in equation.masks] + strings = [mask.to_latex(mask_ctx) for mask in equation.masks] return r"\land{}".join(f"({s})" for s in strings if s != "true") - expr_ctx = _equation_ctx(equation, ctx, "expression") - return to_math_string(equation.expression, expr_ctx) - - -def check_mask_expr_consistency( - name: str, expression: xr.DataArray, mask: xr.DataArray -) -> None: - """ - Check that an evaluated expression is consistent with its mask array. - - Parameters - ---------- - name : str - Name to identify the equation by in error messages. - expression : xr.DataArray - Array of linear expressions or one side of a constraint equation. - mask : xr.DataArray - Boolean mask; there should be a valid expression value wherever it is True. - - Raises - ------ - ValueError - If the expression is indexed over dimensions not present in the mask, or - has missing (NaN) entries where the mask applies. - """ - broadcast_dims_mask = set(expression.dims).difference(set(mask.dims)) - if broadcast_dims_mask: - raise ValueError( - f"{name} | The linear expression array is indexed over dimensions " - f"not present in `foreach`: {broadcast_dims_mask}" - ) - incomplete_constraints = expression.isnull() & mask - if incomplete_constraints.any(): - raise ValueError( - f"{name} | Missing a linear expression for some coordinates selected " - "by 'mask'. Adapting 'mask' might help." - ) + 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 index a17afaef0..0829bf735 100644 --- a/linopy/declarative/schema.py +++ b/linopy/declarative/schema.py @@ -44,6 +44,24 @@ "postprocessed", ] +EQUATION_GROUP_T = Literal["expressions", "constraints", "objectives"] + +EQUATION_GROUPS: tuple[str, ...] = ("expressions", "constraints", "objectives") +"""Component groups whose definitions carry parseable equations.""" + +BUILD_ORDER: tuple[str, ...] = ("variables", "expressions", "constraints", "objectives") +"""Component groups in the order they are built into a linopy model.""" + +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.""" + def _validate_unique_list(v: list) -> list: try: @@ -551,7 +569,4 @@ def find( class ConfigModel(LinopyBaseModel): """Base configuration options used when building a Linopy optimisation problem.""" - model_config = {"title": "Model build configuration"} - - foo: str = "bar" - """A dummy variable to test accessing the config items in declarative math.""" + model_config = {"title": "Model build configuration", "extra": "allow"} diff --git a/test/test_declarative_parsing.py b/test/test_declarative_parsing.py index 8eeb0c937..8b705a083 100644 --- a/test/test_declarative_parsing.py +++ b/test/test_declarative_parsing.py @@ -16,7 +16,7 @@ import xarray as xr import yaml -from linopy.declarative import evaluate, grammar, parsing +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 ( @@ -24,7 +24,7 @@ _escape_text_mode, latex_math_doc, ) -from linopy.declarative.schema import MathModel +from linopy.declarative.schema import ConfigModel, MathModel from linopy.expressions import LinearExpression from linopy.variables import Variable @@ -87,13 +87,13 @@ def _inputs() -> xr.Dataset: ) -def _ctx(builder: DeclarativeModelBuilder, **kwargs) -> evaluate.Context: +def _ctx(builder: DeclarativeModelBuilder, **kwargs) -> nodes.Context: """Build a fresh evaluation context from a builder's validated components.""" - return evaluate.Context( + return nodes.Context( model=builder.model, input_data=builder.input_data, math=builder.math, - config=builder.config, + config=kwargs.pop("config", builder.config), helpers=kwargs.pop("helpers", build_registry()), **kwargs, ) @@ -101,7 +101,7 @@ def _ctx(builder: DeclarativeModelBuilder, **kwargs) -> evaluate.Context: def _first_equation( builder: DeclarativeModelBuilder, group: str, name: str -) -> tuple[parsing.Equation, xr.DataArray, evaluate.Context]: +) -> tuple[parsing.Equation, xr.DataArray, nodes.Context]: """Parse a component and return its first equation, sub-mask, and context.""" ctx = _ctx(builder) definition = getattr(builder.math, group)[name] @@ -171,6 +171,14 @@ def test_find_refs_in_call_kwargs(self): assert grammar.find_refs(tree, grammar.SliceRef) == {"n"} assert grammar.find_refs(tree, grammar.Component) == {"node", "flow"} + def test_node_repr_is_clean(self): + 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 = _math() math["constraints"]["cap"]["equations"][0]["expression"] = "flow <= <=" @@ -208,14 +216,17 @@ def test_mask_comparison_and_subset_and_helper_return_bool(self): ) def test_mask_atoms_return_bool(self, builder_with_flow, mask_string): node = parsing.parse_mask(mask_string, builder_with_flow.math) - result = evaluate.evaluate(node, _ctx(builder_with_flow, route="mask")) + 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): """Bare input references coerce to existence booleans on the mask route.""" node = parsing.parse_mask("cap_max", builder_with_flow.math) - result = evaluate.evaluate(node, _ctx(builder_with_flow, route="mask")) + result = nodes.evaluate(node, _ctx(builder_with_flow, mode="mask")) assert result.values.tolist() == [True, True, True] @@ -289,16 +300,20 @@ def test_undefined_sub_expression_reference_raises(self): ) def test_plain_and_list_slices(self, builder_with_flow): - ctx = _ctx(builder_with_flow, mask=xr.full_like(_inputs()["cost"], True, bool)) + 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 = evaluate.evaluate(scalar_sliced, ctx, expr=True) + result = nodes.evaluate(scalar_sliced, ctx) assert isinstance(result, LinearExpression) list_sliced = arith.parse_string("flow[node=[a, b]]", parse_all=True)[0] - result = evaluate.evaluate(list_sliced, ctx, expr=True) + result = nodes.evaluate(list_sliced, ctx) assert result.data.sizes["node"] == 2 def test_slicer_reference(self, builder_with_flow): @@ -375,13 +390,13 @@ def test_sliced_component_latex(self, builder_with_flow): ctx = _ctx(builder_with_flow) arith = grammar.arithmetic_grammar(frozenset({"flow", "node"})) tree = arith.parse_string("flow[node=a]", parse_all=True)[0] - assert evaluate.to_math_string(tree, ctx) == r"flow_\text{n=a}" + assert nodes.to_math_string(tree, ctx) == r"flow_\text{n=a}" def test_identity_operands_are_skipped(self, builder_with_flow): ctx = _ctx(builder_with_flow) arith = grammar.arithmetic_grammar(frozenset({"flow"})) tree = arith.parse_string("0 + flow", parse_all=True)[0] - assert evaluate.to_math_string(tree, ctx) == "flow" + assert nodes.to_math_string(tree, ctx) == "flow" class _RecordArgs(HelperFunction): @@ -440,25 +455,22 @@ def test_non_subclass_rejected_by_registry(self): with pytest.raises(ValueError, match="must be subclassed"): build_registry([str]) # type: ignore[list-item] - def test_non_subclass_rejected_at_evaluation(self): - """A hand-built registry with an invalid entry is rejected at call time.""" + def test_unknown_helper_rejected(self): math = _math() math["expressions"]["total_cost"]["equations"][0]["expression"] = ( - "not_a_helper(flow)" + "unknown_helper(flow)" ) builder = DeclarativeModelBuilder(math, _inputs(), {}) builder.add_variable("flow", builder.math.variables["flow"]) - registry = build_registry() - registry["expression"]["not_a_helper"] = str # type: ignore[assignment] - ctx = _ctx(builder, helpers=registry) definition = builder.math.expressions["total_cost"] equation = parsing.parse_component( "expressions", "total_cost", definition, builder.math )[0] - with pytest.raises(ValueError, match="must be subclassed"): - parsing.as_expression(equation, ctx) + with pytest.raises(ValueError, match="Invalid helper function"): + parsing.as_expression(equation, _ctx(builder)) - def test_unknown_helper_rejected(self): + def test_eval_error_carries_caret(self): + """Evaluation errors point a caret at the failing node in the source string.""" math = _math() math["expressions"]["total_cost"]["equations"][0]["expression"] = ( "unknown_helper(flow)" @@ -469,8 +481,13 @@ def test_unknown_helper_rejected(self): equation = parsing.parse_component( "expressions", "total_cost", definition, builder.math )[0] - with pytest.raises(ValueError, match="Invalid helper function"): + 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): class _ClashingSum(HelperFunction): @@ -498,7 +515,7 @@ def test_get_val_at_index(self, builder_with_flow): 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 evaluate.evaluate(tree, ctx).item() == "a" + assert nodes.evaluate(tree, ctx).item() == "a" class TestBuilder: From d5d94a59f548e68826cf8d1be8bd5c8492af8d83 Mon Sep 17 00:00:00 2001 From: Bryn Pickering <17178478+brynpickering@users.noreply.github.com> Date: Mon, 27 Jul 2026 10:33:46 +0100 Subject: [PATCH 07/12] Post-refactor cleanup --- linopy/declarative/build.py | 68 ++++---- linopy/declarative/latex.py | 224 ++++++++++++++++++--------- linopy/declarative/nodes.py | 10 +- linopy/declarative/parsing.py | 258 +++++++++++++++++++++++-------- linopy/declarative/schema.py | 31 ++-- test/test_declarative_parsing.py | 150 +++++++++++++++--- 6 files changed, 537 insertions(+), 204 deletions(-) diff --git a/linopy/declarative/build.py b/linopy/declarative/build.py index 53ac77be5..07c603670 100644 --- a/linopy/declarative/build.py +++ b/linopy/declarative/build.py @@ -100,6 +100,7 @@ def __init__( """ 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( @@ -111,13 +112,10 @@ def __init__( math_reprs=math_reprs or {}, ) - def _references( - self, definition: Any, equations: Iterable[parsing.Equation] = () - ) -> list[str]: + def _references(self, parsed: parsing.ParsedComponent) -> list[str]: """Return the sorted names of all math components a component references.""" - mask_node = parsing.parse_mask(getattr(definition, "mask", "True"), self.math) - refs = find_refs(mask_node, Component) - for equation in equations: + refs = find_refs(parsed.mask, Component) + for equation in parsed.equations: refs |= equation.references() return sorted(refs) @@ -198,10 +196,9 @@ def _check_inputs(self) -> None: warn_msgs: list[str] = [] error_msgs: list[str] = [] active = self.input_data.get("active", xr.DataArray(True)) - for name, check in self.math.checks.root.items(): - if not check.active: - continue - mask_node = parsing.parse_mask(check.mask, self.math, name) + 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(): @@ -242,7 +239,10 @@ def _iter_equations( def add_variable(self, name: str, definition: VariableDef) -> None: """Add a decision variable to the model, masked by its math definition.""" - mask = parsing.component_mask("variables", name, definition, self._ctx) + 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 @@ -255,18 +255,22 @@ def add_variable(self, name: str, definition: VariableDef) -> None: integer=definition.domain == "integer", ) # Variable.attrs values are typed Hashable, but a sorted list serializes best. - self.model.variables[name].attrs["references"] = self._references(definition) # type: ignore[assignment] + 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.""" - mask = parsing.component_mask("expressions", name, definition, self._ctx) + 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) - equations = parsing.parse_component("expressions", name, definition, self.math) - for equation, sub_mask in self._iter_equations(equations, "expressions", mask): + 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 " @@ -280,21 +284,23 @@ def add_expression(self, name: str, definition: ExpressionDef) -> None: 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( - definition, equations - ) + 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.""" - mask = parsing.component_mask("constraints", name, definition, self._ctx) + 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) - equations = parsing.parse_component("constraints", name, definition, self.math) - for equation, sub_mask in self._iter_equations(equations, "constraints", 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 " @@ -320,20 +326,22 @@ def add_constraint(self, name: str, definition: ConstraintDef) -> None: rhs=rhs, mask=mask, ) - self.model.constraints[name].attrs["references"] = self._references( - definition, equations - ) + 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.""" - mask = parsing.component_mask("objectives", name, definition, self._ctx) + 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) - equations = parsing.parse_component("objectives", name, definition, self.math) - for equation, sub_mask in self._iter_equations(equations, "objectives", mask): + 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 " @@ -351,9 +359,7 @@ def add_objective(self, name: str, definition: ObjectiveDef) -> None: 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( - definition, equations - ) + self.model.objective.attrs["references"] = self._references(parsed) def build(self) -> Model: """ @@ -374,7 +380,7 @@ def build(self) -> Model: ) for group in BUILD_ORDER: component = group.removesuffix("s") - ordered_items = self._sorted_by_order(self.math[group].root) + ordered_items = self._sorted_by_order(self.math[group]._active) for name, definition in tqdm( ordered_items, desc=f"Building {group}.", colour=TQDM_COLOR ): diff --git a/linopy/declarative/latex.py b/linopy/declarative/latex.py index 8358cc42c..590c32f67 100644 --- a/linopy/declarative/latex.py +++ b/linopy/declarative/latex.py @@ -10,7 +10,6 @@ from __future__ import annotations -import math as pymath import re from collections.abc import Iterable from dataclasses import dataclass, field, replace @@ -23,7 +22,14 @@ 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 DOCUMENTED_GROUPS, EQUATION_GROUPS +from linopy.declarative.schema import ( + ConstraintDef, + ExpressionDef, + LookupDef, + ObjectiveDef, + ParameterDef, + VariableDef, +) FORMAT_T = Literal["md", "rst", "tex"] @@ -34,6 +40,28 @@ "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: @@ -122,22 +150,23 @@ def _build_math_reprs(self) -> dict[str, str]: reprs: dict[str, str] = {} for group, style in _REPR_STYLES.items(): for name, definition in getattr(self.math, group)._active.items(): - dims = getattr(definition, "foreach", None) - if dims is None: + 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: object) -> str: + def _foreach_string(self, definition: DOCUMENTED_LINOPY_OBJ_GROUPS_T) -> str: r"""Return the LaTeX `\forall` line body for a component's `foreach` sets.""" - sets = getattr(definition, "foreach", []) - if not sets: + if not (sets := definition.foreach): return "" instrs = ", ".join( rf"\text{{{dim_iterator(self.math, dim)}}} \in \text{{{dim}}}" @@ -152,79 +181,96 @@ def _mask_string(self, mask_node: Node, name: str) -> str: ) return "" if rendered == "true" else rendered - def _render_metadata(self, definition: object) -> dict[str, str]: + def _render_metadata(self, definition: DOCUMENTED_GROUPS_T) -> dict[str, str]: """Return the documentable metadata (unit, default, ...) of a definition.""" extras: dict[str, str] = {} - unit = getattr(definition, "unit", "") - if unit: + if unit := getattr(definition, "unit", None): extras["Unit"] = unit - default = getattr(definition, "default", None) - if default is not None and pd.notna(default): - if isinstance(default, int | float) and pymath.isinf(default): - extras["Default"] = "inf" if default > 0 else "-inf" - else: - extras["Default"] = str(default) + if pd.notnull(default := getattr(definition, "default", None)): + extras["Default"] = str(default) return extras - def add_component(self, group: str, name: str, definition: object) -> None: + def add_component( + self, group: str, name: str, definition: DOCUMENTED_GROUPS_T + ) -> None: """Render one math component and store it under `self.components`.""" - mask_node = parsing.parse_mask( - getattr(definition, "mask", "True"), self.math, f"{group}:{name}" - ) rendered = RenderedComponent( group=group, name=name, - title=getattr(definition, "title", ""), - description=getattr(definition, "description", ""), - foreach=self._foreach_string(definition), - mask=self._mask_string(mask_node, f"{group}:{name}"), + title=definition.title, + description=definition.description, extras=self._render_metadata(definition), ) - uses = find_refs(mask_node, Component) - - if group in EQUATION_GROUPS: - equations = parsing.parse_component(group, name, definition, self.math) # type: ignore[arg-type] - for equation in equations: - equation_ctx = replace(self._ctx, equation_name=equation.name) - rendered.equations.append( - { - "mask": parsing.as_latex(equation, equation_ctx, what="mask"), - "expression": parsing.as_latex(equation, equation_ctx), - } - ) - uses |= equation.references() - elif group == "variables": - rendered.extras["Domain"] = definition.domain # type: ignore[attr-defined] - rendered.equations.append( - {"mask": "", "expression": self._bounds_string(name, definition)} - ) - uses |= { - bound - for bound in (definition.bounds.lower, definition.bounds.upper) # type: ignore[attr-defined] - if isinstance(bound, str) - } - if group == "objectives": - rendered.extras["Sense"] = ( - "minimise" if definition.sense == "min" else "maximise" # type: ignore[attr-defined] - ) + 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. - rendered.foreach = _escape_text_mode(rendered.foreach) - rendered.mask = _escape_text_mode(rendered.mask) 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) - self.components.setdefault(group, {})[name] = rendered + 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: object) -> str: + def _bounds_string(self, name: str, definition: VariableDef) -> str: """Return the LaTeX bounds equation of a decision variable.""" - bounds = definition.bounds # type: ignore[attr-defined] + bounds = definition.bounds reprs = self._ctx.math_reprs lower, upper = ( reprs.get(bound, rf"\textit{{{bound}}}") @@ -234,6 +280,13 @@ def _bounds_string(self, name: str, definition: object) -> str: ) 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. @@ -372,16 +425,38 @@ def _metadata_lines(format: FORMAT_T, key: str, value: str) -> list[str]: return [f"- **{key}**: {value}"] -def _math_block(format: FORMAT_T, lines: list[str]) -> list[str]: - """Return a display-math block wrapping an `array` of the given LaTeX lines.""" +def _array_block(lines: list[str]) -> str: + """Return a LaTeX `array` environment of the given lines.""" joined = " \\\\\n ".join(lines) - array = f"\\begin{{array}}{{l}}\n {joined}\n\\end{{array}}" + 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 ["$$", array, "$$", ""] + return ["$$", inner, "$$", ""] if format == "rst": - indented = "\n".join(f" {line}" for line in array.split("\n")) + indented = "\n".join(f" {line}" for line in inner.split("\n")) return [".. math::", "", indented, ""] - return [r"\begin{equation}", array, r"\end{equation}", ""] + return [r"\begin{equation}", inner, r"\end{equation}", ""] def _component_doc(format: FORMAT_T, component: RenderedComponent) -> list[str]: @@ -400,13 +475,22 @@ def _component_doc(format: FORMAT_T, component: RenderedComponent) -> list[str]: blocks.extend(_metadata_lines(format, key, value)) if metadata: blocks.append("") - for equation in component.equations: - lines = [] + if component.equations: + header = [] if component.foreach: - lines.append(component.foreach) - for mask in (component.mask, equation["mask"]): - if mask: - lines.append(rf"\text{{if }} {mask}") - lines.append(equation["expression"]) - blocks.extend(_math_block(format, lines)) + 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 index b60666190..8c26b218c 100644 --- a/linopy/declarative/nodes.py +++ b/linopy/declarative/nodes.py @@ -13,6 +13,7 @@ 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 @@ -191,7 +192,7 @@ def latex_number(value: float | int) -> str: @dataclass(frozen=True, kw_only=True) -class Node: +class Node(ABC): """ Base class of all declarative math AST nodes. @@ -207,13 +208,13 @@ class Node: 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).""" - raise NotImplementedError + @abstractmethod def to_latex(self, ctx: Context) -> str: """Render this node as a LaTeX math string.""" - raise NotImplementedError @dataclass(frozen=True, kw_only=True) @@ -671,7 +672,8 @@ def to_latex(self, ctx: Context) -> str: unmasked_ctx = replace(ctx, apply_mask=False) lhs_str = self.lhs.to_latex(unmasked_ctx) rhs_str = self.rhs.to_latex(unmasked_ctx) - if r"\text" not in rhs_str: + # 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) diff --git a/linopy/declarative/parsing.py b/linopy/declarative/parsing.py index f51a5d674..e610841ca 100644 --- a/linopy/declarative/parsing.py +++ b/linopy/declarative/parsing.py @@ -33,6 +33,7 @@ to_linexpr, ) from linopy.declarative.schema import ( + BUILD_ORDER, COMPONENTS_T, EQUATION_GROUP_T, MATH_DEFS_T, @@ -90,9 +91,92 @@ def references(self) -> set[str]: return set().union(*(find_refs(tree, Component) for tree in trees)) -# --------------------------------------------------------------------------- -# Parsing -# --------------------------------------------------------------------------- +@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]: @@ -110,34 +194,49 @@ def _mask_grammar(math: MathModel) -> pp.ParserElement: ) -class _ErrorCollector: - """Collect parse errors with their positions, to raise a single error at the end.""" +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 __init__(self, component_name: str) -> None: - self.component_name = component_name - self.errors: 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, position: str + 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 with a caret marker pointing at the parse position, - for raising later via :meth:`raise_errors`. + 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"{position} (line {excinfo.lineno}, char {excinfo.col}): " - marker_pos = " " * (len(pointer) + 2 * len(_ERR_BULLET) + excinfo.col - 1) - self.errors.append(f"{pointer}{excinfo.line}\n{marker_pos}^") + 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 parse errors as a single bullet-point ValueError.""" + """Raise all collected errors as one ValueError, grouped by component.""" if self.errors: - raise ValueError(f"- {self.component_name}: {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: @@ -153,15 +252,16 @@ def parse_mask(mask_string: str, math: MathModel, name: str = "") -> Node: name : str, optional Name to identify the string by in error messages. """ - collector = _ErrorCollector(name) - parsed = collector.parse(_mask_grammar(math), mask_string, "mask") + 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: _ErrorCollector, + collector: _ParsingCollector, + component_name: str, parser: pp.ParserElement, mask_parser: pp.ParserElement, expression_list: _Equations, @@ -173,9 +273,11 @@ def _parse_variants( equations = [] for idx, item in enumerate(expression_list): position_id = f"{position}[{idx}]" - mask = collector.parse(mask_parser, item.mask, f"{position_id}.mask") + mask = collector.parse( + mask_parser, item.mask, component_name, f"{position_id}.mask" + ) expression = collector.parse( - parser, item.expression, f"{position_id}.expression" + parser, item.expression, component_name, f"{position_id}.expression" ) if expression is not None and mask is not None: equations.append( @@ -194,14 +296,16 @@ def _expand( 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, with the chosen variants' masks and - ASTs merged in. + 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: @@ -213,9 +317,11 @@ def _expand( continue undefined = refs.difference(candidates.keys()) if undefined: - raise KeyError( - f"{component_name}: Undefined {kind} found in equation: {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 = ( @@ -239,15 +345,17 @@ def _expand( def parse_component( - group: EQUATION_GROUP_T, name: str, definition: EQUATION_DEFS_T, math: MathModel + 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 (syntax errors across all of them are - collected and raised together), then every equation is expanded with the - cartesian product of the sub-expression and slicer variants it references. + 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 ---------- @@ -260,6 +368,9 @@ def parse_component( 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 ------- @@ -267,6 +378,8 @@ def parse_component( 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) @@ -274,12 +387,11 @@ def parse_component( else grammar.arithmetic_grammar(names) ) mask_parser = _mask_grammar(math) - # Objectives are adimensional: they carry no `foreach` key. - sets = tuple(getattr(definition, "foreach", ())) - collector = _ErrorCollector(component_name) + sets = tuple(definition.foreach) equations = _parse_variants( collector, + component_name, equation_parser, mask_parser, definition.equations, @@ -290,6 +402,7 @@ def parse_component( sub_expressions = { sub_name: _parse_variants( collector, + component_name, grammar.sub_expression_grammar(names), mask_parser, sub_list, @@ -302,6 +415,7 @@ def parse_component( slices = { slice_name: _parse_variants( collector, + component_name, grammar.slice_grammar(names), mask_parser, slice_list, @@ -311,10 +425,14 @@ def parse_component( ) for slice_name, slice_list in definition.slices.root.items() } - collector.raise_errors() - equations = _expand(component_name, equations, sub_expressions, "sub_expressions") - return _expand(component_name, equations, slices, "slices") + 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 # --------------------------------------------------------------------------- @@ -368,6 +486,7 @@ def component_mask( group: COMPONENTS_T, name: str, definition: MATH_DEFS_T, + mask_node: Node, ctx: Context, *, align_to_foreach_sets: bool = True, @@ -375,8 +494,7 @@ def component_mask( """ Evaluate a component's top-level mask over its `foreach` dimensions. - Combines the `foreach` existence array with the component's (optional) - top-level `mask` string, breaking early if no valid element remains. + Combines the `foreach` existence array with the component's pre-parsed top-level `mask` AST, breaking early if no valid element remains. Parameters ---------- @@ -386,6 +504,9 @@ def component_mask( 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 @@ -393,16 +514,13 @@ def component_mask( (see :func:`drop_dims_not_in_foreach`). """ component_name = f"{group}:{name}" - # Objectives are adimensional: they carry no `foreach` or `mask` keys. - sets = tuple(getattr(definition, "foreach", ())) - mask_string = getattr(definition, "mask", "True") + 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_node = parse_mask(mask_string, ctx.math, component_name) 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"): @@ -413,11 +531,6 @@ def component_mask( return mask -# --------------------------------------------------------------------------- -# Typed evaluation entry points -# --------------------------------------------------------------------------- - - def _equation_ctx( equation: Equation, ctx: Context, @@ -448,8 +561,8 @@ def as_mask( 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`. + Mask to combine (boolean AND) with the equation's own masks. + E.g., the component-level mask from :func:`component_mask`. Returns ------- @@ -506,23 +619,21 @@ def as_constraint( Returns ------- tuple[LinearExpression, xr.DataArray, LinearExpression] - `(lhs, sign, rhs)` for constraint assembly; pure-parameter sides are - coerced to `LinearExpression` and `sign` is an array of the comparison - operator. + `(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( +def as_latex_mask( equation: Equation, ctx: Context, - *, - what: Literal["expression", "mask"] = "expression", ) -> str: """ - Render an equation's expression or mask as a LaTeX math string. + Render an equation's mask as a LaTeX math string. Parameters ---------- @@ -531,17 +642,36 @@ def as_latex( 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. + 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. """ - if what == "mask": - 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") 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 index 0829bf735..e06f74566 100644 --- a/linopy/declarative/schema.py +++ b/linopy/declarative/schema.py @@ -10,7 +10,7 @@ import logging from collections.abc import Hashable, Iterable from functools import cached_property -from typing import Annotated, Any, ClassVar, Literal, Self, TypeVar +from typing import Annotated, Any, ClassVar, Literal, Self, TypeVar, get_args import numpy as np from annotated_types import Len @@ -46,22 +46,11 @@ EQUATION_GROUP_T = Literal["expressions", "constraints", "objectives"] -EQUATION_GROUPS: tuple[str, ...] = ("expressions", "constraints", "objectives") -"""Component groups whose definitions carry parseable equations.""" +BUILD_ORDER_T = Literal["variables", "expressions", "constraints", "objectives"] +BUILD_ORDER: tuple[BUILD_ORDER_T, ...] = get_args(BUILD_ORDER_T) -BUILD_ORDER: tuple[str, ...] = ("variables", "expressions", "constraints", "objectives") """Component groups in the order they are built into a linopy model.""" -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.""" - def _validate_unique_list(v: list) -> list: try: @@ -194,7 +183,7 @@ class DimensionDef(_MathComponent): @property def default(self) -> float: - """Dummy variable to align with lookups and dims.""" + """Dummy field to align with lookups and dims.""" return float("nan") @@ -208,7 +197,7 @@ class ParameterDef(_MathComponent): @property def dtype(self) -> Literal["float"]: - """Dummy variable to align with lookups and dims.""" + """Dummy field to align with lookups and dims.""" return "float" _group: ClassVar[COMPONENTS_T] = "parameters" @@ -371,6 +360,16 @@ class ObjectiveDef(_MathEquationComponent): _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): """ diff --git a/test/test_declarative_parsing.py b/test/test_declarative_parsing.py index 8b705a083..038a55bc3 100644 --- a/test/test_declarative_parsing.py +++ b/test/test_declarative_parsing.py @@ -102,11 +102,12 @@ def _ctx(builder: DeclarativeModelBuilder, **kwargs) -> nodes.Context: def _first_equation( builder: DeclarativeModelBuilder, group: str, name: str ) -> tuple[parsing.Equation, xr.DataArray, nodes.Context]: - """Parse a component and return its first equation, sub-mask, and context.""" + """Return a pre-parsed component's first equation, sub-mask, and context.""" ctx = _ctx(builder) definition = getattr(builder.math, group)[name] - mask = parsing.component_mask(group, name, definition, ctx) - equation = parsing.parse_component(group, name, definition, builder.math)[0] + 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 @@ -182,11 +183,81 @@ def test_node_repr_is_clean(self): def test_parse_error_carries_position_marker(self): math = _math() 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): + """Broken strings in two components raise as one grouped error.""" + math = _math() + 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): + """A syntax error does not short-circuit undefined `$ref` collection.""" + math = _math() + 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): + """Inactive components are neither parsed nor built.""" + math = _math() + math["expressions"]["broken"] = { + "active": False, + "foreach": ["node"], + "equations": [{"expression": "flow * * cost"}], + } builder = DeclarativeModelBuilder(math, _inputs(), {}) - with pytest.raises(ValueError, match="equations\\[0\\].expression"): - parsing.parse_component( - "constraints", "cap", builder.math.constraints["cap"], builder.math - ) + assert "broken" not in builder.parsed["expressions"] + model = builder.build() + assert "broken" not in model.expressions + + def test_check_masks_are_parsed(self): + math = _math() + 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 = _math() + 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): + 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: @@ -290,14 +361,8 @@ def test_undefined_sub_expression_reference_raises(self): math["expressions"]["sub_expr_test"]["sub_expressions"] = { "bar": [{"expression": "flow"}] } - builder = DeclarativeModelBuilder(math, _inputs(), {}) - with pytest.raises(KeyError, match="Undefined sub_expressions"): - parsing.parse_component( - "expressions", - "sub_expr_test", - builder.math.expressions["sub_expr_test"], - builder.math, - ) + with pytest.raises(ValueError, match="Undefined sub_expressions"): + DeclarativeModelBuilder(math, _inputs(), {}) def test_plain_and_list_slices(self, builder_with_flow): ctx = _ctx( @@ -365,12 +430,12 @@ class TestLatexRoute: def test_equation_latex(self, builder_with_flow): equation, _, ctx = _first_equation(builder_with_flow, "constraints", "cap") - assert parsing.as_latex(equation, ctx) == r"flow \leq cap_max" + assert parsing.as_latex_expression(equation, ctx) == r"flow \leq cap_max" def test_sum_latex(self, builder_with_flow): equation, _, ctx = _first_equation(builder_with_flow, "objectives", "obj") assert ( - parsing.as_latex(equation, ctx) + parsing.as_latex_expression(equation, ctx) == r"\sum\limits_{\substack{\text{n} \in \text{node}}} (total_cost)" ) @@ -382,10 +447,22 @@ def test_mask_latex(self): builder = DeclarativeModelBuilder(math, _inputs(), {}) builder.add_variable("flow", builder.math.variables["flow"]) equation, _, ctx = _first_equation(builder, "constraints", "cap") - assert parsing.as_latex(equation, ctx, what="mask") == ( + 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): + # `\infty` (and other bare LaTeX commands) must stay in math mode; only + # plain-text tokens (numbers, coordinate labels, booleans) get `\text{}`. + math = _math() + 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): ctx = _ctx(builder_with_flow) arith = grammar.arithmetic_grammar(frozenset({"flow", "node"})) @@ -658,7 +735,17 @@ def test_equation_masks_render_as_if_conditions(self): doc = latex_math_doc(math, _inputs(), format="md") assert r"\text{if } (\textit{cost}_\text{n}\mathord{>}\text{1})" in doc - def test_sub_expression_variants_render_as_separate_blocks(self): + def test_infinity_not_wrapped_in_text(self): + # `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 = _math() + 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 = _math() math["expressions"]["sub_expr_test"]["sub_expressions"]["foo"] = [ {"mask": "cost > 1", "expression": "flow"}, @@ -671,6 +758,31 @@ def test_sub_expression_variants_render_as_separate_blocks(self): r"\textbf{flow}_\text{n} \times 2 \times \textit{cost}_\text{n}" ) + def test_multiple_equations_render_as_one_block_with_cases(self): + # 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 = _math() + 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): + 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): doc = latex_math_doc(_math(), _inputs(), format="md") assert doc.startswith("# Math formulation") From aa8ac98e1d49ce8640a7772e60913bbc0baaf051 Mon Sep 17 00:00:00 2001 From: Bryn Pickering <17178478+brynpickering@users.noreply.github.com> Date: Mon, 27 Jul 2026 10:48:13 +0100 Subject: [PATCH 08/12] Add attributions --- linopy/declarative/__init__.py | 6 ++++-- linopy/declarative/build.py | 6 ++++-- linopy/declarative/grammar.py | 17 +++++++++-------- linopy/declarative/helpers.py | 11 +++++++---- linopy/declarative/latex.py | 11 ++++++----- linopy/declarative/nodes.py | 15 +++++++++------ linopy/declarative/parsing.py | 11 ++++++----- linopy/declarative/schema.py | 6 ++++-- 8 files changed, 49 insertions(+), 34 deletions(-) diff --git a/linopy/declarative/__init__.py b/linopy/declarative/__init__.py index bbcf650d2..5000cf44f 100644 --- a/linopy/declarative/__init__.py +++ b/linopy/declarative/__init__.py @@ -1,8 +1,10 @@ """ Linopy declarative math interface. -Build a linopy model from a declarative math definition (typically loaded from -YAML) and an xarray dataset of input data, via :func:`declarative_model`. +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 diff --git a/linopy/declarative/build.py b/linopy/declarative/build.py index 07c603670..3fd8647de 100644 --- a/linopy/declarative/build.py +++ b/linopy/declarative/build.py @@ -1,9 +1,11 @@ """ Linopy declarative model-build module. -This module contains the entry point to build a linopy optimisation model from a -declarative math definition (a dictionary, typically loaded from YAML) and an +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 diff --git a/linopy/declarative/grammar.py b/linopy/declarative/grammar.py index 216aa9a25..ea7a13c11 100644 --- a/linopy/declarative/grammar.py +++ b/linopy/declarative/grammar.py @@ -1,14 +1,15 @@ """ 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 `eval_arith.py` -example (https://github.com/pyparsing/pyparsing/blob/master/examples/eval_arith.py, -MIT licensed). +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 diff --git a/linopy/declarative/helpers.py b/linopy/declarative/helpers.py index a6f9244ab..1ef3140e2 100644 --- a/linopy/declarative/helpers.py +++ b/linopy/declarative/helpers.py @@ -1,10 +1,13 @@ """ Linopy declarative math helper-functions module. -This module contains the helper functions that can be called in declarative math -`mask` and `expression` strings (by their `NAME`), the abstract base class from -which users can define their own, and the registry builder that makes them -available to a model build. +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 diff --git a/linopy/declarative/latex.py b/linopy/declarative/latex.py index 590c32f67..c6e897983 100644 --- a/linopy/declarative/latex.py +++ b/linopy/declarative/latex.py @@ -1,11 +1,12 @@ """ Linopy declarative LaTeX math documentation module. -This module builds a human-readable mathematical formulation document from a -declarative math definition, without building an optimisation problem: every -active math component is rendered to LaTeX (equations, mask conditions, foreach -sets, bounds) together with its metadata and cross-references, and the result -can be generated as Markdown, reStructuredText, or LaTeX source. +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 diff --git a/linopy/declarative/nodes.py b/linopy/declarative/nodes.py index 8c26b218c..7b41805f5 100644 --- a/linopy/declarative/nodes.py +++ b/linopy/declarative/nodes.py @@ -1,12 +1,15 @@ """ Linopy declarative math AST module. -This module contains the AST node types produced when parsing declarative math -strings, together with — per node — the pyparsing parse action(s) that build it -(`from_tokens`), the data evaluator (`evaluate`, returning an `xr.DataArray` or -a linopy expression), and the LaTeX renderer (`to_latex`). The evaluation -context (:class:`Context`) and shared utilities live here too; the pyparsing -grammars that produce the nodes live in :mod:`linopy.declarative.grammar`. +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 diff --git a/linopy/declarative/parsing.py b/linopy/declarative/parsing.py index e610841ca..7f37f5970 100644 --- a/linopy/declarative/parsing.py +++ b/linopy/declarative/parsing.py @@ -1,11 +1,12 @@ """ Linopy declarative math parsing module. -This module turns a validated math component definition into a list of -:class:`Equation` objects — pure data holding the parsed expression/mask ASTs -with all `$name` sub-expression and slicer references resolved — and provides -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 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 diff --git a/linopy/declarative/schema.py b/linopy/declarative/schema.py index e06f74566..799cdb3ad 100644 --- a/linopy/declarative/schema.py +++ b/linopy/declarative/schema.py @@ -1,8 +1,10 @@ """ Linopy declarative math schema module. -This module contains the pydantic models that validate declarative math and -build-configuration definitions. +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 From 5ac32c85265cbf170cb60a86bcaf23082a8e6ab9 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 27 Jul 2026 10:25:00 +0000 Subject: [PATCH 09/12] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- linopy/__init__.py | 2 +- linopy/model.py | 3 ++- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/linopy/__init__.py b/linopy/__init__.py index 4df75f513..602e597f5 100644 --- a/linopy/__init__.py +++ b/linopy/__init__.py @@ -77,5 +77,5 @@ "read_netcdf", "segments", "tangent_lines", - "declarative_model" + "declarative_model", ) diff --git a/linopy/model.py b/linopy/model.py index 1781fdd34..6ea2259a2 100644 --- a/linopy/model.py +++ b/linopy/model.py @@ -933,7 +933,8 @@ def add_variables( def add_expressions( self, data: Variable - | LinearExpression | QuadraticExpression + | LinearExpression + | QuadraticExpression | Sequence[tuple[ConstantLike, Variable | str]], name: str | None = None, mask: MaskLike | None = None, From c25178a8c42acc2b7c199d2792faa0f16647d323 Mon Sep 17 00:00:00 2001 From: Bryn Pickering <17178478+brynpickering@users.noreply.github.com> Date: Mon, 27 Jul 2026 13:29:16 +0100 Subject: [PATCH 10/12] Please mypy --- test/test_declarative_parsing.py | 167 ++++++++++++++++++------------- 1 file changed, 95 insertions(+), 72 deletions(-) diff --git a/test/test_declarative_parsing.py b/test/test_declarative_parsing.py index 038a55bc3..4ae1103af 100644 --- a/test/test_declarative_parsing.py +++ b/test/test_declarative_parsing.py @@ -10,6 +10,7 @@ import re from pathlib import Path +from typing import Any import numpy as np import pytest @@ -24,7 +25,7 @@ _escape_text_mode, latex_math_doc, ) -from linopy.declarative.schema import ConfigModel, MathModel +from linopy.declarative.schema import COMPONENTS_T, ConfigModel, MathModel from linopy.expressions import LinearExpression from linopy.variables import Variable @@ -87,7 +88,7 @@ def _inputs() -> xr.Dataset: ) -def _ctx(builder: DeclarativeModelBuilder, **kwargs) -> nodes.Context: +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, @@ -100,7 +101,7 @@ def _ctx(builder: DeclarativeModelBuilder, **kwargs) -> nodes.Context: def _first_equation( - builder: DeclarativeModelBuilder, group: str, name: str + 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) @@ -126,7 +127,7 @@ class TestGrammar: NAMES = frozenset({"flow", "cost", "cap_max", "node"}) - def test_equation_tree_shape(self): + def test_equation_tree_shape(self) -> None: tree = grammar.equation_grammar(self.NAMES).parse_string( "flow <= cap_max", parse_all=True )[0] @@ -134,7 +135,7 @@ def test_equation_tree_shape(self): assert tree.op == "<=" assert isinstance(tree.lhs, grammar.Component) and tree.lhs.name == "flow" - def test_equation_rejects_mask_only_operators(self): + def test_equation_rejects_mask_only_operators(self) -> None: import pyparsing as pp with pytest.raises(pp.ParseException): @@ -142,21 +143,21 @@ def test_equation_rejects_mask_only_operators(self): "flow < cap_max", parse_all=True ) - def test_arithmetic_tree_shape(self): + 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): + 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): + def test_sub_expression_grammar_rejects_refs(self) -> None: import pyparsing as pp with pytest.raises(pp.ParseException): @@ -164,7 +165,7 @@ def test_sub_expression_grammar_rejects_refs(self): "$foo + 1", parse_all=True ) - def test_find_refs_in_call_kwargs(self): + 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] @@ -172,7 +173,7 @@ def test_find_refs_in_call_kwargs(self): assert grammar.find_refs(tree, grammar.SliceRef) == {"n"} assert grammar.find_refs(tree, grammar.Component) == {"node", "flow"} - def test_node_repr_is_clean(self): + def test_node_repr_is_clean(self) -> None: tree = grammar.arithmetic_grammar(self.NAMES).parse_string( "flow * cost + 1", parse_all=True )[0] @@ -180,7 +181,7 @@ def test_node_repr_is_clean(self): assert "instring=" not in rendered assert "loc=" not in rendered - def test_parse_error_carries_position_marker(self): + def test_parse_error_carries_position_marker(self) -> None: math = _math() math["constraints"]["cap"]["equations"][0]["expression"] = "flow <= <=" with pytest.raises(ValueError, match="equations\\[0\\].expression") as excinfo: @@ -193,7 +194,7 @@ def test_parse_error_carries_position_marker(self): class TestParseWalkthrough: """Whole-dict parse walkthrough with aggregated errors.""" - def test_errors_aggregate_across_components(self): + def test_errors_aggregate_across_components(self) -> None: """Broken strings in two components raise as one grouped error.""" math = _math() math["constraints"]["cap"]["equations"][0]["expression"] = "flow <= <=" @@ -207,7 +208,7 @@ def test_errors_aggregate_across_components(self): assert "expressions:total_cost:" in message assert message.count("^") == 2 - def test_undefined_ref_collected_alongside_syntax_errors(self): + def test_undefined_ref_collected_alongside_syntax_errors(self) -> None: """A syntax error does not short-circuit undefined `$ref` collection.""" math = _math() math["constraints"]["cap"]["equations"][0]["expression"] = "flow <= <=" @@ -221,7 +222,7 @@ def test_undefined_ref_collected_alongside_syntax_errors(self): assert "expressions:sub_expr_test:" in message assert "Undefined sub_expressions" in message - def test_inactive_components_are_skipped(self): + def test_inactive_components_are_skipped(self) -> None: """Inactive components are neither parsed nor built.""" math = _math() math["expressions"]["broken"] = { @@ -234,13 +235,13 @@ def test_inactive_components_are_skipped(self): model = builder.build() assert "broken" not in model.expressions - def test_check_masks_are_parsed(self): + def test_check_masks_are_parsed(self) -> None: math = _math() 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): + def test_inactive_check_masks_are_skipped(self) -> None: math = _math() math["checks"] = { "bad": {"mask": "cost > >", "message": "boom", "active": False} @@ -248,7 +249,7 @@ def test_inactive_check_masks_are_skipped(self): builder = DeclarativeModelBuilder(math, _inputs(), {}) assert "bad" not in builder.parsed.checks - def test_parsed_math_shape(self): + def test_parsed_math_shape(self) -> None: builder = DeclarativeModelBuilder(_math(), _inputs(), {}) assert set(builder.parsed.components) == { "variables", @@ -263,13 +264,15 @@ def test_parsed_math_shape(self): class TestMaskRoute: """Mask strings evaluate to boolean arrays.""" - def test_top_level_mask_returns_boolean_dataarray(self, builder_with_flow): + 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): + def test_mask_comparison_and_subset_and_helper_return_bool(self) -> None: math = _math() math["constraints"]["cap"]["equations"][0]["mask"] = ( "cost > 1 and [a, b] in node and any(cap_max, over=node)" @@ -285,7 +288,9 @@ def test_mask_comparison_and_subset_and_helper_return_bool(self): "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, mask_string): + 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( @@ -294,7 +299,9 @@ def test_mask_atoms_return_bool(self, builder_with_flow, mask_string): assert isinstance(result, xr.DataArray) assert result.dtype == bool - def test_existence_coercion(self, builder_with_flow): + 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")) @@ -304,14 +311,18 @@ def test_existence_coercion(self, builder_with_flow): class TestExpressionRoute: """Expression strings evaluate to linopy expressions.""" - def test_expression_with_variable_returns_linexpr(self, builder_with_flow): + 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): + 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" ) @@ -319,7 +330,9 @@ def test_pure_parameter_expression_coerced_to_linexpr(self, builder_with_flow): # 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): + 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" ) @@ -327,7 +340,7 @@ def test_sub_expression_reference_returns_linexpr(self, builder_with_flow): result = parsing.as_expression(equation, ctx, mask=sub_mask) assert isinstance(result, LinearExpression) - def test_sub_expression_variants_expand_to_cartesian_product(self): + def test_sub_expression_variants_expand_to_cartesian_product(self) -> None: """Two variants of one sub-expression yield two equations with merged masks.""" math = _math() math["expressions"]["sub_expr_test"]["sub_expressions"]["foo"] = [ @@ -356,7 +369,7 @@ def test_sub_expression_variants_expand_to_cartesian_product(self): builder.add_expression("sub_expr_test", definition) assert "sub_expr_test" in builder.model.expressions - def test_undefined_sub_expression_reference_raises(self): + def test_undefined_sub_expression_reference_raises(self) -> None: math = _math() math["expressions"]["sub_expr_test"]["sub_expressions"] = { "bar": [{"expression": "flow"}] @@ -364,7 +377,9 @@ def test_undefined_sub_expression_reference_raises(self): with pytest.raises(ValueError, match="Undefined sub_expressions"): DeclarativeModelBuilder(math, _inputs(), {}) - def test_plain_and_list_slices(self, builder_with_flow): + def test_plain_and_list_slices( + self, builder_with_flow: DeclarativeModelBuilder + ) -> None: ctx = _ctx( builder_with_flow, mode="expr", @@ -381,7 +396,7 @@ def test_plain_and_list_slices(self, builder_with_flow): result = nodes.evaluate(list_sliced, ctx) assert result.data.sizes["node"] == 2 - def test_slicer_reference(self, builder_with_flow): + def test_slicer_reference(self, builder_with_flow: DeclarativeModelBuilder) -> None: """`$name` slicer references resolve like sub-expressions (feature parity).""" math = _math() math["expressions"]["sliced"] = { @@ -404,7 +419,9 @@ def test_slicer_reference(self, builder_with_flow): class TestConstraintRoute: """Constraint equations evaluate to (lhs, sign, rhs) tuples.""" - def test_equation_returns_lhs_sign_rhs_tuple(self, builder_with_flow): + 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" ) @@ -414,7 +431,7 @@ def test_equation_returns_lhs_sign_rhs_tuple(self, builder_with_flow): assert isinstance(sign, xr.DataArray) assert set(np.unique(sign.values)) <= {"<="} - def test_foreach_dim_mismatch_raises(self): + def test_foreach_dim_mismatch_raises(self) -> None: math = _math() # `sum` removed: the equation is indexed over `node` but foreach is empty. math["constraints"]["cap"]["foreach"] = [] @@ -428,18 +445,18 @@ def test_foreach_dim_mismatch_raises(self): class TestLatexRoute: """Math strings render as LaTeX.""" - def test_equation_latex(self, builder_with_flow): + 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): + 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): + def test_mask_latex(self) -> None: math = _math() math["constraints"]["cap"]["equations"][0]["mask"] = ( "cost > 1 and [a, b] in node" @@ -451,7 +468,7 @@ def test_mask_latex(self): r"(\textit{cost}\mathord{>}\text{1} \land \text{n} \in \text{[a,b]})" ) - def test_mask_infinity_latex_not_wrapped_in_text(self): + def test_mask_infinity_latex_not_wrapped_in_text(self) -> None: # `\infty` (and other bare LaTeX commands) must stay in math mode; only # plain-text tokens (numbers, coordinate labels, booleans) get `\text{}`. math = _math() @@ -463,13 +480,17 @@ def test_mask_infinity_latex_not_wrapped_in_text(self): assert r"\mathord{==}\infty" in rendered assert r"\text{\infty}" not in rendered - def test_sliced_component_latex(self, builder_with_flow): + 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): + 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] @@ -483,10 +504,10 @@ class _RecordArgs(HelperFunction): ALLOWED_IN = ["expression"] received: list[type] = [] - def as_math_string(self, *args, **kwargs): # noqa: D102 + def as_math_string(self, *args: Any, **kwargs: Any) -> str: # noqa: D102 return "record_args" - def as_raw(self, *args, **kwargs): # noqa: D102 + 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] @@ -498,17 +519,17 @@ class _Double(HelperFunction): NAME = "double" ALLOWED_IN = ["expression"] - def as_math_string(self, array): # noqa: D102 + def as_math_string(self, array: Any) -> str: # noqa: D102 return rf"2 \times {array}" - def as_raw(self, array): # noqa: D102 + 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): + def test_helper_arguments_are_evaluated_raw(self) -> None: """Helper args arrive un-normalised: raw Variable/DataArray, not LinearExpression.""" math = _math() math["expressions"]["total_cost"]["equations"][0]["expression"] = ( @@ -528,11 +549,11 @@ def test_helper_arguments_are_evaluated_raw(self): assert xr.DataArray in _RecordArgs.received assert LinearExpression not in _RecordArgs.received - def test_non_subclass_rejected_by_registry(self): + 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): + def test_unknown_helper_rejected(self) -> None: math = _math() math["expressions"]["total_cost"]["equations"][0]["expression"] = ( "unknown_helper(flow)" @@ -546,7 +567,7 @@ def test_unknown_helper_rejected(self): with pytest.raises(ValueError, match="Invalid helper function"): parsing.as_expression(equation, _ctx(builder)) - def test_eval_error_carries_caret(self): + def test_eval_error_carries_caret(self) -> None: """Evaluation errors point a caret at the failing node in the source string.""" math = _math() math["expressions"]["total_cost"]["equations"][0]["expression"] = ( @@ -566,21 +587,23 @@ def test_eval_error_carries_caret(self): assert caret_line.strip() == "^" assert caret_line.index("^") == source_line.index("unknown_helper") - def test_duplicate_name_rejected(self): + def test_duplicate_name_rejected(self) -> None: class _ClashingSum(HelperFunction): NAME = "sum" ALLOWED_IN = ["expression"] - def as_math_string(self, *args, **kwargs): # noqa: D102 + def as_math_string(self, *args: Any, **kwargs: Any) -> str: # noqa: D102 return "" - def as_raw(self, *args, **kwargs): # noqa: D102 + 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): + def test_custom_helper_end_to_end(self) -> None: math = _math() math["expressions"]["total_cost"]["equations"][0]["expression"] = ( "double(flow) * cost" @@ -588,7 +611,7 @@ def test_custom_helper_end_to_end(self): model = declarative_model(math, _inputs(), {}, helpers=[_Double]) assert "total_cost" in model.expressions - def test_get_val_at_index(self, builder_with_flow): + 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] @@ -598,7 +621,7 @@ def test_get_val_at_index(self, builder_with_flow): class TestBuilder: """Model assembly from parsed math.""" - def test_overlapping_equation_masks_rejected(self): + def test_overlapping_equation_masks_rejected(self) -> None: math = _math() math["constraints"]["cap"]["equations"] = [ {"mask": "cost > 0", "expression": "flow <= cap_max"}, @@ -607,20 +630,20 @@ def test_overlapping_equation_masks_rejected(self): with pytest.raises(ValueError, match="Overlapping 'mask' conditions"): declarative_model(math, _inputs(), {}) - def test_multiple_active_objectives_rejected(self): + def test_multiple_active_objectives_rejected(self) -> None: math = _math() 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): + def test_references_are_sorted_lists(self) -> 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): + def test_dtype_coercion(self) -> None: math = _math() math["lookups"] = { "flag": {"dtype": "bool", "default": False}, @@ -635,7 +658,7 @@ def test_dtype_coercion(self): # Empty strings are coerced to missing values. assert builder.input_data["label"].isnull().sum() == 1 - def test_checks_run_without_active_variable(self): + def test_checks_run_without_active_variable(self) -> None: """Input checks must not require an `active` variable in the input data.""" math = _math() math["checks"] = { @@ -648,7 +671,7 @@ def test_checks_run_without_active_variable(self): # 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): + def test_check_raises_when_triggered_without_active(self) -> None: math = _math() math["checks"] = { "too_expensive": { @@ -660,7 +683,7 @@ def test_check_raises_when_triggered_without_active(self): with pytest.raises(ValueError, match="cost too high"): declarative_model(math, _inputs(), {}) - def test_check_warns(self, caplog): + def test_check_warns(self, caplog: pytest.LogCaptureFixture) -> None: math = _math() math["checks"] = { "pricey": {"mask": "cost > 0", "message": "prices!", "errors": "warn"} @@ -673,7 +696,7 @@ def test_check_warns(self, caplog): class TestLatexDoc: """LaTeX math documentation building.""" - def test_components_render_with_decorated_reprs(self): + def test_components_render_with_decorated_reprs(self) -> None: builder = LatexModelBuilder(_math(), _inputs(), {}).build() cap = builder.components["constraints"]["cap"] assert cap.foreach == r"\forall{} \text{n} \in \text{node}" @@ -684,7 +707,7 @@ def test_components_render_with_decorated_reprs(self): } ] - def test_variable_bounds_equation(self): + def test_variable_bounds_equation(self) -> None: math = _math() math["variables"]["flow"]["bounds"] = {"lower": 0, "upper": "cap_max"} builder = LatexModelBuilder(math, _inputs(), {}).build() @@ -694,7 +717,7 @@ def test_variable_bounds_equation(self): ) assert flow.uses == ["cap_max"] - def test_underscores_escaped_in_text_mode(self): + def test_underscores_escaped_in_text_mode(self) -> 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() @@ -703,7 +726,7 @@ def test_underscores_escaped_in_text_mode(self): # 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): + 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}" @@ -713,14 +736,14 @@ def test_escape_text_mode_escapes_content_only(self): # 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): + def test_no_unescaped_underscore_in_math_text(self) -> 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): + def test_cross_references(self) -> None: builder = LatexModelBuilder(_math(), _inputs(), {}).build() cost = builder.components["parameters"]["cost"] assert set(cost.used_in) == {"cost_plus_one", "sub_expr_test", "total_cost"} @@ -729,13 +752,13 @@ def test_cross_references(self): assert obj.uses == ["total_cost"] assert obj.extras["Sense"] == "minimise" - def test_equation_masks_render_as_if_conditions(self): + def test_equation_masks_render_as_if_conditions(self) -> None: math = _math() 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): + def test_infinity_not_wrapped_in_text(self) -> 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). @@ -745,7 +768,7 @@ def test_infinity_not_wrapped_in_text(self): assert r"\infty" in doc assert r"\text{\infty}" not in doc - def test_sub_expression_variants_produce_multiple_equations(self): + def test_sub_expression_variants_produce_multiple_equations(self) -> None: math = _math() math["expressions"]["sub_expr_test"]["sub_expressions"]["foo"] = [ {"mask": "cost > 1", "expression": "flow"}, @@ -758,7 +781,7 @@ def test_sub_expression_variants_produce_multiple_equations(self): r"\textbf{flow}_\text{n} \times 2 \times \textit{cost}_\text{n}" ) - def test_multiple_equations_render_as_one_block_with_cases(self): + def test_multiple_equations_render_as_one_block_with_cases(self) -> 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. @@ -777,25 +800,25 @@ def test_multiple_equations_render_as_one_block_with_cases(self): r"\text{if } (\neg (\textit{cost}_\text{n}\mathord{>}\text{1}))" in section ) - def test_single_equation_renders_inline_without_cases(self): + def test_single_equation_renders_inline_without_cases(self) -> 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): + def test_markdown_document_structure(self) -> 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): + def test_rst_document_structure(self) -> 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): + def test_tex_document_structure(self) -> None: doc = latex_math_doc(_math(), _inputs(), format="tex") assert r"\section{Math formulation}" in doc assert r"\begin{equation}" in doc @@ -806,7 +829,7 @@ def test_tex_document_structure(self): class TestEndToEnd: """Full builds from math definitions.""" - def test_declarative_model_end_to_end(self): + def test_declarative_model_end_to_end(self) -> None: model = declarative_model(_math(), _inputs(), {}) assert "flow" in model.variables assert "total_cost" in model.expressions @@ -815,7 +838,7 @@ def test_declarative_model_end_to_end(self): # flow is indexed over the node dimension. assert set(model.variables["flow"].dims) == {"node"} - def test_repo_math_yaml_validates_and_parses(self): + def test_repo_math_yaml_validates_and_parses(self) -> None: """The demo math.yaml at the repo root validates and every component parses.""" math_path = Path(__file__).parent.parent / "math.yaml" if not math_path.exists(): From d86b9e21ec38bbb87dab2b42dc9953f3bf092282 Mon Sep 17 00:00:00 2001 From: Bryn Pickering <17178478+brynpickering@users.noreply.github.com> Date: Mon, 27 Jul 2026 14:01:09 +0100 Subject: [PATCH 11/12] Fixes for tests --- linopy/io.py | 2 ++ linopy/objective.py | 3 +++ 2 files changed, 5 insertions(+) diff --git a/linopy/io.py b/linopy/io.py index 462fa5b8f..8fdb93386 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/objective.py b/linopy/objective.py index a51b22076..b729cefac 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 From c7df8dd665f60f76180c26a598bb36bd62f56b4b Mon Sep 17 00:00:00 2001 From: Bryn Pickering <17178478+brynpickering@users.noreply.github.com> Date: Mon, 27 Jul 2026 18:39:10 +0100 Subject: [PATCH 12/12] Add parameters/lookup dims; Update tests --- linopy/declarative/build.py | 22 ++ linopy/declarative/nodes.py | 5 + linopy/declarative/schema.py | 33 +++ test/resources/inputs.nc | Bin 0 -> 550110 bytes test/resources/math.yaml | 442 +++++++++++++++++++++++++++++++ test/test_declarative_parsing.py | 370 +++++++++++++++++--------- 6 files changed, 744 insertions(+), 128 deletions(-) create mode 100644 test/resources/inputs.nc create mode 100644 test/resources/math.yaml diff --git a/linopy/declarative/build.py b/linopy/declarative/build.py index 3fd8647de..140992d0a 100644 --- a/linopy/declarative/build.py +++ b/linopy/declarative/build.py @@ -207,6 +207,28 @@ def _check_inputs(self) -> None: 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( diff --git a/linopy/declarative/nodes.py b/linopy/declarative/nodes.py index 7b41805f5..fc5c4f463 100644 --- a/linopy/declarative/nodes.py +++ b/linopy/declarative/nodes.py @@ -333,6 +333,11 @@ def evaluate(self, ctx: Context) -> Any: # 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 diff --git a/linopy/declarative/schema.py b/linopy/declarative/schema.py index 799cdb3ad..d5a6cfe18 100644 --- a/linopy/declarative/schema.py +++ b/linopy/declarative/schema.py @@ -196,6 +196,14 @@ class ParameterDef(_MathComponent): """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"]: @@ -212,6 +220,13 @@ class LookupDef(_MathComponent): """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.""" @@ -512,6 +527,24 @@ def unique_component_names(self) -> Self: 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]]]: """ diff --git a/test/resources/inputs.nc b/test/resources/inputs.nc new file mode 100644 index 0000000000000000000000000000000000000000..a5a5e9d67de59548912e0afa0df8cb7703bf044a GIT binary patch literal 550110 zcmeEP2|!Lu`=9pGhNQG;)1pX8D3Y`yrA=HZEw7X|?UIU+ERkeMWeKga3rVQ7pe$vV zs6?WLRNDS?-t&%|m-F4r=idAO+?jhd=lRW?nK{3i^SsZ@nI+o9z=(g4*dSJJZh%}r z!&SSGm>MIXO3CPfxru=#E8l5WSvtnHoe?A>M_w!yC4B>V(jUYMScofO3BV3LHWq*q zXZHCFiFj;40B{pg9nAFS`g;4&eEr>Men1EY5jnELe}Ob#KR0hLo$+J5X#UPFOLQiv zYN=|BS?aPxQ%7A@eS+#lFa$;eU%}<_Yj>FIo6Ke#MwDrYv@XnvaMQsAUNE7b4FEnB zq69=4_{cKQGnioD=IQ1F2E*9IGXNk>F9-2lEy)v3|L-KpJ(MUt8_PRJT~A*~$nde% zi6>(NoFw%DUT*#*<-=hhLHSbe0AD|8KQ}LDn)G-zb*-^#>hOQ!HV$|j4|$ujqX*5) z#nG2uSww{qcl!8xd(eFR-JBg)_?+NrMo7TkI^7I>iCh5eX9x`G20e;%f3Kl0H*+>)yAiWB)6uBO3kp zQz!Lp>WK0aOn7;_5R^46&{#lMp#MpE5e3*gc{mrzt=0f;2f+}+A^*>DOIM7HW=SxN zRL!Un^u)k^()-*nIsQs^iM+`7VlEe%Zx$naK2uWD3Bvaaw))>2r|C>9Ye1BhXtyHp zAMq^2Tgu$%3Ax0aOm;N9-f}#p1JtY^CkL;6_k}kSzXn7u#7n^o^B~?WBR3C^d5#_d zH1Msr`&+pTq?4%H@A_ea@c|I0|E2z;D!Jz)5o)Y&Zc2YfhH^pAOa9UUFx5<@pGR`|L35fEJ+Lm!nNl~b4zpO@j2_qjQ{`LCcO#7`sPP=1IfqU{h4J)e`Kzdy}) z1$u12u#djFkMW*-g^Q8NeJ&s3SDt}0=dZxPQGLy2{3Iw+;7Szf5N(@q%;_kT@K@!8 zzXNo~9DfW2K9+$qmtU14{y2*GY7{uMP3R8QE8=Y`MSL9!oZk3h7;`ywDdJD3!1XBb z85Foa1#Uor8&cp#6u2=3ZoTD+#Z&a@K^6CUCDeSzCR;Ab27MuQQ)x@IO-TP*KZs}{2df{JO#d!fiv@V zfFk}u3Ot1ZKSY5arofL-;HeaN8U>zCfoD+QCm1+04<{+`vlMs^1%8eKKhJ~*F$|## z44k>%b1CqP6zN=J;9Lyf6WUB({*IuSAtSrXm*zr8N*FlG7k0PluX-s3UdF(YeY1k* z;o%(w@1l9pL#{Kz2Qhruv0nxaH!oKj3EW`d91P#d!_j#eyuQhhKbOPXgI>UL2L3rc zZx>gZuM=Su(_IB49MvP@^A-hun}MTzko|Rs0>4X<&Qpr`)fD(M3cQ8_e@=nFpulS> z@Hz^-o&tZ(z!9Iw{%NF$|Aqo@qDZHM2}gEwCj&?AirH@N`h)OKjBw_9?f!%NdMNJe zrNH|r@cuu@2aSGEJs;{YGV{*L=vXo1a72{um~nOroPz@Aq`R;!kG6QG6{X9L3jW!cqJwOgM@^l?g}j zi7~o?ZW>d(r3^)`OMy?P!1XBb85Foa14m^+{p-JA^qSkl+u6|{4$*&FMNeul;uCq%TR0<}h)cHt81fJ(56mBy zo9v@A!d02V2^#54Gs4F(g)1|Kt1yKV{0zKbfk{?mim${JPK+ZAEUzJx46F5vaD66O zk4c`vBu{0MiGKXR{D}VhKv|bbp3WrOG08-~ap3(#|8k&A^iK!M3z%dpCYk6j4~$Rr zlLyN4m}FZh2S0jynUjh(f;mxFYz9@DY{Z;t_8X&~8RlS$e z#!vt6t^0v*Mks6<17%`lVxUZnU<{Or5s-m0(H$Bn`!UJHh|IunVkBmu9LOXGG0AaE z@^&Vf7%3Q--c}}g8cIRmm}H`-Jup0tNhW&q1H)69g36&kTRL==F{xqPn(00qJ*r28vAlEM5we2&e;?@ENd2OaS;{uu{fI93#o{Na>_=JEjxvX{UP zMk&0GK;eXMH167JvUvvl)Ni=sVA|*#u8SC%<9^cn`zi~rqCgzSv)?Bk$6 zJ7N6M9Vrr*YhNW4(fc->M8yfx`!e@U}Q3L29S|Mw*oMhs5a;lh~EzA3|EqgtSTdYi)~!y zz_NnFTOgB1<9xG|lioOx`gh?)k2HPe>sn zpL~~yWC*uC{pzL})}}^)W0xB2JQhO@t_oV_hjD<%@Xdwq!SG@ZFd3o}z+P%2ClDSC zfEWl{MdJdCVT!}RhUqW4Krbx43^?8^Ify=jE)GI-Zt{ZDa5P{fnA!+A(t?8ol3@A; zIYE#L6O#rZ87l{awU7}}pw_*3C{TkOD1hRLc|xEaMoDmj}@h6$aZku2KMxp)CcA6C$PvT%a}+1&=_g64(h*1#oDys|qlKg%kl= zpTbl?GkjtRkSVlf40r{3kN|<3;iWKqcxj-SB&7<7&xc~*-UgR(z#JA*1lTOws0Q3& z(Zxa6HgHCLRit6y~M?&d-=S z3FOdOuQ)C_88pDBkOM;=pUtP&-_Q6%uYXR+I_b!VfE1E$7xcv7htUxP(c@o!6mvAhD=I*o=Za?AQSyObq1Ky+c7cnMR#PgCI-gn&L4(K z)8|yt#T`%-h5s>J78E7vGq30XS!E1^12EbP;Z))9U$(5b|I6zYkjut0RBViD4hY@j;b9On3bA6V4zXbccLc z5uYcV5f(ZjqvnDJ&6OTLUQFs^wLi21zcpO?y*ilCEgTu4PNBezCj7$u^G`I{|)>pEyBjcW(n6!!p{#n?!(piC( z<>%$-K6@3a|VJC(P2**xL=}ckh6&7}+V7M4}^)plq3p>&^ zT+ZHJ#OyZI*}2)90x}QSv#=9Cri2slbo6zF^MO4a ze0<%UXlfhU04MARdHXKYW4s+K?0jT2 z!0x}ohjx`18-=4xZLq7VL_|**z;Lv9cMgq*15n`L_ zevD%VFawwY%m8KpGk_Vu4E%2~K=#ctzI&|fmA+Erpy ztW*7u4mg}#M(6tTA3T}xGZpv(Gk_Vu3}6N@1DFBKz%R)FY3)D$j>&}72VVT}gu>&D z<-6G8hiz=2g-`4cH&7<_h5N2?58~#`l2?o9XC{2tw)ALD<2jguoJ?5=eU}XbqB!3g z#QD}b!2h&uT7s3k>gV_Aou2AlvK_AZTNWI03lUkfv9i7-cHM!F__|8Wqx}0u!+si( z-_;t3XC>Fo+wP|v1DML^uK17C#_$NDeE-JUI6++-S92_<{J+rA37g5Cdn5ChjKO7d~TB_;d(9m}sv{#`xynfA@ZP zKRRmmirBtwfWxC^1IWOam;uZHW&ksQ8Ndu+27W;Xu+{vnjcXXYlYUj-VbWhhcryE! zd|(MPfEmCHUgyfkztpY-1v1We0zyQfFU1^ar1{;dpmpk`ERu$-YfWau`^a%0kMI>4nA%)XIhY(AI-tr$3M;uMV6-D1b60l z_V)C2^Y?dj@}RADK%p{>P<|ITKWFa%FMkKP1367*8473YL=Jdp4t@ci#Mayow~}G- zMyfstBXvg7UwhWOzr^sD{xA8!5@rB1fEmCHUJ#5A&9aLg6Ld~&xJh*=I=~1;{{QQ@WL|jcQ?AvJ==Wd$ z>82%nP5v1Pm7nh&8DntTf~sSLx}eOgW=MaD;Eyw3@_{AH0A>I)fEmCH{E7^Gy)+aZ z-iZ94w?O=F`A9(=jr#Krl=-~2nCC26`vfb*a=ho$iSDF0n7kq05gCY_(d7`+JZwC0P=sH1^3w;valsT#S>Zl zyS7^4JwUnDis{2+9B8r#YHXHPdw=_0i|&y$x|{S-Ju~K@sBhBHXz1+gl@ltH|Hl8hK#Rb79fZiBWc!p#a@B()-a5iK-TC^aQ@~A#uND zSQ=qMQV;w6oOO4|Y+MPei>N##he9FY-0L{T140k2ghI&5f;lDhK`lnrmcxL7eBFdH z-SB0E-+(eNaWg|MgO>yEzYb;kcU4CCH7L{RVaUZ$9+-X!lU&Lq*E7j=OmZ!g{G3UC z!6XwawFc&2%_I{Gwg!gRKsor4^Wy!Wr}Ih++trL?u4!BvKz9iHl7Tc3tPhm?nUd*a zl6#osUMSNUV$jq9<$>vSGRa*`_kUsv?}l>lquGlW2j4utOE2Wja2w9`QG7S~n&-vu zQB#}hyv;fI+}xQvS6^A{I97k$8*^UH^^>Kil-1O5adMiiSLHmA#oHX#v$#LDu%)d+ zFzQoCaR?VD7Z(@cyeS?-u=PujE3&`;O+Onr9XnW%n6P?+KXUqDgWH>)3LO+{cvs__ zqi10<_yw!#dj*pxfyX&6-hTDw%)Ro7IW8N1(4`&perag%gU;sZo}?elcSA;q>qOZf zbQz8Nx&{yVF=t_ue`CfEKCgkiVWF`4HT55T=Cp&=-P*dwKbpI+&c!H^_e6vF#`{UO zqtMRQ+Jdcu(sKp5Kz6KvH*7;s)mh-A=fcao=DC`r&$%X`dAa=QFp-(4_5Q!`AJi%x z-eEaw^nd#&-7O&1Fa9SaH2QGm=?u$$zVx3_s?moz+d)h6Tm66cc7F^_N2M02@UQgh zmjGBDz^7It;HtCsR7OqKXrG-T8;F!**355jpI!%C0O)HPG#LJ1!T$jZH2c2(Y>7-m z&2!3h_JrgA>nmmJx{n@C|!WZHV#Iztdh^ z&1eP4UR(p^t=qet$O007ne9ciSRTbVrV>z6`_?W&ksQ z8Ndu+1~3D^1OueKd1}BOO~?X3{?D^FKO1inj6JB3@%BGxh<>-d$p?o{f8Lo#deAsj zaMic?0|Qls81Q2Q0)USgNOLfw`8&HTnPA}N>E;3kLlciax(1{Hu@VEO2)lt3PXFGl z|H$4%+c2R!p6oD8rLsS%K8Xv;pu0&QwK{Bne#r}#FawwY%m8KpGw>@h@K<~izYPic z`%KE8`bnJF*b;pbZ@NrY^ndwDyk|A#C-FNo!9Bu)L^-OukY?BC9J7eQRpOjZH2^mQ zwKOy`v@|r+H*}z%u{;r4jl^sO7e{|bxSuEh#6F)y0xoWzG%sQgP{2)ue>O7V_%{9X zHN=pQHVK6g-F>ggiFW%$_piHY#z8hIweLCXn~$)534@q0f{BOV;36*BlnHq?d6=+! zJGn;Lf*$4f7RCnpDIA|65tR*6$4AhI%x`^fZ(kQTFSx6x-;3u`^J#VrE@XbV^*uS~7KO}3UJ13JL;DHJdY^1pvhP9wY*Sb7uECxm}JSDkC)+BGWZ|q@D*l*XbHF-U=DCQz&a8C zE?`ZFH6hl7SQBDRh&3VBgjf?|O^7ui)`VCSVoiuOA=ZRg6Jkw>H6hl7SQBDRh&3VB zgjf?|O^7ui)`VCSVoiuOA=ZRg6Jkw>H6hl7SQBDRh&3VBgjf?|O^7ui)`VCSVoiuO zA=ZRg6Jkw>H6hl7SQBDRh&3VBgjf?|O^7ui)`VCSVoiuOA=ZRg6Jkw>H6hl7SQBDR zh&3VBgjf?|O^7ui)`VCSVoiuOA=ZRg6Jkw>H6hl7SQBDRh&3VBgjf?|O^7ui)`VCS zVoiuOA=ZRg6Jkw>H6hl7SQBDRh&3VBgjf?|O^7w&f7XQL37)6^@C45x(z^8HX#toH zCh&p@{qS5|4gtt6U{fK_y(P}9>+iLQLa>h+zzkppFawwY%)qb00BP~@{U?i$kFeneLxX@c)3^Z#7jem^OvtOz;p{0_ zNlK6bfvh<8uS^kSn(*lncd&tBfDiry4rVleXO|@t^u4`(UEI7J{b_#mT-oXAN&{AU z7naO~ED3o{9Hq_*lM#T)z_=(G12<1MmoF0__|J_2l6L5f zOiDX+b@qnrf21~SSN&USV-a<)K*8iS>mT$AzGoDeV7MiD!s(>_@B+Erwzu-B>~*RU z+|AftNAzJIGk_Vu3}6N@1OFNZ{t*i?J?h_TA?8q9h`INhKM@w<=a~(D?BsAS>eYPL zpf1sX$u7dAqm*5QMOF)Y;%8XTGy|M6sKsA_lQE=$J`y!ye@X(-kSNSYga#sUrvLy7 zzz|xlgp)K}LDs9xTqSbcUB-3>c`Ym504dm9q`?>zstd)DU@Q_BK~Wsu&2YhVaU2rkpePU2ka!M?azGu4kLi)eBT<+W z02!cx#EDQG0VW{P9E#$Am?1-V?sRb?64yad6ih;55)|dYWF%gOq6E-F;uE^4jl@p6 zI0cErxQNJ8k*G!&b&zNZMNu#fi7s?e7m1-zlm*j~m<&ZlpohdFC@O#%NEC)Iv@{q$ zrS>rcF)}fj0c22x!&(4t9aC7?QG*!l;vMAW;P2*1b8w<9@%E+Z2e8oZWMLo0NP&mu z;1}TO;OXWSwGM^JGr|VZmMn2|cB6SYuW*e;feMU3?j_!y4$j_w{-t|SkRl55a9ctn za@9#iL8u||dpi2M!cTl2^jnKhp>Wi)#J>o4@NuI#(}LXmXb#>!{tGUl*r?%&d>MOr zfTN4AqraQ;>-#7sYLoo%0UbWGQG5@DqejY0ga`Q0xw`!bg`zgg?HSs5}{`(%HAk>zJ!XRHqPag*l_(9U2p8H)`P3VY0;(~DAt(})fAs2?0P%xGIM~v#!wg^sFawwY%m8NK z7iHid8AISrCJh(VJO1~4^BbK7doa*_WA2&O(^U(JF$A*be<7c2aBE5LC$3l^5@g>k z^^)*qmVIeKqW%F9MmV*lv))o!l%iymul`=b7Ufr*U;s0K8Ndu+1~3D^Is^ZRMY$*L z-)d0~IYwzwZnU^APu4~T3k!X~mlZBZSjtPbKs3_FER*)bL1UxHN0#kRT@%t`f|=5*zO;0CF3Ft%`!pq zwYO4bR2HWgip1o)-$Qu)UCRQ9o6A}T4;!u&{*Tz)*DsR< zF^}T@`YW}cQrq0?0*6QcL3{iuHaBWD7noX6E|h!cp8H~Y?;qLRpC@y}dkNQ2v0RJF z1{XoY>15Qe{&Ez22{C_7!1IW|I)fEoCe82C|( zG>EZ%1+s*Gc*b+_JJQO*7Ab621QlrPnwZKwP(I6Oc2#magMPP2x!1yNYoPo2GoHWu zYs+v~+o{BCc@}6xekAN;$RfSk4_$SYRB8<<3`x+DZzT7tt z?&i1Z(VhktcC@jefOjCx*UQaysXu+2x*%n85(Epo2BTea`Z&0Gx%tx+Q2^R>klV+> z(=ph=C*Zvj3PKwY5FP=jXY zFd7!H5clyuy#`@-_^ar3HEf4h_<2E$g zu$dvlF-E#$#?d!oW}Lo#8oYxUA4!pp6a_9#fy+?fvJ|)+1y28YfMLw#{CqwJ<37X# zIvj&JJ{k&T#?eqHGmgd(aZAgI|QKW+g8JW{TW)d@w z^b0d?&6NJq6)wKst~4(jit@~3;KVcq!qGfv&i=me)E560blH~v7w-7X@X_O?x^6{p? zeJJqd6u2(~NBN@n1X=gYxIabwl@xe51s*|xucpA)P~g!N_(lqR69vAR0^dS`$57x~ zDey!JJc$BNrNA#!;8!T{A`1K(1zt>nmr&rP44k<=mr>x?Dbl$|5&u2~PON4p9CNLL@DCJtI|VMuLIxn-XCwtangUm*z{gSGY81FS1wNhv*Py^BP~e&r z_*4p9hXS8Lf$LM?CKUK=3Y;`Tf49R>z09GAZ$^QeQ{Zzca5QGX+#U!kpKu6gFD2Uv zvi1><%E>r=@~ikEOzEKbE17VVekc=;;xo>*|0?}3iulCJrEir#f(b|IuV%tgdDbxD zC_d5q_*VKvFJvH&;xi8C`>OowDdICuY5gibu}9mt(vN1sQTaDA;ix>Dm~a$-GZT*D zZ(+hweDr;txjmq5;}DMGZ(~XarO!B@@T=$DP7$9mO!Za#9Tf5788{b7AMg+Gb#n8f zft?JT%HG_?lpeAxcQbH7MtV+;{{A%I71DnG-oB2mG_Z#e&d&(<@OE^O_FGDGadGo< zg;)Y3oYCzBC|#5Q<50%W{ZT(}4@bh*Ok~`LerJN*yj+MZk{CE~O}Jzxc^{L!pCP0A zLHQqG;8c8&sRM9}aUsp{HCRyy>Hn|bRP#pqzI>dor2EJKWQ5>(BVSR30n7kq05gCY zzzkppFay6715|b+8dsot*hy?(8p_yiq+$x=m;uZHW&ksQ8Ndu+27V0&f*-xT2ETc! z@v*8UI`8#k{b)ViSDtBX_W$vdPFuL-M_tF6^gTc5k~)Bpi);s#%`HknSQYvh&j=?b zSrQJ7d6N?Uiy6QSU0A>I)fEmCH{3!-zSkE*A z904q2Cxc;do}dJHd;fJH_`pY8%A|M|fmk@3ZzM1rO{S*EmQBvWQ@A^v{HUjt=ugg) z9scl;;*~#@Cyw<0hk+T^rbd9{hJ4}*;3xoqCYX_$zY^SrQ#3`u{56W9pcvM_BpBwL z8xDkFtB?jdg)$L90nU|_02>u5BLJ~Ymo(5?AiWlh-m$%5`%E2F^< z7(o=Ae-OJ71i{3_K)d6mjlcqqNXdZvwQQS!FT4~1+-Hw%2AVL7(csNSjV-_!k|YII zN2_lIfouSb1?yKvZv~BT^67A(H=-*J6u~SMz>sODcYr#W!U*8BH+>fnhV4@h*trbe z4JN_jh=5sahKb-9%vcmW*feP$AZFi>1UpRA4**q|m=s_;U3Cz2!o`jf;G)5YgFq7| zCIQy)+&BbY!7|7I{_gEZzy$hJ-05=21Qjp^STfegV_*!7AP(|&Z9WZXFgGbMzjEdo z5Din12F(t+XMs6P7e){rodc5TMQ1gw&H+h~UooIMIrAJC1)oh4?0KGb0r0}h5x{24 z%uB$XUb2y%8!rJ{h)RLHtV@@`AUf;aRlWrv8(t~_nYJ}(G3F08DRRxR%J=W<$C#|` z!N~WA*UYb;OzM$;Y|Z?4@35SBffmeo>YHwbeN;WJKblQDXAJq^sN(Rf2*d>TF$0(Z z%m8KpGk_Vu4E&l5ke0W}cU#`8&cN;mbZZ&3{IKQy*6g3QybXw5b$-UF_XHISOJdN* z2!ynDMaqxvr&MP`j%c9$PJWo-9Q|L0;h%fPBC)cWXvJt>BebHK%Um~;thK*+Pckca zvZjdXGO+Uw{E=fGv^qZS$0zeZWtJ8Ndu+1~3Dd0n7kq;M)xR=!yYb#&KT%^}424@8OOK(EW2e3jDNn zP2asY!85;--Y=@9yX4E6+We80@Bfh%Hgi45TKV^_usL#(a)phV%&dQAb;nNse_w6< z+|?b#;j4QJ#cqQc(}l|ZkvD3m)xWO}$u(=!aw*@awu?Jl$lBohPwS=lh!s17;VK?d zJ8}aKF5;4ne#Gx;?xBPZWn-}@A$2L*L>sLNWdCXc9-5EGr$@gNt{&neRuDOu(fplV zmQ2w1_V$H`U^@EK{OGx|)6MtqJWbeNOfTsB6+_U#eq)d+T0&xEZt2LFy&32jkts_Eu!@o=$?>zLM z`NE-?0n7kq05gCYzzkpp{?{1@eiWJ%?b7?e4(01|I>D;UuoE`QuL(Dd}xi@_N z^nELBx3)^Mu!x@viO94)_o&QfMMFkaUcZ!uhDnXPrmg<@3x!o8b~N>h{VQ!1Svgx? zkIKGO*Eq`d?YWTT%q8?HUH%HVj%07Ai)9OrZ zyEErH`a@WB-Yq&aZ)EQJObxr6FBalOO?Y7vM`6;G(-t)$Yu7$-BLpriw zD&=Rq+kbhqYuLK;2|VR%$LmHvGmI(P<<~B4yZ79ujji)eB;W7!j~43Q>Q`0VbyhOV zsWLCKnkDDK>w7)9@pf}+m5b-gad|t8Xi*C9-jb<%u69oLMnCyy_qJxv?RoANsIf8{nan-^tTKmrWPjxyXRC@osGw=KBXFAPU%%A37IGpuygsrgT zW)s(oa!t3YR@$=dFn$s+HEX2p*4?JN*e*#3a9>egld34;ajdAyBvMN&*eSB#Wm|v0 zgVS_<#r&BH=U?rgE)`kvn&r*AMeTe_1*s|Lwo8O>%M?5xvOe>1M_;~3zTop}i8GM} z=Po?cyXtMaHd)?sIF|%%YV^8rt=8Pqx;+v;Z@VLsIUlP1anOMiR}w?^cL=h)xsv&?Bm1hb!`_DAr9tiABfqrHi?g60xn=)?vGNtwOwQ-e!KteD!`b2Fh>|WeMJM&yCq)wNAq%~`M;N||2wJ{@D zgcI||T&s$+t1k!|ck--n4UOyL>8;M)Ip;U3zY>j@ytJy!LEzZf%VSJ=&G zXSE)vE=W_b?O%N0W=r?`rmJ&beu%S@ALH|&MEH1}d7DZ0E<4^0Yjag!#-uCMva|bZ zXwjy1BuD8~u8l2H@*47D^B9jqLr?s+x>K|;%v4dM%HPA`YI^=+k>mks-` zPCEPIp?*ubk|HmU%2AJt>0_6w@0vVFZf~`v--6oC+2K~j`}!*yeb4J_?uha6Q0|h` zwQUxCCkbj*dCbBD)~}NCi3-;&IZ?H!*kQep;sFDDoevw{)>v?)9+ax9sai7i#mc!k zL(a?REGe^oGQ;|asd0aIWUxw1@{Y)UO~<@(Q!aI{PO)$)(9UOZC|i@?-SuHX)INpf zRU0zzzS<~ap|?I(Fz0yw%EaH^`emz%hV!hKVi)4&+SAKfl($v8Tjh=Crd`u!TzbOd zXJ=kMJ=9-L)`5SXQ&E)Nh_8U4k zg{su>%zq@|6tXd8)mXpI*~2=6My`|PRXN7xVDB@(b6(U;TbZ(#anDnCFVe}(*ge+l zB2Uy5)9hXmKOXt;(6W5p*VcJ6V+4gebyvL3_ZzX5zf9xqk(j>FyQLOEySokNCBNZX z_3)^hvUB)bgNL?8Y}-QTPgy)Z-qz})zv{m57cHgFZZ?HW@AEEjUt{Si9Mn`>bKNeI zYh6pLq|bVdr7eoX!XLi2J1na6Vr856xSrFB4?nmbo|_buAy7K{WZCuLUX7OSMNNGR zj59A(ls577WO9|7cAt2bvUt9Q;@RzU=J8%u-Q(}yIFv<^)!wyUP%GHx>gy9vWR6WL z>i1O-9RKJEzh>mFq%qT;yz03f+~l>SY9zn=#j7Ieb#0UTa=5+JlhZ``-rQ)n&mJFFIZsM!R*k;G_Hl8q@7=1H&DXiRRaUZR zZ>^H^vCF{g@}!WOekrGVDaoFdMYm7Cdm--7x@c_lm8UhQ3?3S!&bD}bd*9wG(?1Q3 zSJAuoQs8>;&dfy$ic$rFjaR)M74XDa_s`MZ%9fOHN`6SccjuKA^RMMp^6s1Rab@}Q zo=I~Tl?_=KU|OznyV_w%2TZQEXjtB-Uzna;^Ibi-XIqI*SWMZtdX; z@a-@^{Ul#PzJIrcNZU!?b!jIu*?ZO)SzLGFUUV!_L#xlPDU{>VV^NMC=B8?~W$# zMX8sDHG5nM**7sagze?d1LNmsq_Bs1o~+`{%Hft+#eS(Z?_|{quK2rrRwlBl{p<7X zMxEW6?kD&rNu{*>*r|gHq(6FOH!eOEAfU4J_0mBN53lGxcQ&s*9DdSxnZ(iQO9JjJ zHWQ7#GdSF=@Nsdbzz1WqUW@C;?>`NG>gCbT?f%ODu(!2OTkO$y0#|STCe+})ooiv{ zZ%5Now)T5xzZNJl;tNr}5OE8xKqKY7KNouWs4);ySXFsfB( z!FZ#j!H=5jwhG&gEqd|h#M|P!mvao-f@1D_zI2aBm(gA$U8-d6cjA$9z`f~9w4Yh$ zX)Rr)Yo!x^S#^2cy0e4yrUhv8k5j8w@LxME%X8ZKt}*k5_hy{9-Pc~a{#c1|cx^z= z^myYd8=56|s_gF6d4Juecl{8JGtuuV=9=g_AJKAmy_S${DX=Q@NZ~!r4^K_~w=JzY zv*Q#Cm&D7D<5y}9IiePEqia}+bcWsp+pDAPHD@heCQIX;(#g5cer;^M(me0(6vG>n7Tt-yQ7x|N4>z^A$=?gy{2(kLwxD#xx`$U% z%qq=NvO6Q1mAUd8jq@IDf9hSpm!V}5TNr%3L}y{^uBmkw9ScA3O_hCA>Mp3!Em^4{ zykXtr&e&JriO}F{3U9YXo++_N3<=);c5%peI=E)DM zNN-M4?K*zhyYr=7Q@zqO9e%5$i^s36EI(Pz|6#azTEn7|dd(Y;tY5Wzih@?~@f{<2 z)_jz&5Wl!Z`GZ=I*0pHC=(>*92CWatN1RJT_PY#HKPSFv)5M%2r95v7&x+!C^))9B zEcN5nQp}tAy6Dc#Pmxc>`1G~(v&Gn|9TY>{YaR2oXZa63B^5m@XShumc;>+Q@z!2p z{=nUR6*rD)3~!Da7aevXsm3l@9u-&^?jIoG23-nLHa>nwI2-xamt-Fc&iI~BGQmgrvCcBTAO{<4S9cv`l& zg&M1*`l*jt)^yoj@wu-^%ktNiiOz;9Nj=MBl~e|$r0dD|u(!{!-*!adO}z3OijhI3E#feOUolF zRUw424$W-$EF1N)ACk4%(v?sFD6xCk!Ag1i&<8O`FIcZ`MC*MRXzQk1>r%*x1?Lj zZ{HdbQ!xF*Q)$bU9%^Dz&kiMziwoXBm1sQ4Lca6zGm9fv4`)>Y;fYg*U6)rzbDHFY(7_5 zlAG8&I6iND_37x$E3b0pOAFgCr!3FP5lWe$_9<8=qbRblW=iE%y6?)p6iKQ4VAomtY>Mq1Zow_EwH!Ut~OE1Pk`fQT~Ba~_UtjU!gg46IA7jX)ii_C##q5s-p@K!^NMtUz^d&koi;;1 zT~Rw(H@@9t(^NNw@RhNxz8Y3CZL|q{>Rdm(@QIgi*sZ!_YSC!(J=ZQ@6~8)l-a)0e z^A^0y7V^#%Z60OcG4{c<-CD=im)eg1%^{-9wB0l)ee9k?@|_}l9a?GDtr`Nli(En` zid%cX;8BheZE2Q|;z%0bWmjIae*d=Ys52flcNaB|iA&r5Y`0cJO-EdsF6WS^HfbUo z(}Y%BZOk)k7k1LzwqT}0`BB?$m3KE=bMFW( z%lII_51A|0r<<`JemQiM5PY{ch(8M{`{g)h@kNs}9po zGT#0^{B4q{2!D!FMMu=$;3)^@A4)5)O41p}YBfJfIk`J#PqmWzYO}G~t@ac4?7Tr6 zYJX<6sjv9iQ4J%esV}>GNtP>G@(U9#@Xwu?kzGK$N7G^ z_`9RVu^O^tZ}o}CYKE=fx{80O?f5+Pyv*J?rWG%yY>Z%?bo`w}K$vgrDPF;ZS?4B( ziLJk&n@~Of$c(|!Q&&V?>re<8l+ko@i&8fjx2%Pu7 z8_;;=h*6x`n)afWoaCmy6;-@DGG3jo)xV>DUVqxgB}Yyj4;4w_f3s-Lc+GC*sZNuI z2wxqsKi|eW%XhB-?sU5#=hmLkWBaFTwKug#nT_H-S-ukC!ac&aPg9z{-p4>57MN5ka%IutIR8>QLZp|;L{ z-R`UA4nCb#VfUsaw@b+0oulyBb#l-C)b}|Z53esQoH1yAE7y?8*Pe&VNyvX1^3iv+ zjb}ioO+e4#laUA8a`e`h))jx`a5S==e&b`A`^kJ8$xB1@-mTbLT_VuZ5jKAL=GRF} ztz$FqsNK7Nyp3(Kv+%ap$0HBAN*)PWy)xfbr{h$cZc9e6^YT^6zbPLR@b_08napZZFPCoG_9-f!`t{{bD=qKuVMK zTAE9Ze8pwASYbLCZ13&56+x^JZ^p$e?ecn zmByl^#FJYkZ`7|1C~ZCSe2HSEeb^i6_TKefeTR}0!jf9c#15P68#nT`wd|xT7qm79 z4eC9}vS)_-feQkIOyp0-hX+J>FS_2R8nJ3sZJMcX^^9tR56W-yX7?r)*ywy3y-RXo z>CB?vI17FgvhXel-tIWwvC%s2_$QB)9IiCQjI+n9o6}Ej6xuSAGxkGQ!{U@1*VHl{ z)3#i_-JX8ruo+FK#3jG@^|`16>hFsD5;7!rq;0dR_T%68WV+3Q$h0|5BIOyUT6_6y zZ&%29sZFt3_v9&u9j|SV{Nr?y>IqeadtZrf(h{D?W3X}Vn6`_NJzZPmw?^z9QE4nR zLTgd$0qq#weUr+To6CNXu2^!qt~y!Sd6Hj7&!|(??-B=pTHQQ(jq-#ANuO3ePPbep zx@DzM=Xma$&SJ0p>{l$`5Lp?%=atv>1uup+HS#yE<1R9c((0S^Xz>e!OGhW!Sa`Z0 zv1yEsot*m2dCDT=k_&#(tF+|Lh}dW4g1p3D&(+2gj03cchmR@F;Sc=Ks=j{Rb+co# zXAkNntoPaH+Zc1fOy};6SK;s z>vr|8Ij7Ajd#`A=NI{P7V%L#AA9h&@WK_O5C9+tnBJ$v==UfL%Y&7hqO;0zDn%`_W zb@`xnW7X_c>lM=S>hy~DiL39*a5jzN46w}+CQ8> zM^U zBU5uFrunZ6wH7<(&S~GiZc9n->Z?f^&MB`ehRJsLUwPrZUGt?%X0+wB%I9mXPB_Rp zWe(3erE~8AZG6mqTM&GS8ni{CKr}^*OSM#qIznr z^fC_TEfPIi%~7BjeR5F6l2?K(1~yy{4^yvXb^~H9f83u}@NUll^9^U*s%T9w)kF z)SEmLv%WptPxqJGm-U1Ne8{|=l_t}buhyks?T}b+ZOD~aUNqclPNM7!@q($R-F@yy zUn~$%+`PEH&Gd*h`?JYo^xGCXYIL_QjosQhtXs9c+;qjxd1fn`1NtVG-EbFscuDiT ztZ|-O-;ucJWm=t~>qBIOL(i9{zF1ZoGk&;Rs8Uo+x&74c%^}OOS|a^jZMn1?@6BxN zy{$CGueG5_V5iC959{Mnj3Q}yt()b|vW16N?ozQ>woh~A1QFe&18tW&E`3rp>$x{k zbf>}DM~BALB_)LJ$~}drFtV_CZiD@b>cepbAHr$jo1BV$mVBDBFEFUnZBokB z#yCx@%iYDd_iDGjwlxb|S(r9B#I^2bY+lWjecB^c8t(t5m2g{G=xJODZ@cx=r;nn0 z*Bj+j&AM%s)^t_<&1B)cP6vBM~oo{&mY>hzpZHlsA-qD z9zT6vS@6=!;N2H>Ego!XX%e~os3^-cY>QKR%DQI_i_h+_D3~uV?2u7weeOQg){`Ih ziR{R=SzK#tRyFZ@X@gMW?l9p(9^Yw6Z<=nfDhfSLt$nvWAV_H|cR+^ymE&Wg_#2ed z_eC!2s?y5wH9By=VehC5o2Ja_)7(9e$2G4z@rj(rdpD~Fj?C0?3EEHQUASW_xysH% zO{i+z$a#aC{702(uk4FWnYmRueCk%kk-qQZO~+`K?kev|Sp6cpDR{SE)Wxdhv8CI! zr@cr&`tqvaq3gld9BpbPF3Ge`86Fy!8lyUeuPey7E3vk+P$TW&M$6`S{~q_4$W)y~ zlQ2o^l^c4bH}q{kwd-->B7U`5H$o5Z@((&U-B0m?t9HV|u+A%zFJ6!4iuyP?-bXj; zXz^gfNXI)4Q=a5rD99+B{c^;Iz9CEc&N~g0OETHHVAb=^TT2dlXvQjswI5CCZ4kAw zT;}E5maC%iuD0{FkbugZz&7s%ACHONut+|9v!Uko;3&h16E9SZXtPNCG_iDXj%4>! zpU!JCb3d@(XvnAyub7-1*R*u1v4zZUr|xb^T&427$8%GGJ@5FA=5)uQ{jS#pf<7GI zqp;V^H7wP-sUmRRib}`(IyZy6TVDh=u3cRjCL>@q*?q3zcmr4EH8O*4*Kg^1pQQZh zP{Hg+S2mhkn!gS3HPkJrVSADjAjO&Mx9kBsts%_Jn)|n@H7y%VdJc28pEii_N$?ej z>$~h|)WBC!ITjrX#Yi}#0 zsg|*(zF64gb54pc*PhpE>_OR@X~xf%zIbpfJ~~H0>BeezU9Ty7MpfSSyEL(DlE6-N zyRP`-8Ifl3+6B*h*t%xxy;N*n`0PZdkXfo;W17(BT}2^l<^&vc4&P#t)SmN{L+0f3 zWg-vU(piJFcgHQxij7cwSLn{;@{zqh+&){h**b`$HHz!ufnnFW3)4GtC%2n=zDq02 zm(nw|++c2)e zKD*?iV7l;D$M}sgZ#fL};_Bd7o_9-<-OhV6s)X~BOz+&)Tdl>puj8(p&JX=CfQU@;2b-FjC0M9a}g%PpO@wl#|96W?%m zjh#g+U0}AwuVC?f*4}HJ!QJW^NlI1ion?CQs+?mF)~pRkf8?~=ZNH^S&Xn#cg-e4i z)N=IwM`x>??kri*%ww!2;CNA~?aq)8j`+?Gg2uFGjn-HDCrb1VYiP-iixV1l*V}EF z0%V_@;Slw?fBW(EJ-=B7n@#k4SM+RnjehH(y^pjPoQ$v-uD>B?@H?f| zoC)Uj&*licUsYbW=Ztj8v`cal0h;=)N3vt?=ebmNPS4uQ_x>oi**zd}~vN?a|R+kyY1uU4U0Ga@`!)`jyecqH{uG@^=KPeH4FJ!YTNuRZ^}czanYl z>9diR=IWV3O*2%7X1iYWIDAvm-~Dcqg7TWUDFBo3dHy=Y~Nc7C)ornn4QYG@}a-j$H?l{iJ=|_6*6;IEFWv^yd(Q| z?ff*co@r|r1%v*Cch>67+6gsc$44b?NZ(*t#q#dAd+C~oHnp8SQ5V?sgy-Y7!uPGC zGo%7;l#W`+-~YIv^q%bEgFWkuO#=d7N~Ycky=G7*mUUu%Pi&b0XU6&9REhlYS?=1Y zd+u+SxwFS@+l!EeNlLRdP0vVbhXl3V7^F5)pJp(k*#X)&8g{3Z-3|2vL-Ma?UK3C` z>UD8u*gECp65hxw#_2-F9a?>sS1VhS=4$S2k&I2VSXPzZu-$A<&&J#8&Wq2#pFBF3 z>&QNH9lh3AtGEsMW8}~B4qb8cmYi6MUi8WNH*aSfa#cqdH?ahqDj6$vV!V0guD<3W zktZHMA6H*`_w~}*T|2IJ#1{k|6OBnWG>c)&n|*MmYkIZ8&XmH7u5Hb4q(g*0)=DJ$ z94$MODJdUvpmkTYT2*x6j44wUO_n#y#PrD~m}Wa{SyO5q@z`q(kAljxIldEyS-ii* zo@;0`#h~6tcjqw9Gu6>2Z8r-{jqu1{Tc)tbZsClgzJRF13tZlAN=j+zW>-_RAARHC zD%xb_k0~_{+)v&JXZ1|$w7mJ8)z_Z3w^a4$WZ`*x&$G84>C2dPc%jQMlbhE26z)&E zZZ~MVu0mV0;Sjgx&i8{;tVdedsSXxTbUt!SdqD@k;t_AP4xX&giAnKCn-cb04^b?; zq!Ap?Zu{y<2M39v)`xY6D_D zlMgrX`DNRW4=Ehk);!_lH4qw_Wtnzpcyn{zum$ylw!QK0>{STNKU~m1(r0}4_~YjU zhR7dY?tRrIf!IOw*bc9NJ@k>1Yy6I4C(yThv9HaUJCewqf~?F+lUskmsI!?f(lv+Lw7`K7uBFwhKDpWk$L z^{TzbZYHvU`zT&oVRyP&XLw6qT#3~R){+5rkQ0jV=t?^MHh^iFp-YG_5Yg6sV4 z68I9Y;-j+T&#TRzH$+po4zt`&7S{;jbP#LU%qrw(3*=hqU zC@<`*BnY!dosH_`pmKoGEEm9_I#VM~z~Zkf{r;+RbcXPOHPqeoO z@&BG=gHDE#>HWQIn0a5SiJMrm9p2too&6Q&R?5MvF68wTJTVhrp9(uHVH6u}zrltn zjDyK6Kq;!olN%}O#Quj){*Y5XdHAlY*fx~P;WZ-pZKvwp;OOHez{Y=Ief^1P4vuE zEVq@TA4t7=l-wh~lCY1=`3Lb{z$my{nFugSU=eD+MeN`V`lb``zXZkVw;5DU0DS_< z8nIaD6gBZs7;S>Wc2TnSn9M$zKF$nsV=*FhiF%^+lBtE{jBo-OF^$Yr9hELoy))!0 zcVwAbTux#oP-RJlmPM|ZN7pElPZQKrtEG){7l^;y#(v3sZzQLh&XIc(6qAlNyBh^s zEJ4MoNpyF@!}frp395JE?eACS)9L6Nm9gY?ctnC?rE>n@FFw5=*3<^1nGc5j4l>Q) z4mrF3lwmRn{JS`iW&`$1grlz|_7hYUCRs?ntu<>~2j6{=k@Jwxo3n$>oqDS?b;UEF zUI?FcpB0ePNBTm4n7vXDBs($<1C_8~vko#dM(cnfWR@9qJ zr_Ttdo6M8}Hw$DgGwFoXV}F-m$pn?pyBk)Tzzj;#MJ|Z07o~oQVJvryaVn8??%0fn z%_n}_@~!D$XasjW=u`>mV0S+9MQEKQY=}rJyAQHNZE^z52t?nkR6m$0&;St67$=_{!YgB`^jo= zVCM&zpPY~<0U++FdU_Qvbv&868W1d(7w3v8baqn$*(m7%&k zeaw1;K#M};qv>$bePDZS5Ism$a=)9ZBO6q1pq@84fGYGT{5ggAXzo-+3-RZ+^q!)y zuB%^hXb#?QBzk(| zVV9lyV*#^Y4VTM}!Qo(bI@zi+(Q{U@$lpvo9f%g$1a{L+>8INmPw@1#(=~b(tG$oD z`wyA3A1r4k>wSS4eB{JOJu_YgI9Z5T?M4;~0?`B2tfXp_Z1oAgwueZ0242z#?}Lf%c%yUaCq?_RAE(G*o!H-0 z_O~B=V=LeBq@%Rm!5Y)BT6aEMh>uIBW;82$(a|T*5oIG`NWG|p4^c171!OIo;-;z>}E79qa-o2kX6QD-5J=gj-#)|G5eA*#+Ufk zT<$#$)nEzQdXPfV%|ZJ@_41-`4CmkVgb5D8#}iZ+`DN3^%f?m}In~qxY+681vz<~1gq3bJk(>=W zW>OnOf=K=0O+(o8R;ucK)nwuzvgBCeBZG=6797e4n@dpd0uhaSzNtTg`P*a|k90w6 zGD?E`Zj9|%qZi2a8o4nBn?BB+q1^Gc)48}0YfOjBS0gJl1tWGdat&ucGL6Q)7Iknf zJYo#i9U=Zx|S9uY^bm<8Vd0iH_Umsuq_^=di_!J96odPo9AmrQ;#Ktuoks1{A$bWjGEWTF!o~b2?|o8Se*EMYlsgEyfH&v2LI`wKmph zk!L~grDUs_*h@HH8)ESlhCR#{s;5AlJ6&rOV)+zEflhu_EHVZ!7S)Tv#0 zi^nNG4FaVK!h+-sK{Iwzo^$Qz)YHohyGf*v*bR)3{5ug`OHgNtrIYk1pT0-$VLtP_ zM5izsJ$4v(+%&Tiq)wnul14UI0IFo7#JsLp{Sc>Kc!}|Hb*{t2)=Va*)#mjg3~X>p?k9&=bWl*BVzrpCA}_8+DdRRLTBk z&Iw>nBjz%Q=*~~QkXLCL^43`CR8KHI6vs{nQ439Dl_gYfOrDhA;wVpFF~5n-uOz!F zELWaic$X_VSyRDqxXSoiAid~7_|CYH025rU@`t=rR(>C6iQ!01;L3SnP}Skt!_~d@on6>vD&gMb zs4%kq5PD;`)mgt36Z`r+!+By`bdZBYSx<7FAjn~-eye3PvYzPU7wDm-(A8S5R8GN? zCmc(@hOaGUlvbeV9oS)-Qx(NHJ|rEJiR7cI{Hw>Ql*Tw-BD?yesoArLP&sLL38E=XNk)rwDad=&bdOyuIr@l-!p532JA|3pRRLphN`QG}*H}qTUen+8r z7Z%o>NG?YVo_D+?*kM|E@wK1voaa$VQc!}#pYJ;3?KS?l1zsEj1Kq=J?N?b+_^F(9 zk~4uHnto3qzBirSny#we_Kwxso6giY>XhZ!N=~|8S6oH-U$p55NH(1a<5|PV)3^#z z@0-hSa<^8pdh%T%Aj4&9h)?k*;kNUfnRYdvxCtHZ2s^l*Y}!b@fuXg7Dsokoe%@TH zoDDYSXHN32oN#lm&YVTLcP0C~l$AHct3=CfX{t8;eCWT?qV9my-(Z8qM0b>fFKJGE zdZ;+R@~40ytl5&Nx*hS*_jE(hLI%)sb9uD>99UIpWVjb)fpln4OJr)9^2pj ztu1+T)e+;JNnm22$`GTO-*QJ~y6*6``HYuE6@8X?&kOUA*m}$HI>{*)*}>e}ZxVHU zZxx?uj%pw4jCTdabq7jTAI9qhUI{Cm?dU^ejs5oUW`J|BhZV$XJJki4qV|zfHqQ9J#W9wVZQPm&ztO$pYcw?<}%b1`o*1kVT;M^ zTk!F+_(mflE?J%U4snpbq_Jc+D9{0KFU83F6=VFt@#jv4a}7Yh6215Y_YP1QTa?t> zPN!inRZ3}g=s0u!3d@VCalk>n3r6E^MU^lJ8*F9IJ1H+d=2X@*O^={1eWvd4rs{ZX zvf`rBpZU|#7+=Bm=fU%i;NJ~6^L&T3BsqID3R|_r*9S7wkwnT3N9C;G_|ScL#ClMC zG+A&V-}bY^^(s5{<&&%^7TYcYIa~AnGQYA;J@c!{N1ez9q9jaGZ!R6}?4_KqlQW;2 z@z{|>UK*B?liKpEhdjw2;ApDru)rkhkFsby&#*J{ZaF#sSI4j_ncs5IR-PIR1PP}r zMNUw@j8SoheCE_|)EgT=xWel~1aQ278@ zlAo;6OTA%Np4tj^=G>C?_JPTK1Z!y2`I$ z<}Byj)?>Bqpv4fCXXMGmlTJ;1h=@OpzxHDn{$_p|j(*z9;m_yD-335_;jmbFGD6ihJr$l$Z|;8<_~#OQy$H0#-0Y#eR3g1@xOAs$#1**x#qw zqsOpRA14;15=k`b>REE$0-i;g21k!k&o>rv^vN>DrUN*WKM0?kh?dbw>ErtxJ*KyNn;J99BY43+GD$V~wxGZ=XT0L3o{`RIB<%k> zih7i)0xLTxFZ%dt?)a9<^(Bx-YR=gX{ypL_-`0%xHM!ywSlBphCi(Yo!_9Q%FDDyn zlA*ez4>jgfc|(aXebE#)7!I7k$)8Q&T0IamQ1w^kltrY&wDvPz2;UV8pWVr)b)DGi z;@Iyg?l_DsX0S7Ij`kmCF7o`AREu9QUOf155>)#amI`rH=Oib72b%RxfFF#7a|Nlh zm-0TI%Z`#Tgca>X#bu|So%;L~5hnnJHQAeJ2aa-?Ew2Lh{iy<~N>CgB`syk71D-pZmeb=72Bq zRHN+g{alrEkpFS+C=c@(0&{!R@I>9){?y@cn~bs@Mx>m>$D>r0koG_MMhC+l?x5x! zQYurL!-C}GK#HSQT*ZD($&w9??x zI`28wZ#9$OGT3EeTF$u*;Vj((s?-G4Vf)9)LLE(wEzf-hlNY+FbI2#0{PvB*Pdgd5 zz8!WwPF0oS&%+!QLX^}jQ+GEaFRWGPT8cQaFG^rOXXLK1oSAs|7Y4_*|H&^kz3Xsb zIa9eEWc|_Eg#u31QO?1&ci7*@@MoEQs#66D+SWJqzI3X4!Ohw!zmIYDa<{2EbWQg* zs3H~k%g#C`ICx*l@HNRK;p+E$K5$s>9!IrbXiz{fClxhsgE|W_-zYYk^$jw5vE-=N zncrj;3)37+c6KzfamFW`Q3u2;RGI5|nArJYN13R`x+k%+F~n~kM++b8psk?Kqp-+1 z#OD`IFX4zYUOp$=??AcB0vkIzt8C)r^DB-Yj35`3aw1+-jk^xB^nSm#C0XbV;Z!@gmhSQPgJs^0G$FZsjl)=f@a&QmDA+KFVD>5q;NRpa~1lix}?TCPN1 zX$O_V4bw{D{PT46UekX&_wI5c?_m??>p_dks)N$NxMpne{)`ahodermC`bnL0}EWi8XMknVbQP7Iz$b$f=fAi*XzJRVU`~JMz2bW0>dwdwjGvyO9-OZHYK4=*;vMV0 zVwhkUcHW>W`cqDnt#rnlXYi#I2%M$t{FJk+e>py0(BSxOJocPI$5@Bk_H?Y#&*?tA z3o=}A7-&zUW$2!l?fl4_RRfjwdfd5s#K{42$3%8D4*#2Psxf`4u8xz{(hWyn&N-uT zN^8k93XpztZW-rqqhLO{csgA1oEA#9JI1e4^Tqy-p@6=qR5(3}2fKr!J}dYKpVJWu43<81W%Evzqyp za8%PLogBZ(@tm`am&ST-IQ(jtQyK1ezAXuaC=Opc$4*OKA!pUbII2c4f2qU2AI#rw zaaF}pan?CiOjpAwDl&33wR4u*U#a)Mck00(xML$(vJaJ9b3V;6JX_md&b z7yPppTC$v|5G7yU^q63J1X=K~Q1x~(S&=-YR>g_wjs~-{$Rw{*<)^8>!3+o6<(yz4 z!_u^zdJdgsnOw@QO#B+J7+1a&#R}! zegQ^94yA(?0!CyhCF>!_|B9G;ItvwS39*`uHW#6?~v*vXoRXO zgw@Y+FmV@m+)w>`4HmltlJNjgW3HED(ki6{Um%S zLUqWyIP)v(WV#eG=uqZ2i1`h1>?hwLzZR6k40?f*RlxHT zh;_%%tvUkdTC4nTo|DTu8042XZl5DQ($Nn(D5fesn^KM*76uE8V&x@>luqpaT4z?H z9j5vdcPs`E-sd~|;D2(~R89u>c2tP%%u0IY!>Ayl@g}Ly#gflC{u}~-Pi8-YndyCC zKzT=t-0sv&fhIy5!k;ISlgFs@Tz@%ATW7;}}bRt(>5$lByL=H&U0wvIMe+vyQ3)>G0WTe@25<`O;yfQO27l zQ5(b)d0{a42&KqnI{wn#aQ+SKTQk063#&|3s$ZZ{)OCH{4KH2;ici6ZuJNx1Ikvj# z_)8vQdo_{Q0zKkOP**C*#g29@nK%xQSOV(4!ph`5;0GOrWrb7YEF<35v1_C8zjLg& zx#PPB9UGiA73XXaJqGnRP31TFm%|n+QtS0~6ax|GqSoA0IWfIJB0y*zdrxuEE)VyEGbE2w^rE(~VC9{n#s!uk&Y}BoA zl$~bi7Gub(-Cf0Ih?7Zj(c%ntjzE3OVt(<8TkJCmsaCO{0sl6W`}Sew2}IUnM|-I1 zcw%0xT!{?-7rXj481}qV3rS2%zxxuo_7mo^np_aezkkHh|Asq0RK_siP;%Y{{61Fk zEAJFiOQWF-0v(IMuMV(V*Hy+I=|q{F!~29>8;$msN#!+yJr$Pmj>9L58O?Ve`MkP1 zgWDAw)N-`BgAS)zNPP4p-fH15TR_4ijuQB^Q?vI4+g~J0Hy~!pgK7y*EndRODYD9) zAlFv-{RVdbNk?n>$jR*Ou}Xb1(~sD1Fc`Mli7nBzq`Q8Edvhbck@|Ip4_SF5hwFXg z^jO|Nd%b}FH6?R@4l;FjRNd}Q&6LT%nhI(Zz}`>5ul_Q+s8+EbG`eyaQFcK6K2BGZ zhr&*MzR{`t+kkq5Sz8c%C>~8(p177S)VEH3R2x(ah0nf(r(Hoa-RD#>BOJ~kZwX%p zU)%vx&CSz6s*`lpsj?&o+>aMaY~@x|L+KTbRfu`oQ34y8NZv`N8v;h0pjSRtowHo) zC_8T(j+qMT%_gVJAP?47eXx-ZGy2YOxSBATaWIAb*dSI#*^3U#kbcE$XANiO_O`{GZzmq%6BlC+U)NR_HVGGYRD-0 zjm|I61}%k!N0V0@s%N?9JN0U!$?@Cq({}LPGoVywr8muXn3iB-DiQq~d-*s@`EC9d z;%G}>8r4AX?1u+|v%i`3@%)5@>e#+?g(-ilLWxVIG_pinyw01qz?9l|g^C}uZ zO}f`T)$?FaImj=X$9gK*V|>GV`1*9?<628w;b2%6s^c(F;1)<7!Fo>{#-qKio>4=O zP)AiqKV5;28Lz6W{Ek1DHyUVdvTk+ossenxvdWjDv?Vy>jb**xm|g7!bBj|MtbwD# zjx^aS2o@9!BE^7uOVkt68yr+VZkX03nA<68#D>IH19k4dj zw;jCCaq_vmRc|%*?KgO7Fng5g;DO9)roo&rBEC2NEa#CDRR>y-wV$)o`(bJCbCsXZ z$`T(VOa$n&wYIwJk7|_{^WU*g!Lvs_;g;s)kCvv}(uwFf3+pRFc8*bc$?rz{(#X&quYVlY_YN_$ zm^sN=139_!nW-1D@SL+sKkY|7FY4v-|G@=BZ_i>yuX92phUznwDlk-?4UngzLQEEy zNx7#7fLB*Z*!~gH1&DdUk8Qde6jsM|n@;)7E6)QtcQsT(eNs;j%l#^Zt$i>N;TNw)aN4%N6$;`SC$jiQ^5O>($z7^#z5r0V_>yAN#--1HliSNi#s65*2Zv;%C z86J_Oe$TK6{NMuP75Ja~`=t5w0_d0@OIBBB0`D`mhRife{fbB`Nb?{VR+#KOkvr@Y zZh|yfs5rAZ?GS|hT7a>sYP`dal90fs<5}DLu*5oy+>U%+g&n-lskg>ZSNsC%9fM!J zKzz!tROVv_lN{D8Z!F(J?y1Umy-y~)OHAj(Bbqo>Z);{BN~w@WQCr?{^de}gxABPWoS2zM?w&xmIZ;K$Z$`6Bg{4JP_ml#gN3_eIXRV4rV+J$;X`uuI@)WZ(!NHREF{0%yhuGi7)z#2WzV}|zo#Aj^hNyTK|FRA zD&0)1Je|t0661*~-q0X_1-cJwnBNij&`thU8I&)gdOM$(_~<}|cAFZzHkl-o{Y@p$ zM5yyL1&qF10XxgDYRrbeY*6pSZ*CZc?t)%tx0=CFx}rS&4%3(4F}cGIws826Wbq!b z@Ve}3Jj(GxkTurmdb-Q~G4@-;e(Z!hULaD+p`IURet}MvItM%NGFaA*3OthcfY+v` ztEyCS;S8%x_N&dT#(=)=5$GZ=>UCC?jtjViPo{=A3yxX-DIe!!m}ruTFkOYY!XD-eSX z4P(@u$64&wD~$XoTx1$v`ZJs_lKF{#AUm^+-O7ZaEFo8Y3f})l^lZYzt~tuiC1yVv zl`c01OHz&PCtv2rbE+y9`H`uo6UlUwiQi&MKW)r-PlGQp4qhF{YDKBQJ|=TM1(QC+ zdaE%5>2DP^%tOYL-%2@5MPH086bqs!s;9Q3YLjfW7hhXKqzr}`-38xb96l*KGoJCD z!Am|P$|K1SAy^}XJEEL=bvdJSf&nCxg|g8jz9jn%#X{cm6_3G7j*>sBW1;cH#}IaA zGwVI!`14wJ{{ZnT{GbZ>C%+N?F*|eAQNg-`2Rn({6;w9Y@V^5@>L;vsprg0XL2=j& zVh&`cO{qCw2VYNRs$GsIzFDIy2#?$b6IaN8FmbO;|N*n#{cP({C5Wn7#1m!90fYw zMsL~zB9-8@=5&0#h{8Xq7oNv{*U5PG@bTsNOIOfV)Ny(0=x3u;O8qgBzYT@)NEbAT znm9|Htq^8e0c+d^x#p4^7h=_Lx(fTxAOcQ+ND=BR>`|j)XVI-Xq*UidpyV6G;7abO>2&s{vfd}LP;=^q>QqY8 z@b;F>uf3yR9We9j#NTEz({v&^S)ELJ-{EU{87~~ZaT8CxLcLOl-H&3t9ZsjFG3$LC z6=Erwq&z79HGfNC2J$YrLPm*+pch?=@lN3PGr^{Gr5A@fdiKYBD!+J^n^OjnNw&hd z!WHY2r=tFJ?6=d57lI}@lvVatQTB|3&+;7ZEPSasSeyrCVgqr0oI3;$o-@cQZ{eN+ zlj(`#5XIG5^_=TAN59W7?4|`Aq!^m-YP7FXXqH)uzjQG+$b#SR0bhEu3vsBX3DjFz z>LfyYgVb5{N!k$sm(UMxvoj6UDMxu$OaAs1<5hxt)xmd*W3?PmS+eBAj=n4EQAO6f z72jA1>Xrgu3K-we@s`GTkKlE)!1h;}^RKY~Y@_vQ{+!Bqn^1+Gq=Zsdb&*w#8n~sFLrPIm>f=eW{cN z5#6Hs2CGxlXPo`$XMAWjRaPvu&{p;PoVyLT)o*E$-G7q#m12HZ*i~V<8=b06l;bsq z!{sL%B+`XXq)slc*89AvIjE}bZQN6I;*kP8AUp+ zW6}Okg5qc4V^^r;B!6^s^pal;N16iV`#ZZj+~8Fjs_7Q)sLXkvQz|Ot{Dzz{eA{I9 zEg;fNzBiDH`Ive$*%70vX=J^IHd37&FHfb6huutb{7>u`ZZM~U`lW~}L~>DLaG_J> z4|F;W58!K=c+PvMBzsVT#GjL#@h%$wTSP`ofr~C7es8KQDg1PSliAZ)ZzAg%NT$mq z3l?_rTe$PLAK1$XVqqanY8hF7wzHS{oJ#m7_HrY7&J_676#V5MCw`+G=Fu3hYYQ@5 zqlUPFHwVL-- zg*@}P58EH6mr$O+RaIK9R7ER{9wBGm7htPR?0zftcm?#_A3UT zeam3U1*L7u`(HLXoun|vOQVXO2`8Kj^ALt7PtE+x$tkKiRER*em*P2c z66OGk_DDbDj(?;x8AvE)X zd>XIvo4iTvj#HCdX1zCI=viE)(oc~S$q$-)_4)RH;FA$>(R%Qm>UhD2+;Q6JM^!Mp z)d4_!WJDfistYrgP|9v?6Ti8d=f<5cZZ;#SQ9i@TF8w(hFphh}RmL9W?DQqG-nDqn zKE5N0XW}vy-hb^>sYMO{D+1dbL_K(&x_g82;>u1`B$-agB%Tng2OsYRQyrtYsIc6Y zP9^*@wqFRI=YW6tVPOv%#Xe?t3 z3qLt1aLrT>X=IWnAWd($&oMQ>6en|DHaTSj*`OeIG{Q^rs#;C z4H1s!mW~f4fotnk)h2HVs_*#oN#=4tS?DYFY7$xXFNbARbC9(MsGI=v4Fp+}`BdI* z)0eA)&YYiQy#-)0-+}U-sl=~2@geW~x^9?i6mOHf4&UuSEX1nZbJVHw*O)n{qVHa% z3tEplZ@S7e#hv^nBhMw`WAN8)?7~*&*TLZ`(;SYTLhjxV3M{3vk?Q@V!%C$Fn`k1r z6#1hUyE>ixyT^S$^-|M63Q)a??O^faFxxYXEHhYXuvl~8t=Q^!_V*Hdl!m3|I4B?| z)9vp#Rb`Y*yz71|96ff6sWh{|#B=;@8gZ1W5H8)B)%PYlN2sTwCbOPFjQ5aHyfohkG@5T3vA-0b z>`7g_R;{Rm!`#H1zGRfDc)>n0NpYCPQHQ0Cb4I?5{Z3-r1XPmgSVI_@Jnwhc(Z^%B z;{__$3LwpSFyNR|?|@>&ML#Ib&d531 zj~#u;drQ4o=wIN+E>KOnBby!7d4aR?n`XTYQ6bL4xnfoC_>Qx`@rEf>B?4AcrDjtB zC90c)yhu6$bz6k~!s_ znITT?vdj5fab{JYd^D8#$*JrcMiJ3n)oKntFTtXl!54`}QRGDVl;<_%ncK&>qbJOx zBFt^Klffd5)~7RFE2C+JF@sL*euT2$2&YR?+Cher_1 z{YsWx=5#oIb&%#ohshL&ABbKjRnfbS*7v-le8w6js{x*KMs@L?ck~T;KlLUP({j#j zI!y4AdQM$B+8;YvNZvFcCu@d)_t(&yX7lM;C%^qrhvJP@~p@CEH7jyx*`O`e%X?b2WQeTswaXB@P6)6^*yRR!M0sYn|-c;DLaHOV9y z;O+($jo!2H*^cHr)o{W{)Vv_&e-Ap^zwp1GO+P?#R0FiMp-N$q-i{!loNSO^U8rL! zk@axuc%{R=>3G;shiS=a`o&=F8dg@B_)T`|%fb$3A2gN71oF`(6&2@9PehN`&B^w| z$k+1L;@2EIPj>S88mHR1jAukR5g+g9P0O83oML$PSSrmm$|FQ2spw?g)=td41H=4> zOi_{j5I!Wod0O7lH`bYdw}(icrz(Q^j>;(-MLoy=(y3A}Qo{~)Dyjuer$N}iym#h9 zP;C~hV78-0{^qD^r=1!xANeB#PaNp@%e#&q(aXs*mswfvS>pZbcNKa$4B!cqC3SpA zr=}eI;2p!mG^_f;;TsQ`%6cu?v9X$QUiS7cR zTCRRIK%KLE&e_XVPR@%meS-=xnNmvO`^U*gFFBQ0bE7Hrg};nogfw5g(X zPEKRpD^aoAI(epy!%9m!*7%M6OM*F0b-cKuliwuU4>Gz>MZ9Ua`nHSCRaYmv<$W!G zfTDpYm7N{`+v=!bA3JqL04hWbJ1xAah7(&~IN3_htDSaJYXVk z0-0o<@)FS^n>qNL&&iz8tTMMYEuj?L;Re_AF38A*OomuaJ$OMO`c5Z{OHQ8S*kvlL zxGeFp%NgaMsaW)Q^4|C>SCERE1wjCr*B8D4aWT$dU!9g?Mlk;9RI^242BR`0lrhX4R$;mjcIlAqS z?8hj20HYbNx})ko;;4CNxFZdI@-A3$7k}yPC_w9+?n5Hubz`sgvEJ^C*W6+Cn;rdr z4tL1!(YJ=h-UYetceIRfM}G-pJb6EGVffQ2X8($_-aw}+uf}-Kl9!^W?Bz`>Rh(LA zm4hz>4M(5Hse@kRr7`L8J%?ou{s(3B~sN}Go4y(z2Q01k=RPTx&oh! zSMf2$$&&L7zEs9vDsx9YD%WBvwu(4jl5SM6NYHy95zvA4#;Jb57tVO!88yBI)}4Wx zoK5#3Rh{DB?r`*foA`~xOF9uB`_T{L6jR;dC{II;KZl}OMuJz3s1cLZDS?0f2Z{!$ zI%F_-P?YZ|fd2_oDB;wry^OLGsebD+cSU#9TO!Lls_vhTKevUyk6}OLmlMW<0X>~u z`<25K5>15GhP9LecW0>cT%#Qpbie5g1QE5tWbEbS&{V~w-*iU4VR&{Fs5c7bxD)k9 zH#OdJM_H7fbxnM?HfyU4UJhhWA9k!JIY6E-`-H!gzJF0D__h0$}iK$sS2pNlZjt3)q4gVwt}b?z1iu`iZQ7OMvrNgyAFu)kk9dPx;W6*_@M zOLJ9*nE4v(20C2+g7Y1@SUg=-q2pnW2grH5ovJt4QSOD<>O{ z8f7~J^jVKKJOj0{o{IfmP9;3kpv3`fkc^k4fjRL?0ov#E2n5XrgMZb~OJ2sx2f!DJ zib{?u8EllSyeL$gn9Di3HwBgA5aRSEx*7C714_wnF-;_=bW>XCLykX>FiK^Vs#02D zgF4JqR2Qin<~up+cjjK3mCs=Rs{F|m-?0WTG$}o(dLq!e##HT$3fyvl&3`?VYZ`?zH6dZ^oVBa+3-TdOi$F^_ncb1 zyOUF7m2E(-0q~PJcK;tot9s1I?9YS6-;$Yr#(qt~u*ptriH29w=%-oiR$uk&3~8+V zM@L`#x0Bg_$E(iaf4Tk8+dPjrO=&^5oGhMf>W|5Ito+tTcU0mQ3SY`QD_?GOW%;e? zit6`qYO#MaocjDVM{`SM#$U6xSolypnzVE{mN`nyR8!-ur7p`)E(?OmOi~P`r=!e@ zraGOhnn9Iy7p9uaGGf&UHOU`OIXPe{UM#Vdd&Vo1Ty|Ee23H&<@JAELZ9swPDB$C$ z7SAgFT-RZO(uHpf1DFBd_yzwzij^}R{A=K7TBq@inW%83iNQq#lw0MzfTY!t*J*P_ndU_s-V#rGvLbduIv;l*Eq$c);Rrz zR}Ba0K$YJRzmHN%K|ZH;o@c5yIVDw=%Igv_(}|ebu2@Er(^)xT@b3~l_5kDEWcIBT zBNP2kp4^Bw`LYzd@H#V{3Ad=No+^9LsaMyV%vk|H?GN8QK)ruT^^=D?OiM5^l87$H zUQR`idzQZ)bM&xR4T=i>b*9!Rh@ESLn6p(Uu8GmoG(IO9ZJ`@nxi%F`w9=bCaJ1`g zW~UqCopdFf12=$8cV>nEF+hVAgxZ zFdmKfYmFM3O}*6%Y~P8$#H%VR+VSU}rbg@yI`$&dt)U{#ukxjw;*wwb6Yng{dY?DD zIt}JFTV=3$jtaZlWUE+M&_s4Cg|#(PdQOmo%6W|vy#lK($A@YYTZP=;vX`qKrV^e_ zZj?9hSH&axs{Y^$4qqE=dIa(CrlDA}5x&+-<;G*qtdb1ogz@h-;m_mn_7Gc^D&g1FsO=)a^GWzeHt;i2hUC-!){RL zeBW@o6jjZqx8bZ?sByB;noB9KD{g92eZKf%u)RC@TLK%r!>6%MmEX}(W~-6cE8&Uq zOUh$lhtfN#<6Qk^m_j0qb_Mw;o4RL-I!Scg$sbRe&cz0zrz)&3FT0hZRMYv6(%anB z5W_)+(x_$e_)9yby0mssv<_ZUlUW7wub!fcYDDDaHHwWs>-?@^eGS;v6c|(_RbWB& z+{6JVUq+frU4Ex4izw?$2YRSF^(QR%nz2C|s$T-*wIaH&;xAX3X_`7yB+Bu4KK+Z@ zpbTAroy>GBF}vNt;_D_V?8i~KV^JcdDD`bL^Lx}`Mw_v74H)AH zGW$zZBFD`hg({`e>amjjdb6S#u-p**d9C6((@if+m|z47!6G=wIFKO%Eu$}0V4Twa z_nY;e1GATcITPrZ>|<7&i0*T$lKa?H9rEj}ht=~2Jy39<(89|D>`RivN90q&~}f z%Q$-bW)z26Am&YG`YJVN12A?g>uu)ngW0Ti8L?QOQT}2t%Mji1M(5H>)&cfo9vLji zXqzR7lyDF|R6QS6)YO-2;Lm5V{aCzdfa=sd%!;0JRQq~FSuUct!9tnjfbKBHn=1c` zGS<|n&dI1Q=QwGX2^W||*3FBzpHok5$y=I*-FymGG$LkJvC2JIHx~OnXI3fvG))VQ0WMhLE-L|DW~zcLy6d`qe(h zn*v_#2cOU4f9;vy6WDK#)1P?{`)wrSt;NT`#b0WHwj)@tyq!GLsNzz8l;Urhuomfp zE}(i(Q~XLaqcvEg3CQ&(x$y^V`UW{v&aGc}Iu}#0#$C94VY0$yFrqOdf8y*%uo*8B z4jc)OC{2_tC9mEFpVQPcvi3Y-BrLoHvl>H9JQJSN3$Ke)vl?meFOuDAMLm6;F6aRw zswcIJesXAq@$44VZ>!a>hQ0@XT7-r-mpkllUx=OMYJ-qf{oNIHWV)IAX8~-o9Vy#d4P9Dbw>xgWs2@ zvKd4sxsG<7q2A0RXRXpr=PV7rCgUALYsqAlo75?WUmSeCV7w^^e<_CF568m0qW7?z z`R#;t?m&-~4wjrkf8HQ<8qwV!)oKOID1rP|NS&pVQ`GXe?~L00I$UKknWQSpXPnBC zQw%GRx0^`4x`b1ho8Z)SLFGj%3ta$;Ky|-x;;kLyeS_CM4z`bA&V6A2QBD;-!PJ9k z^sF24G(sR5Hk~X~LY;v4#o?#&4#W{;4e3q%!LEG{{{3mRdi^}?6q5t;^Y0F^`{(eN zN$M=qGADk2VZGhaUaNrzbIHHGVTbQItfh{FG_Mh>zcQ7*SL07c}3S>{U-iJ)pq2H9g4qqAy(wxE5rh})F8>JsF z&zy}h-ad`$Bob!+mQqCxKHAo$c*1UwSRz9$59F-%8F=C*#r5R30zPyy zaVT850G*mw_^#L3tu4l5^{HLSbdzBeb@1Id6K{#?R9yjwuYGB(n~eP`QpXoTcRZn@ z*8WzC?%bAUE@#kMPLfG>z?RxOT>$CqZ8FvfX1sB%H=6iur@UA!Inz-p-@r?!lgoy( z(Csk&*Po(b{qdIORm_%^FqzY6`d8n2$y=-QMQk> z-WIH<44LjW*?xw@4>Fv;tzs{eiG}*;5gBBPe;od|!^!7@Gdbuv8{t$XIXg$UM2PuX8EV1Y!9TP^9D%$9yX9CpOYPY>EpzQJWnjo>MsE+ z0-0Z^VNklCGSw(Mlfl?9vg9CCk_k#R?d9N=oCh8U%b0duV;B0$Y)p64hQx ze@2u9>FjkR&n1!T_cLBO@amF-?SlMr{&@~qH<9r+)3aFw)0f^;Uk3%GGVIL1y3Dt3 zWPe4klRLsqENG?>%X&{!+of?8Ne$D{iPdWko_-9UOooeAg1h!I99`qf*G6v-P)}5a zu-^7mS)UP&D-0s({CmuBh9y)$32>ieJX^mXB#(nCgG}W*gpPX{=TGHV%Q970MzT;1 zr>5J-igLYkE6Uvp)VzW!YXlpv;{Q_96e2n|cgwG+$Ej08-#GO`SyN4vB@Z4WlI5B2 z8LA4b<{*DVW61~*b2jI5_A_#4#TeT={v2eK;|x@s5_nD^_hzV!C{I4W?sOVjQl)fb zhaO^FcK`qD|g2#2V%2x z9a%~ycidD44zpI*lWT;5s`Q&3&N)&1al@&W;UlD`G!=c>qVAFsZZ5^i2b+VMr-O2m z@uCs@YB$HXA9OOg>zMClhHLzoF`rg9L`Z#ghojHp(U%d0eg%`?6TH^QFf6*hFWNK> zuly~#v>wk|%eMoaDn*Pl=6^9?sh@w1m#^m?C#bzlVjsf{izvuab~bb%It#~nLgC3c zN2=q@n8%|-7w~mk;Au{(I1^#8kKVcR1s4Aa`nn$}j9`2-9sX2&tn5$dPVBx8FK`K& z%_!D1%gGmJ8x$m#Tm}}Mg(jZFJ1$ij?p`N;lo)*}<5r(2WYis0a8h{?v6LswETz9L zz+6HWg+@3Z|TT~aVS(7eZG@bHZ~Zz5VYQmsGE())VpGLILwY@p9ARcR(#o0;HSI! zHr1Kst&Y#$$P;7X4;mpE*+(mU`zga6>X`iu{k6t_e1uoK25H^uu+HB&_It$4w;LLE z3z+MKLZ!05uCK#ion%IF@Hvx+AAd(4&pKJr*A71U4Ey=95m=mgszToKf9QJQ~FKI`i!U2NhgqP^9+H zCFBe9@h?jqdy{&Z_}HtBCoIIy<$mvS6~VHcxixk$Tr=a%x?%wdadpsS-(1>z?7`@5GOtx|hYi(<8;KCYEWYg^ ztrKQ^P4RH~3UR;faDB30CBfLV>`XdIop==ytWcd8p5#O!@j0@yMk3y5_@*&@JI{%4 z4V;|#W5fHzk~wHrQl9vy6VD}63Nd`nQR2dQwZp2TWBX#iS2=9*ZLsV!@QKeDKd$xf z2IsR=MsvE5Z{4VNR7`TT?psIOD~)7Mf@lY+bxAfA6a7BO-#Rss35>mltRly;fDas$ zCLCxdlc$uc3_ROe=WfoJ-!OPhSV2{tAaLxy#86 z8<`l@l>Uw>hDB`EyTwTMU#&Gh`!py#+P&4{l#!sWguG;GB8{l0%)#3hs@Nl3czefZ zCz))g2k~f?vWfwYKWOjR+Y!Tg^u(r5C_nU?gV!QWl|pAbQ;7v3j>Shiy0*iK#5pEU zKF->YQyS6Q8Q;|o2bym0fF7ZYg?@u1*AfX8eZc!etM} z{U*;VBey@S^j;W(w+zD5H&M#Iz87KrW;>|ByWe|)qvOIUWDq}+(ZuT$iH#Z~ zg^)$ffJw`A^y+=bpB6Y6E)TC(&tKiAc1?ci$Wp2@vKv=wHeVs%*=VA0OKqK8W23_f z&Nb2MC{{K^`MQ%13L4>LhD%NB1o1asaU(+=Tm8h*l2u0AYw@;ilr5AymfX?d6CX7A zG>LgkRe5Mz2V0ano}r)7m!<5`ov2uYmz-x4oSdYuvGWUH!pSNcbaMRZ1V;+RhW9B4 zLA6xBCcbmGlSRv&KVwWrkPe?6r950)M@v3;^0)Pdcl?$tuvTSR!i{WjWOb98NoQ)v(#m(%^%qc8IQ z^N9|=`enfQ@wDkrvXdl*;01x+=R~@34#y)mtIDl^GmZDrELeeq6ztoJCxXJ$m9Jat zAc_YayzzI#9H*1_tXEu%NVSQRIUIDd_c`F_dJi{6<>wJjJX+}Jcw57JPbVV{QT{pA z@uy!oHG%?zu9KC`USrUMw*6xc=DNk$+Za}|TFGO#Gou@wHQi=bvL)J?qV!5wY{78P zIq^-<$nWqoLl|G8v-WKru0{0pQ^Q})qq3skI&+U>``%3^Z4J(80|J$|3WcaFaJs{N zg_zX{LEER36Ni9GL)5OJ&Q7d&(upKP(1#sJs5kvR?&SM&GAGpd#1JAx6k6OPTz0a}YlcgPLNWjjyE`{lMeGPzwwUAR4D z(+1}V_i%=MHIwm0D_!gB_?M>*T5n9vwJ}d@LB5?7pVQFE5@L*hNydBUvJX^ldt0C) z`#MMKzBZau&FZ!$r<+6!7_Tg#jpL=}nrv|aQcb`<&Z?8Z6CLKHzZ1{9f*9^)=85Rn z4M=~y6Jc9ANVSqDHscTa@f&&QpZEjWtI@})IsA%Itw##~g#XDvn}#^KZkeOcD~V81 z$m2RjdIdgUg!AmL96b4$vCv`Uux0GbY^`|eI1?MSq<7&VMu8%)B`4~LzKnM=^|$~? zW)xm8jOt+m`Nvi2)S$d;%f59`P}(GRd$ndY0+-!83d^ zoLKr(KD*ArK94%A)$6QK99G#JDf}M)*2|1Qzlq_LK?-MCp{1az=li zBUn@wn%$9E#U_7euKJPhET)5XIrZpIrMBYq#H$HZOLfk_ge*CWndGsTe#i(01 zBGAsujYrq0beqwp;bb<^=zUiZe7K6W9UOhxV)FQ@)O$|z8`s0yCoBJ)X&7*A$(_*R zF!=0~=ujztHPTtD7aV=L4BM9Xk)DH(sAMF<5X3sHsf@G^jfltV{*Krp_kh3TWR;IQ zSwb=+{xf`1OZ4w9<~zaB-CC!=4OnzT{7X+*u0eb|#j$-ksUp138hAyy|LR-fVFk1J zmBWy}?r3%dEBYF(0ocq}tnPZpznpL~&EeD|D*63WAji%4j9Cl@JEUOKcRJkHVn;@g;QyMF zVL!ntZ${H%9ejV8Q_X0CmHmSVelALJcI*E7FC5{7mtQj>^1!?Cja&6wpbhGwuhXBd@!ALLF4 zPj)rzc`mVR7Qb6a42V%Y*8s=AL>qtb0MR8(y`SWIJmDmVIg!@4hL4CR{&v92_aka_ zQKJ&cba!Hk%<^tL*BbDBW2E2K!D}x%k^LHc@o$L=gIK>m;KRxt+iK+CwMjf%Kg*-) z?F^O7Tx{W02M={L`gbF?FYlvYO(i5%t^G+SieF}W9?h8LT_4>+o7Kv{3~?$KI}ARR zQ_1D({W$en*G=u+m;C&4hjaT0UhgaBb^<&U4JUoUFs_x^bxMUE5e)MwT?Irar|*~Y}zKIr;HqvM)O%W-^IE)3LLc!3jmYPkvj=LYfU zN_w_oLnkTZ(T}Y59lUO+qr0NdBMb+ch>T{z1dhYkO;+(z_-EN8C)(Mb`MzzWno2KY zR0Mn4VNz$9XmtVqORQ+B6MJM&`LhP;we^|kHNP3D^WK@?QIl}tOef=;@QQ>sx$IEi~lI+3GwYu8O|+)XXoJs!Wf^NO72D+tXA?! zF&KCOy9shw(?$5i?u?}nFCU`njpI#@xPVzsLWg93#CrUi+z}tHSnTyiU#iIjyls%Rw1~@ep?rKc?UX-Xb%YjlSH)x=4JRXfRSTQs~9L zf;{zB;_@)Y0d=A*6l$|Sw!Fu7U)*<`Nb`d3Cz-p7Ql%42!68T#d5|mnzoJN?! z_-;eK=UMxMM766~=h2L>rIY1f#`s=BzGLxRcVTaTB@>aIsM8$g=^)ZN4MJH9^E4Zz zFYhrqiv-U(cu006-vbKoiBzNc``huX;xT1cagxC=3*Z#OsE!OrOJb42jc8{R#rj=t za>QizzsMUTTY#mOz*QxZ(=}17xMY>ZCi*4fXC|TeFRpq(P%y$6H>{#$~4=N_PXk8zC;sBM&eQJSbH;B%t}^cB~NT~^3`7SasvoJxUb(6^(*=P z^Q@3~0r7;dFz-u=A6K$MvDn9GBr}`&&U5h9c%)yB{R%#io6-yD`8<-j)M3GH!ynv@ z)jm#UGl*4r7E8T^`PMj0)_fR;*?7z&jPw<9&ZhX-GtBof2Os>4`Tho5yq#VoR@7kK z(T3;JOiVlEkxm4gihO0y*c~8?q5Q@Lr$5=H*NIgt0LvAj<5H$ z4cmGStq&s}B`}`__`JqW72u!f-$tIehWVDEb$5{&yv7PW;^6Vc%y$71wl#P|qSo{O zQ@-*a5?ti?mt6YWhQB(Cf1Zc_6)?Uc)^EF0X{cfS4ifRQ(eYjA%WZhu$;?;kwWAC( zCegHv&(4E*qzd|tI%jj4iQ=+9F&r(4W=$`E(k6i?bEwmWtN0<6iw>+tFV-TCsPG*= zVjVs2apdv6nQt=ll~yCHY%CFWHL)y2?Z1*5_%G?@Q8eOOd_@HQ{!9F+?3Df9aN!Nq z8;A!JTYkY>T#nbBPUgIm=-1MbhwK~9$A0Clp#$J9*TUg7;R*Zh3*i%+;CF{$$z#bO zO3CQue*SQU*Df=eTsZvSBD`Y}b;OxiKr8%JxH`-8j>ES{5j6_YzZXG$i?PAIhEdn2 zDo!!qwyco6HR?(-r88*zn~YCTi|iUYz^}>s0?YBDg=B`k(4iZR1of$z_vo(xyitoL z_J`lyf%KE;Pi~#O-9)-3XlGNN7>)O-#f!@O?L*WJ48m<*U^QMQYRJ2KIx_PtW-ffO zScOFP2)=zC{j1MwM-oXoz+LxKUc0ShZ%yd$Zqr|TwB$Uq?4$b2aQw6I+0FUxb$sz< zWQ~on^UHZccJ0b}=iPie053HJ6m*X0a*k*q=Uh7)7Dx7<3c86P$KFKj8G?-R;Qm9@ zUSc`<_Xyv9OzmL@<9mWi;Vk&rtvn%F;5EimOSPZ{dp|Pa555JJmZ}q_!urW)U(;Vb z&-DOUWi*kb3;Mi3MM5qR|WOPH@Td1+$_Mo&8)FmGvPsvuX zTlR;&U^2rHPPPwckNEKh@JtpS^q&9Jr$v^b_H5LSj=duTPGM`NV zfeyf1^i+G9gq3{2j5&lmSMssEkyv*x;`Fr&JIX0`sZ>90bf~48c{H=A&sVpo+Subx zJ>@#%6Z4VBNbr5OLb-CkitNvpOg+qCxOwCi)l}iD$ty;v`MzMX4t-PlPl;+1@H9Qh z5r^VWB{s_bwPPkflG9kJ)B{f7krLQzuu0Xjq&_G+mu2VFBqQHw^_HiVXwzCp3L^|& z6H7UUXC26lx}lx(9PF*fcZc!RC*YcU3EH?=FtjOsPqoQg6Tcc zRcP@^V%aeyeVq8w-f&gA7umDj9`64mykr-APk%Bwi68b|e)^=L8 zoTw{Om}xTMX2`NDGoPsNL7o$T#S_Mx7&=0sC3(|9EprqP@~El0X-jTGl&QuG+>a$! zv3|n6%AIjvI4k)nx^@}8bi|8RVNXACvZ6JP4oStN6pLBKNMB{WrAk}oNc9qvBZjJ7 z5rgrf!^y^WI{x{ujyzV=%a@?~Z}0-Iky%MCaIzCML>{GZBUcjRzrvHmf^N<@`ggnG z$7Ouu)SDt2V#%HO)!Usczs9LF{2X2Dh~{Lm-o0Uh@=PwN`%5$W*BM;A6pf7sDby+{ zh<)5*>Sf|z#-Yz+`Q3r6Mh~^$_629vUNDs@dFyl=)~cP_4L;k!e^U8=(A1RjKs+Lzb2LWE%X;k7TpX?cmv2o_NsMr@@a`F z43&fb2A4OWr;LL9%21Kh|QqI$zHGHgS=OldMSyoc^=bTo1r|)+>b-u}GBC)7M z{;m-ov#UCPBxjetGJJajb)u?*`94IJwFPU;bfggHMB=MKEmO%?`hn4ZiH=K#8{*W< zMw-lZ4z-xBNdGj^;Jn(!B~j>J(g z7#m#Av;9=WZti4qnP$G@$slw2jU;q;wc2%c#>q7AFdTPB^5l`^E<^ddsY;6{JDIg` z6w^VBaUkiAoiEqpa=i`19BBe;2@-Zxg`YvbfIeGL*63EoJh z>nPqkQdP-bbw(w~_X;$&9LsCLO7-T~>Wl~0|9)jWbsljhgf%^ZzGSNyf6n3W-#6oH zL+&D{2CMOjJrpA+`<3O4$y+9e=!f;31&w@zU%Q|fT;X1S>*(`G{OX59p)oM|!@+C9 zGR-r#rFZtYwu$=E+ zz{^(>T?9|Q<7D|^X3WWOEk*E-*`Tk9D$d9|Bjh*UH0;)CDwLC0-6mMn0F`BhI-XVb zC;WogJpeC|L}t^GHT}rJYho!wj3wvc%~I7{&*MN3PdQ#dFqhQOR+-$sH9XB0P{BmC zw^sJ+WgFBi{WSr}48;p1V-;uA*mnc)VLV~1c-*P{os;Ge@)Jw^QK8%}<{mnYgX_HXet+wtC!4j*UbZ^){RSMyzDEI^+(mHP<~ z(qD6;_%2vyxqW$rV_TmZ)r$@v(aQS88tzad^XKSq9v)&lUM(MKHFsF& zRgV3>VCFjveA)$@8KCxFie#R3IHQTo=pACg`^1lv$YWZ-tf7s)8a^jqRjQwJ{EKL;RDC-+HH_7;?4$9C z&pEl@8Hc%><}lIm^fCs|dd!i+(@wtM$r=5d$oF~lS=L2vbd~+%;%(X$o_xUG@QIh9w4?aN++VUE9_;^N5kZ86n{1K9>FDpkn~B= zdWecWN1UwuCyw{&0g9ZBe|gS{jjuWVb#kKpTQGN__{KjO9!}%6?@X0KXFI!z1*aX0 zk2Wq?!nswxx?gndE{#Sl$RRf4`)E#XB{!>zcD~}Es%IR$k!IML@ni`d)eSVg z9WOu0k?&~3b6pQ^nxk&W{$wiECOai$-%^<)%cmUv>R#hF#(`w=)Vs7^#VSM* z=2-Gej{ZGk@M$=!c0%Q$A2=wor6Z4-CZ6}Eg4#o|LhlBAR`QCkjGY(3B+OIVw8`EWw+QW`T-R2<6WhOIh1Lkk7Jlq9GOO`t$jW@hwAsOEa_wJ`l9a*+^ zeD)H?{syd1nKNI}#6m}(+Zp`a4IOWzSO}>oo_74R+_5F+CVDZ;dK}wbdF3e%ADim* zyv#(1bn?L-%3Iv+MDckJ-hJ30!LFc|6AHDYJ686P(_gH?m}l|8p=zCPb8_dOIyyeY zuqlPqHEvdRl;E(YUpnzaDy=C7jfAVcB!WO?|GnJzaL(a)i*rPCBUIypy@68VR<9gO5=1NOR;n#96Sa@;V`Kej6O?6q67Kr6GXabRXhI5iIT#u&IQG;r&a$%Hana*DekMI zgL40FSg<8nd@h{+0{EseiZ?spXx%3U14ppBjj;VuNqmt0o^vwINTWlYl>Lr{X}ce5eA&@r$zdNi)|iY(>I2d$ zV_$fT;^rM~S89`K-ptqURbeP59GDqRAAwtw5nFN=G6`v^eytT+O2dqAwc=HvW z>7cV#mphr^5U}X|=*F+n>=kI|IS0SgI-KfpVnu7LHjCKO9Pb$6WNo4^*(Q%aL4Ozc zjm~7&OBCyO$l*Y<4QlRgY>LY$4_5 zO2AXEBA!>Pd%&-Ayvnbfcytx^mIa@53i(~ne6MhHx0S<}TuN?rITrmb>+?0=c6M@8 zInT7+uyo7t=;z@@rxJgcs94*;U<;iiwqiv~X&uC7USoA1bFh4h!^s>4F*L-YV?mD3 z(nk*m??yRUTR(7zuxee%95PwS%}&00#^I#TW4{Nnk4B7YBl>^2!zfK~#$4|$%GE4> z21?7prmuFmudWV@{RjNt74Rk5tn#z?7qN;r9Ip8~!(>GgHJUM&1lXe+Res*yFp#?5 z{TL&DoUx1`ek7T{*GyZIqeC~C?{37`bwk4v)Tv8}>?@oYbv0v8Wad-xLbH&X#E){v z%48p22(fk+R`e~I*KTAQ?c@f6CqFanc`j>|s&>__W(`_9+?@EAXyXr>5?$uNGEKs( zR4aV-kkfk6XyR$&?@;t=CaAuN!lIJ7ZgFCY%yK$fIvjjohx9KvT-9Hl$o>Mpcs`NI z?(sUSScrCpvC*8Ryr-c)GUcncGgLA2Vuu?(=H!=G69v1YfAx6xTTbDeQT*QtC+fB| zZ2U1aR^Iio1(rBM`Iiw6D&1vx<*r!SIrV;=di@)w_U<2WqIh$I%%hOo8OGNKJ?-UI zq1rkasShZvEqL~8R_c#1v)PVU+3m3B6N#!<5mDX%NezMRyu_zI=MS53Hx1$eN2HFlMGUN z-l}5G{~X_DfGm27TtViQ2{(K* zJ;$nEzHu`Dc}8)p@EiX@7!HBGNRxd}N?=#xo2R1$x zn-1f%c`9D^G}uv}q3pnXYm8LSfZSTD2)4vwVarXl8bFjR1Gme~Jf~It@`U4+W#4rd zyig}}vlSzqLiAhg*urrqepE4|8F=h_X>G)3?{-i@sT1k0K!Pum(bb@(70BqM!)@N< zSaN^xP-F6=kFla6@Rd}cA9gAOj~I@pGg{XfJ*{O&z%a!Z3)69{!+E@cUQdVdIE`fu z=QDeQh33aUGk7-?J?+gpe*=1$s^a-Vr?S$-eETh0ay>at2o>HurF9Z@MTe%D_)$(a zu?U;)j<+pW^=ZlSC1+Y=YD$l>k}cJnlox;wWv}0Z4!;p+kU}+x_6f!^ioHkiifg{Y z(Ym%Kr<;oPGy(NR5}lVS4O?NjFO6BViGr_#`u>c~R5RZzom~8HCMJhs3v!-#JXkq} zxNuZ?kSC2S#ovcAqf+worD*IDvcMbFxruU9d(uAfcH&VwTK5aQNcWVq1R%vvwl} zgz$u%T0Lc$y9jkta4YO{JH5Pvw~#$lanATM$jYl&p-Q}-ywC6;_}~uudpe+eFEQVC zQ#TP^hP{vGi4G1vJL?|Ha*AseMfFLh`HTC_TGdiW;>b4%RCW>hTV<*eubsHjAtEa*nOQPe2(wF zpZM_!Uf>p@!l&rVTGnN>@lwKhECKP%0E?bQ!!E!#_2HCrjH)vJ&16bT=%oWQZ%5sw z9sQNUy+){9Eze{QF`%Fdb^f=M_+cJv9FUahv6lAkr_tP-*t|C zd`^^1!)hr&D!5UR9nG1Phfn)+COdh-7qSdGg#dz zR2ff@6|Gl0Q6-D{hsoa}z$HaQzXfPrZ}lFNBS^5hgNI~CatSDWI8u#<4=%#9CNMt9 zyuuA@ItET*4zcVgRx5AQ8iH4jQQU)g{@JYahakpAME0I!BJUxikLgcX@n=nCK%!qC zuu3CF&>npruW*i>1eTiyu4DzCMuTb%s~ZlNkxv#lM(N7{Gv5SuSv){3p&QJ5E*8}c zkFJ?nyW2mUY;3GLZ?FwkXgTpZ9{XtTWJSg3^Ht2Y$j@HH29sHF;pU|0*OAu4 z^xOk1(g7{)2t}6F8uDjM9$k-m3GWGmJx`3%r%}> z`Wr**0eI|>v48~ip1o2hYm;cT6J0xvrMyB^sK*Sms6%KxInCspF+}Xic*kMrMj{zx zCaoC7-|sP5{$x0}P1wvbkYIPDRe`j^6q6vH@Of~_1mcIhE1*C2F#*XmXTGxM{hvl2 zQPfVxDtzz>{!f^Zdy&iphmTl|KbVfy{))`zP4w>!mKwr*2RnTGY#4{;c+3|VX&pId zDL%FX^Bw2lgW1gYA#AaTUgD7DV5~dZ@LZa2Dn}k|@Q5dn@Ab&{Ht@zQ$A0ZC0Ab+g zE(*&nM8~C4vxgaNbXfboSXmf6RtIFYihN-ho<{NoIkWtVk#7{NO9EW;YOug?@Kq$* zo}l=(%M9NX1#7y1_%WGia5EmW6{Ct#Sy5kuv_jcAa11TZ0ecUHHIO^N5|kH>FwBX( zEq4gkc$T<2nlni&KyHz0udZ-1vTA#=t&M2?9OBWNT+kIv%R3G7lJjH=G=26nJPodhi6bPuVjbrBLZCW69xo$3mpgnbB`y<~`8yNW~BB zHu-r1QL+g&n=xcmH>0PU7=5U^xwNA}zLJ0B@!4vknp8no;j=?kCGL>n%%jkfXx6j{ zlr{>azZ9E^QcPQ0aC%D@Ive0@NVD<`|b`;pntibzN@jp`*=c7eVv(k48L{;MB5GYFqoJU%?`RKwJZC#rsq(0 zdQL^EqZ#Qya8qldX@uID`KB3PC>})iPIpJ!zhMMJSltM<-?5kBB4hC0r9`V^=<^D^ zsJ!1kRP73rQ;jdN8W}{5%aCOxGyf$s7rt2bl1OBi==>`BSIbkbK$k&nfg5fxl{qFK!7lQHPzs#}iS`e2?;N0$yqaC`j(l zX-+hdbFMEL7Dteqpqo}$&&Sx?2y*PRRPjQUReVb84}AL;wTE4d?{4h93z^kJJW=BC zbqo0RFWA->c(tR%(A)4dQkCi9RBApk{-7yX<(PVxeJNg9qU2XjJZfz0Hy%Ex8FTAN z4DGC

N%3ElmE~7|Xklnn+Xl@$qV}lXokjoWpI52E7Z5Jrt=Pz(>d_gKr#sD*HjE z5;deYv5L&$QT$(w;p+8y*pr4ina4^Vhk;Chbs4Vq(+Dg1XEWwdDn#vIAV*@|Gl=44 zYJaxeZy}ZHJB$uZAu=V93k}8lWUJcPPL(BE=IIi~cb4&uWR(TwerIZ9`mTdH2E+9t zdyl|+BJjocsrjavtV7?Fex9fnfu~tcj`&afX`RC&j&f$6OuWoyEnWhjp26qHek9qe zCcEkcqrcDkMWAbSXmL}XT<47CK7-dp(&yDKy&GUUR*RXoXOhM3PH@s24}y4vBNn<;MnEn2b@=9I~?yI^Hq`L$BC{fx@(r#Sv~_GE2Oka_h{?@oHt z>93htjce&|J^I%R$+S?YG}^)U6HFwICjKtN22-$r&T9W;veS!1gAQ>2+ld-Yz+-=3eG*;k&Z;~G#(2|U6>WKfTfU{2$3gY?-~}Fq z(~`SNUUODWJnLk*5vg#z27c*}*X`iwpWFl}9P-;R6!NBsg;?@zY_P~#zrjwW;YD=q zdo<@GSdw^{pm!Y>O!|A+=wF4pzw>{f^&`o={^H=%R3{Ub{Y??fXE2fOTcZ8j4nCOa zd{&GsqnJ+-e8g3ZWRipbgu{Kt*x*?Z&q?x+p|mEeyQyXv7Du0~D?{rZ#WPG}y+1;m za(H5$g9;89P3#B{mxKf-5KY^uT`x~MUPW5<_0>u2s+)-fbD7&Xbv|@~siJ74dK=GP zgM34n`Bi9Dgkx{Na+r=6spu?5LRqYVoP3UQ>?6;ykB{(D{mClt#8+f8z8Hg`bS4sR zn4M#Ii>|EhNcMhoQ>O)KPV#+!s7|LIjBkA|B?mQ>H$xaxlj4lNGpCUsX zKz}dbuX>vdQu9Ny|M{o*x*`1iY-%>^K>7=vS&np=on`!0Uw-vHFUoa$kJ*GQ4ABN0u!eWEqFXhGBVwS*gePwQC*L zs)2)9H-l*M4W~K)eHpLzy2$>&Hbx3EzF5wV$f>~ycCmI-jNEYt)o(R9#7uQd{bcfw z--21>jSW(dn(yfIc7s|5A%(+usTd{SA}6yEZ7N{}LeQ#e^#3Nll~XvkIxN^wgOO%} zGCHf9gX%rzEov{*Ny9Da%KRxT;5U4ihL>+hbeZo+VYOqw@)nVQkR_}^OCDssCC;pM zk6t#J;aT=AHBu~lIX#y;tgh&Iw&^bh4V#4*Si%z>)LElOj{VC0iM{FX z31Z*7XlV=nO0bl8ns$!;PDP(r5#PdCgSPmqLdVMwa3u3NUTQG-ufFE6i8bh?a-pvc zH=_A$sr&t%`0*gCHj9<);Ml7C8)#}kQSd`oG5SHEnDEJF)&i6kq*+Y=3gtNC%kg#B4RId$|1EJ-g`qt5B&l*5m&;fW8( zTwlb~+=};>EK91Ol@8~7webfdmhA!4Jjot?ZR7cY3+j zVBk>jK}T5oWISess;|m<61f9R?kAK={}ou#=VY}#`F4$CThk3|Alej)HyaG=a!&1| zz1L`j#;0N*8|hE>W6sB`J%Y5Za9HO|$9|WZ`6j|uHDOH?)ZWV~TGJd%-GdoTgwNSV z{E%HL_c&S6cBe8RJ+}r$G)FVKI;_D`$MQ-YOei-u-^NM`R=M24SLKG)koXY+LvS-Y zFNUi-Gn+c~g_+K08yM9>;?a7>w~}w83?riNtBP~{^Ec!R&G9dH^6e={OP+Q(3OTuP z7~~NFa$D}m@>(Z%zSUv3UM1h(hfJf<=f;i~NTGFAfa=ICoCWBjSf(fk1)X#JpH%DH zvwnSvl2@5Jj%NJ?AIKhzafVHK0#-6oMI4E>qaA)y_MjygCTkh~pj_<*%W||tZYgMN zY+ClY^&l=h0ZI#1-0;0l6k2SMk(>n1BQAWb&VhS(DK|7$p>HRWnBpJzcP`jSf4KCJP{b9$Fe{Rs+BsI$i94$$7w63%9 zDiTQ+U>_$`Ua`@sddS?0jTev{=X%(@0cw}UE+-#+(D8b+OpIE_imz4-OIs%^N_TQF z;o^ASDxegB(0hk3==<)o_r7PJarV7?eD~gcSz{zKbIq(d=lajGe(TA6vhKX6 z^sEVxf|NJ7_t**SPj>zI%p@ePckUnrnxVX)(~Pe4juJkzNB;~5OllxKY-?40D@OA; z!`D9P^mP?IJo$@cn{Y=(eD1+^!Wzj3EPP@1XK#tX8%}Mfj06jwbghr(?7`kjJUfR|pVu8t&Qulo119`p0 z{UEQ4cZ}z4#6uHx;OzT#cRU+~wWUNL6v|YP3 z9X5s|s`@bWdL)C?ReW{#rImLN`@LHFC{rDpKhpTOuCKHtXA0S7E`M<3 zqD25+U`gTAh@86y25iVDD|-{(V3j>fd5sF9@;L~fqi^C<_MxvLpAiAkO|Tq{2Bstj z(I(5Xc2CUB4^5xFNAP{zegHq@4lJW~<*Tk(cl14py10X?up$`{_A0MxT3>1f+D#>T zvBBoxO|L*q%C??UI%Il!!<9z;!ivxoAPMX#yo;+^^RS_|9@GTO0w}W0Sep+jBE}~4 z6TQo;;H6_n24#!by$VtKb{x4RP>V;6l&;5YP9CK+-G*jy*=>pPY^fsxAU!LF!HM2u z`^GUvsFDD5Aj=_>7xD#e)@sXH{c2M;2g+b{JGz$v*aJ_v2+Iuk?9FR5CM^S&@^CvZr>o}<;ViwV&zj+C?U)6#e@J~xzrT5JoU7t z(SwtzDE1+PVrpK}tadv+7SVF0JvY%$H+E~}zox89;Sa`(hEh#ES?b80C|Rc-;WMy) zUd7wZ+Ahu-!9YegGe|RWF^(!*D^-B8SJy>XdELqtXVQyQJ+c+gN4=~9KNDdlM!hwDXa=5U`LDz zcX;6s+=5Pa5Ea8Ut=1m%(nK=QJUWg){jLX-Ao2Y$8IqVB&UNRc$Z=tmL8$5Zh5zKNwJpdvZo{-!Zk9$a?AQX6bVGv8Bjb@#FW7lm3lei^=qwtW&<|p8l%%x#G{b+`e0)*4M{3yQU#lJ| zr>uUgIWN;|m;2rry!Cu-acQO?47ORH-W$3B7N1n}-P`VH8*SZkJBw%UPT#Ba;R$zN zZNcKP4rKZED5Wt3_>uP~Vb1FZb3{k$`1f03P0BA+ls>FGKdM#l6cl?QTIMLroy)qq z{rF6l&4{j#k%FB;KrD*C=u<3P465DcY+uMW-f_O+UMHo=nc|1g8YCoM7@A|AaH*rl@D>d%j1cI?)#MP#HggvFbN&U*9(i@@ zp++d6M0?A*A;6v}g26ElaEH7Y683RyV)M)4O6J)2-9!sJ3%rr%28oi;^GOxf5$YTF zB~bKEgg!+D4mfA9xlzLO4#W8pGfdqVoQ);f0EEw$ja(}xz1B@$Rx>2eCf6SQ^HNcd z7hgM@rFkQe30XQ}$U`2g(~9=j@2Yn7-_`h{hBFVV|S` z{y$L)gGQUyH zh9~wNZnvvsq5^%s^`&Hjlp&CZDYiqwl(RZk>L2EAImh-l6`(xZ6yM6O>LY9Tma4 zig|h+fE?gcLT88jm<5asxj69u#F+?z8Rn}=kWb*$ZT3_A^jMh`PN zG;@WiKOnezLdgcZOm}VLaE7n7G%&~nd*H#$6ZpgV0AdX@G0Ms1H5?tMwDio)((k9n zv{t_l-F8OUGeVxR0c2aQG2wkP0eJ??S4o)RUW^NC3X=74m7xExKx@Xthd%MOp}~Rh zieO+8Pdt;w6EU?ye*DQ4CEv`!8R6qFrw?z)lB)yo1$buEtT~K#S198WW?PDgofjjP zcJw#TX!8(*#{7io8ABfqrxdEt4wU}*BHv{?MRDu$a|wh>KteFkTpfJ}kOg;OW-hRJwjH z)VYP*K$-qck;I&2Kh-X&OWz#RlM?#%@o|OxxE^(kSlB?JaQ#mtnZS07iHC0UVZtoO z$h9JMj|?`og@Df8G--GjUPt15QgUG>=Yo|1j^NpliUQ-^2xhb0))F&F*LEYIo0oSc>Zez> zs7ift`O@CbVU?Jfefe&pv~S36;_6|6Q;VTlMM4 zU;6;eARNqE6?@0J@7YuHG%VkTh)_{j*l8zqtMrEyjlJSe)4xZEaY9&ceBdosFK-Df zF%zGdcaN9c3&i^AbPXngCBUX$ojFUZ6jth1-n$SE*UbJR%5?F3Ub~4X<~u9gjG~C) z{VDyI2Y^Veiw=`kygR;3xKNLKow%#6jv3WhOk}k{O>g9%=hnL-WTey?_{8;aE zC6`cgz^ZWc>tuU%xjUqpbIv#!mK~O@TK)AhAdwHrRFMv25LaLE$!^Q4mKs)24%bK$xW4w zL6=8rUS|9Id~aa)RjH$+(rKg;id%KBV{7X$6H9FMhZ#s|tXI7sUjA zwern0**AS>dUB*tb^Q+1c=x0{?f^hs{4Y7BJ^!_ndCmc^;4E2J0Rk3g76bk2>4M_i>9x z<+kj&?-!HC0hNej#YP-OiX>-`aJvFZjJ;;@&D3@^GoG7$45i&+CW)GjKtLAU@FWF? zv`35=MQ~?W2)6FAxhky1b+qJqiuk^I2s7K9Ah0)-K8;}R zKn;8B&PN17f;oMeLb@BO$6>q;>$bs9f{xhn&phQjzpP7pWNxi}5q}%_RmZEu?7m;p zc!-VsC!dN>0f$3`BV2u3z6pI9?KQ0>er*TEufq&V)C<0#Jqt(G$TSv0Ab zG)<-%7_TI6@MOG>GBmXG@-^WvK(FC@U$!nsUBRZ8(=QIf+0_YdBF24U&5g^^X}M_l zfQJ>p^g~fH%RSQf89jO4cYHIlY^z8s%ey;Y!H3z}BV5ffLGbDwm`s|a*S@?=EaB$- zK>ZOMxNU!XrDkB9W{YK9hLSYJ3<9Zo?o0r~kMr)J_FwE@tYXXG zHEXZWuT@f5m2^b`@0;=YS-SDgV3&9koYP`Q2}vPe@nFvf;Dmp_2g6n{+FZUy(#Lz7 z+9Rrp_1Ye?Na5oBW()`D_93KEUzMgj*R7a#v*M5fh23u^ma9tf zGRB*DpyaDZd65?O$|8#^wMX^+nO8&Ng9yMm-*kr?U44w&N8*vD+cXX61p``^4_c=a zc5H1`aO!~zc-9JR$j2$d5;Nhi+BREm!`H=*Qyo=by<_8wUzx}AuYJ0_AOI@cSyQx# zTOXSQ&NpN*iRUK9!$&sPfZIFLL9NrK73pF&Iw*r z#C~=ztV%ft=J?_*bETa+i|!@g*#zHIMS2w)&unj!1|^%r5AWoaR*$ zLEMfEdFGx@qn%3|S=2S1he&&P^{kqp&{>G@QIcb>E+RSHQNVE5U~IWnl{l1}Lr69) zuy8yI?)(4%*U#Prn!TNd{b^I?oRPp20IapQczv(OtvJS!Wl-%^tp#kncsJ1> zaeVu1GQ_L$N7zvMSwnGePYt)f3gZqdeTJipnz5zO9;mbMy|P$-?3Q_LA*WVP(h1v4 zlmGb$ebjUKnM)zU8}rq6!KUs_{us@QNM!(EisHMrA$C=Ob+R!5Lo9K9VeAlPA9)l) z5d+`;Q+L_i?x)4bcz4T=q37AR<>FBsh_&Ui1#y~8W9*h;5gxFdX@S$i{vsL$I%<62|7Hh- zY0B2$hr4*cX?z4C(%M=ZoS{CR?5&wTYF<-iU?IKu*pUSRIq!Q_Nf%=tn1-I)WAHm- zli1CA4a->~U*SG6V%rCXY6lJ}0!UNHv&1=KTm)U9^OM-S72O61zIH99WUuTrZn7a4 z2xtraOZNkd-IM!%t61Wnp1)$jM|WnPWtr7Fxy*b>)iZPMR7$nRfnmD}N$&*-v+})I z_xP=|@p2BeN}}>N=Dp|HR9>l+FLXjiyL8U^re41zc0;J81mJ59JRe`q@DpFG(0f8J znUF8G!qdhCID4{c-{3OT)m{LB^x3kv`j91g$ik$Rri(MYBQ4|4jG@YnIk6;)fW~>; zZ@!stCF1qy%@C)e`+~B#Pw?{R%!wDdzQlaSQbz`*$)}lS^g|glzqaQh!o1n9U2w=X z*<#O1cv`%u+-_|sU&cCSyZxF>C8K9cBsN$Y%x~yQb@Z8a8_Bo&Nyt> zJin)yUmv2fE7Ses&((9UV2ZK2}|8t<6Yd$Zqmu1}v)G{k&2Ts?Sv2AHK%O5Z_Dr@vva zvLgWMNkMO2IMaT2pTb%cFQW7qS7DoMrtq`OkiX7i=+1ivD?}Jz_EC(bKZZuXsB>IlYiC{cPK6E^FS~k7%U91$ z?WL`^%-D$*14jDyjpglDfJFqEe2<@4+(JGL7pjwC3HZCMDQ7W?3LO| z5q)Ba7fe$360WPghwj)lH3zAITe{+vq8(-4MRG4cEw{G4IY$)~D{GNZmd-b@ePYov zmYB`~gf_-Q4CV|unkRMX7oX22ZIHTXx>tKbxF?IV?F$g>7IkEu^W{-={o|TZ&(`|t z6)f(naTz-RY#Y8>+ zEKt>Jo5t=@Idt3uH`%g)PU*>gj92FNB~xwFY_=ph40em-(7x6)qd1Tl@N?C|Pqc03&u;%~e&$?fdH*;@5%)?loZ~emJ*`1Lv!)@fSRl z2YgJ$aKeR*zB(t)Exf1l*0t=#aKM(zXT7$t<`Du;%CWb&R&l^QU94n7JiiuK^-=8> z8+>wan|k`gcx2jYoK|l>?b#2B?GfRlsdJQrzv$tKUDKXT{B&=gN*lhtj9iF~V_zZW zilb8`{-hk^N9?%rp$AH#^trcjO-cPY6r10kP#VQjya5le>B@0xzw@TN+aZkG)H#3e z`S>=&MB?j>N|xK7xJio*LK?89_iCo=j*RPcwAgwMRMnF?7CisA4|S zD3QXpG-?y~=+z-bbQP1>gMp^GTSB==F zRKmu^XQ-`N_o%4=icAmf6El8zsPHhb?^qC2g{ir zr&po-tEEVcW5Kp~DkVPU+~KMWQS-7S)2weGjt|cfYh}#{1K$heo=}nR#q=4@0tNajMjsF_ ze4EezFby<%RDB6^d8t=-jhg z@Hx%thPKh2k%oZ3)c(Jj$A46Q!?y}?fjX$cT{yWY5IsNJWhm;C`c$x&CidT?TjBkV-uh$jjlW5~id-ph zrNET}R|;GyaHYVN0#^$B->!g~f{Z5Cok~nh;>$C^!o7vO#AS$yaTA^U4Wi33u(7bp39ns# z!niEr@?a85jLRL^IR68_{zT zeQ&c}RtFp7F1kYKQva-wjJ3VB+22;>_YoMkl`pG@je(DG@3Lr_XDYI2Gq;+mk@7Pw zZHx!#Z{0(`$k4A#MgFs*FRRc#{ktNf=zvM3Ms+kx>~{-JVq!i*r=yj}#Qb}P|NHpA zuthKHi*8YDPu$xf*$B zCI5P)1)?K=9e;NjL2ny9iJWRr6_oxp6Y?JyMGOga?x$n$)L%vXb+N7JJ>}ZXyK)N8 zq%j`S|4RS$#=-UbbIR|bjUns#yMo?bzeD0{zo%3((C^T}?vE*!{ddn<)%YF$d5QdI z!$YExHVll{6KieR1Q?HRw|tlp%?0Ow{&bzp9Ao)*+J98^|MX9Y$@EU4#Wa`R|K&H1 z4Q+Gh|MhGx>(T8=!+$s6A6eyC|AD4V9e86T`6#-z?A}53S230rNF)}lT+R|{SE(rCSCEupGsVXR|;GyaHYVN0#^$B`zr8y!t@0#jo0B8_S-jE z)ORt;Zr*kKM+`yw4`dLwyJCocAkEe8D+R6;xKiLsfhz_6Z56m=h%AiLLriH}bt8Of zNem2iVw|aemLlG`UQxuqZ5O#Zu`30x6u460N`Wf{{(%Dj8j5K9W0~2XgS1O z2>*}pK$eHASix0}8#_txsyeYU`5dHsx z*Vu1o5xK74Oa1>7ACzDmo!NiAH-j{~kGz~THUHJ*mm~epX+RwP*|UR8?Nj|`@%o=@4Njk^=*szyO|{t9dE6XfA&Y`_@4;=IV%Fe zKPch$A1+aUCSO?*{#4>Byi(vwfhz^B6u460-&cX(t0GU|1b93_``$L+E&NA|8Q!CD z=)c+kiGRDO;qU)Q;mR%Zzf$ii?Mi_w1+EmhQs7E~D+R6;_`gp9^q-?N+SjH6lJ{x4 zw9|qm`MbJb+@g#2$0C2gYU+g%xGF`g_2RWV-F6NE=0lSHlv`M;80F*y#F`q-ai4?b z;z$P*rxNW9(^5ESpJe;rF6Zrt20<9~BnK`EJI@x}P9x(@9dGqrC`iC^?bki0Yz;=; z(-+;#s@BJ3R(&s8PGEIsz?_+sA;+&9HEqYR`T zY%!RvE&VqROz&*&;4g?lqJ`|fbRV{qz`)Hi4Wo&djma

  • hw4an(EDE>*4sy<^HO zJ4nab;U+sTy*sLR4N>$7F*s$?zi4&6-S{@^<6ifr|?qKTLh&)G28++Yc4m{5ba*G1m{R zo#hlWr#g>6rVN_t4N9U6fU>_8%?U4inxJ8jwFnL$yMai5^W=WOh;*%@foU#A`M|qz zE-QWxNnp5DYk`NV&=ah3ZrpVxSfs@%Woy0gY4x28$7^lk*|{P#KIQ5fj*Po;$I|2J zLF1up;#~fS++=Gu3!>vk-#zFj5ZoPc)wARVWjXh`wKfQXp3tJ7>m`J_BgZdZ>LiS3 z@$S%?a-AZGiF^RkaHGvfeAkiv-U3l*A&w*_L|)q<_^dBzED6t4QM4h8o)1f>^V1aR zXP0mm-W2>?IWYxWWWGkOV&a<&(zq`igDYSCeeH_~t)Fg$Y{Y^nbzs~|&}DYnu$xIp zZa~=E^t{Ji{U_H%H0h-Cm{+}HeUoxF7cn-V1d`p{-r?Z|u`YIc5)h3MFhD!4wJXXa+BrAZyh z0(jovoR6ZD&7RrU@R~Q{9iNdDZpH4W# zd92Q;IXzr%n78W_7oqU%+uPjts{+Z(={NCGtbNO3bf_eoObr(+at%p-^vl6W;>5Jc zwjqT>ew7AYUQDVJO84|sBlGL+l8eOmlfX|DaXchEDu zs5B|}0lAJ5U+IlBjSQ6u*8;=9T8`yV_;x@b>H&%Nz>p^bKC)a<5t9rmXWN{_>RF)M zOC|J(5Bap%@o6TXqm9C-u$dbK` z(ptYgDlNT&?g#r<;%(<4knK%x#^HxANAw47uoX4%DlbtoTk*b3-ivj4eg+m2gh*o(A9h-QvuE<^*^lbnBCksTg1k{{E>DoNQ=MHpp!c7Fh2 zYb#1Ne5S%iPfwDhJ`c(Z=2H@Dj`RSyXqO#NC>b$;-c1X};d~v`u(dj4Uizq^AMweX zjAb%>ML&}HAUQWBhqsyoL=G1a)aR+ICri;i9ObLLiGM0vO!|oc%>OLy>@Ys+?_tMw0w%E*3 z6T&C1Q8!x)l7vmUmoSAkT=;$4e~@Wgl0r%qilX4^286S%cg!IzlcrsF+lNO}b61T^ zfx;j``L672{j3tvk}Qzv?Tob7i%G)HM4=({$8k!wc@fi z@Z#Gx@!X^-pk;W*7n0>*lOgE^+g^AFEFUD-(=onRTTgCnjt}QyfkE8#q0Pz*N%pdh z!ryXpM1#}=PwSH`WQF_r8F2UtLX)2mS_g0=Ps}4u-T|{K#3w1&0J`O31rwiiI)}O( zAwxUyM6nmK1B+Bz=1+8NbIa)1mt0zqovcuNO`V@HcPSU^D}M6HtkMD0!r+4%3(8hF zQ0RfM(cW&-NJa#TigB1DwO^Mu$c7P4NZxQ034~MEW>NAM*d=|N5PG~nH<81y-O{1I zH|W))x!cQDmOGwOR^lja!0(*Pd!R{3R8-`XRe$@TOU6ivuH;C1V-HSz1WgeUt(q62%f3Rs2*~XOGXRww6wZNnJ%< zWuJc+K4K%Io40ZFKm%c2@N+=UJiI9*e)=?eQ9$0{Lm#+)ZF*iJE~Ihliv5UE_;5lyk75kjwc)yG(tr>9FdY+5+@2Ux)=odd|>{-RjR)T z6_nmc@Aey$aRgdTY3hOLf0)S~?XV-sUCkw|A-PSh!K|zkY)iJw!d-$^Dz{;`!j`pq z2Xxnp$Wx)yY_0Q@fXq9sX*U-_jy|P5Jmby@?bG2tgnZyOCw843yAgiR+5i-zc_Tdg z?YV(LO!TX6O>n1`{iEIusM8r=XJBZx3!gw&Xv3T^?;S9_M}J` z{UZnc^7LNpj%%pEx4@G+5Svw1TIS?)jg6`-u-_~79-Xgz)WO@4)ANh{p}CPLX7p_0 z=uHGE3279bPDCDTk+>35fcLH~i=D{rah!(tSO3gMd_KsX zq%fR`ba>!N{+TNg>gKVKy*=PN71+rn1JZX{KfjnFa=cvDx-Idg#)2Y~h)>=Nj`Y1*3?Reg}rHw(> zS?3*d3{%j1wysidn3E5DR`pyXT!on_e1*M7#|PP8zu_1@C1%6yaWbv5y@?{n zb%Ko%&)sRvD9h=Djxbz%h@S;JzwZ9P=N$`jrcRjCJ>nVpLy6J2X__kV+a*|5=2vM? zM-*x_AL$ZiTLpn`yIo)K#P8grXDT2ndWx5#Ctg4uLJzVv@Q$UsuazIC(H`%kRLIdX zOoL+1d?c3&)Gk<@Dk?Qn_7=OaxVttFu@&GP9-{?Hu3WTLt-t5oL6#Ch_!T+{o+js# z1noE6bITKSqC&C|fee=Ji2o?fRH4y1UlfT^Sj&qSV;{beTSL8r9iQe&Y>Mg~M%Lt~ z6&%F+sPe%q_DN3Qtd>o32GVuR+|6;@mXMYKF^`uE;oDvs@lFB};7 z&ov)&!Y8s}p>@~jwfzwEq#1zq_irx}N5yM?)U-5tyo_g^CouOStq4b9iIthk-1SNI zf3k#r43l43z+Bs^;(j#U8Z(j^5dl)aS*V)%Oa&cuM!M^9@+9}x8B(2N9h zF58u`A6Sv9-&sBDLU0~wQBaVivZNKXauk_BZKkN=6j3^=9nQQWHgOMEcVP)$S)9+` z1anFWN`B}w0W^YNr> z*P6?izK4u+vxtK+@qga-3+2o0QxcAo6NOFRv5s}1mB?z0$x~T@fVCdmJOiA&HbN@C zB^fg0+i~7$qyr}oJ@2;pb`tWAM7F2ILd^hxE2eidhqJUtF#k}Qy{7nMr2dCm3wrnb zXA6C@OT5<4x3z6kNS@W=B8y7vq1+L@6}`&(TY-mCK!fa@pBjo&$w{-IJ3f7TjC};5QNZP5nPkusYpJmh=;bp>&j8*08MA?0R(#vGlvHAw#Gy8gL@to5*bL z4l@tE2~BOrvq>#G8x{}C?<^q(4cMxV9di(i9gnC;ROmYkD?RUlYPy&D?gG|B6@Ni~ zh=&{wV0U+#3kI%-+U1pyG~7gRe9|XHIzm0MzlzfbY`6JwceCea$_8D z&34x{Hox2f&87tHAb4Tzi2sQIQirpca2djuJSdbSxKpY#38^sgt5er`rQl_hu5~7m%Pj#HVCfDFvNe>&YP0v+OqEG(=%GvpQybGfg1B1) zB7te4&Li)5`b+BN!g344@j~>Z{;3 z5nk^K>S^AFzcYtc6z598H`??Y_v9*2~++#f=1^2f^(LJJRkTJS4Eb>ISTwu=fX z)eaO`1X_n5gh}Cdxc6C|SAR!SJ_N#sLeT?|3=a*{v*$22Xn3$%=wd6VTlg0)L?QGE(m(g4%y3x8qVLtQ`%AP2w)09wlc@&$) z=v?BbNeM2M#2UiCk}~9?AwioA-C+!6DNLz<>(po|GIz()9uDjHgA? zJ_gE_bW}?zu~?{L6UIoHiL62o1k@>Wm~|LwX=_zhAG*ekyxydk0rb&IS7l~;biq*C zH>qWwG%&FQy>aFlZ4a)fM+KP0yXRrI3qebgn?M@cd)mW_^6{0WB<#>>Tj@PlHCr{+ z#xM2t<}B$r&P7YR#r16t22~_5Tu$cpa25*2nrckbnL$!mHl5s0Na9cWvLyozL;1>h z^^rtP9gf|!X)oIMr8ZMcBWW6lQW%=<%L^TDHTq@Oq|ToUyog+45jnPt8$a>CR>g#!Y=0qp-33|}KyA}SMWL`zzA8|bJ_y)6u#BB&3F(6HfC4K* zPalAvTj4mVQT_1mb=66xZQOX}Fd3jrize7oLOrGogOMe6me4lvFbD-9SWo!H^=Y?C zS{P)b;k~YkR1A=8mEaa<+T&=6`Ox8m7BE0wGg*Se6Aj5T1O>^p7iPHj5dvuhpmq!%^0= zFYRSZHM3H(QdVu3tYHB8OQv5gn$aE+?P4+A-7DGs86;ZsuKopJPs9_FOz6iNrjNZ% z?=?gB8Xx_#hF5JH+dK+do1~ud;VmOY582euSr&(gk+6mx5|}afdbC>4j9qMMa)&_Nf%=>&^t zIr8Q&&1BRC!0oN7Ucj_;ZU{8 z0mY;|`ZBa;KKI>#SsNbfO@w=-zO9BI^v|^9S54ILH9Tt|uJRuc?)GoXbg2v#wuhsH zviEq3gH=iV??qy0s3KE*JcyII}gJN=^Vyy^AdM zy$DA62&n@!TOVErY21;rv>O220C)(?+nWdFEl+g0EWPEIepg?g2-jrF--=b2j~tO5 zB&S6Ibeu#BNbHMZWbZVQfrxE9{Trj2NocPbf^?qGCg*n-1bdKbEvf z>9Gm4zSb;Tgj%sez&E8ju#1WjQiks`jTQFjMCii#N<0Ii_j6HBSZYwA?xUXRO(r759@k0r)!{-KKSD~yU5lxo=n?)W zr_lrMlaT0aPoH{7(u&6=uj`Eyx{vM}GMAmw7E(*n^*86c-8Wb$+>&gX93nUBA5dVm z%nDvAI#^%|HRDrMz|q{xr4pR?(`;Jc?=yu&dd~^71g@1oJ>3C>T!L?yPveMA-ceondUAQ5PXPJxjzd! zKG1X#o*n-rS|~WdjV#bb*28T&)XK6~0PzcLyR7Hb2_rty^*d!fulUx_WWE&8S;$v1 zK_Oq%gz_9c3XuVHUW^HoT88qR&-E%}J`PrFTTav|ktc{Eb1ZORarxx&T9=mhab)sV z<#Gs(>o7@|74-r#J%pw7NMtX(*9T;<7|tVDFE0sf8R2K*cXIik=} zjHi*qo`p3i!)_V~Q`$gZTT7=f3kB!x0hLX3 zWqI{<#ScJQM(~0g-+*X)`|!aQK@EdJ4A4g%&{vcdI_>u5Q9Rx~ie;PD0v7`-!54C8 z1to?IG>4S%Y9^S%+PZ(H5OSD6amngYyf*`%gcq)?l4bnqhiYfPw2GB#6FMPNl?4S3 zDhcOT&Y^Efz7o%nDO4r&GnL)-tNfWE_grcU7Ru9X@!bj@S9WytvT5?Y>f@E-vNUyu zSS=l%vfq~WW9yUhPx>9KB99QI+3c{=Ub@EeT!60+D_DHfy^}YzIo>+aMQ7#q8tabZ z^ZGJPAp5%2R3X(*jy-N~+z0NL+#L%{D{vmkqlEaa{(8QP>!tv<_nEDkBg*@gUQ69>U@%#5Ow6qwR%dhK7*^Cx#$l); z#Bi@Fu}Lv*KPVv{2jDu~^%v+C%eAX4aM9IVf3J;(Zr}*|lux4A91c6bYR@fiSSg={ z7bR?0#R;1_fnQ1M%q$>b`2gE_-reCdTID>X2t#W`sYzrer3|+=l5&^GAcVUB`D2A8 z>=96Hw_)AstI~lIK@5PV7{VFi>-RQ^b!$JV z;!E7I$1`($CzvbTAu{f%Zk?p6`*W$P>*<5?A9?d5^$oep^VD>QJmR!V<<(?LD}C4M zv;T^?9oa0%kiUootEVeLyM!RsO>t?6#*4Zix;2oMgnnL># zxXVm}i^I6f%sUFWkuLJOyxvuGxhR(0_+EsfGD%79w5`FOtDHO1;Q_7n9TILHokpMDyV-79?o?la2D(Ou237qo9&CQf z^$>98a~J)y`jipt&nV;2?XeBGx<$Ry>bPnlBeedI)iM_lp?v!&4|}n_EM!AkNdUog~|KQ zzed1O^cU2RSYkqXlKVGi2^tF6SD0}-p`Z885?PR)?%;})ODlRWqALT58r9SZ_CkHE z@Q(zpM)uy0c&#R;Wb(^~&XN9Fo*Ex!z)-C!+5jC$RwuS!DXKzG>)+O{3zx!oglL|J zoI%CVCiIWOirP_BlNTJx2Y)0 zF02AFn8!3N7UubA>Fg9`hYyTa4yMvLinE!s|KRMB-<&P<2WQJ)a`v?PC1-Qk zTyl2p|G?R1@V{_&!6j$cmF`uP{K47ug}*sF>`Ra9FU~$wO8V1e9@l?qGDjOD&eE!L zSU2C+KjHowSj(yj>Xt!d0rZJ&pP9p3h`RmsAw?wq->5x^78_=*M_tu+BF?LF)|Hi@ z4w~|#2wxd6f})7Ec=jH)J>=bHnaPRp$GBb%l}meh${+S}=%u~9)Ar*qGA0)s1HT^Z z5wArVRSX?bCnXN3Q)XVwsc|p|p9NTxL@nu5fi?RIxgKXK1Gy+-nV31l$9+Fnd7O zyf1L~_zI7|jMizKU99VCxkT5>$|sAPOM*;0ra-7hXlsUv9Z<+v#WK6} z6jB17nYrKKsSVuSZ!DYt0p`1F11e`}>-G>ck{!~9Dcz8!mee@Av(0mQTML44qDm}i zdlV6!X{)#q6)G&f3+(T+$Ziy*+FKb!)gj>^u=bHiz;(f2J=xs$thZHoq9%V zA3RGuvKjqob4Sy;MLibeWQD+UlJ_7=R8Nl`MM*nu-gG;&J5Je_m{*XFj|YdXh?NU31hE?sU^ zmpzhtD)z=BD?n-Ev2&N&q@R}{+ESdAs{!nlB+$^@=vE^Rv5kwqm4}iSap=I#_kvnw zv};qHG}%j&|9|X#2S60bx^B-5%#eeK;O6Mnxq^Qqquf4w7NWhzhKVfJzbt z1PLk#A~1s@BA|dI1z|=d36gVuJq)vFci(!?+jIBseP_>&;5?m6AG< z5X0H0?qd2rI%;R|ahXNPq?e`1&1>J~_}2ND9&HWL2dcdWitT=7gIxyWW5I1&?qOTX zOT0WE`i*-(cD`qyGCvY%-xO%bb`dt# z@3a&ubRM#Q*WqO3CGK^rgKeAbzJO=D*Qs?ki*%4C!}7Nch3z`$8ik=OU4;b&*cul5~Gih>VBcK65^?p z%JOR+bQvpih|&+P@XLS{hMha@IDY2Gi~!gq)lE99(fA-QwmsLYOjEzivtIu@@*?J9zv86u9 z`-SIC-&bkNwSfHgD&65(yqJZu=$OLLLtERdeZGl;s;cN!jj*bMhmKr{VX9I^UGW}~ z9u@b6+pt3I##isHim7*@%XqFVX8#1z+XJm(xyfp^IZYwSw+dR%Wt~tK@!9A^?T9XPk5W$4vl)duq#=(~m9bhS1PI#qQy z=lD4Tvt3(IN(qB66B4$u0SgOFP29c4-yHfmg<^}n!eeGjBb;bu?tg>oBY~pt;fDL~8!W6Y@p?B&9ElI~xu~KUK$lIm;w$ks-Y3ghlS;4zl&)M0(@UJM z&8+y?CEVF&FXCLXPWB?z35MLKErNcaWsg^`d|CTAS3uKe%n`zz9Y+6N+s!SX{@l53 z;mWngcB@Y6rd@ibBI#ktslw+^)$}|Z`&Qk}t(aYt<90#aKAclnHl!fuQJ=2s5H|9R zQzs+HYTwlApPo{F&n|3bQyjQ)hf8HDV^;Q7(Zk2J@l{UQBlFU*G#`WUUQ+)K3j(~g>EOh=(?0g^0=;kufD4< z&~|z6y>AaRTm+YdM!)CjC#Q7D&3AnNI_7TE$@I`UWUl1Fsw}yp&=I(w07z=LIBb*J z1-S)>bMU10!EX;2@uap?S@&uRr*%`KVDH@6(dfX1>^$xpk{KJfIUKuAfKA17{AX=&oy1r{V!_eJSi45?$C99P^eLPo7%_Jam0IH>8>@Sp!c#J z%y^~Bd&*bOl;-)u8krwVU7DL46O!FH9kr?cJm#+U>5o-mjLCK&p{+t7wDmdz*^GWo zXkT$s+hQlb>VI}T;1r(FuE!JFy-n)3$EXLp7MumCn;f*(j^1?OaWqq?e){-j@m)=Y zB-hoC;m40C7$h?fjSiX}FpwKzemysq^wdC`DON<-s|(ge)@6fK)pjo0)x;mts;#n- zduLep99vF_{r#hggU^z?$O{XwX-02v(CS|pWKZJxytaq(HQ3syA0E92@hZ9%FT+)I zt@fxq%V_P2*eHmG-+7&*1rG@dl`q($c*NqXWM{(~#=xmh6!rT9pQw!AJY2v*dBC|T z^LByMXjABWeKWA>z9Uc>`-ec~rmx|60eCCdoI7*VAmk#mw{R0vP`HfhszCORKBf9F zm8|C#Et|zDh7~t$MayfO@6bUfpW8mD>MWDqHL}!I8G7qp=BjD;2M58)r#^O7WAMGO zXWF{y30048vVIaM`D&_PJB~a#GsReGKI%~J?p&1Bw+|; ztK(EQ`qF9Va(D2B%f(C1KTA^-@X{1-!*N?9#&lsaK2C)bs&gswY{vcFF==mv`0pD$ zTTKyvL!`5d@pLvwKY+>1>T((%Cogbhh3$o&9)+&K4!m z*}ZzSpOLQ?efADxNZv@%w>upL?+SUx;&kun_pNQx zk2$I?#FlvHte#%DI&kS;S>9&U<tym$4@a18VGfDdfvx^kK-c+qkXbp%D5AB zKsx)uKj>^>M*^J<5)cmviWIOHd2cW-`$qYfwosNY3G@4li)HWM?TQ}2F-@+1oL3(T zt7W>>sb&08)j;6NK%mXGASJ-pUURJPoK{EBDa8Y~43DhEhZrYc@5-N>3TI-Hzl&v` zsCf6if)4!rR^c{TaWP6{ihSuRB}DTQpgFaqxSudq4dX52-7Y zb4f=!R*TXt*NeYbbkCctEts?}njD=snTn)%&NjJd(%6OIlctMcMYHJ)olZiBja9rG zt)PopS18BcYIol##je_*7yQxX)!3>cY!&SA{NA`YUxw@L-&+n(b6qD7WOqkK7Ve99i;~tIQ1nm>TDdczx8j{;a`aWM8 z0CNOdFILT>W3$2zC=IG9e1JZ>2p`08B`ILTKUQ%yh&9FMtrcC}=)S(u?Vfc_4eie~ zzbVTancs?>st}#57@!QPrFZhu&oKuM9_P&XDyy_ov1pPBrjbj{9Kd`KyqNcr=2~$G zV0>W7WJ!hgs>hO~Qn7RRh&+8-^;z_S?s89wGTkI?<0?f~Ytcrxtw-{1o7v;joNG#( z=1_+{SC{+v!g#06FdO+IxJvdQ_`TspY@yRT}b5>M#N5L*b+8;ktb^o?#5r zAMz)+0P!H|`bi$#rMqb>9_WQVN-kv8WdSn_e$YhB8&Yq7b~92-T=2%Xvz9IDpckBS9xxD>_Vp+UYBM>Kqea%p#=v z7Ti~sdMOpLbd1Z%nl|+XALc&)*~hTy5n0}oj82__7Ci{)xCl!2BTUDZLn+2PYfK~O zX#rN2O8BlxSbT*MF(-6REv_#2$2pAur z7xZAnbv+W#RfsE1mPOzG?AGp{7)g^3FBqk)^zx}dzTA>CT~V$dri&aoi+b9sdWIYF zHJJ93wX}XchvlOE%edEB@2V*r9v#4aOBXL}r#as1&h^oKf=)0{zJ3F4A-`0&#StF( z?2ve(g7W1QoDhBVS}jxAq@9JaV`G;9YGaG){A9bhJIhO~fraQ4N>GN2|NF(|u(PZw z>EVT@G>@LMSiF%{Y8gkp&Y)N?p-Z;e$7@HUkuu=NLl!*Bo9`2#Ls_9Zg@hUeC{!vO z1)$oo#uwQwMp95|DUf9CsFU}Sqiqka@Wr~2?VoMGS5x{w?P9`bwHM<3@wpo+As>;9N`HOe2?$x4eUbo~Tv zTsf>}`SMLJB0rb=RpNR_g0FEy$HU?$59d;WJ=EZl+FJGwRlP_KzBBF>@lZ? zEBk!3;hB!&csV}lj2Ke|I{q(CZ43sDkDD=<0;qIs*;`i5#g#`YrQe}i&scO)0GO{9^F$ZkXex0RRU%c3uPYq#ejY3|5o2Dpv!7?uN&ed*YZ~oj~2}Hi!kM7S>sG?z9M}4XvAqR zx|TgPBc@c=RY2+sbB0zqnC2@!jw#eJ1-Qn9dTQd4?BW4<<#QluTn$vN-^0S*jp$Bu zCIXN5C{Z7ZvhSmJnuR2FbcD}M(Y5;s$ge2AErZUuLL=X>U0y*Z58s}K& zHv^{qdR{u*A(v2d4!f+gWX&+~#M0DDTqOu<{4nx)4Mkcom(BVDYTuH|ejXq{SRkvD z$(XT1Y5fptq=v4t6*@XcDgMPhYP=sF@{ETSW27Ww3UOH6gA02Nl?`Rkp9&CXN5UT7 zhE!dYsIPNYj;*Dy#0;TS(J`!8;EAS)ooIweilt9UxM@F6YI|VZNFJY}SC*nIim^;i z=*uV)#d^2jivGdjPYC^PDF&koRcidqU)E*3H z2#e$w-K3eIy5LT}U5U{%a_p0yl35=o#dFC=9l0Jqq7I&bD9y7#UP{A|a!?(oBh!Tf zydqDd9eUZy+VmLCK2e6GQ%D?#HQadjt`{1TTjI(Mv?_zj?bj4T7VKV@yZJqNaKdGU zi=ImN@_55qA+;AF^SE=YBX6f&USp26qHcLfT6_;&c$t7|)PiZW0Gqoa&!o|@AtN0W zx$bV9j7!Nk0`IVo)Q98Wy^>)|ovIaIOB77JpMK64hgnjQ?+DHHwpiyz(Hvu1Rgrr( z8GG=7_78r!a7ktMr0v%*y2JrT}(W zI^k8u#;&=#Qy-Xy*F)n*SSxiaUb^!@s`*FmQ6^8G)KThrxsVW}E1;ji??P74;3|GZ zOyy2-Y;X{T)wS9^K2j|K_a08gZ28gF&si-`NT8xppl%0ag6`1O^g@pXKx`a+S_=V!AN&XjEpSz`WrX0 zh7YjKb*9v2>l>~bjBB!^U9y;z4Zg=>Yv6cDN7xhvc5i;u5Qy|le$yG_ar2Bhr2g6jf0~*G#d=Yni{!mBG z=c!^?k*8dJ*Mp?OrI+XTg^!q03YN=VnPiW9Lv3MtE{VEKZ<7MYc){YP(1Ju>UECBU zqbFBOA2hC?TpUh)1>0KP#p}$Gv80YckB6Uj1zsKIXo!E;)(_0Qj|`HBN`&Sm97|tTD}0a; zekYA?B}lL{-X7w58KdDA)158eQ5NCPtu~T?%Ig94yTvf(wVVcEHMFyS+*b^k`8}lv zyN6KEeI9pj6a-C@+4ysoIvQELGK*@eLB>X2;&oRi`lA>yMmn&cX|$!`$Xq8>}4S!liGC|$K3}*-`!X{ zP@w70*zk-NZC282!6g2T90NR!bYkdVI_2?rSUE!x+WMSrY~=#OzSk<5tr0Gt*ffsL z`ejD0d;t8zxGjRB#LBQMv8a#QN*e7c_qY+YK?-VS^k3!=VzOBWSKt=Ugi>>8QG5F8 zXxzfbla)SmLypYtYH8^}sk)}7iK=~=VR*izGUZ4vk)Fi5fsRf}u~)fqhtgUbN$W50 z>8oab{n3||KB$^=8fLVIWJ{bNHUsHde4~T4oHo`lyF$7--SBnR=&hdUf*HBduKNs^ zMj$LjP^z|-ynAVjQH;Gjp+ALFDciGWQY3tyOA_jNxy}uf$^hS>(;f8NZ8vQ(+O^+9 z?d|}U;+W47-jgu9X`6+c8L?Hqd(`T$qD&dlt4wOr6SWF^X05|gD@HHfgaq1Kc8;Ew zl&O5(3b=~JO57J4y3^;r!I>Ix1A>Gm8<4kLpjdEFWSucn{&<78JGbuqp=eJTE-UL3 zqu+ci$d^=$-zPL4ng-u~^0b?Li{&dm4oSB0i~EvKr9jcd$DVx~C-f}@sClrmx21dE z^kAaWpw5|Ex{Oh+M($pSIIaJejWP9e@4nMG=o-DMpq!MZogF=y4|%O}Gabf!Q_Chn z!{I27S67siZYiBEWRu!sh2!(oO~GB zyM$4Kba_!94qu);bo+pX?a>YE79UcOyP&1+r^PdHAMoX}1E zorcq2w#6yxx{=b^OUH&2E$ffMs2+wF-T@R8xSZ%iSOSb#uQrDZs8cL)jn>iNj@v7% z$-!JHVi3>=KJ&vV&QU(7ozUYCJP4_!dNtM1*BXWhh zsPt~~`6bv`Kdomu;>xjfb=7#ua|bJwGVGOJAK>f0i8>xXIdc)WbeV6nIQBp<&8mZJ z%j-QJEW32BPXVutP!~cdL!X;()Gweg9g;5=$x8P1GEax(!;y$dZ`t`WSQ?8&#i&g$ z%<|#r8gr|vMmm}AIE@8>cxk{LBNz%~JccbeaNCShi-)qnjGnR0tH2-6oU@=xEHv?? zwR!a>slUv4d0!lbY^;y8I7$pa$T0hWY63Y8`#rJ6^>OQL17KdVa`bBBk#-Di*d>p+j7chHvXIb^r0nhUBxf~%kbUin%|r&SDC zT-|lWU{}UETk3{M_M%I~hpWPG_}U`hq8RBhaTQsM2g&ZGQM1S1c|@hAcY}uq9;$@g zcO2?9EyH^CWmfJdh|xzN^*q98mcMdVQfVq)mZ?m3T-^Vz9^`c68#Gmc2Ah-o%1KYw zl-K2Xc_@>Ml2@=DxI3Sy&6@=nE_Gyw1*z7i!$StR0@bnl#}Q*=Vi+@jF&l0APBPCI zy7>&Zv)&H@$WfDbd9cUtdum5-KkrjPhQiqxhBi$;U)uAM^HBJjhFo75jBShG;vyyM z*H8wd9#n3#ib55vVV<+?LrUF3cU|hRk~R;8bs(7yCIM3>tz^nY#q7d_gS?#=XzDIR z?KuHki#J_#n0S#1D0i!xWz`(jExGGWdlATZLOE@_yP(Tu6l;YkJcN->)E^n7$ahCV zbJ29WmG@+<^C^v)p`v*J7dP3Kb}4bRoI}z5vB!WgW%_eXEPV|uRE|n&ihQ)o%A#P9 z-W6dS%!O-%>6q10o`0-xrRz`=V-kI; zYPxW6jJh*0d3JJ=dKjgUV>LSoXrJy!Iz2+E^ZquGi=B%AN?%oTSuNdw!BGa=f3S>Zea^tCO-5y}fh=lXDDt zK3YEPcGGH@Amf)Okg#J=FFP6i!jIxkiK_JPfKYacq8?1#6f|v?qHOdwFFO3#vh<+L zGap4rGDC>bIF>K*CfFfn;^&8JOLBQhu@}h`+VU#fcdMiAT&1|1+_P6Vp z!<*&UD#)11;2Smf-L0`s3~yg#)-kRy)^}}s3?0Ze>WkRVH*!f!musqG5*p$I_do0o zJCUkIyT+Y`aYP-*LXvqt*BM(xxJMocq}E@HL{;SRdWlAU|48f8L=P-b=_KA2T%oZs z*FSZ)TMlGy1|*3BD8y?SBcWhhzKW-JPk8F^-NcYZ?Sv$Yd(RRXvG zEX|X0MEzW5o!BXi*y#Fqhp6MUS|+g66Z6$!WLK_^y$Dl+e}ejFl123~dwi9M%hj$k zxP`HiO)5+QQdZg4`!_JjNCcG^+x*xuoTaUY9>4jlcxNIP?|v+`@wx!YG6)kK3_Y(L zbxE#{X#)BXrm7R89A_8OLM6E#@wdBWE`TL1iY7W2aV)}qOgwIX`d@m_BKph5W% zu>ZD+vRi^kr<09(GpMcE>IBJV?b?M~RbO4F$9-3G7d3>r$EOale3qeu?(;e#KAFx7 z{_-3Q+MPO|rz}(ih)y?$j-U0Y=?kyOgmP=MO|H4MDTGBY2f=kD*|SHat1agv*dCU*;l6A1hslTt5GgywJ62n+E zteN#$zN}xER&Rygd?cXuA#Bm4y|1xEC90p6bu9qrQK2H4PR+XLjm!I*w1$k`Jd413 z3YA8LKvr^0C!j`-D$K?}^*0YpTx8bjThs_^q^Un_-3+Xxa0(XNkFUcIanX$~^Lbc` zMQT>_REox;8Uk2KKk-18WPk=+4v%xKF*3JcsWq@lL;6VLJMg6q-u$TBDB2UcXGEn# zY$ez`!Xgd^DdOB<1(R$H6ILfKh4l2n%Qk|5w~Ko8I4)G{ z300{O>gj8K$~QP^j*pv(hlYjwuRsJBn6sP@04s{L<^1nWu!uu4h&u^~J)T?;zh>)U z=@;3%OpD8NkIL9YHr~KS^gZF8f31qLq|JZlPJW|A_0^`lP3!A$?|xuTg6`QgJMO8s zQmdNmkh@NvP<$Bq%APK4B9vu$uV6$hO5qNz)O@p!Di+bSRm3pC0$Wq$WPLRi&Q=uh z$e5yaEeLlhFMRk8q$-Ge6om$iAChZTiVgN*Bej`6W4~poA`u5?4>r~C*m0s8S zRIJ>OX7F%8S~x0$E(^Om8Q`DLk<~krU7U;2OQm($8#r5N z@lFxpTLqMVPejou#<*M=YhlMVi9yWMY2xSiP8VDxJ6H=(eVrrzg;|I0xzg4oe4&~BHeCh*(OAXg8qYII%j7;3iI!{6X92)3F7Q(mAoLhHS8juTV`UwcsJ@|4U) zA)IJUo;uvP8QAv~p64Y$ zIO^rG%MB{J8dBNhhC33Bu$%U`D6{pyn{N3IO{F=hgDH3ARh@#oLo?{i2l@}VE8hs- z1O5c@VI5A)(Y2vf!`aNJo`jK$$M32SK{!pkL|u1>V(kw!I`&F!SYy5*BWGpV$5x7T zXy3)OZ%9n9c|%J0IaiHztPw92GEnwGa#z}{EVdAs6r8QLrG)G`>bm49c zV}j1uZSh5!wm5p&+7_$LsF{xEsI%Z}$W3|b;yU{9kznT3eIb`T7#%*GB-a=puZ7Dj zL)EP44vQ7MX}|phsnY36@4aH=apNYmcAb0^%bW$Pg>TNudZ=*&37dKzo*86!KJkoo zUt)MQr}RdhQNw;Ro&3|%dacvcPxe9s$>pKM$le7Ixl=HQ5%-P zLvi8aQMK5072O;En?Eo3tJiMLI=`C^l2@b%#WyMZ55WKINHo=)wedn~APQ2u0H8 zV{u^KiU(^{hdBg5teg0 zuYCfT`zdC1#u4*DkGD6+clQ=WD`VHBpGxh0rNRxmjEPgJ%?GmSYMEG@Vu3Aw(NzTI z+7{T;D_;K0FyZ6YB*&{)kcDwaX$DhuR`kH^v@zJTFbxhz^jq-`k??*~+TdBuQK6#e z&mf2A6>x7SU;!EA`@MxaQD~L@NQA!*3b{_3ikte3Jdqu$wzPbp%@!H%f9j>_Bz@W} zgG1GIzQ;p?@M+4{K*SUso2+;=3{KZN&KT}6R7cans+eiV207dy+7%aZ+h)L z7Q+K!J`W$Lr9NgJs{+4)d$Yy8doHDJZkes~hQrfE@*DBOR!n!lh2ZvYu`O@B!<21O zHmtZP;Ywi6&vs2N9?Tzg&CXqCEWoB?g2fOE0Rp9E_Na;zy|%QSIin#1(c+?sk<9@D z4)s*LTf&$xiZ7Vjd!!X9uYB3-|MYCu(+ct$hPP7&;vMM7cK;*!>;0dby!q&k-=#B# z6x8)#(g(=LHXV~7XXCy{hBM34&CCX3X13s4yajKtz{ISbj(QlGV#MMKmhRZf0i_si z<@8z}f6DZp#}t(Vn)!=m1K-WEIJFRaeCRyJ;sVsBqNad~gY4yDh>^KtjnzvqXbEae zWu&JD&@N-0{b~&MN?wN>GKaBo{*+_!^diXNFnY}H%Gb>#E3cs$c4b+9C9m^Z0fYXE zDrsj!A5cMxM)#(|A3jGc_x5BLaM1_MbF{IS4IG%2R-}?xp~3Z1OfTx+zcO)aE%%9{ErpN zZ%q*8o{WRtQ-C@ae&sGTNR&+PD=_>CyQzS?)cSJn{pQA;()I5&Gt~@ZKKezAG5IEx zGX9eFZD2$dvk-G1>Z>ef}EQ{a@XWYh%M*5RHDA zwdAyh>U$!lwkRzw`Wl+K&M{qe4D|Oy25iOVb7sHgi_bXA?00P`Y?KPP3 zO}4}C-1=+dyjW3Ky*y=Sy#!7|BW}0MCP(;vS6Ke+!B4LHubkGW9WDHT0KV3;3D0n2 zGI^p|0~zi-MMfNDc{S(2X2erMuR@%t0~*hPKS} zA$`hj>PP1ESov002V1k>eWIwgL3%I>grXsnr;%gpwDrAf*_)xzPhE=2YHA7!WCe2> z;_f5R!!1zlar)^S0!cI4Dp4U^rE@(DmRLA?QX3_v%XMWyCzQr%U>_>44|*#d&iIsa z+UJBt))a(XhQ5F&q<>8nCvA?92-rVAWtMd|r*RWu!A1?RNkeY6r=fzv69W$8Cd+Zw z?_dx6C`%{R;@%}f6T;Y+8_dP8-A3uUz*Zj7P0z!ot2?qaLm7qztmxh0BgkIFs-WBnANtGeh zaNB!zJ>a7NS$22C9y3o5)fTTZr=Inbso$W?s^GtR1af9o#sg8v@R1|@NDJBP1zOBy zIfm9P$TMbxnM9>fdHBi&9wRbq)5K7ZpeDoM43F4wsiSx z5;L>I#sP<``*dK^D9sNP1z$OrCjv)>_=d%l-b~PP4vypQWzb3C%yf+0BQJyiH--mv z_9;@#adffuYA+cpZ8Ca@H)K&w0DdMk2iMDS6xs;4r_`ev_SuL75yf$K3xy~dCh=S} z(6UM0*)V4zx<%c#CT!t-R*At_Y?{Mj+#4sn7MU`0nzgXzb5oh(RHBaJahl^IAwzD= zo%sjqLY~7q7Z(`L76Lydz~~-x-syDZOVm1dXa&}WE`fzg(a+AFRN6e*pANb6SwMX4 zIk1!lxidY5YkCa`pO6}A&e~rb!&t=-oY}|AzX*#t2g_9-xiE>C2}X40u;F6gLhh|d z>qmy5!jBgi$Qy?<6~X-BPD?4u3Wp*z*QG|WWeY4N`;&van{0 zsn}L=>CAIB*4^fN-p9ZQim;wt^+PE>gO@r|JB(DQO`nIymN* zl+6PlZbz|K9l|x&VU$HWF=B@_eFQy_W6SbV@bjuA0NXTRa+d-&wI;DUN)!fX|bf0Rt zyD3%U^nMR96BNRr?V2J&ZB@u2#T5mShpk+c$a>Y@79_iGmP5R7Ng0_yZWBUXiN?Z$ zH^uVN2T(_X$mdIdE~gWz%{>r~0E)*}$Lkr2WLAVNzSLrpD`L&Lpcx!12zS_Ku%O`T z1x$$>^x~YG(hC#dXhw{z_p#{?1nTqzUEHRC&aEX>q$eZq=d`wt<3awL#Uoowt>+y2 z``r)4cE^W0UqJh^lxeILm4_E~c_ikS=VG1UPwyfhMj85GC0yGV^20?r6}#ocu~m5;tI_YzgA37&I)U^=LHT{vF7sfK%D;t ziA(+;+2-Q zT%rrE=?eb}X1DUx*xGe^RnEWqVfVT#9jBPH)aL%k)^m6Or-nab=ws~oLKgP z!@`YesZ6XvNz+Py)y&bfz`=Ffg0Hsa4Q)2dr@bvg#r?>c}~Ck!Q`Yh))f(o>vVs>}&H__(FTU(m_uJg|4~hR+_Vi z<^!E!D>X=WSKNZ(c<@8E4?M=?f^8HJc?!#%B+H}<4#i&^YU_}?&WdK%I~?k~-?OHu z9sS@^iWKa1vv3h}l(3((v#rIT|9TNO>jj0d9vAZ9w3Y>xjsCil3aM$^@-Nitr&J4@ zsfO&Dp)#~-{VKnz>WuB-nJ>#KjE@eBn|ITE@;Y97(IS+msRYD5Z!plgEUY>AUYNO}yt$y} z-Djp$KXtTWMd_=(zB03;!BU1zx8Up4P(z5Z-q6w_m(;L&-AtD*UxVYd6#@gi_0 z4Ra}v5<&BuG@N4hXeXn(uZ8E>9G*Gc$Uwy>!qWlY0!84p8EZxZ=w#SZV)4iY#=x4k z@~a7m(};v?_ZjwO&v^OqhGm?oK6u;ec$rK5jML!_j`aW|vhKN}W+Z%R96lmeo9Q^u z`~_i>;$AJg()_rHK`e3&{D1v=dq?(oj|Spmw0Zgrr?J@7xfHkQr4y&krO}hgv)XyX z8|`DK=39-_&y-^&m)CMqTCrT!!?X#AT5i$KDbDozh1{yJNaR%UDfyD|uC|5i&Cqe( z(Y0?wO-K;eu5bx%S?X;6(m>p%YyScxSWh}1+MnN(Z;uDPE=5W<^lErv!!4ll6;_3d?jGyjA>8p|%`nibnQH~t z9=yRe6dpE*W5p)d@(}2ZrPi5!^TC;I3xTu?(nby&_eZ#bw^lm8^O$nPJ=gQDY4g9p zv0U8@Pq^0RPub5?sxo8y4YL~TG}c~7iR{*rsI5L1EGTInyw%bazHsw3wn?eguV5y0 z;)bbv!P=6hnftW=h2!G2{?$>l{mih}nNwC?VHFaT)tR^EZmYD4^s|qYL}#$5eH|G& zONSPM}L{Gar*FTwiSLdfi#~r2@QoD4XckA}8xuJ^5B^UVC@p_wgh3*A< zYR?px1IV*Sf`Y2g38y zQ1cDLh8XcGrnju2+UE3Cd!A_JhHY+s`lwi>va}EtDs2qgL^^D#7(Id+}_+MG2yPKD=Zux)SitSpT zez;=$O_4G!{3_iPO!%4xgF>m~e~l822pA3iqXTFFu+fgxvh}gHv6EDDbak`=7{E0r z2`zbE0J>9EY0lu|U;sV90PF?;wNq!+we(JDYa!2^($fdnz`Yp2KNj#0zhEU40tRr4 z55L?tY9qGePJ*AJ4gqhrjhzp)g}T-MN)76T{Tdr)&h4uQW4C?~Xq53~hJirP;K3lG zr|RzRW#i~(>0|5tYYcw}4SWVPYIt;D05zBiJ{t{37Z*cI7hhX|kymjWM*vU)B&h+( zEpUy{-@_E23ji#K{b6@9;JpRR_x-R2aafj#;A_yISH@d8N z%O_kRaUFHy_Q&hYlhfg|f{FZNfz;(N>OJ^4C=?2AAlnGyw-A65%>CzkA@dZWa{nhi zG5(y}_w@k5c`D$hM6_(|>0bEa1cwq2#Xpbm0e|Rr()OX3Er^Gn zA|wNFFx|)3`+){AL31a8&esF!gM|%gKXXeBN@Oio& z$gF(5#gW;>=(E3wmcZR5Mr-Xv(^v;syEs}SZERgE1KvC$CTs5`lW)fxJt4;H?8L)e zY~AdA9QXx@3FmeapdJz^PhzCrP9&X=t(TXjqZ`uB(%Q$}%WCJx079m}vjfC7TZwSc zP6S~CE@jv7h!zm1d084fS)s5e?fD?!#UM!l`B!`Xe`>7vx}|%+lNq`5-~ISBbIs*X z%IO>T8}^TiKceI|{|S;$7=P_25+(Of03#)n5FjByLV$z-2>}uUBn1A=5dfRyb=hEV zjRKYxGTEo`+?9*i)%pMb!cgvn-sbO@G@s`8;|a0Ju&y~C*&nr<8Gmr-XZByWmWd?9 zKPZMD97u{LAwWWaga8Qv5&|Ry{%sNP@3pqzO*{UYl_q!@6sWp_?HM9?27j90dWY=C>Y{CZ!yi%?5mj2mEs;3; z^VB;urXYf#R_Xs;|3xedaf>e#QU6s5e*dNY3=bF>uJ%+x@v5UAtA766aS};Hpvea45{oeo z${U#hqPPYq2UD4UevAN3pveaiMKr)TXzm6Kh{rO6PY4dckZ6j7=6=A4Xc~ehCvcu< zdV%Hv-~!Q1+dkHqXqIi8CPcFpG+6=g1}20N&=dg7h$a>IaN-2aiRS)o6G=2xwoMD7 zX})b*63xJE(~4-`1Wi`JnrIe-rXXNLG`m4_FCgZAy<6dMWb%1?He2P{O`*y*&vW@87l+0H{o`p_+8m0{thBO%JA~CMV!RG#O~|IlB@~8PMbc+=!+jXz~H>MAHv6*#Qrt`4ebz z1D-^)1T}uUe@6s%#A?ekusG$1tE+*$@JE1y6GHz-gWh<400PTS7loKi zYDW-}AZ|XXOL;em)2$K*IWQ5E zkRG;P_&-N?8x~UnD*+)Rj(woCMLN1U`Z!v;+&3ep5eGu-{yFUzSrpwN?nWFS0hYnl z(jV#JE0;k`A`TG1CxKI)dic&{5|fAnC&+D)Ztkv24~g-_5flvD@km>LA6qvYODh*! zSSfKA;;0I+%;2H;la|&H(}+VXh-r>)$I-+z;&=;Un!8O4hL~otQ+0OI+&z3=G!fG* zf1T!S>!Z^`Od}4<051cKA9zcQCyvixcD3}f2dArc*`Dkc=?zZxYH4pv6UGKL5zOH$2#X?lj|PKu@|J8k5K1T4U z;XWh2?y%K^5$1=0kP#K8 ze80crFZiUiJ;Y;(GE^p8J3B{fM_V`R0Hlqhw}+*VwS$`*p}h(rBdSvwez6yxU9sKk zLEOb@=PKZiZgwv2ekzK@GdS-gP~!P22X|jD?$ZL;YMJuxw0r>Nw1c)~1y^B-+~LKp9k&)#(O7$NWfw)E0plyGGv z=_TkRJgGtyb2t6QZr<@>ZSgJ%g4a&|0qG@y#WCHP#4KtjuYXx~S)lT_@fscqhTYe% zAlKW@l>mjIr{5<1*N228Z%Nk5dBq2aTfP?@E0TKcv1V{+{ z+ao|4U-w7G*NN>6042H|F9g8{()c<84x~>K0we@T2#^pUA@FxZ;E#^4`)BN1*mV;7 z_IIo&q-!D}Kth0o00{vS0we@T2>b{E68rXh>>K+J;40b~enm_sv2R2qNLvyDBm_tZ zkPsjt@b^W4AkF-XXCGskoD=W5`03Q(~Qra>6W3Fpkhn_&3sq~mAsRSX7DgUcO**M`&phX)y2y7<}xzjrD4 z3nx?`LdC=v^_Q2V4%&Zcx@dd`gojj2GvOidI}a&?CtJq|4#@pe9e|%pi0FXQzlQ_x z(;)p}2mHA+gJ^*j#9qI)3Gdn;grDEOeT6YVltV&l6d1s2Rp>~o_#^c@+YfSDDkI!C z#AM>#NYe2T5g~0!2#^pUAwWWaguveq0YbCzubfd&N0`v!_{--SBQ%rI7SL@k633Ws z_YVQ!S!`P1_F7%!qrDg80W+{I%>jHLkWv8rx4)i4$teIS;F0XWTO6PUbW;PsUcjtZ zP93-jeq{xQPrlLs#wh`?6EJdX_X(gBY{l^ay-DLIfjV#sL0+IdZtWyc4K9zx1DMJ= zp8|Nme0hPll?$hV>)_XYK+Z&(7QhUajt6*F%c2cbg0TVsRFdZ$&;;ho2PmfJ8USlx z4y=H5nY|%!1q6}M4Ag^7H(nt1gQgiE4UQIP1JoL1EC9j+U}(-1 z@GIEH;{-14eq;rBZs)N1s>uc*tS;8Vc+LU9Pl+G^Fe=M90Dyf|-wu(a6KzBC^o8Y}}VkjXdf1Ehgv-~i6M zMqLE#KnJh@H%0Bn|0nK5!o5KAUv-ga@YRI@EHStN0|h{iz<-4atNTM;M3{&2}$R%ht!*#!gbz-QCLuoW<0~)_eP`glu`Y`_F_D z5KBTh5x(bwXHV$BY`}eyJK3l?x;onYIs?M~J6*>iD2O%~ey8jBK!_iLU&D-W^Vsbt zlpt#7@IP-hMtp{)bPaig{CBbke#;~vvwxkAyZdYV9^zRa!Ha(}H+t;_p&)pVOhIW2^@Kaa^AoLdbEx8a2o<4seRq zm`9HDJD9QK=k4e;UyWt-yo+Xcy!`cieY^O^%WABe5jWW%Gdzb4uGGJO>58SN8^)x! zVsC!IjGH?bm_PS1Jm>enb=GG~=QLH#7x@2}26fSYr^0)`OU;q5Tu;+HPux#eB66P@ z^v6!|{hxND$h}jvX_Y?RV`Ff<^b6BQy-_}@u_k6*Uy@~bJ~h)I*(4 zd(M7x^)rUGN$(wX?RTX*Hdq+r=!b(nZd?jAn=k(-Q~s4buBpLfUVTWKGOLpO2e({RphJY zUrge%iC8&r$Cj^UYcy8bgx9&#OMtMG|uZ5aa%SR^3fMh^2AQ_MhNCw`_47`>)_(!WF zU#-tN@R2v2%N9}6UxsacEnDNpw#d3QF642!H5$zfonmc`h>y2_-Cd{XmpNo@nK$5M zZ*RArqQA@hF;@Lbx)HRSALk0`r}#MjPB&R^Z{3u=oNP!$TS$S@LUK&0JeCYd1|$QL z0m(q5$iU0i0?QP&DsQiR*QLA#yPJ}CyXyo^+g5shV;$( zcyornQ=X3J&DuEY`3baRobKP|-nP^s^FU@0Vp7wGH&A|rDNzS)0li=7Vc)=;YXvS))z1`=CwblO3JH8KOd+W~Cf30uVts@sS zm{_&d%kJ?*io0o#?=jBmCYL?l+?yTu_+w7-akaafQq2IQpF^rRItVXc@@=tG(Cylj@ z&U$!cR;5wQG4Pz*OwO@&D}wL_(0l2dJ=8ubt@Cit{4|@zw$}PtGmg+ugsQ23aq*xU zCWg^EY@D{RhhhIzK7M1PtCqH_A@Se;#NNV|aTK2CQyxhMBmpcg)MJx2ig5zWoSQ)YQkdE~)Su<+2+8Q-?)f)&}^OhcCOr zR*6haT=x!Nrbbonc(HZk#6ABLEv{kbq{S6(i|d@Zog9ykerv`N`XZ|*@2l{Gy4uJv z7L>DEOR7*!J(k?-=&Hr_S7!~cacX%vx65TDJ1fWXuVg?nAQ_MhNCqSWl7TlR1OG(J z`=Xf~-SW;}=Z$E2Uot;ZgevT6|FBx#y{Efud6)A_mg%*$y!ZQB?@3zTb!z+8tk$5R zdCxv`dzt?yGyhj@p73tl<=BE#_cwZQXT}}tOky(#n?fg~{cjxmd~iit;^DT$&zsxH zaXWrz#u55rt0#X`vfKTz?VbLqd51TLsou3ZU%6_DCv{lP@!_h?a_rMFAM#i-AQ_Mh zNCqSWks$+C1K?jiyEJA%L0BB5A4!5`&s>@9y}?L^OHjS zLu=pGOZ8+a?>1p8>Es{OJ~SZeT|K)Ov*WZ0?h+E>8{*$4G_XrhtP)llJfiktezZW} zpe~(@7`nVK+&&rT|GkD7=x`+S1qz6EbGi3X^}anL}JN+WI!??8ITOT@flDfRJgt# zDZ;(Qcjgw_AgeUOyS7VmX?Uy4r+jPshqh_ku27}G&Vg->_~rp9t9qY4>==eR7NPRoYJvw+2}Y~&Pf+vOCQMMJdv-} zZS6isT;Fe?Yw+)1{}X$K-r9z-sD-s}>yl=bdzX%A3+=SD&~*Al$$(@)G9Vd{3`Cj? zyq5anJT9S&Oef7*^~E8skD1-|X%*wOY>g;0B5Upa*C*B|59bKxs1&KjzLu@geOYAP z8ugyK+!~wK70><;w59%mRoYu*&MmZ+v+eV<-3d4p(H7e+X|Y9`zs~YOB?FQH$$(@) zGN3bHKL^Xa(q{RJdsjW84rPAyP`|Feo&9<>uA)!m8D^Uk7<8kNo->>|vDI^e{JVz+ z`*qsrr{|7fZVWXyB(SZ2r8aujNZYJX|E|Z{>RF6LZ$$5*PABrsKQTnl8Li)kHK#!E zKs{$HbE2B38U>FZqKn5-9J5WA&Rx3%`3Hrj9;WBL%e+J(L4I9B0=k6y_Vf>IAD~XZ z3JHDJT8c#k$-o-bTBUWGMWQdPkMk;Lb(jCK=Y3Z# z`xbxEvah^FTJ~?iA87dqk^#wpWI!??8ITM}2L5&iDm1L=V|eNDMvSfINGa(IJ-&!> z)(q1dz4h8I$1#lLh8|tScx28^ZsaHfO>#L$8dY4A4#WTBy)#CZkdKu$iyDyLR@(0=2^?CWQkwFdP z=jcU}gnbvHb&MW!FFVneS-JP=J!VVBBZjrNm;I8l!yKYDwb7+x>C487DC!GQ{rzZHjJ4*| zQX5;>`}|~lY%ZI`DBx4`x>3_?;wLvc_wl`HREVt>XuoFTZ6ny+F4>GsT@TzgE}8?8 zrZ)05eR$uPYc8A3STcOYL!+j--o(b7tRFu%@|&fpjhbs>JTWGlZzro!sP6QqhV_l8 zyQA_mV}rS}R7SVlt1IbjL}pb(KSm;z0naE=jvO?TXHM#CU|yRYu>oV%6Wu%5mHG|d z=jSu|#|Exi`mRPcTDUc;dOLqi`g^TDluVKV$$(@)G9Vd{47`~ccrC}L>K(V&p8mB; zzPDME^)Nh**k&ErY|{Ap1a|7w*soJ}|8U153h#~&)?V=2zzC18W3?}axjsJCKJ)I_ zuVrgoJ{wuL#(CFU5S)!(Y?I+3lUD%z|YKG~vvm^ZTVm3v4 z^*xqmP4%*_T7Y3qSL?&s{@p1pKwCZ}kPJu$Bm3^n`v--Xw|X#Q=z6YUnB&EoNA2>n zBWsIirvUs*+T#XQ?D3u^_IOEcUooW_x@;yFGri-@d+HbZm(+vA`h z`|_E)+T#jk%_El%)86*>xZFMa^3})K=TDqtk58_#$7ek3?T7xkXsuE9++A*wCzvipZoOF|Yl<~~P-oy~7Twa_5qY zUAhFf4Gc1Gq!?nqe~H;1B?lazQBSu{p;O0Zz0B6NS%kg2<9^1{kz|0+33DChxLM(Yt3kG%^~}i1dyJfLH4$- zZu{zOuG_vk9@~l*i~HW0-mlep)P;-Z4RYU_^=HmuzWaDX6C3VZvwqKc%=diY)XH-1 zUvoio%~>aPjk?3$Li{q*)fy3PAzqahB73_a8ITM}1|$QLfyk2q_f<}Gg({O-p<0o_ zFU0+8&TOvv>YU>qYR!Kf%)mdGKQ-|3KUoI2Z_L5ar{*y0Rz+8FOVmF7nFMs;O6fxIaG26h@*1K2#X6UFs z?DNU?G1t$qX8bdm&1bH>ccU#$-2XAv5zt4~9;=;^j^V^Q;UE)=a@Qm@&4-}7ST znwPh0+t&SSE^e;bEB@u9?q72WbIs?bWU4{e0m*=5Kr$d1h%_0{6&d$c>(nWUWrG4z>ovQtTKBfO=G}8AMp0}2qs3UegLO*U zKhv&hZocGF?}_!@|6%;iH7Bp}$wb?S>5;{Lk(qlBHK`*(KU}=n!Tldboo1S{$NK8- zUvpP;&HDmUv~d5L)fHtov-{LC2|M2DzE{wVwVY;?YFf!W#oWKpNd{AoBP*X&|GtYbV)qjzvjZ` znzwGAH`)DbE^V$kTK@{u)qed~EylWs>@CQSF;~!8t;h)ZH4tW~(ne+{E&2ZW9r@6S|QoWR#^c9yX zcHe%c+msc})%s7UYVQ$xHppt!xVJ}0zhqVOCC`l8@g!^q3>&+}9Coz-)PM3NYnU(j z$)V}-+@HT%=9-IanG)>&HPsPKy#4A@{XLTwrt6%^ z9La!WKr$d1kPJu$Bm@7?8E{{%RJR^4ndx|Pee{~J$~zZUe%97%t)(Yb5xM|WT~asf z;@qe{4H{H7I%Kg{p)R9*rkY0{=IkeG5f8I@@HmcoU>>Y~Xbx0ls2g@XP}N%x^KINR zpI~^LjAG_cE{)vZ-daxotH%b? z%O6lzxH(MOM>_=O!(;E+V?CCTe*U86_W6e|+T#h0^paNjlt6nt*w-FApUqscyM6xY z3HEs8aC@xBHqxaBX6Z-P>z7_)kM-C`x>S#Sq+>nyk&gGB)Q_z79X(@@_1H+dRF93M zV?8#Kj`i3`I@V(&>9}W3dJssfe>Qup$4=7A>#>t`oT8RhmkJY5(V%6=>?=!?6>#>~l@{TQ(DUu^-EP>FZtj z2zu-%{dx&Q?Cam$(jNORu*Vk`K9Ak^OQyRa>Y8tBfB$8h-M41_8(H4`{ExQx`_cVt zu4JycZq_x8+`ncYbIln-z8+`w(S79Mky&lTDCYGb&V2^XvDF(lsi@V{&_F%SoFaOU zq!{W76!ZQwo(Jw(D^f3``=4~zUR2jzdzHhU9+4~c^>Xkqh5n-@ zgSzM~65bnJYTnn$^)$`%#Qk(7T=kl%OF&*Ut+;FLf_fR}J>IJA=((a~Kr$d1kPJu$ zBmcP9PF z@X0Ey3 z_0oR~WEtH!eBSz8$1EZ?!lLyhsgeD+9|{kg{<6Jjm)+Hi_AHutn8ci<7tKt9+B0=8 z+J2|_fElal`zoh4TaOIG>^0kC^_rzG`)lV9t|yH*sGO@2<@Iy^J9e#Eb#mrvfA7Dm zz9awcqSYp!&$DILx^b>L)&A;KG!OPjoha#4)0q$@1CjyBfMh^2AQ_MhyqOtzVW(Pd zFT>ZZooX#eqEJGmTBMN{9Tf3s6fy{%Gvf<-$wDD>7+A<7&g476=|(I>0W&>9o=1{#ul2HywMtc+eZC;jag$? zQmd2A+u?FjeOa6K`Q@b6^I{m$qO0B@Pc^QPF<2kt$;|jG`e}dTbhGm2^UXh4Zn!aw zIzcS1d4Cymd--yInM#431KZl~ETdk|oUC*0&TO&{5jC?4)m>fmw6o#(u#r`*1-~t3 z&Cy4A+can%Va%=`~aBJ*(BHlwLFGHBLuyc;3sV`@UX%xM$s;QrmpReFtv0bpM*`m}}m& zXU6Yac6~X^>O~v!qF%Jbs+%W4`m%ov%2p`0QRtc3AsOL|JoRW_d(&Cn8qODZM*C9- ziS8(I*R7ZEa#YjH@N;YB$d~P+mCBLhjw+*D$mkaK%8-n1 z@w{@NzA2Hpa-=@L=NfxH);Qhu!0}#v?yvHwQDPINdOP5GJ`8ffl3|B4Ihk7B$FNDM6HRbHBUpev6Avy9;G9Vd{3`hnf1CjyBz#Eu>7xr9q z@-o`EwdcZ*3~Ht4!oEvM&xO@o(Tx^4@z8#(T`t$Px7^>P%XQT+!HZ5jRP!TpY+Wty z*mI%cT?g&?t9!^EpFeExS-5=49xpv-ozcxvPEsRnk>yX`QH8HNtnM z=or=N$rz?Ss5wRWq(yyaigS&v6}Zv?zJYD6kGW^Lr@o?{g1cSfGt&Gx>ggO6PtxSIM9y_mIsC8p7*;52 z^+xC($>tl+AG3N8^ZzTCI;*RFqZs*%Tm7QyTrM@~kp=&cu5&iu^|7U%Yr@|C3;&TG zL#yLWdJOe7`z7?3&&7U|I8b^FU8F{eB?FQH$$(@)G9Vd{47`yUcwvv>I4{HBtv!Zq z$e~tx4DCCW^cYHyq2nGQJ%;ukGU+k2=i^Q9F&x&w-a9d^=edej4`QwepM#i+doXw6 z3-;3IAeQTGy{XrC4q`5K-;=+2%M$g2I_~{SQyqU+93F@D7@Et(@C>`pZs37S9ut?^ zPi%Hr`akz@e9V_TGq7OamAhZ^OvXcH-tmYU@V%>Z^uos13@1_LWE`q+vLZ#20m*=5 zKr$d1c(XI`!k(7dUPe2&_O!GmnOf;-vF}~d(;_`Bj(ddkwAgzVq^HH6kI2^3vbVkc z+^tz3cU$eM#rxFpixke|USsXOI5Lh^I7yNs$$(@)G9Vd{3`hnf18-0UUfA=o!pmsy z)}9Z4GN+ZE5BshpJs;BZ;kZXg&xgIGFFhajd_?x1k8<~{d|1C4hS912TN$l$P>cwl zgYw|MwXRn_2W7>)*L)6&^$XJU!QGUOza`;6HqSBHxghKulllE`o%Pw$IZf5@7yJ*F z4(g(}dw6edsdsp#x@?-~iTmkFylk)8s=3l@rZXu@1|$QL0m*=5Kr$d1hzuEcVXxU% zFC)OMy=Lu5saAT;?E9GXnn|yj;~pWsX7*kp={2+GLwe2Zy=K+NSiKp#H|0{4(UV@b z7wzl?>zjClFXT&K$#rR+^`>6im0Z>bC|qw<3qFbi(re~n-gQxrW2F77bImTVwKl0v zq?~J}Ga*U_BmT}a-#0Q$>qK5s=jc+72j7rsNRj%EjTAkXy zHLEpfXrwTgFmG*WBs2e4tB>MgPNr&dbblS|ZL~AX+f>+JFl>Qv|B+s^m+3XD`<3*X zg}XaRp=3ZZAQ_MhNCqSWl7Tlg15#^#nOgIRvr=n*L%$I6u_ObM0m*=5Kr$d1kPL*I z0jV{IS8FbzJMP2HnPF`GQEJUDQX|EZ0m*=5Kr$d1kPJu$-pCA?wdTFn7}|25_P=}I z_7~r${jWw6%6m_0+;8MpNj{WhKr$d1kPJu$BmhSuh6S~M&gKhS6;?=3| zYEtK?%On|)3`hnf1CjyBz?+$Y7hXZ0JG#-utyhqDCW%_Pg517W$ra?*=@j~;ihU>b zBP&lw&)8$%bN1!uowmovlSc`AqJGxuEcUooW_#Q-rxok1;-AeP-_LGe{?UH>>!oOE zUw(R7REfn;fFEwAH;HSF_?RkuIC&k%dOx1W9aO4;l!ocdAhEsmf; z_UGTPWM97D0Q>sJwz5BenU8(>gdz6%n_JprX@T4GAuVwG{`0z9;I(ht-_Q2l$5;&V>%NvHUfQ~ZrneA+2KW5yA!;3TUf6&y2dGD!v`1CjyBfMg(2 zWgy%N&VDKds+!|i!TH50{@E$M=M>*}if=l_x18eJPVpV5_^ugOsaU6GSoJHu8JLq* z6|>%_ic8flUq-1%T+m~Xt97eMk96Q%uFoTOF4yPxI+yGF138!L>XdW2j1W!LiKi*nK(wm+yvbe{E@IXZv3T#r=e{CU%`ex1LN z&@RFLA%>Bba{am6`v>_4`H#_B;98Sw+nqz&^ zoG!|9xhT)=qCAg_@_a7J^SdZ7;G(>si}FG)$_u+FFXE!SsEcy5-4-_K{OY_FcTrx# zMR`dV<)vMehy4hd^PFF|yo=>4xhVH>QC`JGc~uwX)m)UxG3-H zqP&}n@?aO`>bFWw&f7iID*jV>cNgV7T$Jl8R-M=1%f<4&U6l86QQp@@c|RBB{aus~ zpj{9dPWz5G6>a=rY1r*ggg z0jF}k{6VL3z5EwW<$C!;PUU*}!%pRT`6Eu{digJ%%JuSJIhE_>k2;m><&Qa)>*bF- zmFwldb}HA)pKvPI%d5_GH93FZPg%u(lf~3GF4xPe z&VDuNa=rdQn*_44PP%JuT*X&xmFwj%IhE_>zo%StSwFxZ zoZ=sy;ww(^Pfqbwr}&yve4W_U-;IbK;P|b*{)=U;I2P;I`Ez=C2_*xP0m*=5Kr$d1 zcmp!9UWW7Dkn8_Zm$!x`Xe;d@z9=HDE zgdpo>7nHNlPyf`@_I{Ezx37Ou3wykBwmmL6*S@~H^X&1jE9}dAuCd25fTulQ|K$NZ z)s7OYgg)?W815g?v!nS6y{>lF$L!3=c>vF;IiJ72dOw|N=3o8FT?pexNf&}jrc9Cn z$$(@)G9VdH6*HiE3?)!oDUb%~*?K0_NP|1`@ zG9Vd{3`hnf1Cc2M5xr2d`Yx$2M5g^iKC5IvG9Vd{3`ho424W<5GaA~kQ_Iq z3KH3NJNdkl0m*=5Kr$d1P~%|4kFWBuak*8!2q7unxeRl(Do-P>;jQi_<6GN5v`yP~ zg(?Mh4s2`0H*eH!4HlB#i1Iv(Wm|q2`pQ+38^1`EA zGqGWgU6>$!S;MSom{U2UP;8^nGxM%>5vrD`@el9JEwsTTZSE!K)tV0P4!H5kRg~>d zrHZ1GDU)PCG9Vd{3`hndQwH2OdMNc9ql{TqX)@)5@58b<*nXhLbB@pHzg=Z1%D|Wp=U7kG|2kD!#t*SqVYZ*Fknc}b7{|Btj43-S^qw8V7~$5tr3_=uTDe<_ z-gcI9w-mK!8>-(iY}$8HKeFPZXY8@>Is5YSPTS+-$*lvU`dO>9*yCE6?QzeXR?N#A z{@LvD{p|MTAMLllUW%6X<(K%{@$B)*wDx%BX?y;&pS8!eZ`RISWx=&#OM5I;D|_?CZPH$R0=ZN~DP)_U8u=w8zF@^bR>!t?rGu)Uel?n>!1b zKX@@b=g--@ZeY$DHxgB~{YaX@pnPWiXz4qKG5?mazgyv;cxIWdEln~XJJz1mefiWO zFFk*P`Pkbbe%E}gj)m#@i_8ixrb1N%R8s~EBQs{hoT#fO`H2f+5iE|S zuq;-<%BVL$4dU8Z4;x|=Y=$kc6}CZL&7wMGbfUZq(miQ(C+>xPaR3g+VK@@iHc^w> zwrW!QfSRUQ81LZ>oQ?BvAuhtDxExpET2$wmscDmi@hNV@9k?6!;sHE_U*a)5f#2X+ z3_GFCe0-7kGG4)JcoXm7eSCnw<70e=QOxh3nxbJ0jE!+I0Vc*|m=e=qdd!4bF$dB%tcwk>F*e2K=!>n<9|N!>24Oc0#h%y)`{N)SiX(6| zj>8E!8K>cNoP~370e*-}@FQG_Yj8bo#Lc)Bx8pAS9O-d1z92q=NAYVsg}TS|TjC4& zJ^qMS@dn<;d-x0fhL7+mdQhE?iqSC^#=-cQ2$Nz8OpWO)F9aS~3& z_i-l9!TIq@8bjf9Ur4Q&t6Ts zTOgXM$eL3OjE!+I0Vc*|m=e=qdd!4bF$d898SQ=I1Q)cES!rA@IzdJAK^+|gX?i4 zZpN*+9e3g9xF5g3BX|_Q##49(zr_pqJ^qMS@dk1v!?;KM3;u?W@F{X6#fVC?GdjjX zedHoOaUx8LDKIsr!;F{(vtusIiv_R{7Dau;r8IFltcX5X4QpZ@tdF{R{xv#+A;?MXiKEx*&)*)g3j(XDY_rjRy zjqxxcCc)&G3e#c+%#7JEC+5NYSP+X~aV&*pu>w}cs#pVSV?At$O|TiZz*g7>+hHJf z!Y&w$-LV(;#Q``Nhv7)%h^+B0@kE@0@8Jxbjq`9JF2bd_99Q95+<=?#Q{09-a5wJ7 z19%9(#AA2@zrnM39xvi$yn@&8Cf>pO_yB*$$M_7Rs17Z2iiR;THpaySm>82`N=$?4 zF%xFR9GDyP;agZ3i(yGDgXOUjR>A663+rM7Y>Z8@Ir?I2^v3|~h(Xv5L$N3J!TvZ1 zhvEnvjpJ|vPR40C9cSTOT!0_q68s2P;u>6!8&Tiib1U(7)VB`*oOnNefk*HtevPN_ z41S9j@O%6bui_27jrULw2=N>7Bh*)6dZ>uhv0A= zg=29%PQt19KF-8BI3GX2#kdSt;A&ilALA#u1wX@`xCi&)K|G9K;c+~Pr|}$qhnMgN z{0Xn)Exe0A0e|1e0SbOp6&XGiJk_mq@8bjf9UtQ} zj1tTGy^Mx2FgC`;1eh3;VMi*^EQ95-5>~{1im<6DmcY_j4lAM$R>PWD2kT=ad>h}vmgt9Vu|0Oc z&e#=0um|?WemD?^;BXv;V{trA!m0Q^&cr!5A3wmwxC~d|YFvjO<0rTUKf|542lwGY zJd9uAaXg8q@f?1Km+%Mt39sWVyo*2MulNw3pgzFwsZKmFCohbN-WU%PViHV_sW2^O zz|5Epb7CIMj|H&^7RORp7As(7tco?THrB(2*aVwl3v7jLupI_sC+vd3*d2ReUmSpg zaTt!oG59V{#3}e5&cN9?4;SJhT#Cza6|Th%xCuYSZMXw><6b;~hww{0h9~eFJd5Y? zB3{NTcnxpj9lVbZ@OONS&oGKQ;l-SyVGN9oaWMfV#$=ch(_nhcgjq2M=Ei*Z78b^0 zSQ5)%d8~w0usYVly4V04V^eI7zStW5F#tPa5O%{*?1_D_KMum7I08rGIGli!aT-p? zSvVIL;D@*bKf;x`2G`?8+>BdsJMO~IaX)^6NAM_qji>Moev23Id;Afv;tjlw_wX0| z4IklC^oVQir%^FF#=qLrUMzrxuqc+m(pU~Fq7PQXnpg+x zVSKw+~ zhacl7xCKAMowx`0;Xyo%U*T~)iKp=#eutOv2mA@I<1M_4KjW|X5TBqvvCoqe{JbzG zdSg6Hh)FOxroyzC0W)JZ%!zq0KNiFySR6}XS*(DSu`1TU+E@=8ViRnJEwB}~!FCvk zov;fAV|VO@eQ^K|#$h-T$Kbm-5vSmLI0I+nJY0y2a49awRk#*6;3oVOx8V-ljeGF` z9>Op27@ojy@GPFki+CBY;5EF7ckn(wz~AvPKEo*L1Xy#5hA}WU#>E7f7?WX2OoQn$ z6K2I6m>cuqTUZ#2VM#25<*^c0!RlBG>tX|Jj7_mQ`eJMJ#{lezLD&sLu_yMy{x}GS z;s_j#<8T5_#%VYmXW?92fFI%#{0LX#8eESXaWihk?YIj+$Nl&P9>JsdHJ-vV_$^+* z@9{^xiZ}2!-osz;H++Oo(IbJipGL*#7z^WId`yH%F$Jc^beIvdV0O%fd9eT%!lGCL zOJg~#h(1^iYhoR&kB#tcdiAGXEz*a16ZR}8@(*cK0cj6x0hX?U6euc;JB%a1|_#IxtAMhu?DxW$=#B9(Atu4(m7>D6V9E0!TM4W=};S8LO^Kc<9!lk$zSK(UR zfSd4B+=e@FH}1s)cnH75V|W6;!LxWCFXCmqg4ggS-og9$0Ds5F_za^YvVK>iVGN9o zaWMfV#$=ch(_nhcgjq2M=Ei*Z78b^0SQ5)%d8~w0usYVly4V04V^eI7zStW5F#tPa z5O%{*?1_D_KMum7I08rGIGli!aT-p?SvVIL;D@*bKf;x`2G`?8+>BdsJMO~IaX)^6 zNAM_qji>Moev23Id;Afv;tjlw_wX0|4IklC^hnI#e~gZ?Fb>AYM3@v)U}{W<88HiH z$6S~f3t%BEiY2f#mcxqZgVnGm*1`JN2;atcuqFCoTWpUVurqeW5bS}yu^$e^AvhdI z;aD7xlW;1&k27%&&c_dMF)qUuxEj~t$M^|u!Ow6f?!kR{5D(*5cpOjSX*`GD;U)Y5 zf5Pi{3-993_$xldC#Wyv^W=g)FN}%a7!MO-5=@S%FfC@l%$N;xVjj$o1+fSg$5L1p zD_~`;iZ!q{*29L_1e;+CY=v#G9R^}2?1I7A9eZJ49Dsvy7>>j-_%2SwDfk}Fz}Ywt z7vdsZipy~ouEh~C-56Qi|6qoUdAhU4R7KdypIp?cYKV` zFiKMX{$mV`jd3vnCdOo#64PLM%!FAn2j<3n_!bt%VptN(V0o;BRj@kN!n)W18)H*! zj=tC${V@PLVi0!2Q0$3)us;sMp*R9Z<2amvlW`hO$5}WR7vP7u1V6%+xCYnbM%;{B zaXaq9&v8F~fk*HtevPN_41S9j@O%6bui_27jrZ^u{0$%BQ}js2-+zpbu`mwC$3&PE zQ($ULhZ!*oX2)EZ7Ykq^EQ%$tG?v4P=!4a;Cf33F*a+Xocd#Y;VOwmE9k4TY#SrX) zy|Et-#348wN8wl;kCSjJzK=6;4$j99a4{~!6}TGL;m7z1Zo$uRC+@+0cn}ZcS9lyx z;%PjG-{B?v0e`~lcnk01&-g1o#3!gP;P>Q0elLuP-WU%PViHV_sW2^Oz|5Epb7CIM zj|H&^7RORp7As(7tco?THrB(2*aVwl3v7jLupI_sC+vd3*d2ReUmSpgaTt!oG59V{ z#3}e5&cN9?4;SJhT#Cza6|Th%xCuYSZMXw><6b;~hww{0h9~eFJd5Y?B3{NTcnxpj z9lVbZ@OONS&oD{~{{CYOjE!+I0Vc*|m=e=qdd!4bF$dvLz8vQWo0fRk|=PRCg|7Z>1%xCB4K zmAD4i<3`+!TX8$?!q0I(et}2uD1MEn@C<&77w~)h5wGG6yp8wp7yJz$;ZyWT$=`pB zjYgM6jNYoOotgU3ueb$m=_CRAuNg|ur!v#is*yYuqM{Q`q&8H#&@tK`e9pa zj~%cxcEu3vfxWRG4#Xii97o|;9FLQ5D!z|1aSqPM4{$Lq!xgw1*Wt(b32wp9a3}7; zeRvQL<5zebPvU7jhu`5P`~iQ$>v#+A;?MXiKEx+zFmR72dSOiT#(0?$`_a;s6|s z!*C>y!FO>YPQmwZ2F}KLxDXfNQe2L!a4l}YP53Ep!yULA_u>IOgkRz@Jb~ZfSv-#y z@iJb)Yj_jy;C+06zvE+khEY=U_a9?mY>bNuFfk^>l$Zw7VQQ z2Fqh5tb*0C7S_cE*ch8)bM(d5=#K%|5reQBhGI|bgZ*(34#g2T8pq)ToQ%_OI?lqm zxBx%ICHN7p#5K4cH{xd8iraA)evbR`3p|2H@oPMVXYgCRfZyYfcolEpZM=uS;BWW{ zpQ48v8rPhnVswm!aWFn6!lal2Q)4>Jh*>Z@=EA&K01IJJEPov|y1U=Qq#{cs=-!QnUx$KrULgj4Z-oQZRAK7N3UaT%__)wm8n z#!qkyeug`75AMT*co@IJ<9HHJ<2n2eFX0dP6JEz#co%=hU-2P6L4yJNJkbkdqBq9F zgqQ@AV=7FG889t@caSXnT6LAW@hcj?C&clVc2$$k=T!m|K18%}kaU1Ty-MAMI;351H zkKqaY2G8Poyoi_a3SPsTcn9y}1Nt+h7>oaXcz-yV_Zyti7^?b#59;5GhtTD zfw?gszJ-Oc7?#8`SRN~36|9c6ur4;h#@G~_qc65be+X2ObKH+#;1N8EU*jn}gWuu>{2qVA zt9S!%<30QZf5S)k6g|?bz0~~spFRvDI>y2{7#|a1QcQuVF&$>aESMd0VO}hNg|H}= zz|vR_E20lp!tiE)8{ff}=!b2wJ$As(*cC&t2lmE(I1q>6a2$nWaXe1KsrWw5 z#5p)0KfuMf3|HW4T!$ayC%6Sa!=1PX_u)Z2j9=k#Jc+0A9DawF@CW<}uj4Jei$CM9 z_z<6<9;nwdgAq1)VNCSKc$g5AU~)`_X)yz4#%!1q^I(1~h()kCmcp`F0V`uwtbw($ z9yY`#*bG}>D{O=9Fc3Ro7YxSk*bDpO033|Na3qewcX1+4!S`?m&c=DT5EtQ6T#l=7 zEpEU~_$h9~9k?6!;sHE_U*a)5f#2X+JdYRgGG4)JcoXm7eSCnw<70e=Q8HTlX*7(1 zu`w 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 index 4ae1103af..bc64ada38 100644 --- a/test/test_declarative_parsing.py +++ b/test/test_declarative_parsing.py @@ -10,9 +10,10 @@ import re from pathlib import Path -from typing import Any +from typing import Any, Literal import numpy as np +import pandas as pd import pytest import xarray as xr import yaml @@ -32,19 +33,29 @@ NODES = ["a", "b", "c"] -def _math() -> dict: +@pytest.fixture +def math() -> dict: """Minimal but representative math definition exercising every route.""" return { "dimensions": {"node": {"dtype": "string", "iterator": "n"}}, "parameters": { - "cost": {"default": 0}, + "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": { @@ -78,7 +89,8 @@ def _math() -> dict: } -def _inputs() -> xr.Dataset: +@pytest.fixture +def inputs() -> xr.Dataset: return xr.Dataset( { "cost": ("node", [1.0, 2.0, 3.0]), @@ -115,13 +127,76 @@ def _first_equation( @pytest.fixture -def builder_with_flow() -> DeclarativeModelBuilder: +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 = 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.""" @@ -181,11 +256,12 @@ def test_node_repr_is_clean(self) -> None: assert "instring=" not in rendered assert "loc=" not in rendered - def test_parse_error_carries_position_marker(self) -> None: - math = _math() + 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(), {}) + DeclarativeModelBuilder(math, inputs, {}) message = str(excinfo.value) assert "constraints:cap:" in message assert "^" in message @@ -194,63 +270,66 @@ def test_parse_error_carries_position_marker(self) -> None: class TestParseWalkthrough: """Whole-dict parse walkthrough with aggregated errors.""" - def test_errors_aggregate_across_components(self) -> None: + def test_errors_aggregate_across_components( + self, math: dict, inputs: xr.Dataset + ) -> None: """Broken strings in two components raise as one grouped error.""" - math = _math() 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(), {}) + 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) -> None: + 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 = _math() 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(), {}) + 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) -> None: + def test_inactive_components_are_skipped( + self, math: dict, inputs: xr.Dataset + ) -> None: """Inactive components are neither parsed nor built.""" - math = _math() math["expressions"]["broken"] = { "active": False, "foreach": ["node"], "equations": [{"expression": "flow * * cost"}], } - builder = DeclarativeModelBuilder(math, _inputs(), {}) + 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) -> None: - math = _math() + 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(), {}) + DeclarativeModelBuilder(math, inputs, {}) - def test_inactive_check_masks_are_skipped(self) -> None: - math = _math() + 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(), {}) + builder = DeclarativeModelBuilder(math, inputs, {}) assert "bad" not in builder.parsed.checks - def test_parsed_math_shape(self) -> None: - builder = DeclarativeModelBuilder(_math(), _inputs(), {}) + def test_parsed_math_shape(self, math: dict, inputs: xr.Dataset) -> None: + builder = DeclarativeModelBuilder(math, inputs, {}) assert set(builder.parsed.components) == { "variables", "expressions", @@ -272,12 +351,13 @@ def test_top_level_mask_returns_boolean_dataarray( assert sub_mask.dtype == bool assert bool(sub_mask.all()) - def test_mask_comparison_and_subset_and_helper_return_bool(self) -> None: - math = _math() + 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 = DeclarativeModelBuilder(math, inputs, {}) builder.add_variable("flow", builder.math.variables["flow"]) _, sub_mask, _ = _first_equation(builder, "constraints", "cap") assert sub_mask.dtype == bool @@ -340,14 +420,15 @@ def test_sub_expression_reference_returns_linexpr( result = parsing.as_expression(equation, ctx, mask=sub_mask) assert isinstance(result, LinearExpression) - def test_sub_expression_variants_expand_to_cartesian_product(self) -> None: + 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 = _math() 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 = DeclarativeModelBuilder(math, inputs, {}) builder.add_variable("flow", builder.math.variables["flow"]) definition = builder.math.expressions["sub_expr_test"] equations = parsing.parse_component( @@ -369,21 +450,22 @@ def test_sub_expression_variants_expand_to_cartesian_product(self) -> None: builder.add_expression("sub_expr_test", definition) assert "sub_expr_test" in builder.model.expressions - def test_undefined_sub_expression_reference_raises(self) -> None: - math = _math() + 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(), {}) + DeclarativeModelBuilder(math, inputs, {}) def test_plain_and_list_slices( - self, builder_with_flow: DeclarativeModelBuilder + self, inputs: xr.Dataset, builder_with_flow: DeclarativeModelBuilder ) -> None: ctx = _ctx( builder_with_flow, mode="expr", - mask=xr.full_like(_inputs()["cost"], True, bool), + mask=xr.full_like(inputs["cost"], True, bool), ) arith = grammar.arithmetic_grammar( frozenset({"flow", "cost", "cap_max", "node"}) @@ -396,15 +478,14 @@ def test_plain_and_list_slices( result = nodes.evaluate(list_sliced, ctx) assert result.data.sizes["node"] == 2 - def test_slicer_reference(self, builder_with_flow: DeclarativeModelBuilder) -> None: + def test_slicer_reference(self, math: dict, inputs: xr.Dataset) -> None: """`$name` slicer references resolve like sub-expressions (feature parity).""" - math = _math() math["expressions"]["sliced"] = { "foreach": ["node"], "equations": [{"expression": "flow[node=$n] * cost"}], "slices": {"n": [{"expression": "a"}]}, } - builder = DeclarativeModelBuilder(math, _inputs(), {}) + builder = DeclarativeModelBuilder(math, inputs, {}) builder.add_variable("flow", builder.math.variables["flow"]) definition = builder.math.expressions["sliced"] equations = parsing.parse_component( @@ -431,11 +512,10 @@ def test_equation_returns_lhs_sign_rhs_tuple( assert isinstance(sign, xr.DataArray) assert set(np.unique(sign.values)) <= {"<="} - def test_foreach_dim_mismatch_raises(self) -> None: - math = _math() + 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 = 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`"): @@ -456,24 +536,24 @@ def test_sum_latex(self, builder_with_flow: DeclarativeModelBuilder) -> None: == r"\sum\limits_{\substack{\text{n} \in \text{node}}} (total_cost)" ) - def test_mask_latex(self) -> None: - math = _math() + 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 = 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) -> None: + 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 = _math() math["constraints"]["cap"]["equations"][0]["mask"] = "cap_max == inf" - builder = DeclarativeModelBuilder(math, _inputs(), {}) + 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) @@ -529,13 +609,14 @@ def as_raw(self, array: Any) -> LinearExpression | xr.DataArray: # noqa: D102 class TestHelpers: """Helper-function registration and argument evaluation.""" - def test_helper_arguments_are_evaluated_raw(self) -> None: + 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 = _math() math["expressions"]["total_cost"]["equations"][0]["expression"] = ( "record_args(flow, cost)" ) - builder = DeclarativeModelBuilder(math, _inputs(), {}) + 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"] @@ -553,12 +634,11 @@ 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) -> None: - math = _math() + 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 = DeclarativeModelBuilder(math, inputs, {}) builder.add_variable("flow", builder.math.variables["flow"]) definition = builder.math.expressions["total_cost"] equation = parsing.parse_component( @@ -567,13 +647,12 @@ def test_unknown_helper_rejected(self) -> None: with pytest.raises(ValueError, match="Invalid helper function"): parsing.as_expression(equation, _ctx(builder)) - def test_eval_error_carries_caret(self) -> None: + 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 = _math() math["expressions"]["total_cost"]["equations"][0]["expression"] = ( "unknown_helper(flow)" ) - builder = DeclarativeModelBuilder(math, _inputs(), {}) + builder = DeclarativeModelBuilder(math, inputs, {}) builder.add_variable("flow", builder.math.variables["flow"]) definition = builder.math.expressions["total_cost"] equation = parsing.parse_component( @@ -603,12 +682,11 @@ def as_raw( with pytest.raises(ValueError, match="already exists"): build_registry([_ClashingSum]) - def test_custom_helper_end_to_end(self) -> None: - math = _math() + 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]) + 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: @@ -621,35 +699,35 @@ def test_get_val_at_index(self, builder_with_flow: DeclarativeModelBuilder) -> N class TestBuilder: """Model assembly from parsed math.""" - def test_overlapping_equation_masks_rejected(self) -> None: - math = _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(), {}) + declarative_model(math, inputs, {}) - def test_multiple_active_objectives_rejected(self) -> None: - math = _math() + 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(), {}) + declarative_model(math, inputs, {}) - def test_references_are_sorted_lists(self) -> None: - model = 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) -> None: - math = _math() + def test_dtype_coercion(self, math: dict, inputs: xr.Dataset) -> None: math["lookups"] = { "flag": {"dtype": "bool", "default": False}, "label": {"dtype": "string"}, } - inputs = _inputs() inputs["flag"] = ("node", [1.0, float("nan"), 0.0]) inputs["label"] = ("node", ["x", "", "z"]) builder = DeclarativeModelBuilder(math, inputs, {}) @@ -658,9 +736,10 @@ def test_dtype_coercion(self) -> None: # Empty strings are coerced to missing values. assert builder.input_data["label"].isnull().sum() == 1 - def test_checks_run_without_active_variable(self) -> None: + 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 = _math() math["checks"] = { "too_expensive": { "mask": "cost > 100", @@ -669,10 +748,11 @@ def test_checks_run_without_active_variable(self) -> None: } } # No `active` variable in the inputs, and the check does not trigger. - declarative_model(math, _inputs(), {}) + declarative_model(math, inputs, {}) - def test_check_raises_when_triggered_without_active(self) -> None: - math = _math() + def test_check_raises_when_triggered_without_active( + self, math: dict, inputs: xr.Dataset + ) -> None: math["checks"] = { "too_expensive": { "mask": "cost > 0", @@ -681,23 +761,33 @@ def test_check_raises_when_triggered_without_active(self) -> None: } } with pytest.raises(ValueError, match="cost too high"): - declarative_model(math, _inputs(), {}) + declarative_model(math, inputs, {}) - def test_check_warns(self, caplog: pytest.LogCaptureFixture) -> None: - math = _math() + 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(), {}) + 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) -> None: - builder = LatexModelBuilder(_math(), _inputs(), {}).build() + 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 == [ @@ -707,20 +797,21 @@ def test_components_render_with_decorated_reprs(self) -> None: } ] - def test_variable_bounds_equation(self) -> None: - math = _math() + 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() + 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) -> None: + 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() + 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. @@ -736,15 +827,17 @@ def test_escape_text_mode_escapes_content_only(self) -> None: # 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) -> None: + 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") + 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) -> None: - builder = LatexModelBuilder(_math(), _inputs(), {}).build() + 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. @@ -752,45 +845,47 @@ def test_cross_references(self) -> None: assert obj.uses == ["total_cost"] assert obj.extras["Sense"] == "minimise" - def test_equation_masks_render_as_if_conditions(self) -> None: - math = _math() + 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") + 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) -> None: + 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 = _math() math["constraints"]["cap"]["equations"][0]["mask"] = "cap_max == inf" - doc = latex_math_doc(math, _inputs(), format="md") + 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) -> None: - math = _math() + 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() + 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) -> None: + 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 = _math() 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") + 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 @@ -800,26 +895,28 @@ def test_multiple_equations_render_as_one_block_with_cases(self) -> None: r"\text{if } (\neg (\textit{cost}_\text{n}\mathord{>}\text{1}))" in section ) - def test_single_equation_renders_inline_without_cases(self) -> None: - doc = latex_math_doc(_math(), _inputs(), format="md") + 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) -> None: - doc = latex_math_doc(_math(), _inputs(), format="md") + 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) -> None: - doc = latex_math_doc(_math(), _inputs(), format="rst") + 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) -> None: - doc = latex_math_doc(_math(), _inputs(), format="tex") + 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. @@ -829,8 +926,10 @@ def test_tex_document_structure(self) -> None: class TestEndToEnd: """Full builds from math definitions.""" - def test_declarative_model_end_to_end(self) -> None: - model = declarative_model(_math(), _inputs(), {}) + 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 @@ -838,18 +937,33 @@ def test_declarative_model_end_to_end(self) -> None: # flow is indexed over the node dimension. assert set(model.variables["flow"].dims) == {"node"} - def test_repo_math_yaml_validates_and_parses(self) -> None: - """The demo math.yaml at the repo root validates and every component parses.""" - math_path = Path(__file__).parent.parent / "math.yaml" - if not math_path.exists(): - pytest.skip("repo-root math.yaml not present") - math = MathModel.model_validate(yaml.safe_load(math_path.read_text())) - for group in ("expressions", "constraints", "objectives"): - for name, definition in getattr(math, group).root.items(): - equations = parsing.parse_component(group, name, definition, math) - assert equations, f"{group}:{name} produced no equations" + @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. - for fmt in ("md", "rst", "tex"): - doc = latex_math_doc(yaml.safe_load(math_path.read_text()), format=fmt) - name = r"storage\_balance" if fmt == "tex" else "storage_balance" - assert name in doc + doc = latex_math_doc(larger_math, larger_inputs, format=fmt) + name = r"storage\_balance" if fmt == "tex" else "storage_balance" + assert name in doc