Tidy lexico code

Author: tsmbland

src/muse/decisions.py

@@ -38,15 +38,19 @@ def weighted_sum(objectives: Dataset, parameters: Any, **kwargs) -> DataArray:
     "weighted_sum",
 ]
 
-from collections.abc import Mapping, MutableMapping, Sequence
+from collections.abc import Hashable, Mapping, MutableMapping, Sequence
 from typing import (
     Any,
     Callable,
 )
 
+import numpy as np
+import xarray as xr
 from xarray import DataArray, Dataset
 
 from muse.registration import registrator
+from muse.timeslices import drop_timeslice
+from muse.utilities import tupled_dimension
 
 PARAMS_TYPE = Sequence[tuple[str, bool, float]]
 """Standard decision parameter type.
@@ -170,63 +174,180 @@ def weighted_sum(objectives: Dataset, parameters: Mapping[str, float]) -> DataAr
 def lexical_comparison(
     objectives: Dataset, parameters: PARAMS_TYPE | Sequence[tuple[str, float]]
 ) -> DataArray:
-    """Lexical comparison over the objectives.
+    """Lexicographic comparison using the best available replacements as reference.
 
     Lexical comparison operates by binning the objectives into bins of width
-    w_i = min_j(p_i o_i^j). Once binned, dimensions other than `asset` and
-    `technology` are reduced by taking the max, e.g. the largest constraint.
-    Finally, the objectives are ranked lexographically, in the order given by the
-    parameters.
+
+        w_i = min_j(p_i o_i^j),
+
+    where o_i^j are the objective values of the candidate replacements and the
+    minimum is taken over the replacement dimension. Once binned, dimensions
+    other than ``asset`` and ``replacement`` are reduced by taking the maximum
+    (e.g. the largest constraint). Finally, the objectives are ranked
+    lexicographically in the order given by ``parameters``.
 
     The result is an array of tuples which can subsequently be compared
     lexicographically.
     """
-    from muse.utilities import lexical_comparison
-
     assert len(parameters) > 0
-    if len(parameters[0]) == 3:
-        parameters = [(u[0], coeff_sign(u[1], u[2])) for u in parameters]
-    assert set(objectives.data_vars).issuperset([u[0] for u in parameters])
-    order = [u[0] for u in parameters]
 
-    binsize = objectives.copy()
-    for obj_name, weight in parameters:
-        binsize[obj_name] = binsize[obj_name] * weight
+    # Convert (name, min/max, coefficient) specifications to
+    # (name, signed coefficient), where the sign encodes whether the
+    # objective should be minimised or maximised.
+    if len(parameters[0]) == 3:
+        parameters = [
+            (name, coeff_sign(minmax, coeff)) for name, minmax, coeff in parameters
+        ]
+
+    # Lexicographic priority of the objectives.
+    order = tuple(name for name, _ in parameters)
+
+    # All requested objectives must be present.
+    assert set(objectives.data_vars).issuperset(order)
+
+    # Define the bin widths as
+    #
+    #     w_i = min_j(p_i o_i^j),
+    #
+    # i.e. the weighted objective values of the candidate replacements,
+    # taking the minimum over the replacement dimension.
+    binsize = objectives.copy(deep=True)
+
+    # Temporarily flatten the timeslice MultiIndex to avoid xarray's
+    # deprecated Dataset assignment path when modifying variables.
+    if "timeslice" in binsize.indexes:
+        index_names = binsize.indexes["timeslice"].names
+        binsize = binsize.reset_index("timeslice")
+
+    # Apply the objective weights (including minimisation/maximisation
+    # direction encoded in the sign).
+    for name, weight in parameters:
+        binsize[name] *= weight
+
+    # Restore the original timeslice MultiIndex structure.
+    if "timeslice" in objectives.indexes:
+        binsize = binsize.set_index(timeslice=index_names)
+
+    # Use the smallest weighted objective across replacements as the
+    # bin width for each objective.
     binsize = binsize.min("replacement")
 
-    return lexical_comparison(objectives, binsize, order=order, bin_last=False).rank(
-        "replacement"
-    )
+    # Construct lexicographically comparable tuples and rank the
+    # replacement options accordingly.
+    return _lexical_comparison(
+        objectives,
+        binsize,
+        order=order,
+        keep_last_continuous=True,
+    ).rank("replacement")
 
 
 @register_decision(name="retro_lexo")
 def retro_lexical_comparison(
-    objectives: Dataset, parameters: PARAMS_TYPE | Sequence[tuple[str, float]]
+    objectives: Dataset,
+    parameters: PARAMS_TYPE | Sequence[tuple[str, float]],
 ) -> DataArray:
-    """Lexical comparison over the objectives.
+    """Lexicographic comparison using current assets as reference.
 
     Lexical comparison operates by binning the objectives into bins of width
-    w_i = p_i o_i, where i are the current assets. Once binned, dimensions other
-    than `asset` and `replacement` are reduced by taking the max, e.g. the
-    largest constraint.  Finally, the objectives are ranked lexographically, in
-    the order given by the parameters.
+
+        w_i = p_i o_i,
+
+    where o_i are the objective values of the current assets. Once binned,
+    dimensions other than ``asset`` and ``replacement`` are reduced by taking
+    the maximum (e.g. the largest constraint). Finally, the objectives are
+    ranked lexicographically in the order given by ``parameters``.
 
     The result is an array of tuples which can subsequently be compared
     lexicographically.
     """
-    from muse.utilities import lexical_comparison
-
     assert len(parameters) > 0
+
+    # Convert (name, min/max, coefficient) specifications to
+    # (name, signed coefficient), where the sign encodes whether the
+    # objective should be minimised or maximised.
     if len(parameters[0]) == 3:
-        parameters = [(u[0], coeff_sign(u[1], u[2])) for u in parameters]
+        parameters = [
+            (name, coeff_sign(minmax, coeff)) for name, minmax, coeff in parameters
+        ]
+
+    # Retrofitting compares candidate replacements against the current
+    # asset, so every asset must also appear amongst the replacements.
     assert objectives.asset.isin(objectives.replacement).all()
-    assert set(objectives.data_vars).issuperset([u[0] for u in parameters])
 
-    order = [u[0] for u in parameters]
-    binsize = Dataset(dict(parameters)) * objectives.sel(replacement=objectives.asset)
-    return lexical_comparison(objectives, binsize, order=order, bin_last=False).rank(
-        "replacement"
-    )
+    # Lexicographic priority of the objectives.
+    order = tuple(name for name, _ in parameters)
+
+    # All requested objectives must be present.
+    assert set(objectives.data_vars).issuperset(order)
+
+    # Define the bin widths as
+    #
+    #     w_i = p_i o_i,
+    #
+    # where o_i are the objective values of the current assets. The
+    # objective values are selected by matching each asset to itself
+    # along the replacement dimension.
+    binwidths = Dataset(dict(parameters)) * objectives.sel(replacement=objectives.asset)
+
+    # Construct lexicographically comparable tuples and rank the
+    # replacement options accordingly.
+    return _lexical_comparison(
+        objectives,
+        binwidths,
+        order=order,
+        keep_last_continuous=True,
+    ).rank("replacement")
+
+
+def _lexical_comparison(
+    objectives: xr.Dataset,
+    binwidths: xr.Dataset,
+    order: Sequence[Hashable],
+    *,
+    keep_last_continuous: bool = False,
+) -> xr.DataArray:
+    """Lexical comparison over the objectives.
+
+    Lexical comparison operates by binning the objectives into bins of width
+    ``binwidths``. Once binned, dimensions other than ``asset`` and
+    ``replacement`` are reduced by taking the maximum (e.g. the largest
+    constraint). Finally, the objectives are ranked lexicographically in the
+    order given by ``order``.
+
+    Arguments:
+        objectives: xr.Dataset containing the objectives to rank.
+        binwidths: Bin widths used to discretise the objectives.
+        order: Order in which objectives are compared lexicographically.
+        keep_last_continuous: Whether the final objective should be left as a
+            continuous value, rather than being discretised.
+
+    Result:
+        An array of tuples which can subsequently be compared
+        lexicographically.
+    """
+    # Restrict to the objectives participating in the comparison and
+    # temporarily flatten the timeslice MultiIndex to avoid Dataset
+    # assignment issues in xarray.
+    result = drop_timeslice(objectives[list(order)]).copy()
+
+    # All objectives are discretised except, optionally, the final
+    # tie-breaking objective.
+    discretized = order[:-1] if keep_last_continuous else order
+
+    # Convert objectives to integer-valued bins.
+    for name in discretized:
+        result[name] = np.floor(result[name] / binwidths[name]).astype(np.int64)
+
+    # Preserve the final objective as a continuous quantity for
+    # tie-breaking, while still normalising by its bin width.
+    if keep_last_continuous:
+        name = order[-1]
+        result[name] = result[name] / binwidths[name]
+
+    # Combine the ordered objectives into tuples that can be compared
+    # lexicographically.
+    return result.to_array(dim="objective").reduce(tupled_dimension, dim="objective")
 
 
 def _epsilon_constraints(

src/muse/utilities.py

@@ -310,47 +310,6 @@ def tupled_dimension(array: np.ndarray, axis: int):
     return result.reshape(*shape[:-1])
 
 
-def lexical_comparison(
-    objectives: xr.Dataset,
-    binsize: xr.Dataset,
-    order: Sequence[Hashable] | None = None,
-    bin_last: bool = True,
-) -> xr.DataArray:
-    """Lexical comparison over the objectives.
-
-    Lexical comparison operates by binning the objectives into bins of width
-    `binsize`. Once binned, dimensions other than `asset` and `technology` are
-    reduced by taking the max, e.g. the largest constraint. Finally, the
-    objectives are ranked lexographically, in the order given by the parameters.
-
-    Arguments:
-        objectives: xr.Dataset containing the objectives to rank
-        binsize: bin size, minimization direction
-            (+ -> minimize, - -> maximize), and (optionally) order of
-            lexicographical comparison. The order is the one given
-            `binsize.data_vars` if the argument `order` is None.
-        order: Optional array indicating the order in which to rank the tuples.
-        bin_last: Whether the last metric should be binned, or whether it
-            should be left as a the type it already is (e.g. no flooring and
-            no turning to integer.)
-
-    Result:
-        An array of tuples which can subsequently be compared lexicographically.
-    """
-    if order is None:
-        order = [u for u in binsize.data_vars]
-
-    assert set(order) == set(binsize.data_vars)
-    assert set(order).issuperset(objectives)
-
-    result = objectives[order]
-    for name in order if bin_last else order[:-1]:
-        result[name] = np.floor(result[name] / binsize[name]).astype(int)
-    if not bin_last:
-        result[order[-1]] = result[order[-1]] / binsize[order[-1]]
-    return result.to_array(dim="variable").reduce(tupled_dimension, dim="variable")
-
-
 def merge_assets(
     capa_a: xr.DataArray,
     capa_b: xr.DataArray,

tests/test_decisions.py

@@ -55,17 +55,38 @@ def normalize(objective):
 def test_lexical():
     """Test lexical comparison against hand-constructed tuples."""
     shape = (5, 10)
+
+    # Three objectives evaluated for each asset/replacement pair.
     a = rand(*shape) * 10 - 5
     b = rand(*shape) * 10 - 5
     c = rand(*shape) * 10 - 5
 
-    parameters = [("b", -rand() * 0.1), ("a", rand() * 0.1), ("c", rand())]
+    # Objectives are compared lexicographically in the order
+    # b -> a -> c. Negative weights indicate maximisation.
+    parameters = [
+        ("b", -rand() * 0.1),
+        ("a", rand() * 0.1),
+        ("c", rand()),
+    ]
     param_dict = dict(parameters)
 
+    # Lexo defines bin widths as
+    #
+    #     w_i = min_j(p_i o_i^j),
+    #
+    # over the replacement dimension. For maximisation objectives,
+    # the sign convention means this becomes the negative of the
+    # largest weighted objective.
     mina = (a * param_dict["a"]).min(1)
     minb = -(b * abs(param_dict["b"])).max(1)
     minc = (c * param_dict["c"]).min(1)
 
+    # Construct the expected lexicographic tuples by hand:
+    #
+    #   1. Bin the first two objectives by flooring after
+    #      normalisation by their bin widths.
+    #   2. Leave the final objective continuous so that it acts
+    #      as a tie-breaker.
     expected = np.zeros(shape=shape, dtype=object)
     for i in range(shape[0]):
         for j in range(shape[1]):
@@ -82,10 +103,19 @@ def test_lexical():
             "c": (("asset", "replacement"), c),
         }
     )
-    objectives["asset"] = choice(objectives.replacement, shape[0], replace=False)
+
+    # Associate each asset with one of the replacement options.
+    objectives["asset"] = choice(
+        objectives.replacement,
+        shape[0],
+        replace=False,
+    )
 
     actual = lexical_comparison(objectives, parameters)
     assert actual.shape == expected.shape
+
+    # The decision function returns the ranks of the lexicographic
+    # tuples along the replacement dimension.
     for i in range(shape[0]):
         assert actual.values[i] == approx(rankdata(expected[i]))