diff --git a/src/calliope/backend/backend_model.py b/src/calliope/backend/backend_model.py index 53667465b..9b2b03ce8 100644 --- a/src/calliope/backend/backend_model.py +++ b/src/calliope/backend/backend_model.py @@ -59,7 +59,6 @@ def __init__(self, inputs: xr.Dataset, **config_overrides): self.inputs.attrs["config"].build.union( AttrDict(config_overrides), allow_override=True ) - self.valid_math_element_names: set = set() self._solve_logger = logging.getLogger(__name__ + ".") @abstractmethod @@ -197,10 +196,6 @@ def _build(self) -> None: "objectives", ]: component = components.removesuffix("s") - if components in ["variables", "global_expressions"]: - self.valid_math_element_names.update( - self.inputs.math[components].keys() - ) for name in self.inputs.math[components]: getattr(self, f"add_{component}")(name) LOGGER.info( @@ -275,8 +270,7 @@ def _add_component( ) top_level_where = parsed_component.generate_top_level_where_array( - self.inputs, - self._dataset, + self, align_to_foreach_sets=False, break_early=break_early, ) @@ -285,7 +279,7 @@ def _add_component( self._create_obj_list(name, component_type) - equations = parsed_component.parse_equations(self.valid_math_element_names) + equations = parsed_component.parse_equations(self.valid_component_names) if not equations: component_da = component_setter( parsed_component.drop_dims_not_in_foreach(top_level_where) @@ -297,9 +291,7 @@ def _add_component( .astype(np.dtype("O")) ) for element in equations: - where = element.evaluate_where( - self.inputs, self._dataset, initial_where=top_level_where - ) + where = element.evaluate_where(self, initial_where=top_level_where) if break_early and not where.any(): continue @@ -462,6 +454,7 @@ def _apply_func( kwargs=kwargs, vectorize=True, keep_attrs=True, + dask="parallelized", output_dtypes=[np.dtype("O")], output_core_dims=output_core_dims, ) @@ -519,6 +512,19 @@ def objectives(self): "Slice of backend dataset to show only built objectives" return self._dataset.filter_by_attrs(obj_type="objectives") + @property + def valid_component_names(self): + def _filter(val): + return val in ["variables", "parameters", "global_expressions"] + + in_data = set(self._dataset.filter_by_attrs(obj_type=_filter).data_vars.keys()) + in_math = set( + name + for component in ["variables", "global_expressions"] + for name in self.inputs.math[component].keys() + ) + return in_data.union(in_math) + class BackendModel(BackendModelGenerator, Generic[T]): def __init__(self, inputs: xr.Dataset, instance: T, **config_overrides) -> None: diff --git a/src/calliope/backend/expression_parser.py b/src/calliope/backend/expression_parser.py index 94e5fc357..c63180947 100644 --- a/src/calliope/backend/expression_parser.py +++ b/src/calliope/backend/expression_parser.py @@ -31,7 +31,7 @@ from __future__ import annotations import re -from abc import ABC +from abc import ABC, abstractmethod from typing import ( TYPE_CHECKING, Any, @@ -44,8 +44,11 @@ overload, ) +import numpy as np +import pandas as pd import pyparsing as pp import xarray as xr +from typing_extensions import NotRequired, TypedDict, Unpack from calliope.backend.helper_functions import ParsingHelperFunction from calliope.exceptions import BackendError @@ -58,12 +61,120 @@ SUB_EXPRESSION_CLASSIFIER = "$" +class EvalAttrs(TypedDict): + equation_name: str + where_array: xr.DataArray + slice_dict: dict + sub_expression_dict: dict + backend_interface: BackendModel + input_data: xr.DataArray + references: set[str] + helper_functions: dict[str, Callable] + as_values: NotRequired[bool] + + +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 + + def __eq__(self, other): + return self.__repr__() == other + + def __repr__(self) -> str: + "Return string representation of the parsed grammar" + + +class EvalToArrayStr(EvalString): + @abstractmethod + def as_math_string(self) -> str: + """Return evaluated expression as a LaTex math string""" + + @abstractmethod + def as_array(self) -> xr.DataArray | list[str, float]: + """Return evaluated expression as an array""" + + @overload + def eval(self, return_type: Literal["math_string"], **eval_kwargs) -> str: + """math strings evaluate to strings""" + + @overload + def eval( + self, return_type: Literal["array"], **eval_kwargs + ) -> xr.DataArray | list[str, float]: + """arrays evaluate to arrays""" + + def eval( + self, return_type: RETURN_T, **eval_kwargs + ) -> Union[str, list[str | float], xr.DataArray]: + """Evaluate math string expression. + + Args: + return_type (Literal[math_string, input, array]): + Dictates how the expression should be evaluated (see `Returns` section). + Keyword Args: + equation_name (str): Name of math component in which expression is defined. + slice_dict (dict): Dictionary mapping the index slice name to a parsed equation expression. + sub_expression_dict (dict): Dictionary mapping the sub-expression name to a parsed equation expression. + backend_interface (backend_model.BackendModel): Interface to optimisation backend. + input_data (xr.Dataset): Input parameter arrays. + where_array (xr.DataArray): boolean array with which to mask evaluated expressions. + references (set): any references in the math string to other model components. + helper_functions (dict[str, type[ParsingHelperFunction]]): Dictionary of allowed helper functions. + as_values (bool, optional): Return array as numeric values, not backend objects. Defaults to False. + Returns: + Union[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_kwargs + evaluated: Union[str, list[str, float], 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): + @abstractmethod + def as_callable(self) -> Callable: + """""" + + def eval( + self, + return_type: RETURN_T, + **eval_kwargs: Unpack[EvalAttrs], # type: ignore + ) -> Callable: + """Evaluate math string expression. + Args and kwargs are passed on directly to helper function. + + Args: + return_type (str): Whether to return a math string or xarray DataArray. + Keyword Args: + equation_name (str): Name of math component in which expression is defined. + slice_dict (dict): Dictionary mapping the index slice name to a parsed equation expression. + sub_expression_dict (dict): Dictionary mapping the sub-expression name to a parsed equation expression. + backend_interface (backend_model.BackendModel): Interface to optimisation backend. + input_data (xr.Dataset): Input parameter arrays. + where_array (xr.DataArray): boolean array with which to mask evaluated expressions. + references (set): any references in the math string to other model components. + helper_functions (dict[str, type[ParsingHelperFunction]]): Dictionary of allowed helper functions. + as_values (bool, optional): Return array as numeric values, not backend objects. Defaults to False. + Returns: + Callable: returns helper function. + """ + self.eval_attrs = eval_kwargs + evaluated = self.as_callable(return_type) + return evaluated -class EvalOperatorOperand(EvalString): + +class EvalOperatorOperand(EvalToArrayStr): LATEX_OPERATOR_LOOKUP: dict[str, str] = { "**": "{val}^{{{operand}}}", "*": r"{val} \times {operand}", @@ -71,6 +182,7 @@ class EvalOperatorOperand(EvalString): "+": "{val} + {operand}", "-": "{val} - {operand}", } + SKIP_IF: list[str] = ["+", "-"] def __init__(self, tokens: pp.ParseResults) -> None: """ @@ -86,65 +198,47 @@ def __init__(self, tokens: pp.ParseResults) -> None: self.values = tokens def __repr__(self) -> str: - "Return string representation of the parsed grammar" first_operand = self.value[0].__repr__() operand_operator_pairs = " ".join( op + " " + val.__repr__() - for op, val in self.operatorOperands(self.value[1:]) + for op, val in self._operator_operands(self.value[1:]) ) arithmetic_string = f"({first_operand} {operand_operator_pairs})" return arithmetic_string - def operatorOperands( - self, tokenlist: list + def _operator_operands( + self, token_list: list ) -> Iterator[tuple[str, pp.ParseResults]]: "Generator to extract operators and operands in pairs" - - it = iter(tokenlist) + it = iter(token_list) while 1: try: yield (next(it), next(it)) except StopIteration: break - def as_latex( - self, val: str, operand: str, operator_: str, val_type: Any, operand_type: Any - ) -> str: - """Add sign to stringified data for use in a LaTex math formula""" - # We ignore zeros that do nothing - if operand == "0" and operator_ in ["-", "+"]: - return val - if val_type == type(self): - val = "(" + val + ")" - if operand_type == type(self): - operand = "(" + operand + ")" - if val == "0" and operator_ in ["-", "+"]: - return operand - - return self.LATEX_OPERATOR_LOOKUP[operator_].format(val=val, operand=operand) - - def _eval( - self, - to_eval: pp.ParseResults, - as_latex: bool = False, - **eval_kwargs, - ) -> Any: - evaluated = to_eval.eval(as_latex=as_latex, **eval_kwargs) - if not as_latex: - evaluated = xr.DataArray(evaluated) - if eval_kwargs.get("apply_where", True): - try: - evaluated = evaluated.where(eval_kwargs["where"]) - except AttributeError: - evaluated = evaluated.broadcast_like(eval_kwargs["where"]).where( - eval_kwargs["where"] - ) + def _apply_where_array(self, evaluated: xr.DataArray) -> xr.DataArray: + "eval util function to apply where arrays to non-latex strings" + where_array = self.eval_attrs["where_array"] + try: + evaluated = evaluated.where(where_array) + except AttributeError: + evaluated = evaluated.broadcast_like(where_array).where(where_array) return evaluated - def operate( - self, val: xr.DataArray, evaluated_operand: xr.DataArray, operator_: str + 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" if operator_ == "**": val = val**evaluated_operand elif operator_ == "*": @@ -157,44 +251,42 @@ def operate( val = val - evaluated_operand return val - @overload # noqa: F811 - def eval( # noqa: F811 - self, as_latex: Literal[False] = False, **eval_kwargs - ) -> xr.DataArray: - "Expecting array if not requesting latex string" - - @overload # noqa: F811 - def eval(self, as_latex: Literal[True], **eval_kwargs) -> str: # noqa: F811 - "Expecting string if requesting latex string" + def as_math_string(self) -> str: + val = self.value[0].eval(return_type="math_string", **self.eval_attrs) - def eval( # noqa: F811 - self, as_latex: bool = False, **eval_kwargs - ) -> Union[str, xr.DataArray]: - """ - Returns: - Any: - If all operands are numeric, returns float, otherwise returns an - expression to use in an optimisation model constraint. - """ - val = self._eval(self.value[0], as_latex, **eval_kwargs) - - for operator_, operand in self.operatorOperands(self.value[1:]): - evaluated_operand = self._eval(operand, as_latex, **eval_kwargs) - if as_latex: - val = self.as_latex( - val, - evaluated_operand, - operator_, - type(self.value[0]), - type(operand), - ) + for operator_, operand in self._operator_operands(self.value[1:]): + evaluated_operand = operand.eval( + return_type="math_string", **self.eval_attrs + ) + # We ignore zeros that do nothing + if self._skip_component_on_conditional(evaluated_operand, operator_): + continue + if type(self.value[0]) == type(self): + val = "(" + val + ")" + if type(operand) == type(self): + evaluated_operand = "(" + evaluated_operand + ")" + if self._skip_component_on_conditional(val, operator_): + val = evaluated_operand else: - val = self.operate(val, evaluated_operand, operator_) + val = self.LATEX_OPERATOR_LOOKUP[operator_].format( + val=val, operand=evaluated_operand + ) + return val + + def as_array(self) -> xr.DataArray: + val = self._apply_where_array( + self.value[0].eval(return_type="array", **self.eval_attrs) + ) + for operator_, operand in self._operator_operands(self.value[1:]): + evaluated_operand = self._apply_where_array( + operand.eval(return_type="array", **self.eval_attrs) + ) + val = self._operate(val, evaluated_operand, operator_) return val -class EvalSignOp(EvalString): +class EvalSignOp(EvalToArrayStr): def __init__(self, tokens: pp.ParseResults) -> None: """ Parse action to process successfully parsed expressions with a leading + or - sign. @@ -207,24 +299,31 @@ def __init__(self, tokens: pp.ParseResults) -> None: self.values = tokens def __repr__(self) -> str: - "Return string representation of the parsed grammar" return str(f"({self.sign}){self.value.__repr__()}") - def as_latex(self, val: str) -> str: - """Add sign to stringified data for use in a LaTex math formula""" - return self.sign + val + @overload + def _eval(self, return_type: Literal["math_string"]) -> str: + "string return" + + @overload + def _eval(self, return_type: Literal["array"]) -> xr.DataArray: + "array return" - def eval(self, **eval_kwargs) -> Any: - val = self.value.eval(**eval_kwargs) - if eval_kwargs.get("as_latex", False): - return self.as_latex(val) - elif self.sign == "+": - return val - elif self.sign == "-": - return -1 * val + def _eval(self, return_type: RETURN_T) -> Union[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: + return self.sign + self._eval("math_string") -class EvalComparisonOp(EvalString): + def as_array(self) -> xr.DataArray: + evaluated = self._eval("array") + if self.sign == "-": + evaluated = -1 * evaluated + return evaluated + + +class EvalComparisonOp(EvalToArrayStr): OP_TRANSLATOR = {"<=": r" \leq ", ">=": r" \geq ", "==": " = "} def __init__(self, tokens: pp.ParseResults) -> None: @@ -239,30 +338,48 @@ def __init__(self, tokens: pp.ParseResults) -> None: self.values = tokens def __repr__(self) -> str: - "Return string representation of the parsed grammar" return f"{self.lhs.__repr__()} {self.op} {self.rhs.__repr__()}" - def as_latex(self, lhs: str, rhs: str) -> str: - """Add operator between two sets of stringified data for use in a LaTex math formula""" - return lhs + self.OP_TRANSLATOR[self.op] + rhs + @overload + def _eval(self, return_type: Literal["math_string"]) -> tuple[str, str]: + "string return" - def eval(self, **eval_kwargs) -> Any: - """ - Returns: - Any: - If LHS and RHS are numeric, returns bool, otherwise returns an equation - to use as an optimisation model constraint. - """ - lhs = self.lhs.eval(**eval_kwargs) - rhs = self.rhs.eval(**eval_kwargs) + @overload + def _eval(self, return_type: Literal["array"]) -> tuple[xr.DataArray, xr.DataArray]: + "array return" - if eval_kwargs.get("as_latex", False): - return self.as_latex(lhs, rhs) - else: - return xr.DataArray(lhs), self.op, xr.DataArray(rhs) + def _eval( + self, return_type: RETURN_T + ) -> Union[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 _compare_bitwise(self, where: bool, lhs: Any, rhs: Any) -> Any: + "Comparison function for application to individual elements of the array" + if not where or pd.isnull(lhs) or pd.isnull(rhs): + return np.nan + elif self.op == "==": + constraint = lhs == rhs + elif self.op == "<=": + constraint = lhs <= rhs + elif self.op == ">=": + constraint = lhs >= rhs + return constraint + + def as_math_string(self) -> str: + lhs, rhs = self._eval("math_string") + return lhs + self.OP_TRANSLATOR[self.op] + rhs + + def as_array(self) -> xr.DataArray: + lhs, rhs = self._eval("array") + return self.eval_attrs["backend_interface"]._apply_func( + self._compare_bitwise, self.eval_attrs["where_array"], lhs, rhs + ) -class EvalFunction(EvalString): +class EvalFunction(EvalToArrayStr): def __init__(self, tokens: pp.ParseResults) -> None: """ Parse action to process successfully parsed helper function strings of the form @@ -280,55 +397,66 @@ def __init__(self, tokens: pp.ParseResults) -> None: self.values = tokens def __repr__(self) -> str: - "Return string representation of the parsed grammar" - return f"{str(self.func_name)}(args={self.args}, kwargs={self.kwargs})" - - def arg_eval(self, arg, **eval_kwargs): + _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: + "string return" + + @overload + def _arg_eval( + self, return_type: Literal["array"], arg: Any + ) -> xr.DataArray | list[str, float]: + "array return" + + def _arg_eval( + self, return_type: RETURN_T, arg: Any + ) -> Union[str, xr.DataArray, list[str, float]]: + "Evaluate the arguments of the helper function" if isinstance(arg, pp.ParseResults): - evaluated_ = arg[0].eval(**eval_kwargs) + evaluated = arg[0].eval(return_type, **self.eval_attrs) elif isinstance(arg, list): - evaluated_ = [self.arg_eval(arg_) for arg_ in arg] + evaluated = [self._arg_eval(return_type, arg_) for arg_ in arg] else: - evaluated_ = arg.eval(**eval_kwargs) - return evaluated_ + evaluated = arg.eval(return_type, **self.eval_attrs) + if isinstance(evaluated, xr.DataArray) and isinstance(arg, EvalGenericString): + evaluated = evaluated.item() + return evaluated - def eval(self, **eval_kwargs) -> Any: - """ + @overload + def _eval(self, return_type: Literal["math_string"]) -> str: + "string return" - Args: - test (bool, optional): - If True, return a dictionary with parsed components rather than - calling the helper function with the defined args and kwargs. - Defaults to False. + @overload + def _eval(self, return_type: Literal["array"]) -> xr.DataArray: + "array return" - Returns: - Any: - Either the defined helper function is called, or only a dictionary with - parsed components is returned (if test=True). - """ - eval_kwargs["apply_where"] = False + def _eval(self, return_type: RETURN_T) -> Union[str, xr.DataArray]: + "Pass evaluated arguments to evaluated helper function" + self.eval_attrs["where_array"] = xr.DataArray(True) args_ = [] for arg in self.args: - args_.append(self.arg_eval(arg, **eval_kwargs)) + args_.append(self._arg_eval(return_type, arg)) kwargs_ = {} for kwarg_name, kwarg_val in self.kwargs.items(): - kwargs_[kwarg_name] = self.arg_eval(kwarg_val, **eval_kwargs) - - helper_function = self.func_name.eval(**eval_kwargs) - if eval_kwargs.get("as_dict"): - return { - "function": helper_function, - "args": args_, - "kwargs": kwargs_, - } - else: - eval_func = helper_function(*args_, **kwargs_) - return eval_func + kwargs_[kwarg_name] = self._arg_eval(return_type, kwarg_val) + helper_function = self.func_name.eval(return_type, **self.eval_attrs) -class EvalHelperFuncName(EvalString): + evaluated = helper_function(*args_, **kwargs_) + return evaluated + + def as_math_string(self) -> str: + return self._eval("math_string") + + def as_array(self) -> xr.DataArray: + return self._eval("array") + + +class EvalHelperFuncName(EvalToCallable): def __init__(self, instring: str, loc: int, tokens: pp.ParseResults) -> None: """ Parse action to process successfully parsed helper function names. @@ -349,49 +477,25 @@ def __init__(self, instring: str, loc: int, tokens: pp.ParseResults) -> None: self.values = tokens def __repr__(self) -> str: - "Return string representation of the parsed grammar" return str(self.name) - def eval( - self, - helper_functions: dict[str, type[ParsingHelperFunction]], - as_dict: bool = False, - **eval_kwargs, - ) -> Optional[Union[str, Callable]]: - """ - - Args: - helper_functions (dict[str, type[ParsingHelperFunction]]): Allowed helper functions. - test (bool, optional): - If True, return a string with the helper function name rather than - collecting the helper function from the dictionary of functions. - Defaults to False. - - Returns: - str, Callable: - Helper functions are expected to be two-tiered, with the first level - taking the generic eval kwargs (e.g. model_data) and the second level - taking the user-defined input arguments. - If test=True, only the helper function name is returned. - """ - + def as_callable(self, return_type: RETURN_T) -> Callable: + helper_functions = self.eval_attrs["helper_functions"] + equation_name = self.eval_attrs["equation_name"] if self.name not in helper_functions.keys(): raise BackendError( - f"({eval_kwargs['equation_name']}, {self.instring}): Invalid helper function defined: {self.name}" + f"({equation_name}, {self.instring}): Invalid helper function defined: {self.name}" ) elif not isinstance(helper_functions[self.name], type(ParsingHelperFunction)): raise TypeError( - f"({eval_kwargs['equation_name']}, {self.instring}): Helper function must be " + f"({equation_name}, {self.instring}): Helper function must be " f"subclassed from calliope.backend.helper_functions.ParsingHelperFunction: {self.name}" ) else: - if as_dict: - return str(self.name) - else: - return helper_functions[self.name](**eval_kwargs) + return helper_functions[self.name](return_type, **self.eval_attrs) -class EvalSlicedParameterOrVariable(EvalString): +class EvalSlicedComponent(EvalToArrayStr): def __init__(self, tokens: pp.ParseResults) -> None: """ Parse action to process successfully parsed sliced parameters or decision variables @@ -411,13 +515,17 @@ def __init__(self, tokens: pp.ParseResults) -> None: self.values = tokens def __repr__(self) -> str: - "Return string representation of the parsed grammar" 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): - def _replace(term): + 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: @@ -429,11 +537,28 @@ def _replace(term): + term[3] ) - return _replace + return __replace + + @overload + def _eval(self, return_type: Literal["math_string"]) -> tuple[str, dict]: + "string return" - def as_latex(self, evaluated_obj: str, index_slices: dict[str, str]) -> str: - """Stingify evaluated dataarray for use in a LaTex math formula""" - singular_slice_refs = {k.removesuffix("s"): v for k, v in index_slices.items()} + @overload + def _eval(self, return_type: Literal["array"]) -> tuple[xr.DataArray, dict]: + "array return" + + def _eval(self, return_type: RETURN_T) -> tuple[Union[str, xr.DataArray], dict]: + "Evaluate the slice dim and vals of each slice element" + slices: dict[str, Any] = { + k: v.eval(return_type, **self.eval_attrs) 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: + evaluated, slices = self._eval("math_string") + singular_slice_refs = {k.removesuffix("s"): v for k, v in slices.items()} id_ = pp.Combine( pp.Word(pp.alphas, pp.alphanums) + pp.ZeroOrMore("_" + pp.Word(pp.alphanums)) @@ -443,34 +568,15 @@ def as_latex(self, evaluated_obj: str, index_slices: dict[str, str]) -> str: obj_parser = id_formatted + pp.Opt( r"_\text{" + pp.Group(pp.delimited_list(id_)) + "}" ) - obj_parser.set_parse_action(self.replace_rule(singular_slice_refs)) - return obj_parser.parse_string(evaluated_obj, parse_all=True)[0] - - def eval(self, **eval_kwargs) -> Optional[Union[str, dict, xr.DataArray]]: - """ - Returns: - Optional[Union[dict, xr.DataArray]]: - If `eval_kwargs["as_dict"]` is True, returns separated key:val pairs for parameter/variable name and index items; - else, `eval_kwargs` has a backend dataset, returns sliced xarray object; - else, returns None. - """ - slices: dict[str, Any] = { - k: v.eval(**eval_kwargs) for k, v in self.slices.items() - } + obj_parser.set_parse_action(self._replace_rule(singular_slice_refs)) + return obj_parser.parse_string(evaluated, parse_all=True)[0] - if eval_kwargs.get("as_dict", False): - return {"dimensions": slices, **self.obj_name.eval(**eval_kwargs)} - elif eval_kwargs.get("backend_dataset", None) is not None: - evaluated_obj = self.obj_name.eval(**eval_kwargs) - if eval_kwargs.get("as_latex", False): - return self.as_latex(evaluated_obj, slices) - else: - return evaluated_obj.sel(**slices) - else: - return None + def as_array(self) -> xr.DataArray: + evaluated, slices = self._eval("array") + return evaluated.sel(**slices) -class EvalIndexSlice(EvalString): +class EvalIndexSlice(EvalToArrayStr): def __init__(self, tokens: pp.ParseResults) -> None: """ Parse action to process successfully parsed expression index slice references @@ -484,34 +590,38 @@ def __init__(self, tokens: pp.ParseResults) -> None: self.values = tokens def __repr__(self) -> str: - "Return string representation of the parsed grammar" return "REFERENCE:" + str(self.name) - def eval( - self, - slice_dict: Optional[dict[str, pp.ParseResults]] = None, - **eval_kwargs, - ) -> Any: - """ - Args: - slice_dict (Optional[dict[str, pp.ParseResults]]): - Dictionary mapping the index slice name to a parsed equation expression. - Default is None. + @overload + def _eval(self, return_type: Literal["math_string"], as_values: bool) -> str: + "string return" - Returns: - Any: If eval_kwargs["as_dict"] is True, returns a dictionary, - otherwise attempts to evaluate the referenced index slice. - """ - if eval_kwargs.get("as_dict"): - return {"slice_reference": self.name} - elif slice_dict is not None: - slicer = slice_dict[self.name][0].eval(as_values=True, **eval_kwargs) - if isinstance(slicer, xr.DataArray) and slicer.isnull().any(): - slicer = slicer.notnull() - return slicer + @overload + def _eval( + self, return_type: Literal["array"], as_values: bool + ) -> xr.DataArray | list[str, float]: + "array return" + def _eval( + self, return_type: RETURN_T, as_values: bool + ) -> Union[str, xr.DataArray, list[str, float]]: + "Evaluate the referenced `slice`." + self.eval_attrs["as_values"] = as_values + return self.eval_attrs["slice_dict"][self.name][0].eval( + return_type, **self.eval_attrs + ) -class EvalSubExpressions(EvalString): + def as_math_string(self) -> str: + return self._eval("math_string", False) + + def as_array(self) -> xr.DataArray | list[str, float]: + evaluated = self._eval("array", True) + if isinstance(evaluated, xr.DataArray) and evaluated.isnull().any(): + evaluated = evaluated.notnull() + return evaluated + + +class EvalSubExpressions(EvalToArrayStr): def __init__(self, tokens: pp.ParseResults) -> None: """ Parse action to process successfully parsed sub-expressions of the form @@ -525,106 +635,30 @@ def __init__(self, tokens: pp.ParseResults) -> None: self.values = tokens def __repr__(self) -> str: - "Return string representation of the parsed grammar" return "SUB_EXPRESSION:" + str(self.name) - def eval( - self, - sub_expression_dict: Optional[dict[str, pp.ParseResults]] = None, - **eval_kwargs, - ) -> Any: - """ - Args: - sub_expression_dict (Optional[dict[str, pp.ParseResults]]): - Dictionary mapping the sub-expression name to a parsed equation expression. - Default is None. - - Returns: - Any: If sub-expression dictionary is given, find the expression matching - the sub-expression name and evaluate it. - If not given, return a dictionary giving the sub-expression name. - """ - if eval_kwargs.get("as_dict"): - return {"sub_expression": self.name} - elif sub_expression_dict is not None: - return sub_expression_dict[self.name][0].eval(**eval_kwargs) - - -class EvalUnslicedParameterOrVariable(EvalString): - def __init__(self, tokens: pp.ParseResults) -> None: - """ - Parse action to process successfully parsed unsliced parameters or decision variables - of the form `param_or_var`. - - Args: - tokens (pp.ParseResults): - Has one parsed element containing the paramater/variable name (str). - """ - self.name: str = tokens[0] - self.values = tokens - - def __repr__(self) -> str: - "Return string representation of the parsed grammar" - return "PARAM_OR_VAR:" + str(self.name) - - def as_latex(self, evaluated: Optional[xr.DataArray] = None) -> str: - """Stingify evaluated dataarray for use in a LaTex math formula""" - if evaluated is None: - return rf"\text{{{self.name}}}" + @overload + def _eval(self, return_type: Literal["math_string"]) -> str: + "string return" - if evaluated.shape: - dims = rf"_\text{{{','.join(str(i).removesuffix('s') for i in evaluated.dims)}}}" - else: - dims = "" - if evaluated.attrs["obj_type"] in ["global_expressions", "variables"]: - formatted_name = rf"\textbf{{{self.name}}}" - elif evaluated.attrs["obj_type"] == "parameters": - formatted_name = rf"\textit{{{self.name}}}" - return formatted_name + dims + @overload + def _eval(self, return_type: Literal["array"]) -> xr.DataArray: + "array return" - def eval( - self, - references: set, - as_dict: bool = False, - as_values: bool = False, - backend_dataset: Optional[xr.Dataset] = None, - backend_interface: Optional[BackendModel] = None, - **eval_kwargs, - ) -> Optional[Union[dict, xr.DataArray, str]]: - """ - Args: - references (set): - as_dict (bool, optional): - as_values(bool, optional): - If True, return values rather than backend objects - Returns: - Optional[Union[dict, xr.DataArray]]: - If `eval_kwargs["as_dict"]` is True, returns separated key:val pairs for parameter/variable name and index items; - else, `eval_kwargs` has a backend dataset, returns sliced xarray object; - else, returns None. - """ - references.add(self.name) - evaluated: Optional[Union[dict, xr.DataArray, str]] - as_latex = eval_kwargs.get("as_latex", False) - if as_dict: - evaluated = {"param_or_var_name": self.name} - elif backend_interface is not None and backend_dataset is not None: - if as_values and not as_latex: - evaluated = backend_interface.get_parameter( - self.name, as_backend_objs=False - ) - else: - evaluated = backend_dataset[self.name] + def _eval(self, return_type: RETURN_T) -> Union[str, xr.DataArray]: + "Evaluate the referenced sub_expression" + return self.eval_attrs["sub_expression_dict"][self.name][0].eval( + return_type, **self.eval_attrs + ) - if as_latex: - evaluated = self.as_latex(evaluated) - else: - evaluated = None + def as_math_string(self) -> str: + return self._eval("math_string") - return evaluated + def as_array(self) -> xr.DataArray: + return self._eval("array") -class EvalNumber(EvalString): +class EvalNumber(EvalToArrayStr): def __init__(self, tokens: pp.ParseResults) -> None: """ Parse action to process successfully parsed numbers described as integers (1), @@ -638,31 +672,20 @@ def __init__(self, tokens: pp.ParseResults) -> None: self.values = tokens def __repr__(self) -> str: - "Return string representation of the parsed grammar" return "NUM:" + str(self.value) - def as_latex(self, evaluated): - """Stingify evaluated float to 6 significant figures for use in a LaTex math formula""" + def as_math_string(self) -> str: return re.sub( r"([\d]+?)e([+-])([\d]+)", r"\1\\mathord{\\times}10^{\2\3}", - f"{evaluated:.6g}", + f"{float(self.value):.6g}", ) - def eval(self, **eval_kwargs) -> float: - """ - Returns: - float: Input string as a float, even if given as an integer. - """ + def as_array(self) -> xr.DataArray: + return xr.DataArray(float(self.value), attrs={"obj_type": "number"}) - evaluated = float(self.value) - if eval_kwargs.get("as_latex", False): - return self.as_latex(evaluated) - else: - return evaluated - -class StringListParser(EvalString): +class ListParser(EvalToArrayStr): def __init__(self, tokens: pp.ParseResults) -> None: """ Parse action to process successfully parsed lists of generic strings. @@ -675,23 +698,65 @@ def __init__(self, tokens: pp.ParseResults) -> None: self.val = tokens def __repr__(self) -> str: - "Return string representation of the parsed grammar" return f"{self.val}" - def as_latex(self, evaluated): - """Stingify evaluated object for use in a LaTex math formula""" - return evaluated + def as_math_string(self) -> str: + input_list = self.as_array() + return "[" + ",".join(str(i) for i in input_list) + "]" + + def as_array(self) -> list[str, float]: + 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 [da.item() if da.size == 1 else da.name for da in values] + + +class EvalUnslicedComponent(EvalToArrayStr): + def __init__(self, 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: + tokens (pp.ParseResults): Has one parsed element: string name (str). + """ + self.val = tokens[0] + self.name = str(self.val) + self.values = tokens + + def __repr__(self) -> str: + return f"COMPONENT:{self.name}" + + def as_math_string(self) -> str: + self.eval_attrs["as_values"] = False + evaluated = self.as_array() + + if evaluated.shape: + dims = rf"_\text{{{','.join(str(i).removesuffix('s') for i in evaluated.dims)}}}" + else: + dims = "" + if evaluated.attrs["obj_type"] in ["global_expressions", "variables"]: + formatted_name = rf"\textbf{{{self.name}}}" + elif evaluated.attrs["obj_type"] == "parameters": + formatted_name = rf"\textit{{{self.name}}}" + return formatted_name + dims + + def as_array(self) -> xr.DataArray: + backend_interface = self.eval_attrs["backend_interface"] - def eval(self, **eval_kwargs) -> list[str]: - "Return input as list of strings." - evaluated = [val.eval() for val in self.val] - if eval_kwargs.get("as_latex", False): - return self.as_latex(evaluated) + if self.eval_attrs.get("as_values", False): + evaluated = backend_interface.get_parameter( + self.name, as_backend_objs=False + ) else: - return evaluated + evaluated = backend_interface._dataset[self.name] + + self.eval_attrs["references"].add(self.name) + return evaluated -class GenericStringParser(EvalString): +class EvalGenericString(EvalToArrayStr): def __init__(self, tokens: pp.ParseResults) -> None: """ Parse action to process successfully parsed generic strings. @@ -704,20 +769,13 @@ def __init__(self, tokens: pp.ParseResults) -> None: self.val = tokens[0] def __repr__(self) -> str: - "Return string representation of the parsed grammar" return f"STRING:{self.val}" - def as_latex(self, evaluated): - """Stingify evaluated string for use in a LaTex math formula""" - return evaluated + def as_math_string(self): + return str(self.val) - def eval(self, **eval_kwargs) -> str: - "Return input as string." - evaluated = str(self.val) - if eval_kwargs.get("as_latex", False): - return self.as_latex(evaluated) - else: - return evaluated + def as_array(self) -> xr.DataArray: + return xr.DataArray(str(self.val), attrs={"obj_type": "string"}) def helper_function_parser( @@ -776,8 +834,8 @@ def helper_function_parser( # define function keyword arguments key = generic_identifier + pp.Suppress("=") - kwarglist = pp.delimited_list(pp.dict_of(key, arg_values)) - kwargs_ = pp.Group(kwarglist).set_results_name("kwargs") + kwarg_list = pp.delimited_list(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( @@ -849,14 +907,14 @@ def sliced_param_or_var_parser( 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(EvalSlicedParameterOrVariable) + 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 preppended with the YAML constraint sub-expression classifier `$`. E.g. "$my_sub_expr" + Parse strings prepended with the YAML constraint sub-expression classifier `$`. E.g. "$my_sub_expr" Args: generic_identifier (pp.ParserElement): @@ -876,13 +934,13 @@ def sub_expression_parser(generic_identifier: pp.ParserElement) -> pp.ParserElem return sub_expression -def unsliced_object_parser(valid_math_element_names: Iterable[str]) -> pp.ParserElement: +def unsliced_object_parser(valid_component_names: Iterable[str]) -> pp.ParserElement: """ Create a copy of the generic identifier and set a parse action to find the string in the list of input paramaters or optimisation decision variables. Args: - valid_math_element_names (Iterable[str]): A + valid_component_names (Iterable[str]): A All backend object names, to ensure they are captured by this parser function. Returns: @@ -891,43 +949,58 @@ def unsliced_object_parser(valid_math_element_names: Iterable[str]) -> pp.Parser parameter/variable value """ - unsliced_param_or_var = pp.one_of(valid_math_element_names, as_keyword=True) - unsliced_param_or_var.set_parse_action(EvalUnslicedParameterOrVariable) + 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_math_element_names: Iterable -) -> tuple[pp.ParserElement, pp.ParserElement]: + identifier: pp.ParserElement, valid_component_names: Iterable +) -> pp.ParserElement: """ - Create an evaluatable copy of the generic identifier that will return a string or a - list of strings. + 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_math_element_names (Iterable[str]): A + valid_component_names (Iterable[str]): A All backend object names, to ensure they are *not* captured by this parser function. Returns: - evaluatable_identifier (pp.ParserElement): + pp.ParserElement: Parser for valid python variables without leading underscore and not called "inf". - Evaluates to a string. - id_list (pp.ParserElement): - Parser for lists of "evaluatable_identifier", bound by "[]" parentheses + Evaluates to a string or an array (if it is a model component). """ evaluatable_identifier = ( - ~pp.one_of(valid_math_element_names, as_keyword=True) + identifier - ).set_parse_action(GenericStringParser) + ~pp.one_of(valid_component_names, as_keyword=True) + identifier + ).set_parse_action(EvalGenericString) - id_list = ( - pp.Suppress("[") + pp.delimited_list(evaluatable_identifier) + pp.Suppress("]") - ) - id_list.set_parse_action(StringListParser) + return evaluatable_identifier + + +def list_parser( + number: pp.ParserElement, evaluatable_identifier: pp.ParserElement +) -> pp.ParserElement: + """ + Parse strings which define a list of other strings or numbers. + + Lists are defined as anything wrapped in square brackets (`[]`). - return evaluatable_identifier, id_list + Strings that could be evaluated as model component arrays will still be evaluated as strings (using the model component name). + + Args: + number (pp.ParserElement): Parser for numbers (integer, float, scientific notation, "inf"/".inf"). + evaluatable_identifier (pp.ParserElement): Parser for valid python variables without leading underscore and not called "inf". + + Returns: + pp.ParserElement: Parser for valid lists of strings and/or numbers. + """ + list_elements = pp.MatchFirst([evaluatable_identifier, number]) + id_list = pp.Suppress("[") + pp.delimited_list(list_elements) + pp.Suppress("]") + id_list.set_parse_action(ListParser) + return id_list def setup_base_parser_elements() -> tuple[pp.ParserElement, pp.ParserElement]: @@ -1023,7 +1096,7 @@ def equation_comparison_parser(arithmetic: pp.ParserElement) -> pp.ParserElement return equation_comparison -def generate_slice_parser(valid_math_element_names: Iterable) -> pp.ParserElement: +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]`. @@ -1031,7 +1104,7 @@ def generate_slice_parser(valid_math_element_names: Iterable) -> pp.ParserElemen nor references to sub expressions. Args: - valid_math_element_names (Iterable): + valid_component_names (Iterable): Allowed names for optimisation problem components (parameters, decision variables, expressions), to allow the parser to separate these from generic strings. @@ -1039,10 +1112,11 @@ def generate_slice_parser(valid_math_element_names: Iterable) -> pp.ParserElemen pp.ParserElement: Parser for expression strings under the constraint key "slices". """ number, identifier = setup_base_parser_elements() - evaluatable_identifier, id_list = evaluatable_identifier_parser( - identifier, valid_math_element_names + evaluatable_identifier = evaluatable_identifier_parser( + identifier, valid_component_names ) - unsliced_param = unsliced_object_parser(valid_math_element_names) + id_list = list_parser(number, evaluatable_identifier) + unsliced_param = unsliced_object_parser(valid_component_names) sliced_param = sliced_param_or_var_parser( number, identifier, @@ -1071,7 +1145,7 @@ def generate_slice_parser(valid_math_element_names: Iterable) -> pp.ParserElemen ) -def generate_sub_expression_parser(valid_math_element_names: Iterable) -> pp.Forward: +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`. @@ -1079,7 +1153,7 @@ def generate_sub_expression_parser(valid_math_element_names: Iterable) -> pp.For and reference to index slice expressions. Args: - valid_math_element_names (Iterable): + valid_component_names (Iterable): Allowed names for optimisation problem components (parameters, decision variables, expressions), to allow the parser to separate these from generic strings. @@ -1087,10 +1161,11 @@ def generate_sub_expression_parser(valid_math_element_names: Iterable) -> pp.For pp.ParserElement: Parser for expression strings under the constraint key "sub_expressions". """ number, identifier = setup_base_parser_elements() - evaluatable_identifier, id_list = evaluatable_identifier_parser( - identifier, valid_math_element_names + evaluatable_identifier = evaluatable_identifier_parser( + identifier, valid_component_names ) - unsliced_param = unsliced_object_parser(valid_math_element_names) + id_list = list_parser(number, evaluatable_identifier) + unsliced_param = unsliced_object_parser(valid_component_names) sliced_param = sliced_param_or_var_parser( number, identifier, evaluatable_identifier, unsliced_param ) @@ -1112,14 +1187,14 @@ def generate_sub_expression_parser(valid_math_element_names: Iterable) -> pp.For return arithmetic -def generate_arithmetic_parser(valid_math_element_names: Iterable) -> pp.ParserElement: +def generate_arithmetic_parser(valid_component_names: Iterable) -> pp.ParserElement: """ Create parser for left-/right-hand side (LHS/RHS) of 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_math_element_names (Iterable): + 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. @@ -1127,10 +1202,11 @@ def generate_arithmetic_parser(valid_math_element_names: Iterable) -> pp.ParserE pp.ParserElement: Partial parser for expression strings under the constraint key "equation/equations". """ number, identifier = setup_base_parser_elements() - evaluatable_identifier, id_list = evaluatable_identifier_parser( - identifier, valid_math_element_names + evaluatable_identifier = evaluatable_identifier_parser( + identifier, valid_component_names ) - unsliced_param = unsliced_object_parser(valid_math_element_names) + id_list = list_parser(number, evaluatable_identifier) + unsliced_param = unsliced_object_parser(valid_component_names) sliced_param = sliced_param_or_var_parser( number, identifier, evaluatable_identifier, unsliced_param ) @@ -1152,14 +1228,14 @@ def generate_arithmetic_parser(valid_math_element_names: Iterable) -> pp.ParserE return arithmetic -def generate_equation_parser(valid_math_element_names: Iterable) -> pp.ParserElement: +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_math_element_names (Iterable): + 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. @@ -1167,7 +1243,7 @@ def generate_equation_parser(valid_math_element_names: Iterable) -> pp.ParserEle pp.ParserElement: Parser for expression strings under the constraint key "equation/equations". """ - arithmetic = generate_arithmetic_parser(valid_math_element_names) + arithmetic = generate_arithmetic_parser(valid_component_names) equation_comparison = equation_comparison_parser(arithmetic) return equation_comparison diff --git a/src/calliope/backend/helper_functions.py b/src/calliope/backend/helper_functions.py index 74bff5523..8b3ad6d6b 100644 --- a/src/calliope/backend/helper_functions.py +++ b/src/calliope/backend/helper_functions.py @@ -10,11 +10,12 @@ import functools import re from abc import ABC, abstractmethod -from typing import Any, Literal, Mapping, Union, overload +from typing import Any, Literal, Mapping, Optional, Union, overload import numpy as np import xarray as xr +from calliope.backend.backend_model import BackendModel from calliope.exceptions import BackendError _registry: dict[ @@ -25,20 +26,22 @@ class ParsingHelperFunction(ABC): def __init__( self, - as_latex: bool = False, + return_type: Literal["array", "math_string"], + *, + equation_name: str, + input_data: xr.Dataset, + backend_interface: Optional[type[BackendModel]] = None, **kwargs, ) -> 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. - Args: - as_latex (bool, optional): - If True, will return a LaTeX math string on calling the class. - Defaults to False. """ - self._kwargs = kwargs - self._as_latex = as_latex + self._equation_name = equation_name + self._input_data = input_data + self._backend_interface = backend_interface + self._return_type = return_type @property @abstractmethod @@ -51,27 +54,27 @@ def NAME(self) -> str: "Helper function name that is used in the math expression/where string." @abstractmethod - def as_latex(self, *args, **kwargs) -> str: + def as_math_string(self, *args, **kwargs) -> 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 ``as_latex=True``. + This method is called when the class is initialised with ``return_type=math_string``. """ @abstractmethod def as_array(self, *args, **kwargs) -> 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 ``as_latex=False``. + This method is called when the class is initialised with ``return_type=array``. """ def __call__(self, *args, **kwargs) -> Any: """ When a helper function is accessed by evaluating a parsing string, this method is called. - The value of `as_latex` on initialisation of the class defines whether this method returns a string (``as_latex=True``) or :meth:xr.DataArray (``as_latex=False``) + The value of `return_type` on initialisation of the class defines whether this method returns a string (``return_type=math_string``) or :meth:xr.DataArray (``return_type=array``) """ - if self._as_latex: - return self.as_latex(*args, **kwargs) - else: + 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): @@ -170,7 +173,7 @@ def _expand_link_techs(self, vals: list[str]) -> list[str]: to_remove = [] to_add = [] for val in vals: - link_techs = self._kwargs["model_data"].techs.str.startswith(val + ":") + link_techs = self._input_data.techs.str.startswith(val + ":") if link_techs.any(): to_add.extend(list(link_techs[link_techs].techs.data)) to_remove.append(val) @@ -186,7 +189,7 @@ class Inheritance(ParsingHelperFunction): #: NAME = "inheritance" - def as_latex(self, tech_group: str) -> str: + def as_math_string(self, tech_group: str) -> str: return rf"\text{{tech_group={tech_group}}}" def as_array(self, tech_group: str) -> xr.DataArray: @@ -197,9 +200,7 @@ def as_array(self, tech_group: str) -> xr.DataArray: Args: model_data (xr.Dataset): Calliope model data """ - inheritance_lists = ( - self._kwargs["model_data"].inheritance.to_series().str.split(".") - ) + inheritance_lists = self._input_data.inheritance.to_series().str.split(".") return inheritance_lists.apply(lambda x: tech_group in x).to_xarray() @@ -210,7 +211,7 @@ class WhereAny(ParsingHelperFunction): #: ALLOWED_IN = ["where"] - def as_latex(self, array: str, *, over: Union[str, list[str]]) -> str: + def as_math_string(self, array: str, *, over: Union[str, list[str]]) -> str: if isinstance(over, str): overstring = self._instr(over) else: @@ -231,15 +232,18 @@ def as_array(self, parameter: str, *, over: Union[str, list[str]]) -> xr.DataArr If the parameter 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 parameter does not exist, returns a dimensionless False array. """ - if parameter in self._kwargs["model_data"].data_vars: - parameter_da = self._kwargs["model_data"][parameter] + if parameter in self._input_data.data_vars: + parameter_da = self._input_data[parameter] bool_parameter_da = ( parameter_da.notnull() & (parameter_da != np.inf) & (parameter_da != -np.inf) ) - elif parameter in self._kwargs.get("backend_dataset", xr.Dataset()).data_vars: - bool_parameter_da = self._kwargs["backend_dataset"][parameter].notnull() + elif ( + self._backend_interface is not None + and parameter in self._backend_interface._dataset + ): + bool_parameter_da = self._backend_interface._dataset[parameter].notnull() else: bool_parameter_da = xr.DataArray(False) over = self._listify(over) @@ -254,7 +258,7 @@ class Defined(ParsingHelperFunction): #: ALLOWED_IN = ["where"] - def as_latex(self, *, within: str, how: Literal["all", "any"], **dims) -> str: + def as_math_string(self, *, within: str, how: Literal["all", "any"], **dims) -> str: substrings = [] for name, vals in dims.items(): substrings.append(self._latex_substring(how, name, vals, within)) @@ -324,7 +328,7 @@ def as_array( else self._listify(vals, expand_link_techs=False) for dim, vals in dims.items() } - definition_matrix = self._kwargs["model_data"].definition_matrix + definition_matrix = self._input_data.definition_matrix dim_within_da = definition_matrix.any(self._dims_to_remove(dim_names, within)) within_da = getattr(dim_within_da.sel(**dims_with_list_vals), how)(dim_names) @@ -345,7 +349,7 @@ def _dims_to_remove(self, dim_names: list[str], within: str) -> set: Returns: set: Undefined dimensions to remove from the definition matrix. """ - definition_matrix = self._kwargs["model_data"].definition_matrix + definition_matrix = self._input_data.definition_matrix missing_dims = set([*dim_names, within]).difference(definition_matrix.dims) if missing_dims: raise ValueError( @@ -383,7 +387,7 @@ class Sum(ParsingHelperFunction): #: ALLOWED_IN = ["expression"] - def as_latex(self, array: str, *, over: Union[str, list[str]]) -> str: + def as_math_string(self, array: str, *, over: Union[str, list[str]]) -> str: if isinstance(over, str): overstring = self._instr(over) else: @@ -415,7 +419,7 @@ class ReduceCarrierDim(ParsingHelperFunction): #: ALLOWED_IN = ["expression"] - def as_latex( + def as_math_string( self, array: str, carrier_tier: Literal["in", "out", "in_2", "out_2", "in_3", "out_3"], @@ -436,11 +440,14 @@ def as_array( Returns: xr.DataArray: `array` reduced by the `carriers` dimension. """ - return Sum(as_latex=self._as_latex, **self._kwargs)( + sum_helper = Sum( + return_type=self._return_type, + equation_name=self._equation_name, + input_data=self._input_data, + ) + return sum_helper( array.where( - self._kwargs["model_data"].definition_matrix.sel( - carrier_tiers=carrier_tier - ) + self._input_data.definition_matrix.sel(carrier_tiers=carrier_tier) ), over="carriers", ) @@ -452,7 +459,7 @@ class ReducePrimaryCarrierDim(ParsingHelperFunction): #: ALLOWED_IN = ["expression"] - def as_latex(self, array: str, carrier_tier: Literal["in", "out"]) -> str: + def as_math_string(self, array: str, carrier_tier: Literal["in", "out"]) -> str: return rf"\sum\limits_{{\text{{carrier=primary_carrier_{carrier_tier}}}}} ({array})" def as_array( @@ -471,11 +478,14 @@ def as_array( Returns: xr.DataArray: `array` reduced by the `carriers` dimension. """ - return Sum(as_latex=self._as_latex, **self._kwargs)( + sum_helper = Sum( + return_type=self._return_type, + equation_name=self._equation_name, + input_data=self._input_data, + ) + return sum_helper( array.where( - getattr( - self._kwargs["model_data"], f"primary_carrier_{carrier_tier}" - ).notnull() + getattr(self._input_data, f"primary_carrier_{carrier_tier}").notnull() ), over="carriers", ) @@ -487,7 +497,7 @@ class SelectFromLookupArrays(ParsingHelperFunction): #: ALLOWED_IN = ["expression"] - def as_latex(self, array: str, **lookup_arrays: str) -> str: + def as_math_string(self, array: str, **lookup_arrays: str) -> str: new_strings = { (iterator := dim.removesuffix("s")): rf"={array}[{iterator}]" for dim, array in lookup_arrays.items() @@ -556,12 +566,10 @@ def as_array( # 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._kwargs["model_data"][index.name].stack({dim: dims}) + stacked_lookup = self._input_data[index.name].stack({dim: dims}) ix = array.indexes[index_dim].get_indexer(stacked_lookup) if (ix == -1).all(): - received_lookup = ( - self._kwargs["model_data"][index.name].to_series().dropna() - ) + received_lookup = self._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}" ) @@ -588,7 +596,7 @@ class GetValAtIndex(ParsingHelperFunction): #: ALLOWED_IN = ["expression", "where"] - def as_latex(self, **dim_idx_mapping: str) -> str: + def as_math_string(self, **dim_idx_mapping: str) -> str: dim, idx = self._mapping_to_dim_idx(**dim_idx_mapping) return f"{dim}[{idx}]" @@ -623,7 +631,7 @@ def as_array(self, **dim_idx_mapping: int) -> xr.DataArray: timesteps tuple[str, int]: @overload @staticmethod def _mapping_to_dim_idx(**dim_idx_mapping: str) -> tuple[str, str]: - "used in as_latex" + "used in as_math_string" @staticmethod def _mapping_to_dim_idx(**dim_idx_mapping) -> tuple[str, Union[str, int]]: @@ -648,7 +656,7 @@ class Roll(ParsingHelperFunction): #: ALLOWED_IN = ["expression"] - def as_latex(self, array: str, **roll_kwargs: str) -> str: + def as_math_string(self, array: str, **roll_kwargs: str) -> str: new_strings = { k.removesuffix("s"): f"{-1 * int(v):+d}" for k, v in roll_kwargs.items() } @@ -688,7 +696,7 @@ class GetTransmissionTechs(ParsingHelperFunction): #: ALLOWED_IN = ["expression"] - def as_latex(self, vals: Union[str, list[str]]) -> str: + def as_math_string(self, vals: Union[str, list[str]]) -> str: expanded_vals = self._listify(vals, expand_link_techs=True) return f"techs=[{','.join(expanded_vals)}]" diff --git a/src/calliope/backend/latex_backend_model.py b/src/calliope/backend/latex_backend_model.py index 36b109550..41827844a 100644 --- a/src/calliope/backend/latex_backend_model.py +++ b/src/calliope/backend/latex_backend_model.py @@ -287,7 +287,6 @@ def add_parameter( use_inf_as_na: bool = False, ) -> None: self._add_to_dataset(parameter_name, parameter_values, "parameters", {}) - self.valid_math_element_names.add(parameter_name) def add_constraint( self, @@ -321,8 +320,6 @@ def add_global_expression( name: str, expression_dict: Optional[parsing.UnparsedExpressionDict] = None, ) -> None: - self.valid_math_element_names.add(name) - equation_strings: list = [] def _expression_setter( @@ -352,8 +349,6 @@ def add_variable( ) -> None: domain_dict = {"real": r"\mathbb{R}\;", "integer": r"\mathbb{Z}\;"} - self.valid_math_element_names.add(name) - def _variable_setter(where: xr.DataArray) -> xr.DataArray: return where.where(where) @@ -454,8 +449,8 @@ def generate_math_doc(self, format: _ALLOWED_MATH_FILE_FORMATS = "tex") -> str: return self._render(doc_template, components=components) def _add_latex_strings(self, where, element, equation_strings): - expr = element.evaluate_expression(self.inputs, self, as_latex=True) - where_latex = element.evaluate_where(self.inputs, self._dataset, as_latex=True) + expr = element.evaluate_expression(self, return_type="math_string") + where_latex = element.evaluate_where(self, return_type="math_string") if self.include == "all" or (self.include == "valid" and where.any()): equation_strings.append({"expression": expr, "where": where_latex}) @@ -470,9 +465,7 @@ def _generate_math_string( sets: Optional[list[str]] = None, ) -> None: if parsed_component is not None: - where = parsed_component.evaluate_where( - self.inputs, self._dataset, as_latex=True - ) + where = parsed_component.evaluate_where(self, return_type="math_string") sets = parsed_component.sets if self.include == "all" or ( @@ -506,10 +499,8 @@ def _get_capacity_bounds( ], } parsed_bounds = parsing.ParsedBackendComponent("constraints", name, bound_dict) - equations = parsed_bounds.parse_equations( - self.valid_math_element_names, - ) + equations = parsed_bounds.parse_equations(self.valid_component_names) return tuple( - {"expression": eq.evaluate_expression(self.inputs, self, as_latex=True)} + {"expression": eq.evaluate_expression(self, return_type="math_string")} for eq in equations ) diff --git a/src/calliope/backend/parsing.py b/src/calliope/backend/parsing.py index 11eb2063d..c6da9efe1 100644 --- a/src/calliope/backend/parsing.py +++ b/src/calliope/backend/parsing.py @@ -9,7 +9,6 @@ import operator from typing import ( TYPE_CHECKING, - Any, Callable, Iterable, Literal, @@ -156,7 +155,7 @@ def find_slices(self) -> set[str]: [ expression_parser.EvalOperatorOperand, expression_parser.EvalFunction, - expression_parser.EvalSlicedParameterOrVariable, + expression_parser.EvalSlicedComponent, ] ) to_find = expression_parser.EvalIndexSlice @@ -247,9 +246,9 @@ def add_expression_group_combination( @overload # noqa: F811 def evaluate_where( # noqa: F811 self, - model_data: xr.Dataset, - backend_dataset: Optional[xr.Dataset], - as_latex: Literal[False] = False, + backend_interface: backend_model.BackendModel, + *, + return_type: Literal["array"] = "array", initial_where: xr.DataArray = TRUE_ARRAY, ) -> xr.DataArray: "Expecting array if not requesting latex string" @@ -257,42 +256,49 @@ def evaluate_where( # noqa: F811 @overload # noqa: F811 def evaluate_where( # noqa: F811 self, - model_data: xr.Dataset, - backend_dataset: Optional[xr.Dataset], - as_latex: Literal[True], + backend_interface: backend_model.BackendModel, + *, + return_type: Literal["math_string"], ) -> str: "Expecting string if requesting latex string" def evaluate_where( # noqa: F811 self, - model_data: xr.Dataset, - backend_dataset: Optional[xr.Dataset] = None, - as_latex: bool = False, + backend_interface: backend_model.BackendModel, + *, + return_type: str = "array", initial_where: xr.DataArray = TRUE_ARRAY, ) -> Union[xr.DataArray, str]: """Evaluate parsed backend object dictionary `where` string. Args: - model_data (xr.Dataset): Calliope model dataset. + backend_interface (calliope.backend.backend_model.BackendModel): Interface to a optimisation backend. + + Keyword Args: + return_type (str, optional): + If "array", return xarray.DataArray. If "math_string", return LaTex math string. + Defaults to "array". initial_where (xr.DataArray, optional): If given, the where array resulting from evaluation will be further where'd by this array. Defaults to xr.DataArray(True) (i.e., no effect). Returns: Union[xr.DataArray, str]: - If `as_latex` is False: Boolean array defining on which index items a parsed component should be built. - If `as_latex` is True: Valid LaTeX math string defining the "where" conditions using logic notation. + 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 "where" conditions using logic notation. """ evaluated_wheres = [ where[0].eval( + return_type, + equation_name=self.name, helper_functions=helper_functions._registry["where"], - as_latex=as_latex, - model_data=model_data, - backend_dataset=backend_dataset, + input_data=backend_interface.inputs, + backend_interface=backend_interface, + apply_where=True, ) for where in self.where ] - if as_latex: + if return_type == "math_string": return r"\land{}".join(f"({i})" for i in evaluated_wheres if i != "true") else: where = xr.DataArray( @@ -319,50 +325,91 @@ def drop_dims_not_in_foreach(self, where: xr.DataArray) -> xr.DataArray: @overload # noqa: F811 def evaluate_expression( # noqa: F811 self, - model_data: xr.Dataset, backend_interface: backend_model.BackendModel, - as_latex: Literal[False] = False, + *, + return_type: Literal["array"] = "array", references: Optional[set] = None, where: Optional[xr.DataArray] = None, - ) -> Any: + ) -> xr.DataArray: "Expecting anything (most likely an array) if not requesting latex string" @overload # noqa: F811 def evaluate_expression( # noqa: F811 self, - model_data: xr.Dataset, backend_interface: backend_model.BackendModel, - as_latex: Literal[True], + *, + return_type: Literal["math_string"], references: Optional[set] = None, ) -> str: "Expecting string if requesting latex string" def evaluate_expression( # noqa: F811 self, - model_data: xr.Dataset, backend_interface: backend_model.BackendModel, - as_latex: bool = False, + *, + return_type: str = "array", references: Optional[set] = None, - where: Optional[xr.DataArray] = None, - ) -> Any: - if where is None: - apply_where = False - else: - apply_where = True - return self.expression[0].eval( + where: xr.DataArray = TRUE_ARRAY, + ) -> Union[xr.DataArray, str]: + """Evaluate a math string to produce either an array of backend objects or a LaTex math string. + + Args: + backend_interface (calliope.backend.backend_model.BackendModel): Interface to a optimisation backend. + + Keyword Args: + return_type (str, optional): + If "array", return xarray.DataArray. If "math_string", return LaTex math string. + Defaults to "array". + references (Optional[set], optional): If given, any references in the math string to other model components will be logged here. Defaults to None. + where (Optional[xr.DataArray], optional): If given, should be a boolean array with which to mask any produced arrays. Defaults to xr.DataArray(True). + + Returns: + Union[xr.DataArray, str]: + If return_type == `array`: array of backend expression objects. + If return_type == `math_string`: Valid LaTeX math string defining the "where" conditions using logic notation. + """ + evaluated = self.expression[0].eval( + return_type, equation_name=self.name, slice_dict=self.slices, sub_expression_dict=self.sub_expressions, backend_interface=backend_interface, - backend_dataset=backend_interface._dataset, - model_data=model_data, - where=where, + input_data=backend_interface.inputs, + where_array=where, references=references if references is not None else set(), - as_dict=False, helper_functions=helper_functions._registry["expression"], - as_latex=as_latex, - apply_where=apply_where, ) + if not return_type == "math_string": + self.raise_error_on_where_expr_mismatch(evaluated, where) + return evaluated + + def raise_error_on_where_expr_mismatch( + self, expression: xr.DataArray, where: xr.DataArray + ) -> None: + """ + Checks if an evaluated expression is consistent with the `where` array. + + Args: + expression (xr.DataArray): array of linear expressions or one side of a constraint equation. + where (xr.DataArray): where 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 where. + BackendError: + Raised if the expression has any NaN where the where applies. + """ + broadcast_dims_where = set(expression.dims).difference(set(where.dims)) + if broadcast_dims_where: + raise exceptions.BackendError( + f"{self.name} | The linear expression array is indexed over dimensions not present in `foreach`: {broadcast_dims_where}" + ) + # Check whether expression has NaN values in elements where the expression should be valid. + incomplete_constraints = expression.isnull() & where + if incomplete_constraints.any(): + raise exceptions.BackendError( + f"{self.name} | Missing a linear expression for some coordinates selected by 'where'. Adapting 'where' might help." + ) def log_not_added( self, @@ -459,13 +506,13 @@ def parse_top_level_where( def parse_equations( self, - valid_math_element_names: Iterable[str], + valid_component_names: Iterable[str], errors: Literal["raise", "ignore"] = "raise", ) -> list[ParsedBackendEquation]: """Parse `expression` and `where` strings of math component dictionary. Args: - valid_math_element_names (Iterable[str]): + valid_component_names (Iterable[str]): strings referring to valid backend objects to allow the parser to differentiate between them and generic strings. errors (Literal["raise", "ignore"], optional): Collected parsing errors can be raised directly or ignored. @@ -480,7 +527,7 @@ def parse_equations( equation_expression_list = self._unparsed.get("equations", []) equations = self.generate_expression_list( - expression_parser=self.equation_expression_parser(valid_math_element_names), + expression_parser=self.equation_expression_parser(valid_component_names), expression_list=equation_expression_list, expression_group="equations", id_prefix=self.name, @@ -489,7 +536,7 @@ def parse_equations( sub_expression_dict = { c_name: self.generate_expression_list( expression_parser=expression_parser.generate_sub_expression_parser( - valid_math_element_names + valid_component_names ), expression_list=c_list, expression_group="sub_expressions", @@ -500,7 +547,7 @@ def parse_equations( slice_dict = { idx_name: self.generate_expression_list( expression_parser=expression_parser.generate_slice_parser( - valid_math_element_names + valid_component_names ), expression_list=idx_list, expression_group="slices", @@ -682,7 +729,7 @@ def extend_equation_list_with_expression_group( ] def combine_definition_matrix_and_foreach( - self, model_data: xr.Dataset + self, input_data: xr.Dataset ) -> xr.DataArray: """Generate a multi-dimensional boolean array based on the sets over which the constraint is to be built (defined by "foreach") and the model `exists` array. @@ -690,27 +737,27 @@ def combine_definition_matrix_and_foreach( It is indexed over ["nodes", "techs", "carriers", "carrier_tiers"]. Args: - model_data (xr.Dataset): Calliope model dataset. + input_data (xr.Dataset): Calliope model dataset. Returns: xr.DataArray: boolean array indexed over ["nodes", "techs", "carriers", "carrier_tiers"] + any additional dimensions provided by `foreach`. """ # Start with (carriers, carrier_tiers, nodes, techs) and go from there - exists = model_data.definition_matrix + exists = input_data.definition_matrix # Add other dimensions (costs, timesteps, etc.) add_dims = set(self.sets).difference(exists.dims) - if add_dims.difference(model_data.dims): + if add_dims.difference(input_data.dims): self.log_not_added( - f"indexed over unidentified set names: `{add_dims.difference(model_data.dims)}`." + f"indexed over unidentified set names: `{add_dims.difference(input_data.dims)}`." ) return xr.DataArray(False) - exists_and_foreach = [exists, *[model_data[i].notnull() for i in add_dims]] + exists_and_foreach = [exists, *[input_data[i].notnull() for i in add_dims]] return functools.reduce(operator.and_, exists_and_foreach) def generate_top_level_where_array( self, - model_data: xr.Dataset, - backend_dataset: Optional[xr.Dataset] = None, + backend_interface: backend_model.BackendModel, + *, align_to_foreach_sets: bool = True, break_early: bool = True, ) -> xr.DataArray: @@ -719,7 +766,7 @@ def generate_top_level_where_array( and apply the component top-level where to the array. Args: - model_data (xr.Dataset): Calliope model input data. + backend_interface (calliope.backend.backend_model.BackendModel): Interface to a optimisation backend. align_to_foreach_sets (bool, optional): By default, all foreach arrays have the dimensions ("nodes", "techs", "carriers", "carrier_tiers") 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. @@ -731,7 +778,8 @@ def generate_top_level_where_array( Returns: xr.DataArray: Boolean array defining on which index items a parsed component should be built. """ - foreach_where = self.combine_definition_matrix_and_foreach(model_data) + input_data = backend_interface.inputs + foreach_where = self.combine_definition_matrix_and_foreach(input_data) if not foreach_where.any(): self.log_not_added("'foreach' does not apply anywhere.") @@ -740,9 +788,7 @@ def generate_top_level_where_array( return foreach_where self.parse_top_level_where() - where = self.evaluate_where( - model_data, backend_dataset, initial_where=foreach_where - ) + where = self.evaluate_where(backend_interface, initial_where=foreach_where) if break_early and not where.any(): return where diff --git a/src/calliope/backend/pyomo_backend_model.py b/src/calliope/backend/pyomo_backend_model.py index 95ceab400..7d7b65f10 100644 --- a/src/calliope/backend/pyomo_backend_model.py +++ b/src/calliope/backend/pyomo_backend_model.py @@ -90,7 +90,6 @@ def add_parameter( parameter_da.attrs["original_dtype"] = parameter_values.dtype self._add_to_dataset(parameter_name, parameter_da, "parameters", {}) - self.valid_math_element_names.add(parameter_name) def add_constraint( self, @@ -100,21 +99,12 @@ def add_constraint( def _constraint_setter( element: parsing.ParsedBackendEquation, where: xr.DataArray, references: set ) -> xr.DataArray: - lhs, op, rhs = element.evaluate_expression( - self.inputs, self, where=where, references=references - ) - lhs = lhs.squeeze(drop=True) - rhs = rhs.squeeze(drop=True) - - self._check_expr_where_consistency(lhs, where, f"(constraints, {name})") - self._check_expr_where_consistency(rhs, where, f"(constraints, {name})") + expr = element.evaluate_expression(self, where=where, references=references) to_fill = self._apply_func( self._to_pyomo_constraint, where, - lhs, - rhs, - op=op, + expr, name=name, ) return to_fill @@ -129,13 +119,9 @@ def add_global_expression( def _expression_setter( element: parsing.ParsedBackendEquation, where: xr.DataArray, references: set ) -> xr.DataArray: - expr = element.evaluate_expression( - self.inputs, self, where=where, references=references - ) + expr = element.evaluate_expression(self, where=where, references=references) expr = expr.squeeze(drop=True) - self._check_expr_where_consistency(expr, where, f"(expressions, {name})") - to_fill = self._apply_func( self._to_pyomo_expression, where, @@ -145,8 +131,6 @@ def _expression_setter( self._clean_arrays(expr) return to_fill - self.valid_math_element_names.add(name) - self._add_component( name, expression_dict, _expression_setter, "global_expressions" ) @@ -160,7 +144,6 @@ def add_variable( if variable_dict is None: variable_dict = self.inputs.attrs["math"]["variables"][name] - self.valid_math_element_names.add(name) def _variable_setter(where): domain_type = domain_dict[variable_dict.get("domain", "real")] @@ -174,8 +157,6 @@ def _variable_setter(where): domain_type=domain_type, ) - self.valid_math_element_names.add(name) - self._add_component(name, variable_dict, _variable_setter, "variables") def add_objective( @@ -192,8 +173,8 @@ def add_objective( def _objective_setter( element: parsing.ParsedBackendEquation, where: xr.DataArray, references: set ) -> xr.DataArray: - expr = element.evaluate_expression(self.inputs, self, references=references) - objective = pmo.objective(xr.DataArray(expr).item(), sense=sense) + expr = element.evaluate_expression(self, references=references) + objective = pmo.objective(expr.item(), sense=sense) if name == self.inputs.attrs["config"].build.objective: text = "activated" objective.activate() @@ -490,37 +471,6 @@ def unfix_variable(self, name: str, where: Optional[xr.DataArray] = None) -> Non variable_da = variable_da.where(where.fillna(0)) self._apply_func(self._unfix_pyomo_variable, variable_da) - @staticmethod - def _check_expr_where_consistency( - expression: xr.DataArray, where: xr.DataArray, description: str - ) -> None: - """ - Checks if a given constraint or global expression is consistent with the where. - - Parameters: - expression (xr.DataArray): array of linear expressions from a global expression or one side of a constraint equation. - where (xr.DataArray): where array. - description (str): Description to prefix the error message. - - Raises: - BackendError: - Raised if there is a dimension in the expression that is not in the where. - BackendError: - Raised if the expression has any NaN where the where applies. - """ - # Check whether expression has a dim that does not exist in where. - broadcast_dims_where = set(expression.dims).difference(set(where.dims)) - if broadcast_dims_where: - raise BackendError( - f"{description}: The linear expression array is indexed over dimensions not present in `foreach`: {broadcast_dims_where}" - ) - - incomplete_constraints = expression.isnull() & where - if incomplete_constraints.any(): - raise BackendError( - f"{description}: Missing a linear expression for some coordinates selected by 'where'. Adapting 'where' might help." - ) - def _get_capacity_bound(self, bound: Any, name: str) -> xr.DataArray: """ Generate array for the upper/lower bound of a decision variable. @@ -639,10 +589,8 @@ def _unfix_pyomo_variable(self, orig: ObjVariable) -> None: def _to_pyomo_constraint( self, mask: Union[bool, np.bool_], - lhs: Any, - rhs: Any, + expr: Any, *, - op: Literal["==", ">=", "<="], name: str, ) -> Union[type[ObjConstraint], float]: """ @@ -653,11 +601,9 @@ def _to_pyomo_constraint( Args: mask (Union[bool, np.bool_]): If True, add constraint, otherwise return np.nan - lhs (Any): Equation left-hand-side linear expression - rhs (Any): Equation right-hand-side linear expression + expr (Any): Equation expression. Kwargs: - op (Literal[, optional): Operator to compare `lhs` and `rhs`. Defaults to =", ">=", "<="]. name (str): Name of constraint Returns: @@ -666,16 +612,12 @@ def _to_pyomo_constraint( Otherwise return pmo_constraint(expr=lhs op rhs). """ - if not mask: + if mask: + constraint = ObjConstraint(expr=expr) + self._instance.constraints[name].append(constraint) + return constraint + else: return np.nan - elif op == "==": - constraint = ObjConstraint(expr=lhs == rhs) - elif op == "<=": - constraint = ObjConstraint(expr=lhs <= rhs) - elif op == ">=": - constraint = ObjConstraint(expr=lhs >= rhs) - self._instance.constraints[name].append(constraint) - return constraint def _to_pyomo_expression( self, mask: Union[bool, np.bool_], expr: Any, *, name: str diff --git a/src/calliope/backend/where_parser.py b/src/calliope/backend/where_parser.py index 75887d558..86638b22c 100644 --- a/src/calliope/backend/where_parser.py +++ b/src/calliope/backend/where_parser.py @@ -4,37 +4,52 @@ from __future__ import annotations import operator -from typing import Any, Optional, Union +from typing import TYPE_CHECKING, Callable, Union import numpy as np import pyparsing as pp import xarray as xr +from typing_extensions import NotRequired, TypedDict from calliope.backend import expression_parser from calliope.exceptions import BackendError +if TYPE_CHECKING: + from calliope.backend.backend_model import BackendModel + + pp.ParserElement.enablePackrat() BOOLEANTYPE = Union[np.bool_, np.typing.NDArray[np.bool_]] -class EvalNot(expression_parser.EvalSignOp): - "Parse action to process successfully parsed expressions with a leading `not`" +class EvalAttrs(TypedDict): + equation_name: str + backend_interface: BackendModel + input_data: xr.DataArray + helper_functions: dict[str, Callable] + apply_where: NotRequired[bool] + + +class EvalWhere(expression_parser.EvalToArrayStr): + "Update type reference for `eval_attrs` to match `where` evaluation kwargs" + eval_attrs: EvalAttrs = {} + + +class EvalNot(EvalWhere, expression_parser.EvalSignOp): + "Parse action to process successfully parsed expressions with a leading `not`." - def as_latex(self, val: str) -> str: + def as_math_string(self) -> str: """Add sign to stringified data for use in a LaTex math formula""" - return rf"\neg ({val})" + evaluated = self.value.eval("math_string", **self.eval_attrs) + return rf"\neg ({evaluated})" - def eval(self, **kwargs) -> Union[BOOLEANTYPE, str]: - "Return inverted bool / boolean array" - evaluated = self.value.eval(**kwargs) - if kwargs.get("as_latex", False): - return self.as_latex(evaluated) - else: - return ~evaluated + def as_array(self) -> xr.DataArray: + evaluated = self.value.eval("array", **self.eval_attrs) + return ~evaluated -class EvalAndOr(expression_parser.EvalOperatorOperand): +class EvalAndOr(EvalWhere, expression_parser.EvalOperatorOperand): """ Parse action to process successfully parsed expressions with operands separated by an and/or operator (OPERAND OPERATOR OPERAND OPERATOR OPERAND ...) @@ -44,44 +59,34 @@ class EvalAndOr(expression_parser.EvalOperatorOperand): "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 - def bool_operate( - self, val: BOOLEANTYPE, evaluated_operand: BOOLEANTYPE, operator_: str - ) -> BOOLEANTYPE: + @staticmethod + def _operate( + val: xr.DataArray, evaluated_operand: xr.DataArray, operator_: str + ) -> xr.DataArray: + "Apply bitwise comparison between boolean xarray dataarrays." if operator_ == "and": val = operator.and_(val, evaluated_operand) elif operator_ == "or": val = operator.or_(val, evaluated_operand) return val - def _as_latex( - self, val: str, operand: str, operator_: str, val_type: Any, operand_type: Any - ) -> str: - if val == "true": - val = operand - elif operand != "true": - val = self.as_latex(val, operand, operator_, val_type, operand_type) - return val + def _apply_where_array(self, evaluated: xr.DataArray) -> xr.DataArray: + "Override func from parent class to effectively do nothing." + return evaluated - def eval(self, as_latex: bool = False, **kwargs) -> Any: - "Return combined bools / boolean arrays" - val = self.value[0].eval(as_latex=as_latex, **kwargs) - for operator_, operand in self.operatorOperands(self.value[1:]): - evaluated_operand = operand.eval(as_latex=as_latex, **kwargs) - if as_latex: - val = self._as_latex( - val, - evaluated_operand, - operator_, - type(self.value[0]), - type(operand), - ) - else: - val = self.bool_operate(val, evaluated_operand, operator_) - return val + def as_math_string(self) -> str: + return super().as_math_string() + def as_array(self) -> xr.DataArray: + return super().as_array() -class ConfigOptionParser(expression_parser.EvalString): + +class ConfigOptionParser(EvalWhere): def __init__(self, instring: str, loc: int, tokens: pp.ParseResults) -> None: """ Parse action to process successfully parsed configuration option names. @@ -99,42 +104,26 @@ def __init__(self, instring: str, loc: int, tokens: pp.ParseResults) -> None: self.loc = loc def __repr__(self): - "Return string representation of the parsed grammar" return f"CONFIG:{self.config_option}" - def as_latex(self) -> str: - """Add return input string for use in a LaTex math formula""" + def as_math_string(self) -> str: return rf"\text{{config.{self.config_option}}}" - def eval( - self, model_data: xr.Dataset, **kwargs - ) -> Union[int, float, str, bool, np.bool_]: - """ - If the parsed configuration group and configuration option are valid then - return the option value, otherwise add to provided errors list inplace. - - Args: - model_data (xr.Dataset): Calliope model data. + def as_array(self) -> xr.DataArray: + config_val = ( + self.eval_attrs["input_data"].attrs["config"].build[self.config_option] + ) - Returns: - Optional[Union[int, float, str, bool, np.bool_]]: Configuration option value. - """ - - if kwargs.get("as_latex", False): - return self.as_latex() + if not isinstance(config_val, (int, float, str, bool, np.bool_)): + raise BackendError( + f"(where, {self.instring}): Configuration option resolves to invalid " + f"type `{type(config_val).__name__}`, expected a number, string, or boolean." + ) else: - config_val = model_data.attrs["config"].build[self.config_option] - - if not isinstance(config_val, (int, float, str, bool, np.bool_)): - raise BackendError( - f"(where, {self.instring}): Configuration option resolves to invalid " - f"type `{type(config_val).__name__}`, expected a number, string, or boolean." - ) - else: - return config_val + return xr.DataArray(config_val) -class DataVarParser(expression_parser.EvalString): +class DataVarParser(EvalWhere): def __init__(self, instring: str, loc: int, tokens: pp.ParseResults) -> None: """ Parse action to process successfully parsed model data variable names. @@ -152,69 +141,18 @@ def __init__(self, instring: str, loc: int, tokens: pp.ParseResults) -> None: self.loc = loc def __repr__(self): - "Return string representation of the parsed grammar" return f"DATA_VAR:{self.data_var}" - def as_latex( - self, data: xr.Dataset, data_var_type: str, apply_where: bool = True - ) -> str: - """stringify conditional for use in a LaTex math formula""" - # TODO: add dims from a YAML schema of params that includes default dims - if data_var_type == "parameters": - data_var_string = rf"\textit{{{self.data_var}}}" - else: - data_var_string = rf"\textbf{{{self.data_var}}}" - - var = data.get(self.data_var, None) - if var is not None and var.shape: - data_var_string += ( - rf"_\text{{{','.join(str(i).removesuffix('s') for i in var.dims)}}}" - ) - if apply_where: - data_var_string = rf"\exists ({data_var_string})" - return data_var_string - - def _data_var_exists( - self, model_data: xr.Dataset, data_var_type: str - ) -> xr.DataArray: - "mask by setting all (NaN | INF/-INF) to False, otherwise True" - var = model_data.get(self.data_var, xr.DataArray(np.nan)) - if data_var_type == "parameters": - return var.notnull() & (var != np.inf) & (var != -np.inf) - else: - return var.notnull() + def _preprocess(self) -> tuple[xr.DataArray, str]: + """Get data variable from the optimisation problem dataset. - def _data_var_with_default(self, model_data: xr.Dataset) -> xr.DataArray: - "Access data var and fill with default values. Return default value as an array if var does not exist" - default = model_data.attrs["defaults"].get(self.data_var) - return model_data.get(self.data_var, xr.DataArray(default)).fillna(default) - - def eval( - self, - model_data: xr.Dataset, - backend_dataset: Optional[xr.Dataset] = None, - apply_where: bool = True, - **kwargs, - ) -> Union[str, np.bool_, xr.DataArray]: + Raises: + TypeError: Cannot work with math components of type `constraint` or `objective`. + TypeError: Cannot check array contents (`apply_where=False`) of `variable` or `global_expression` math components. """ - Get parsed model data variable from the Calliope model dataset. - If it isn't there, return False. - - Args: - model_data (xr.Dataset): Calliope model dataset. - apply_where (bool, optional): - If True, return boolean array corresponding to whether there is data or - not in each element of the array. If False, return original array. - Defaults to True. - - Returns: - Union[np.bool_, xr.DataArray]: - False if data variable not in model data, array otherwise. - """ - if backend_dataset is None: - backend_dataset = xr.Dataset() - if self.data_var in backend_dataset.data_vars.keys(): - data_var_type = backend_dataset[self.data_var].attrs["obj_type"] + backend_interface = self.eval_attrs["backend_interface"] + if self.data_var in backend_interface._dataset.data_vars.keys(): + data_var_type = backend_interface._dataset[self.data_var].attrs["obj_type"] else: data_var_type = "parameters" @@ -223,6 +161,7 @@ def eval( f"Cannot check values in {data_var_type.removesuffix('s')} arrays in math `where` strings. " f"Received {data_var_type.removesuffix('s')}: `{self.data_var}`." ) + apply_where = self.eval_attrs.get("apply_where", True) if data_var_type != "parameters" and not apply_where: raise TypeError( f"Can only check for existence of values in {data_var_type.removesuffix('s')} arrays in math `where` strings. " @@ -231,23 +170,60 @@ def eval( ) if data_var_type == "parameters": - source_array = model_data + source_array = self.eval_attrs["input_data"] else: - source_array = backend_dataset + source_array = backend_interface._dataset + + return source_array, data_var_type - if kwargs.get("as_latex", False): - return self.as_latex(source_array, data_var_type, apply_where) + def _data_var_exists( + self, source_array: xr.DataArray, data_var_type: str + ) -> xr.DataArray: + "mask by setting all (NaN | INF/-INF) to False, otherwise True" + var = source_array.get(self.data_var, xr.DataArray(np.nan)) + if data_var_type == "parameters": + return var.notnull() & (var != np.inf) & (var != -np.inf) + else: + return var.notnull() + + def _data_var_with_default(self, source_array: xr.Dataset) -> xr.DataArray: + "Access data var and fill with default values. Return default value as an array if var does not exist" + default = source_array.attrs["defaults"].get(self.data_var) + return source_array.get(self.data_var, xr.DataArray(default)).fillna(default) - if data_var_type == "parameters" and self.data_var not in model_data: - return np.False_ + def as_math_string(self) -> str: + # TODO: add dims from a YAML schema of params that includes default dims + source_array, data_var_type = self._preprocess() + if data_var_type == "parameters": + data_var_string = rf"\textit{{{self.data_var}}}" + else: + data_var_string = rf"\textbf{{{self.data_var}}}" - if apply_where: + var = source_array.get(self.data_var, None) + if var is not None and var.shape: + data_var_string += ( + rf"_\text{{{','.join(str(i).removesuffix('s') for i in var.dims)}}}" + ) + if self.eval_attrs.get("apply_where", True): + data_var_string = rf"\exists ({data_var_string})" + return data_var_string + + def as_array(self) -> xr.DataArray: + source_array, data_var_type = self._preprocess() + + if ( + data_var_type == "parameters" + and self.data_var not in self.eval_attrs["input_data"] + ): + return xr.DataArray(np.False_) + + if self.eval_attrs.get("apply_where", True): return self._data_var_exists(source_array, data_var_type) else: return self._data_var_with_default(source_array) -class ComparisonParser(expression_parser.EvalComparisonOp): +class ComparisonParser(EvalWhere, expression_parser.EvalComparisonOp): "Parse action to process successfully parsed strings of the form x=y" OP_TRANSLATOR = { "<=": r"\mathord{\leq}", @@ -261,41 +237,30 @@ def __repr__(self): "Return string representation of the parsed grammar" return f"{self.lhs}{self.op}{self.rhs}" - def eval(self, **kwargs) -> Union[str, BOOLEANTYPE]: - """ - Compare LHS (any) and RHS (numeric, string, bool) and return a bool/boolean array - - Returns: - BOOLEANTYPE: Same shape as LHS. - str: latex representation of the comparison. - """ - kwargs["apply_where"] = False - lhs = self.lhs.eval(**kwargs) - rhs = self.rhs.eval(**kwargs) - - if kwargs.get("as_latex", False): - if r"\text" not in rhs: - rhs = rf"\text{{{rhs}}}" - comparison = self.as_latex(lhs, rhs) - else: - if self.op == "<=": - comparison = lhs <= rhs - elif self.op == ">=": - comparison = lhs >= rhs - if self.op == "<": - comparison = lhs < rhs - elif self.op == ">": - comparison = lhs > rhs - elif self.op == "=": - comparison = lhs == rhs - - if isinstance(comparison, bool): - # enables the "~" operator to later invert `comparison` if required. - comparison = np.bool_(comparison) - return comparison - - -class SubsetParser(expression_parser.EvalString): + def as_math_string(self) -> str: + self.eval_attrs["apply_where"] = 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: + self.eval_attrs["apply_where"] = False + lhs, rhs = self._eval("array") + if self.op == "<=": + comparison = lhs <= rhs + elif self.op == ">=": + comparison = lhs >= rhs + if self.op == "<": + comparison = lhs < rhs + elif self.op == ">": + comparison = lhs > rhs + elif self.op == "=": + comparison = lhs == rhs + return xr.DataArray(comparison) + + +class SubsetParser(EvalWhere): def __init__(self, instring: str, loc: int, tokens: pp.ParseResults) -> None: """ Parse action to process successfully parsed dimension subsetting. @@ -308,30 +273,31 @@ def __init__(self, instring: str, loc: int, tokens: pp.ParseResults) -> None: tokens (pp.ParseResults): Has two parsed elements: model set name (str), set items (Any). """ - self.subset, self.set_name = tokens + 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.subset}" + 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_latex(self, subset: list) -> str: - """stringify subset for use in a LaTex math formula""" + def as_math_string(self) -> str: + subset = self._eval() set_singular = self.set_name.removesuffix("s") subset_string = "[" + ",".join(str(i) for i in subset) + "]" return rf"\text{{{set_singular}}} \in \text{{{subset_string}}}" - def eval(self, model_data: xr.Dataset, **kwargs) -> Union[str, xr.DataArray]: - subset = [i.eval(**kwargs) for i in self.subset] - if kwargs.get("as_latex", False): - set_item_in_subset = self.as_latex(subset) - else: - set_item_in_subset = model_data[self.set_name].isin(subset) + def as_array(self) -> xr.DataArray: + subset = self._eval() + set_item_in_subset = self.eval_attrs["input_data"][self.set_name].isin(subset) return set_item_in_subset -class BoolOperandParser: +class BoolOperandParser(EvalWhere): def __init__(self, tokens: pp.ParseResults) -> None: """ Parse action to process successfully parsed boolean strings. @@ -342,25 +308,38 @@ def __init__(self, tokens: pp.ParseResults) -> None: self.val = tokens[0].lower() def __repr__(self): - "Return string representation of the parsed grammar" return f"BOOL:{self.val}" - def as_latex(self): - "Return boolean as a string in the domain {true, false}" + def as_math_string(self): return self.val - def eval(self, **kwargs) -> np.bool_: - "evaluate string to numpy boolean object." - if kwargs.get("as_latex", False): - bool_val = self.as_latex() - else: - if self.val == "true": - bool_val = np.True_ - elif self.val == "false": - bool_val = np.False_ + def as_array(self) -> xr.DataArray: + 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): + def __init__(self, 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: + tokens (pp.ParseResults): Has one parsed element: string name (str). + """ + self.val = tokens[0] + + def __repr__(self) -> str: + return f"STRING:{self.val}" + + def eval(self, *args, **eval_kwargs) -> str: + return str(self.val) + + def data_var_parser(generic_identifier: pp.ParserElement) -> pp.ParserElement: """ Parsing grammar to process model data variables which can be any valid python @@ -422,7 +401,7 @@ def bool_parser() -> pp.ParserElement: def evaluatable_string_parser(generic_identifier: pp.ParserElement) -> pp.ParserElement: "Parsing grammar to make generic strings used in comparison operations evaluatable" evaluatable_identifier = generic_identifier.copy() - evaluatable_identifier.set_parse_action(expression_parser.GenericStringParser) + evaluatable_identifier.set_parse_action(GenericStringParser) return evaluatable_identifier diff --git a/src/calliope/core/model.py b/src/calliope/core/model.py index 4a900b371..be360fee9 100644 --- a/src/calliope/core/model.py +++ b/src/calliope/core/model.py @@ -546,7 +546,7 @@ def validate_math_strings(self, math_dict: dict) -> None: importlib.resources.files("calliope") / "config" / "math_schema.yaml" ) validate_dict(math_dict, math_schema, "math") - valid_math_element_names = [ + valid_component_names = [ *self.math["variables"].keys(), *self.math["global_expressions"].keys(), *math_dict.get("variables", {}).keys(), @@ -561,7 +561,7 @@ def validate_math_strings(self, math_dict: dict) -> None: component_group, name, component_dict ) parsed.parse_top_level_where(errors="ignore") - parsed.parse_equations(set(valid_math_element_names), errors="ignore") + parsed.parse_equations(set(valid_component_names), errors="ignore") if not parsed._is_valid: collected_errors[f"{component_group}:{name}"] = parsed._errors diff --git a/tests/conftest.py b/tests/conftest.py index 0d2176441..fd287c6dc 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -227,6 +227,9 @@ def dummy_model_data(): model_data.attrs["defaults"] = AttrDict( {"all_inf": np.inf, "all_nan": np.nan, "with_inf": 100, "only_techs": 5} ) + model_data.attrs["math"] = AttrDict( + {"constraints": {}, "variables": {}, "global_expressions": {}, "objectives": {}} + ) return model_data diff --git a/tests/test_backend_expression_parser.py b/tests/test_backend_expression_parser.py index 16dd52b06..c78802aff 100644 --- a/tests/test_backend_expression_parser.py +++ b/tests/test_backend_expression_parser.py @@ -4,6 +4,7 @@ import numpy as np import pyparsing as pp import pytest +import xarray as xr from calliope import exceptions from calliope.backend import expression_parser, helper_functions @@ -16,7 +17,7 @@ class DummyFunc1(helper_functions.ParsingHelperFunction): NAME = "dummy_func_1" ALLOWED_IN = ["expression"] - def as_latex(self, x): + def as_math_string(self, x): return f"{x} * 10" def as_array(self, x): @@ -27,18 +28,17 @@ class DummyFunc2(helper_functions.ParsingHelperFunction): NAME = "dummy_func_2" ALLOWED_IN = ["expression"] - def as_latex(self, x, y): + def as_math_string(self, x, y): return f"{x} + {y}" - def as_array(self, x): - return x * 10 + def as_array(self, x, y): + return x * 10 + y @pytest.fixture -def valid_math_element_names(): +def valid_component_names(): return [ "foo", - "foo_bar", "with_inf", "only_techs", "no_dims", @@ -64,33 +64,28 @@ def identifier(base_parser_elements): @pytest.fixture -def evaluatable_identifier_elements(identifier, valid_math_element_names): +def evaluatable_identifier(identifier, valid_component_names): return expression_parser.evaluatable_identifier_parser( - identifier, valid_math_element_names + identifier, valid_component_names ) @pytest.fixture -def evaluatable_identifier(evaluatable_identifier_elements): - return evaluatable_identifier_elements[0] - - -@pytest.fixture -def id_list(evaluatable_identifier_elements): - return evaluatable_identifier_elements[1] +def id_list(number, evaluatable_identifier): + return expression_parser.list_parser(number, evaluatable_identifier) @pytest.fixture def unsliced_param(): - def _unsliced_param(valid_math_element_names): - return expression_parser.unsliced_object_parser(valid_math_element_names) + def _unsliced_param(valid_component_names): + return expression_parser.unsliced_object_parser(valid_component_names) return _unsliced_param @pytest.fixture -def unsliced_param_with_obj_names(unsliced_param, valid_math_element_names): - return unsliced_param(valid_math_element_names) +def unsliced_param_with_obj_names(unsliced_param, valid_component_names): + return unsliced_param(valid_component_names) @pytest.fixture @@ -169,15 +164,17 @@ def helper_function_one_parser_in_args(identifier, request): @pytest.fixture(scope="function") -def eval_kwargs(): +def eval_kwargs(dummy_pyomo_backend_model): return { "helper_functions": helper_functions._registry["expression"], - "as_dict": True, "slice_dict": {}, "sub_expression_dict": {}, "equation_name": "foobar", - "apply_where": False, + "where_array": xr.DataArray(True), "references": set(), + "backend_interface": dummy_pyomo_backend_model, + "input_data": dummy_pyomo_backend_model.inputs, + "return_type": "array", } @@ -226,18 +223,18 @@ def equation_comparison(arithmetic): @pytest.fixture -def generate_equation(valid_math_element_names): - return expression_parser.generate_equation_parser(valid_math_element_names) +def generate_equation(valid_component_names): + return expression_parser.generate_equation_parser(valid_component_names) @pytest.fixture -def generate_slice(valid_math_element_names): - return expression_parser.generate_slice_parser(valid_math_element_names) +def generate_slice(valid_component_names): + return expression_parser.generate_slice_parser(valid_component_names) @pytest.fixture -def generate_sub_expression(valid_math_element_names): - return expression_parser.generate_sub_expression_parser(valid_math_element_names) +def generate_sub_expression(valid_component_names): + return expression_parser.generate_sub_expression_parser(valid_component_names) class TestEquationParserElements: @@ -259,7 +256,7 @@ class TestEquationParserElements: ) def test_numbers(self, number, string_val, expected): parsed_ = number.parse_string(string_val, parse_all=True) - assert parsed_[0].eval() == expected + assert parsed_[0].eval(return_type="array") == expected @pytest.mark.parametrize( "string_val", @@ -337,25 +334,33 @@ def test_id_list(self, id_list, eval_kwargs): parsed_ = id_list.parse_string("[hello, there]", parse_all=True) assert parsed_[0].eval(**eval_kwargs) == ["hello", "there"] - @pytest.mark.parametrize( - "string_val", ["1", "inf", "foo", "$foo", "dummy_func_1(1)"] - ) + def test_id_list_with_numeric(self, id_list, eval_kwargs): + parsed_ = id_list.parse_string("[hello, 1, 1.0, there]", parse_all=True) + assert parsed_[0].eval(**eval_kwargs) == ["hello", 1.0, 1.0, "there"] + + @pytest.mark.parametrize("string_val", ["foo", "$foo", "dummy_func_1(1)"]) def test_id_list_fail(self, id_list, string_val): with pytest.raises(pp.ParseException): id_list.parse_string(f"[{string_val}, there]", parse_all=True) - @pytest.mark.parametrize("string_val", ["foo", "foo_bar"]) - def test_unsliced_param(self, unsliced_param_with_obj_names, string_val): + @pytest.mark.parametrize("string_val", ["with_inf", "no_dims"]) + def test_unsliced_param( + self, unsliced_param_with_obj_names, eval_kwargs, string_val + ): parsed_ = unsliced_param_with_obj_names.parse_string(string_val, parse_all=True) - assert parsed_[0].eval(references=set(), as_dict=True) == { - "param_or_var_name": string_val - } + assert ( + parsed_[0] + .eval(**eval_kwargs) + .equals(eval_kwargs["backend_interface"]._dataset[string_val]) + ) - def test_unsliced_param_references(self, unsliced_param_with_obj_names): - references = set() - parsed_ = unsliced_param_with_obj_names.parse_string("foo", parse_all=True) - parsed_[0].eval(references=references, as_dict=True) - assert references == {"foo"} + def test_unsliced_param_references( + self, unsliced_param_with_obj_names, eval_kwargs + ): + references = eval_kwargs.pop("references") + parsed_ = unsliced_param_with_obj_names.parse_string("with_inf", parse_all=True) + parsed_[0].eval(references=references, **eval_kwargs) + assert references == {"with_inf"} @pytest.mark.parametrize("string_val", ["Foo", "foobar"]) def test_unsliced_param_fail(self, unsliced_param_with_obj_names, string_val): @@ -365,31 +370,28 @@ def test_unsliced_param_fail(self, unsliced_param_with_obj_names, string_val): @pytest.mark.parametrize( ["string_val", "expected"], [ - ("foo[techs=tech]", ["foo", {"techs": "tech"}]), + ("foo[techs=tech]", "SLICED_COMPONENT:foo[techs=STRING:tech]"), ( f"foo[techs={SUB_EXPRESSION_CLASSIFIER}tech]", - ["foo", {"techs": {"slice_reference": "tech"}}], + "SLICED_COMPONENT:foo[techs=REFERENCE:tech]", ), ( f"foo[techs=tech,bars={SUB_EXPRESSION_CLASSIFIER}bar]", - ["foo", {"techs": "tech", "bars": {"slice_reference": "bar"}}], + "SLICED_COMPONENT:foo[techs=STRING:tech,bars=REFERENCE:bar]", ), ( f"foo[ bars={SUB_EXPRESSION_CLASSIFIER}bar, techs=tech ]", - ["foo", {"bars": {"slice_reference": "bar"}, "techs": "tech"}], + "SLICED_COMPONENT:foo[bars=REFERENCE:bar,techs=STRING:tech]", ), ( - "foo_bar[techs=tech, nodes=node]", - ["foo_bar", {"techs": "tech", "nodes": "node"}], + "with_inf[techs=tech, nodes=node]", + "SLICED_COMPONENT:with_inf[techs=STRING:tech,nodes=STRING:node]", ), ], ) - def test_sliced_param(self, sliced_param, eval_kwargs, string_val, expected): + def test_sliced_param(self, sliced_param, string_val, expected): parsed_ = sliced_param.parse_string(string_val, parse_all=True) - assert parsed_[0].eval(**eval_kwargs) == { - "param_or_var_name": expected[0], - "dimensions": expected[1], - } + parsed_[0] == expected @pytest.mark.parametrize( "string_val", @@ -410,18 +412,21 @@ def test_fail_string_issues_sliced_param(self, sliced_param, string_val): sliced_param.parse_string(string_val, parse_all=True) @pytest.mark.parametrize( - ["string_val", "expected"], + "string_val", [ # keeping explicit reference to "$" to ensure something weird doesn't happen # with the constant (e.g. is accidentally overwritten) - ("$foo", "foo"), - (f"{SUB_EXPRESSION_CLASSIFIER}foo", "foo"), - (f"{SUB_EXPRESSION_CLASSIFIER}Foo_Bar_1", "Foo_Bar_1"), + "$foo", + f"{SUB_EXPRESSION_CLASSIFIER}foo", + f"{SUB_EXPRESSION_CLASSIFIER}Foo_Bar_1", ], ) - def test_sub_expression(self, sub_expression, string_val, expected): + def test_sub_expression(self, sub_expression, string_val): parsed_ = sub_expression.parse_string(string_val, parse_all=True) - assert parsed_[0].eval(as_dict=True) == {"sub_expression": expected} + assert ( + str(parsed_[0]) + == f"SUB_EXPRESSION:{string_val.removeprefix(SUB_EXPRESSION_CLASSIFIER)}" + ) @pytest.mark.parametrize( "string_val", @@ -443,172 +448,49 @@ def test_fail_sub_expression(self, sub_expression, string_val): @pytest.mark.parametrize( ["string_val", "expected"], [ - ( - "dummy_func_1(1)", - {"function": "dummy_func_1", "args": [1], "kwargs": {}}, - ), - ( - "dummy_func_1(x=1)", - {"function": "dummy_func_1", "args": [], "kwargs": {"x": 1}}, - ), - ( - "dummy_func_2(1, y=2)", - {"function": "dummy_func_2", "args": [1], "kwargs": {"y": 2}}, - ), - ( - "dummy_func_2(1, 2)", - {"function": "dummy_func_2", "args": [1, 2], "kwargs": {}}, - ), - ( - "dummy_func_2(y=1, x=2)", - {"function": "dummy_func_2", "args": [], "kwargs": {"x": 2, "y": 1}}, - ), - ( - "dummy_func_1(dummy_func_2(1, 2))", - { - "function": "dummy_func_1", - "args": [ - {"function": "dummy_func_2", "args": [1, 2], "kwargs": {}} - ], - "kwargs": {}, - }, - ), - ( - "dummy_func_1(x=dummy_func_2(1, y=2))", - { - "function": "dummy_func_1", - "args": [], - "kwargs": { - "x": { - "function": "dummy_func_2", - "args": [1], - "kwargs": {"y": 2}, - } - }, - }, - ), - ( - "dummy_func_1(1, foo[bars=bar])", - { - "function": "dummy_func_1", - "args": [ - 1, - {"param_or_var_name": "foo", "dimensions": {"bars": "bar"}}, - ], - "kwargs": {}, - }, - ), - ( - "dummy_func_1(1, dummy_func_2(1))", - { - "function": "dummy_func_1", - "args": [ - 1, - {"function": "dummy_func_2", "args": [1], "kwargs": {}}, - ], - "kwargs": {}, - }, - ), - ( - "dummy_func_1(foo)", - { - "function": "dummy_func_1", - "args": [{"param_or_var_name": "foo"}], - "kwargs": {}, - }, - ), ( f"dummy_func_1(foo[bars={SUB_EXPRESSION_CLASSIFIER}bar])", - { - "function": "dummy_func_1", - "args": [ - { - "param_or_var_name": "foo", - "dimensions": {"bars": {"slice_reference": "bar"}}, - } - ], - "kwargs": {}, - }, + "dummy_func_1(args=[SLICED_COMPONENT:foo[bars=REFERENCE:bar]], kwargs={})", ), ( "dummy_func_1(1, x=foo[bars=bar1])", - { - "function": "dummy_func_1", - "args": [1], - "kwargs": { - "x": { - "param_or_var_name": "foo", - "dimensions": {"bars": "bar1"}, - } - }, - }, + "dummy_func_1(args=[NUM:1], kwargs={x=SLICED_COMPONENT:foo[bars=STRING:bar1]})", ), ( "dummy_func_1($foo)", - { - "function": "dummy_func_1", - "args": [{"sub_expression": "foo"}], - "kwargs": {}, - }, + "dummy_func_1(args=[SUB_EXPRESSION:foo], kwargs={})", ), ( "dummy_func_1($foo, x=$foo)", - { - "function": "dummy_func_1", - "args": [{"sub_expression": "foo"}], - "kwargs": {"x": {"sub_expression": "foo"}}, - }, + "dummy_func_1(args=[SUB_EXPRESSION:foo], kwargs={x=SUB_EXPRESSION:foo})", ), ( "dummy_func_1(dummy_func_2(foo[bars=bar1], y=$foo))", - { - "function": "dummy_func_1", - "args": [ - { - "function": "dummy_func_2", - "args": [ - { - "param_or_var_name": "foo", - "dimensions": {"bars": "bar1"}, - } - ], - "kwargs": {"y": {"sub_expression": "foo"}}, - } - ], - "kwargs": {}, - }, + "dummy_func_1(args=[dummy_func_2(args=[SLICED_COMPONENT:foo[bars=STRING:bar1]], kwargs={y=SUB_EXPRESSION:foo})], kwargs={})", ), ( - "dummy_func_1(1, dummy_func_2(1, foo[bars=$bar]), foo[foos=foo1, bars=bar1], $foo, 1, foo)", - { - "function": "dummy_func_1", - "args": [ - 1, - { - "function": "dummy_func_2", - "args": [ - 1, - { - "param_or_var_name": "foo", - "dimensions": {"bars": {"slice_reference": "bar"}}, - }, - ], - "kwargs": {}, - }, - { - "param_or_var_name": "foo", - "dimensions": {"foos": "foo1", "bars": "bar1"}, - }, - {"sub_expression": "foo"}, - 1, - {"param_or_var_name": "foo"}, - ], - "kwargs": {}, - }, + "dummy_func_1(1, dummy_func_2(1, foo[bars=$bar]), $foo, 1, foo, x=foo[foos=foo1, bars=bar1])", + "dummy_func_1(args=[NUM:1, dummy_func_2(args=[NUM:1, SLICED_COMPONENT:foo[bars=REFERENCE:bar]], kwargs={}), SUB_EXPRESSION:foo, NUM:1, COMPONENT:foo], kwargs={x=SLICED_COMPONENT:foo[foos=STRING:foo1, bars=STRING:bar1]})", ), ], ) - def test_function(self, helper_function, string_val, expected, eval_kwargs): + def test_function_parse(self, helper_function, string_val, expected): + parsed_ = helper_function.parse_string(string_val, parse_all=True) + assert parsed_[0] == expected + + @pytest.mark.parametrize( + ["string_val", "expected"], + [ + ("dummy_func_1(1)", 10), + ("dummy_func_1(x=1)", 10), + ("dummy_func_2(1, y=2)", 12), + ("dummy_func_2(1, 2)", 12), + ("dummy_func_2(y=1, x=2)", 21), + ("dummy_func_1(dummy_func_2(1, 2))", 120), + ("dummy_func_1(x=dummy_func_2(1, y=2))", 120), + ], + ) + def test_function_eval(self, helper_function, string_val, expected, eval_kwargs): parsed_ = helper_function.parse_string(string_val, parse_all=True) assert parsed_[0].eval(**eval_kwargs) == expected @@ -767,12 +649,12 @@ def test_mashup(self, equation_string, expected, eval_kwargs, func_string, reque @pytest.mark.parametrize("number_", numbers) @pytest.mark.parametrize("sub_expr_", ["$foo", "$bar1"]) - @pytest.mark.parametrize("unsliced_param_", ["foo", "foo_bar"]) + @pytest.mark.parametrize("unsliced_param_", ["foo", "with_inf"]) @pytest.mark.parametrize( - "sliced_param_", ["foo[bars=bar1]", "foo_bar[foos=foo1, bars=$bar]"] + "sliced_param_", ["foo[bars=bar1]", "with_inf[foos=foo1, bars=$bar]"] ) @pytest.mark.parametrize( - "helper_function_", ["foo(1)", "bar1(foo, $foo, foo_bar[foos=foo1], x=1)"] + "helper_function_", ["foo(1)", "bar1(foo, $foo, with_inf[foos=foo1], x=1)"] ) @pytest.mark.parametrize( "func_string", ["arithmetic", "helper_function_allow_arithmetic"] @@ -811,15 +693,7 @@ def test_helper_function_no_arithmetic(self, helper_function, string_): @pytest.mark.parametrize( ["string_", "expected"], - [ - ("1 + 2", {"args": [3], "kwargs": {}}), - ("1 * 2", {"args": [2], "kwargs": {}}), - ("x=1/2", {"args": [], "kwargs": {"x": 0.5}}), - ( - "foo, x=1+1", - {"args": [{"param_or_var_name": "foo"}], "kwargs": {"x": 2}}, - ), - ], + [("1 + 2", 30), ("1 * 2", 20), ("x=1/2", 5)], ) def test_helper_function_allow_arithmetic( self, helper_function_allow_arithmetic, eval_kwargs, string_, expected @@ -829,13 +703,13 @@ def test_helper_function_allow_arithmetic( helper_func_string, parse_all=True ) evaluated = parsed[0].eval(**eval_kwargs) - assert evaluated == {"function": "dummy_func_1", **expected} + assert evaluated == expected def test_repr(self, arithmetic): parse_string = "1 + foo - foo[foos=foo1, bars=$bar] + (foo / $foo) ** -2" expected = ( - "(NUM:1 + PARAM_OR_VAR:foo - SLICED_PARAM_OR_VAR:foo[foos=STRING:foo1, bars=REFERENCE:bar]" - " + ((PARAM_OR_VAR:foo / SUB_EXPRESSION:foo) ** (-)NUM:2))" + "(NUM:1 + COMPONENT:foo - SLICED_COMPONENT:foo[foos=STRING:foo1, bars=REFERENCE:bar]" + " + ((COMPONENT:foo / SUB_EXPRESSION:foo) ** (-)NUM:2))" ) parsed_ = arithmetic.parse_string(parse_string, parse_all=True) assert str(parsed_[0]) == expected @@ -886,34 +760,16 @@ def test_sub_expression_parser_fail(self, generate_sub_expression, instring): class TestEquationParserComparison: EXPR_PARAMS_AND_EXPECTED_EVAL = { - 0: 0.0, - -1: -1.0, - 1e2: 100, - "1/2": 0.5, - "2**2": 4, - ".inf": np.inf, - "foo_bar": {"param_or_var_name": "foo_bar"}, - "foo[foos=foo1, bars=bar1]": { - "param_or_var_name": "foo", - "dimensions": {"foos": "foo1", "bars": "bar1"}, - }, - "foo[bars=bar1]": {"param_or_var_name": "foo", "dimensions": {"bars": "bar1"}}, - "$foo": {"sub_expression": "foo"}, - "dummy_func_1(1, foo[bars=$bar], $foo, x=dummy_func_2(1, y=2), foo=foo_bar)": { - "function": "dummy_func_1", - "args": [ - 1, - { - "param_or_var_name": "foo", - "dimensions": {"bars": {"slice_reference": "bar"}}, - }, - {"sub_expression": "foo"}, - ], - "kwargs": { - "x": {"function": "dummy_func_2", "args": [1], "kwargs": {"y": 2}}, - "foo": {"param_or_var_name": "foo_bar"}, - }, - }, + 0: "NUM:0", + -1: "(-)NUM:1", + 1e2: "NUM:100.0", + "1/2": "(NUM:1 / NUM:2)", + "2**2.0": "(NUM:2 ** NUM:2.0)", + ".inf": "NUM:inf", + "with_inf": "COMPONENT:with_inf", + "foo[foos=foo1, bars=bar1]": "SLICED_COMPONENT:foo[foos=STRING:foo1, bars=STRING:bar1]", + "$foo": "SUB_EXPRESSION:foo", + "dummy_func_1(1, y=foo[bars=$bar])": "dummy_func_1(args=[NUM:1], kwargs={y=SLICED_COMPONENT:foo[bars=REFERENCE:bar]})", } @pytest.fixture(params=EXPR_PARAMS_AND_EXPECTED_EVAL.keys()) @@ -947,15 +803,14 @@ def test_simple_equation( expected_right, operator, equation_comparison, - eval_kwargs, ): parsed_constraint = equation_comparison.parse_string( single_equation_simple, parse_all=True ) evaluated_expression = parsed_constraint[0] - assert evaluated_expression.lhs.eval(**eval_kwargs) == expected_left - assert evaluated_expression.rhs.eval(**eval_kwargs) == expected_right + assert evaluated_expression.lhs == expected_left + assert evaluated_expression.rhs == expected_right assert evaluated_expression.op == operator @pytest.mark.parametrize( @@ -982,13 +837,8 @@ def test_evaluation( ): parser_func = request.getfixturevalue(func_string) parsed_equation = parser_func.parse_string(equation_string, parse_all=True) - lhs, op, rhs = parsed_equation[0].eval(**eval_kwargs) - comparison_dict = { - "==": lhs == rhs, - ">=": lhs >= rhs, - "<=": lhs <= rhs, - } - assert comparison_dict[op] if expected else not comparison_dict[op] + evaluated = parsed_equation[0].eval(**eval_kwargs) + assert evaluated if expected else not evaluated @pytest.mark.parametrize( "equation_string", @@ -1015,28 +865,26 @@ def test_fail_evaluation(self, equation_string, func_string, request): def test_repr(self, equation_comparison): parse_string = "1 + foo - foo[foos=foo1, bars=$bar] >= (foo / $foo) ** -2" expected = ( - "(NUM:1 + PARAM_OR_VAR:foo - SLICED_PARAM_OR_VAR:foo[foos=STRING:foo1, bars=REFERENCE:bar])" - " >= ((PARAM_OR_VAR:foo / SUB_EXPRESSION:foo) ** (-)NUM:2)" + "(NUM:1 + COMPONENT:foo - SLICED_COMPONENT:foo[foos=STRING:foo1, bars=REFERENCE:bar])" + " >= ((COMPONENT:foo / SUB_EXPRESSION:foo) ** (-)NUM:2)" ) parsed_ = equation_comparison.parse_string(parse_string, parse_all=True) assert str(parsed_[0]) == expected -class TestAsLatex: +class TestAsMathString: @pytest.fixture - def latex_eval_kwargs(self, dummy_latex_backend_model, dummy_model_data): + def latex_eval_kwargs(self, dummy_latex_backend_model): return { "helper_functions": helper_functions._registry["expression"], - "as_dict": False, - "as_latex": True, + "return_type": "math_string", "index_slice_dict": {}, "component_dict": {}, "equation_name": "foobar", - "apply_where": False, + "where_array": None, "references": set(), "backend_interface": dummy_latex_backend_model, - "backend_dataset": dummy_latex_backend_model._dataset, - "model_data": dummy_model_data, + "input_data": dummy_latex_backend_model.inputs, } @pytest.mark.parametrize( @@ -1049,7 +897,7 @@ def latex_eval_kwargs(self, dummy_latex_backend_model, dummy_model_data): ("number", "-1", "-1"), ("number", "2000000", "2\\mathord{\\times}10^{+06}"), ("evaluatable_identifier", "hello_there", "hello_there"), - ("id_list", "[hello, hello_there]", ["hello", "hello_there"]), + ("id_list", "[hello, hello_there]", "[hello,hello_there]"), ("unsliced_param_with_obj_names", "no_dims", r"\textit{no_dims}"), ( "unsliced_param_with_obj_names", diff --git a/tests/test_backend_helper_functions.py b/tests/test_backend_helper_functions.py index 083f61270..c4dd82b0f 100644 --- a/tests/test_backend_helper_functions.py +++ b/tests/test_backend_helper_functions.py @@ -70,7 +70,11 @@ def expression_get_transmission_techs(expression, parsing_kwargs): class TestAsArray: @pytest.fixture(scope="class") def parsing_kwargs(self, dummy_model_data): - return {"model_data": dummy_model_data} + return { + "input_data": dummy_model_data, + "equation_name": "foo", + "return_type": "array", + } @pytest.fixture(scope="function") def is_defined_any(self, dummy_model_data): @@ -326,7 +330,11 @@ class TestAsArrayGetTransmission: @pytest.fixture(scope="class") def parsing_kwargs(self): model = build_test_model(scenario="simple_supply,two_hours") - return {"model_data": model._model_data} + return { + "input_data": model._model_data, + "equation_name": "foo", + "return_type": "array", + } def test_expression_get_transmission_one_tech( self, expression_get_transmission_techs @@ -352,10 +360,14 @@ def test_expression_get_transmission_multi_tech( ) -class TestAsLatex: +class TestAsMathString: @pytest.fixture(scope="class") def parsing_kwargs(self, dummy_model_data): - return {"model_data": dummy_model_data, "as_latex": True} + return { + "input_data": dummy_model_data, + "return_type": "math_string", + "equation_name": "foo", + } def test_inheritance(self, where_inheritance): assert where_inheritance("boo") == r"\text{tech_group=boo}" @@ -473,11 +485,15 @@ def test_roll(self, expression_roll, instring, expected_substring): assert rolled_string == rf"\textit{{foo}}_\text{{{expected_substring}}}" -class TestAsLatexGetTransmission: +class TestAsMathStringGetTransmission: @pytest.fixture(scope="class") def parsing_kwargs(self): model = build_test_model(scenario="simple_supply,two_hours") - return {"model_data": model._model_data, "as_latex": True} + return { + "input_data": model._model_data, + "return_type": "math_string", + "equation_name": "foo", + } def test_expression_get_transmission_one_tech( self, expression_get_transmission_techs diff --git a/tests/test_backend_latex_backend.py b/tests/test_backend_latex_backend.py index 2cd69960c..1f069597f 100644 --- a/tests/test_backend_latex_backend.py +++ b/tests/test_backend_latex_backend.py @@ -73,7 +73,7 @@ def test_add_parameter(self, request, backend_obj): latex_backend_model = request.getfixturevalue(backend_obj) latex_backend_model.add_parameter("param", xr.DataArray(1)) assert latex_backend_model.parameters["param"] == xr.DataArray(1) - assert "param" in latex_backend_model.valid_math_element_names + assert "param" in latex_backend_model.valid_component_names @pytest.mark.parametrize( "backend_obj", ["valid_latex_backend", "dummy_latex_backend_model"] @@ -93,7 +93,7 @@ def test_add_variable(self, request, dummy_model_data, backend_obj): latex_backend_model.variables["var"].sum() <= dummy_model_data.with_inf_as_bool.sum() ) - assert "var" in latex_backend_model.valid_math_element_names + assert "var" in latex_backend_model.valid_component_names assert "math_string" in latex_backend_model.variables["var"].attrs def test_add_variable_not_valid(self, valid_latex_backend): @@ -107,7 +107,7 @@ def test_add_variable_not_valid(self, valid_latex_backend): ) # some null values might be introduced by the foreach array, so we just check the upper bound assert not valid_latex_backend.variables["invalid_var"].sum() - assert "invalid_var" in valid_latex_backend.valid_math_element_names + assert "invalid_var" in valid_latex_backend.valid_component_names assert "math_string" not in valid_latex_backend.variables["invalid_var"].attrs @pytest.mark.parametrize( @@ -128,7 +128,7 @@ def test_add_expression(self, request, dummy_model_data, backend_obj): latex_backend_model.global_expressions["expr"].sum() <= dummy_model_data.with_inf_as_bool.sum() ) - assert "expr" in latex_backend_model.valid_math_element_names + assert "expr" in latex_backend_model.valid_component_names assert "math_string" in latex_backend_model.global_expressions["expr"].attrs @pytest.mark.parametrize( @@ -151,7 +151,7 @@ def test_add_expression_with_variable_in_where( latex_backend_model.global_expressions["var_init_expr"].sum() <= dummy_model_data.with_inf_as_bool.sum() ) - assert "var_init_expr" in latex_backend_model.valid_math_element_names + assert "var_init_expr" in latex_backend_model.valid_component_names assert ( "math_string" in latex_backend_model.global_expressions["var_init_expr"].attrs @@ -175,7 +175,7 @@ def test_add_constraint(self, request, dummy_model_data, backend_obj): latex_backend_model.constraints["constr"].sum() <= dummy_model_data.with_inf_as_bool.sum() ) - assert "constr" not in latex_backend_model.valid_math_element_names + assert "constr" not in latex_backend_model.valid_component_names assert "math_string" in latex_backend_model.constraints["constr"].attrs @pytest.mark.parametrize( @@ -198,7 +198,7 @@ def test_add_constraint_with_variable_and_expression_in_where( latex_backend_model.constraints["var_init_constr"].sum() <= dummy_model_data.with_inf_as_bool.sum() ) - assert "var_init_constr" not in latex_backend_model.valid_math_element_names + assert "var_init_constr" not in latex_backend_model.valid_component_names assert "math_string" in latex_backend_model.constraints["var_init_constr"].attrs def test_add_constraint_not_valid(self, valid_latex_backend): @@ -244,7 +244,7 @@ def test_add_objective(self, dummy_latex_backend_model): }, ) assert dummy_latex_backend_model.objectives["obj"].isnull().all() - assert "obj" not in dummy_latex_backend_model.valid_math_element_names + assert "obj" not in dummy_latex_backend_model.valid_component_names assert len(dummy_latex_backend_model.objectives.data_vars) == 1 def test_create_obj_list(self, dummy_latex_backend_model): diff --git a/tests/test_backend_parsing.py b/tests/test_backend_parsing.py index 4639b11ac..da3a459b8 100644 --- a/tests/test_backend_parsing.py +++ b/tests/test_backend_parsing.py @@ -6,6 +6,7 @@ import pyparsing as pp import pytest import ruamel.yaml as yaml +import xarray as xr from calliope.backend import backend_model, expression_parser, parsing, where_parser from .common.util import check_error_or_warning @@ -37,28 +38,28 @@ def exists_array(component_obj, dummy_model_data): @pytest.fixture -def valid_math_element_names(dummy_model_data): +def valid_component_names(dummy_model_data): return ["foo", "bar", "baz", "foobar", *dummy_model_data.data_vars.keys()] @pytest.fixture -def expression_string_parser(valid_math_element_names): - return expression_parser.generate_equation_parser(valid_math_element_names) +def expression_string_parser(valid_component_names): + return expression_parser.generate_equation_parser(valid_component_names) @pytest.fixture -def arithmetic_string_parser(valid_math_element_names): - return expression_parser.generate_arithmetic_parser(valid_math_element_names) +def arithmetic_string_parser(valid_component_names): + return expression_parser.generate_arithmetic_parser(valid_component_names) @pytest.fixture -def slice_parser(valid_math_element_names): - return expression_parser.generate_slice_parser(valid_math_element_names) +def slice_parser(valid_component_names): + return expression_parser.generate_slice_parser(valid_component_names) @pytest.fixture -def sub_expression_parser(valid_math_element_names): - return expression_parser.generate_sub_expression_parser(valid_math_element_names) +def sub_expression_parser(valid_component_names): + return expression_parser.generate_sub_expression_parser(valid_component_names) @pytest.fixture @@ -239,7 +240,7 @@ def __init__(self): @pytest.fixture(scope="function") -def evaluatable_component_obj(valid_math_element_names): +def evaluatable_component_obj(valid_component_names): def _evaluatable_component_obj(equation_expressions): setup_string = f""" foreach: [techs, nodes] @@ -259,7 +260,7 @@ def __init__(self, dict_): self, "constraints", "foo", dict_ ) self.parse_top_level_where() - self.equations = self.parse_equations(valid_math_element_names) + self.equations = self.parse_equations(valid_component_names) return DummyParsedBackendComponent(sub_expression_dict) @@ -276,42 +277,32 @@ def __init__(self, dict_): ("only_techs + with_inf[techs=$tech] == 2", 1), ] ) -def evaluate_component_where(evaluatable_component_obj, dummy_model_data, request): +def evaluate_component_where( + evaluatable_component_obj, dummy_pyomo_backend_model, request +): component_obj = evaluatable_component_obj(request.param[0]) top_level_where = component_obj.generate_top_level_where_array( - dummy_model_data, break_early=False, align_to_foreach_sets=False + dummy_pyomo_backend_model, break_early=False, align_to_foreach_sets=False ) equation_where = component_obj.equations[0].evaluate_where( - dummy_model_data, initial_where=top_level_where + dummy_pyomo_backend_model, initial_where=top_level_where ) equation_where_aligned = component_obj.drop_dims_not_in_foreach(equation_where) return component_obj, equation_where_aligned, request.param[1] @pytest.fixture -def evaluate_component_expression( - evaluate_component_where, dummy_model_data, dummy_backend_interface -): +def evaluate_component_expression(evaluate_component_where, dummy_backend_interface): component_obj, equation_where, n_true = evaluate_component_where return ( component_obj.equations[0].evaluate_expression( - dummy_model_data, dummy_backend_interface, where=equation_where + dummy_backend_interface, where=equation_where ), n_true, ) -def apply_comparison(comparison_tuple): - lhs, op, rhs = comparison_tuple - if op == "==": - return lhs == rhs - if op == ">=": - return lhs >= rhs - if op == "<=": - return lhs <= rhs - - class TestParsedComponent: @pytest.mark.parametrize( "parse_string", @@ -363,7 +354,9 @@ def test_generate_expression_list( expression_string_parser, [expression_dict], "equations", id_prefix="foo" ) - assert parsed_list[0].where[0][0].eval() == expected_where_eval + assert ( + parsed_list[0].where[0][0].eval(return_type="array") == expected_where_eval + ) assert isinstance(parsed_list[0].expression, pp.ParseResults) @pytest.mark.parametrize("n_dicts", [1, 2, 3]) @@ -490,6 +483,7 @@ def test_extend_equation_list_with_expression_group_missing_sub_expression( ) def test_add_exprs_to_equation_data_multi( self, + dummy_backend_interface, component_obj, generate_expression_list, parsed_sub_expression_dict, @@ -509,11 +503,14 @@ def test_add_exprs_to_equation_data_multi( for constraint_eq in expression_list: component_sub_dict = constraint_eq.sub_expressions assert set(component_sub_dict.keys()) == {"foo", "bar"} - comparison_tuple = constraint_eq.expression[0].eval( - sub_expression_dict=component_sub_dict, apply_where=False + comparison_expr = constraint_eq.expression[0].eval( + sub_expression_dict=component_sub_dict, + backend_interface=dummy_backend_interface, + where_array=xr.DataArray(True), + return_type="array", ) - assert apply_comparison(comparison_tuple) == expected + assert comparison_expr == expected @pytest.mark.parametrize("n_1", [0, 1, 2]) @pytest.mark.parametrize("n_2", [0, 1, 2]) @@ -581,12 +578,12 @@ def test_extend_equation_list_with_expression_group_missing_slices( def test_parse_equations( self, obj_with_sub_expressions_and_slices, - valid_math_element_names, + valid_component_names, eq_string, expected_n_equations, ): component_obj = obj_with_sub_expressions_and_slices(eq_string) - parsed_equations = component_obj.parse_equations(valid_math_element_names) + parsed_equations = component_obj.parse_equations(valid_component_names) assert len(parsed_equations) == expected_n_equations assert len(set(eq.name for eq in parsed_equations)) == expected_n_equations @@ -602,11 +599,11 @@ def test_raise_caught_errors(self, component_obj, is_valid): assert check_error_or_warning(excinfo, ["\n * constraints:foo:"]) def test_parse_equations_fail( - self, obj_with_sub_expressions_and_slices, valid_math_element_names + self, obj_with_sub_expressions_and_slices, valid_component_names ): component_obj = obj_with_sub_expressions_and_slices("bar = 1") with pytest.raises(calliope.exceptions.ModelError) as excinfo: - component_obj.parse_equations(valid_math_element_names, errors="raise") + component_obj.parse_equations(valid_component_names, errors="raise") expected_err_string = """ * constraints:my_constraint: * equations[0].expression (line 1, char 5): bar = 1 @@ -614,10 +611,10 @@ def test_parse_equations_fail( assert check_error_or_warning(excinfo, expected_err_string) def test_parse_equations_fail_no_raise( - self, obj_with_sub_expressions_and_slices, valid_math_element_names + self, obj_with_sub_expressions_and_slices, valid_component_names ): component_obj = obj_with_sub_expressions_and_slices("bar = 1") - component_obj.parse_equations(valid_math_element_names, errors="ignore") + component_obj.parse_equations(valid_component_names, errors="ignore") expected_err_string = """\ equations[0].expression (line 1, char 5): bar = 1 @@ -640,9 +637,9 @@ def test_foreach_unidentified_name(self, caplog, dummy_model_data, component_obj component_obj.combine_definition_matrix_and_foreach(dummy_model_data) assert "indexed over unidentified set names" in caplog.text - def test_evaluate_where_to_false(self, dummy_model_data, component_obj): + def test_evaluate_where_to_false(self, dummy_pyomo_backend_model, component_obj): component_obj.parse_top_level_where() - where = component_obj.evaluate_where(dummy_model_data) + where = component_obj.evaluate_where(dummy_pyomo_backend_model) assert where.item() is True def test_parse_top_level_where_fail(self, component_obj): @@ -653,11 +650,13 @@ def test_parse_top_level_where_fail(self, component_obj): assert check_error_or_warning(excinfo, "Errors during math string parsing") def test_generate_top_level_where_array_break_at_foreach( - self, caplog, dummy_model_data, component_obj + self, caplog, dummy_pyomo_backend_model, component_obj ): component_obj.sets = ["nodes", "techs", "foos"] caplog.set_level(logging.DEBUG) - where_array = component_obj.generate_top_level_where_array(dummy_model_data) + where_array = component_obj.generate_top_level_where_array( + dummy_pyomo_backend_model + ) assert "indexed over unidentified set names: `{'foos'}`" in caplog.text assert "'foreach' does not apply anywhere." in caplog.text @@ -666,23 +665,25 @@ def test_generate_top_level_where_array_break_at_foreach( assert not where_array.shape def test_generate_top_level_where_array_break_at_top_level_where( - self, dummy_model_data, component_obj + self, dummy_pyomo_backend_model, component_obj ): component_obj.sets = ["nodes", "techs", "timesteps"] component_obj._unparsed["where"] = "all_nan" - where_array = component_obj.generate_top_level_where_array(dummy_model_data) + where_array = component_obj.generate_top_level_where_array( + dummy_pyomo_backend_model + ) assert not where_array.any() assert not set(component_obj.sets).difference(where_array.dims) def test_generate_top_level_where_array_no_break_no_align( - self, caplog, dummy_model_data, component_obj + self, caplog, dummy_pyomo_backend_model, component_obj ): component_obj.sets = ["nodes", "techs", "foos"] component_obj._unparsed["where"] = "all_nan" caplog.set_level(logging.DEBUG) where_array = component_obj.generate_top_level_where_array( - dummy_model_data, break_early=False, align_to_foreach_sets=False + dummy_pyomo_backend_model, break_early=False, align_to_foreach_sets=False ) assert "indexed over unidentified set names: `{'foos'}`" in caplog.text assert "'foreach' does not apply anywhere." in caplog.text @@ -692,12 +693,12 @@ def test_generate_top_level_where_array_no_break_no_align( assert set(component_obj.sets).difference(where_array.dims) == {"foos"} def test_generate_top_level_where_array_no_break_align( - self, dummy_model_data, component_obj + self, dummy_pyomo_backend_model, component_obj ): component_obj.sets = ["nodes", "techs"] component_obj._unparsed["where"] = "all_nan AND all_true_carriers" where_array = component_obj.generate_top_level_where_array( - dummy_model_data, break_early=False, align_to_foreach_sets=True + dummy_pyomo_backend_model, break_early=False, align_to_foreach_sets=True ) assert not where_array.any() assert not set(component_obj.sets).difference(where_array.dims) @@ -907,13 +908,17 @@ def test_add_slices_after_sub_expressions( @pytest.mark.parametrize("false_location", [0, -1]) def test_create_subset_from_where_definitely_empty( - self, dummy_model_data, equation_obj, where_string_parser, false_location + self, + dummy_pyomo_backend_model, + equation_obj, + where_string_parser, + false_location, ): equation_obj.sets = ["nodes", "techs"] equation_obj.where.insert( false_location, where_string_parser.parse_string("False", parse_all=True) ) - where = equation_obj.evaluate_where(dummy_model_data) + where = equation_obj.evaluate_where(dummy_pyomo_backend_model) assert not where.any() @@ -935,6 +940,7 @@ def test_create_subset_from_where_definitely_empty( def test_create_subset_from_where_one_level_where( self, dummy_model_data, + dummy_pyomo_backend_model, equation_obj, where_string_parser, where_string, @@ -946,17 +952,17 @@ def test_create_subset_from_where_one_level_where( equation_obj.where = [ where_string_parser.parse_string(where_string, parse_all=True) ] - where = equation_obj.evaluate_where(dummy_model_data) + where = equation_obj.evaluate_where(dummy_pyomo_backend_model) if level_ == "initial_where": equation_obj.where = [ where_string_parser.parse_string(where_string, parse_all=True) ] - initial_where = equation_obj.evaluate_where(dummy_model_data) + initial_where = equation_obj.evaluate_where(dummy_pyomo_backend_model) equation_obj.where = [ where_string_parser.parse_string("True", parse_all=True) ] where = equation_obj.evaluate_where( - dummy_model_data, initial_where=initial_where + dummy_pyomo_backend_model, initial_where=initial_where ) expected = dummy_model_data[expected_where_array] @@ -966,7 +972,7 @@ def test_create_subset_from_where_one_level_where( ) def test_create_subset_from_where_trim_dimension( - self, dummy_model_data, where_string_parser, equation_obj, exists_array + self, dummy_pyomo_backend_model, where_string_parser, equation_obj, exists_array ): equation_obj.sets = ["nodes", "techs"] @@ -974,19 +980,19 @@ def test_create_subset_from_where_trim_dimension( where_string_parser.parse_string("[foo] in carrier_tiers", parse_all=True) ] where = equation_obj.evaluate_where( - dummy_model_data, initial_where=exists_array + dummy_pyomo_backend_model, initial_where=exists_array ) assert where.sel(carrier_tiers="foo").any() assert not where.sel(carrier_tiers="bar").any() def test_create_subset_align_dims_with_sets( - self, dummy_model_data, where_string_parser, equation_obj, exists_array + self, dummy_pyomo_backend_model, where_string_parser, equation_obj, exists_array ): equation_obj.sets = ["nodes", "techs"] equation_obj.where = [where_string_parser.parse_string("True", parse_all=True)] where = equation_obj.evaluate_where( - dummy_model_data, initial_where=exists_array + dummy_pyomo_backend_model, initial_where=exists_array ) aligned_where = equation_obj.drop_dims_not_in_foreach(where) @@ -994,10 +1000,10 @@ def test_create_subset_align_dims_with_sets( assert not set(aligned_where.dims).difference(["nodes", "techs"]) def test_evaluate_expression(self, evaluate_component_expression): - comparison_tuple, n_true = evaluate_component_expression + comparison_expr, n_true = evaluate_component_expression # we can't check for equality since the random generation of NaNs in dummy_model_data carrier/node_tech # might nullify an otherwise valid item. - assert apply_comparison(comparison_tuple).sum() <= n_true + assert comparison_expr.sum() <= n_true class TestParsedConstraint: @@ -1025,27 +1031,32 @@ def test_parse_constraint_dict_sets(self, constraint_obj): def test_parse_constraint_dict_n_equations(self, constraint_obj): assert len(constraint_obj.equations) == 2 - def test_parse_constraint_dict_empty_eq1(self, constraint_obj, dummy_model_data): - assert not constraint_obj.equations[0].evaluate_where(dummy_model_data).any() + def test_parse_constraint_dict_empty_eq1( + self, constraint_obj, dummy_pyomo_backend_model + ): + assert ( + not constraint_obj.equations[0] + .evaluate_where(dummy_pyomo_backend_model) + .any() + ) def test_parse_constraint_dict_evaluate_eq2( - self, constraint_obj, dummy_model_data, dummy_backend_interface + self, constraint_obj, dummy_pyomo_backend_model, dummy_backend_interface ): # We ignore foreach here so we can do "== 1" below. With foreach, there is # a random element that might create a where array that masks the only valid index item - top_level_where = constraint_obj.evaluate_where(dummy_model_data) + top_level_where = constraint_obj.evaluate_where(dummy_pyomo_backend_model) valid_where = constraint_obj.equations[1].evaluate_where( - dummy_model_data, initial_where=top_level_where + dummy_pyomo_backend_model, initial_where=top_level_where ) aligned_where = constraint_obj.drop_dims_not_in_foreach(valid_where) references = set() - comparison_tuple = constraint_obj.equations[1].evaluate_expression( - dummy_model_data, + comparison_expr = constraint_obj.equations[1].evaluate_expression( dummy_backend_interface, where=aligned_where, references=references, ) - assert apply_comparison(comparison_tuple).sum() == 1 + assert comparison_expr.sum() == 1 assert references == {"only_techs"} @@ -1062,9 +1073,11 @@ def test_parse_variable_dict_sets(self, variable_obj): def test_parse_variable_dict_n_equations(self, variable_obj): assert len(variable_obj.equations) == 0 - def test_parse_variable_dict_empty_eq1(self, variable_obj, dummy_model_data): + def test_parse_variable_dict_empty_eq1( + self, variable_obj, dummy_pyomo_backend_model + ): top_level_where_where = variable_obj.generate_top_level_where_array( - dummy_model_data, break_early=False, align_to_foreach_sets=False + dummy_pyomo_backend_model, break_early=False, align_to_foreach_sets=False ) assert not top_level_where_where.any() @@ -1089,14 +1102,22 @@ def test_parse_objective_dict_sets(self, objective_obj): def test_parse_objective_dict_n_equations(self, objective_obj): assert len(objective_obj.equations) == 2 - def test_parse_objective_dict_empty_eq1(self, objective_obj, dummy_model_data): - assert not objective_obj.equations[0].evaluate_where(dummy_model_data).any() + def test_parse_objective_dict_empty_eq1( + self, objective_obj, dummy_pyomo_backend_model + ): + assert ( + not objective_obj.equations[0] + .evaluate_where(dummy_pyomo_backend_model) + .any() + ) def test_parse_objective_dict_evaluate_eq2( - self, objective_obj, dummy_model_data, dummy_backend_interface + self, objective_obj, dummy_pyomo_backend_model, dummy_backend_interface ): - valid_where = objective_obj.equations[1].evaluate_where(dummy_model_data) + valid_where = objective_obj.equations[1].evaluate_where( + dummy_pyomo_backend_model + ) objective_expression = objective_obj.equations[1].evaluate_expression( - dummy_model_data, dummy_backend_interface, where=valid_where + dummy_backend_interface, where=valid_where ) assert objective_expression.sum() == 12 diff --git a/tests/test_backend_pyomo.py b/tests/test_backend_pyomo.py index acef7f429..9320b658f 100755 --- a/tests/test_backend_pyomo.py +++ b/tests/test_backend_pyomo.py @@ -2020,7 +2020,7 @@ def test_raise_error_on_constraint_with_nan(self, simple_supply): assert check_error_or_warning( error, - f"(constraints, {constraint_name}): Missing a linear expression for some coordinates selected by 'where'. Adapting 'where' might help.", + f"constraints:{constraint_name}:0 | Missing a linear expression for some coordinates selected by 'where'. Adapting 'where' might help.", ) def test_raise_error_on_expression_with_nan(self, simple_supply): @@ -2063,7 +2063,7 @@ def test_raise_error_on_expression_with_nan(self, simple_supply): assert check_error_or_warning( error, - f"(expressions, {expression_name}): Missing a linear expression for some coordinates selected by 'where'. Adapting 'where' might help.", + f"global_expressions:{expression_name}:0 | Missing a linear expression for some coordinates selected by 'where'. Adapting 'where' might help.", ) def test_raise_error_on_excess_dimensions(self, simple_supply): @@ -2109,7 +2109,7 @@ def test_raise_error_on_excess_dimensions(self, simple_supply): assert check_error_or_warning( error, - f"(constraints, {constraint_name}): The linear expression array is indexed over dimensions not present in `foreach`: {{'nodes'}}", + f"constraints:{constraint_name}:0 | The linear expression array is indexed over dimensions not present in `foreach`: {{'nodes'}}", ) @pytest.mark.parametrize( diff --git a/tests/test_backend_where_parser.py b/tests/test_backend_where_parser.py index 5cfedf691..97006f844 100644 --- a/tests/test_backend_where_parser.py +++ b/tests/test_backend_where_parser.py @@ -87,14 +87,13 @@ def where(bool_operand, helper_function, data_var, comparison, subset): @pytest.fixture(scope="function") -def eval_kwargs(dummy_model_data, dummy_pyomo_backend_model): +def eval_kwargs(dummy_pyomo_backend_model): return { - "model_data": dummy_model_data, - "backend_dataset": dummy_pyomo_backend_model._dataset, + "input_data": dummy_pyomo_backend_model.inputs, + "backend_interface": dummy_pyomo_backend_model, "helper_functions": helper_functions._registry["where"], - "test": True, - "errors": set(), - "warnings": set(), + "equation_name": "foo", + "return_type": "array", } @@ -276,7 +275,8 @@ def test_config_fail_datatype( ) def test_boolean_parser(self, bool_operand, bool_string, expected_true): parsed_ = bool_operand.parse_string(bool_string, parse_all=True) - assert parsed_[0].eval() if expected_true else not parsed_[0].eval() + evaluated = parsed_[0].eval(return_type="array") + assert evaluated if expected_true else not evaluated @pytest.mark.parametrize( "bool_string", ["tru e", "_TRUE", "True_", "false1", "1false", "1", "foo"] @@ -289,7 +289,7 @@ def test_boolean_parser_malformed(self, bool_operand, bool_string): @pytest.mark.parametrize("instring", ["foo", "foo_bar", "FOO", "foo10", "foo_10"]) def test_evaluatable_string_parser(self, evaluatable_string, instring): parsed_ = evaluatable_string.parse_string(instring, parse_all=True) - parsed_[0].eval() == instring + parsed_[0].eval(return_type="array") == instring @pytest.mark.parametrize( "instring", ["_foo", "1foo", ".foo", "$foo", "__foo__", "foo bar", "foo-bar"] @@ -395,7 +395,7 @@ def test_comparison_malformed_string(self, comparison, comparison_string): def test_subsetting_parser(self, subset, subset_string, expected_subset): parsed_ = subset.parse_string(f"{subset_string} in foo", parse_all=True) assert parsed_[0].set_name == "foo" - assert [i.eval() for i in parsed_[0].subset] == expected_subset + assert [i.eval(return_type="array") for i in parsed_[0].val] == expected_subset @pytest.mark.parametrize( "subset_string", @@ -579,10 +579,10 @@ def test_where_malformed(self, where, instring): assert check_error_or_warning(excinfo, "Expected") -class TestAsLatex: +class TestAsMathString: @pytest.fixture def latex_eval_kwargs(self, eval_kwargs): - eval_kwargs["as_latex"] = True + eval_kwargs["return_type"] = "math_string" return eval_kwargs @pytest.mark.parametrize(