SubspaceDiscrete.from_simplex convenience constructor (#117)

This PR adds a new `SubspaceDiscrete.from_simplex` convenience
constructor that is useful for mixture use cases. In particular, it
enables such cases when an API is involved, as the two existing
alternatives will not work here:
* `from_product` in combination with sum constraints suffers from the
exponential blowup of parameter configurations
* `from_dataframe` requires sending excessively large requests to the
API (since the requests need to contain the full enumerated search
space)

Author: AdrianSosic

CHANGELOG.md

@@ -10,6 +10,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - `mypy` for campaign, constraints and telemetry
 - Top-level example summaries
 - `RecommenderProtocol` as common interface for `Strategy` and `Recommender`
+- `SubspaceDiscrete.from_simplex` convenience constructor
 
 ### Changed
 - Order of README sections

baybe/searchspace/discrete.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+from itertools import zip_longest
 from typing import Any, Collection, Iterable, List, Optional, Tuple, cast
 
 import numpy as np
@@ -230,6 +231,133 @@ def discrete_parameter_factory(
 
         return cls(parameters=parameters, exp_rep=df, empty_encoding=empty_encoding)
 
+    @classmethod
+    def from_simplex(
+        cls,
+        max_sum: float,
+        simplex_parameters: List[NumericalDiscreteParameter],
+        product_parameters: Optional[List[DiscreteParameter]] = None,
+        boundary_only: bool = False,
+        tolerance: float = 1e-6,
+    ) -> SubspaceDiscrete:
+        """Efficiently create discrete simplex subspaces.
+
+        The same result can be achieved using
+        :meth:`baybe.searchspace.discrete.SubspaceDiscrete.from_product` in combination
+        with appropriate sum constraints. However, such an approach is inefficient
+        because the Cartesian product involved creates an exponentially large set of
+        candidates, most of which do not satisfy the simplex constraints and must be
+        subsequently be filtered out by the method.
+
+        By contrast, this method uses a shortcut that removes invalid candidates
+        already during the creation of parameter combinations, resulting in a
+        significantly faster construction.
+
+        Args:
+            max_sum: The maximum sum of the parameter values defining the simplex size.
+            simplex_parameters: The parameters to be used for the simplex construction.
+            product_parameters: Optional parameters that enter in form of a Cartesian
+                product.
+            boundary_only: Flag determining whether to keep only parameter
+                configurations on the simplex boundary.
+            tolerance: Numerical tolerance used to validate the simplex constraint.
+
+        Raises:
+            ValueError: If the passed parameters are not suitable for a simplex
+                construction.
+
+        Returns:
+            The created simplex subspace.
+
+        Note:
+            The achieved efficiency gains can vary depending on the particular order in
+            which the parameters are passed to this method, as the configuration space
+            is built up incrementally from the parameter sequence.
+        """
+        if product_parameters is None:
+            product_parameters = []
+
+        # Validate parameter types
+        if not (
+            all(isinstance(p, NumericalDiscreteParameter) for p in simplex_parameters)
+        ):
+            raise ValueError(
+                f"All parameters passed via 'simplex_parameters' "
+                f"must be of type '{NumericalDiscreteParameter.__name__}'."
+            )
+        if not (all(isinstance(p, DiscreteParameter) for p in product_parameters)):
+            raise ValueError(
+                f"All parameters passed via 'product_parameters' "
+                f"must be of subclasses of '{DiscreteParameter.__name__}'."
+            )
+
+        # Construct the product part of the space
+        product_space = parameter_cartesian_prod_to_df(product_parameters)
+        if not simplex_parameters:
+            return cls(parameters=product_parameters, exp_rep=product_space)
+
+        # Validate non-negativity
+        min_values = [min(p.values) for p in simplex_parameters]
+        if not (min(min_values) >= 0.0):
+            raise ValueError(
+                f"All parameters passed to '{cls.from_simplex.__name__}' "
+                f"must have non-negative values only."
+            )
+
+        def drop_invalid(df: pd.DataFrame, max_sum: float, boundary_only: bool) -> None:
+            """Drop rows that violate a specified simplex constraint.
+
+            Args:
+                df: The dataframe whose rows should satisfy the simplex constraint.
+                max_sum: The maximum row sum defining the simplex size.
+                boundary_only: Flag to control if the points represented by the rows
+                    may lie inside the simplex or on its boundary only.
+            """
+            row_sums = df.sum(axis=1)
+            if boundary_only:
+                locs_to_drop = row_sums[
+                    (row_sums < max_sum - tolerance) | (row_sums > max_sum + tolerance)
+                ].index
+            else:
+                locs_to_drop = row_sums[row_sums > max_sum + tolerance].index
+            df.drop(locs_to_drop, inplace=True)
+
+        # Get the minimum sum contributions to come in the upcoming joins (the first
+        # item is the minimum possible sum of all parameters starting from the
+        # second parameter, the second item is the minimum possible sum starting from
+        # the third parameter, and so on ...)
+        min_upcoming = np.cumsum(min_values[:0:-1])[::-1]
+
+        # Incrementally build up the space, dropping invalid configuration along the
+        # way. More specifically: after having cross-joined a new parameter, there must
+        # be enough "room" left for the remaining parameters to fit. Hence,
+        # configurations of the current parameter subset that exceed the desired
+        # total value minus the minimum contribution to come from the yet to be added
+        # parameters can be already discarded.
+        for i, (param, min_to_go) in enumerate(
+            zip_longest(simplex_parameters, min_upcoming, fillvalue=0)
+        ):
+            if i == 0:
+                exp_rep = pd.DataFrame({param.name: param.values})
+            else:
+                exp_rep = pd.merge(
+                    exp_rep, pd.DataFrame({param.name: param.values}), how="cross"
+                )
+            drop_invalid(exp_rep, max_sum - min_to_go, boundary_only=False)
+
+        # If requested, keep only the boundary values
+        if boundary_only:
+            drop_invalid(exp_rep, max_sum, boundary_only=True)
+
+        # Augment the Cartesian product created from all other parameter types
+        if product_parameters:
+            exp_rep = pd.merge(exp_rep, product_space, how="cross")
+
+        # Reset the index
+        exp_rep.reset_index(drop=True, inplace=True)
+
+        return cls(parameters=simplex_parameters, exp_rep=exp_rep)
+
     @property
     def is_empty(self) -> bool:
         """Return whether this subspace is empty."""

tests/hypothesis_strategies/alternative_creation/test_searchspace.py

@@ -1,16 +1,21 @@
 """Test alternative ways of creation not considered in the strategies."""
 
+import hypothesis.strategies as st
+import numpy as np
 import pandas as pd
 import pytest
+from hypothesis import given
 from pytest import param
 
 from baybe.parameters import (
     CategoricalParameter,
     NumericalContinuousParameter,
     NumericalDiscreteParameter,
 )
+from baybe.parameters.categorical import TaskParameter
 from baybe.searchspace import SearchSpace, SubspaceContinuous
 from baybe.searchspace.discrete import SubspaceDiscrete
+from tests.hypothesis_strategies.parameters import numerical_discrete_parameter
 
 # Discrete inputs for testing
 s_x = pd.Series([1, 2, 3], name="x")
@@ -97,3 +102,63 @@ def test_searchspace_creation_from_dataframe(df, parameters, expected):
     else:
         with pytest.raises(expected):
             SearchSpace.from_dataframe(df, parameters)
+
+
+@pytest.mark.parametrize("boundary_only", (False, True))
+@given(
+    parameters=st.lists(
+        numerical_discrete_parameter(min_value=0.0, max_value=1.0),
+        min_size=1,
+        max_size=5,
+        unique_by=lambda x: x.name,
+    )
+)
+def test_discrete_space_creation_from_simplex_inner(parameters, boundary_only):
+    """Candidates from a simplex space satisfy the simplex constraint."""
+    tolerance = 1e-6
+    max_possible = sum(max(p.values) for p in parameters)
+    min_possible = sum(min(p.values) for p in parameters)
+
+    if boundary_only:
+        # Ensure there exists configurations both inside and outside the simplex
+        max_sum = (max_possible + min_possible) / 2
+    else:
+        # We use the maximum parameter sum because it can be exactly achieved (for other
+        # values, except for the minimum, it's not guaranteed there actually exists
+        # a parameter combination that can exactly hit it)
+        max_sum = max_possible
+
+    subspace = SubspaceDiscrete.from_simplex(
+        max_sum, parameters, boundary_only=boundary_only, tolerance=tolerance
+    )
+
+    if boundary_only:
+        assert np.allclose(subspace.exp_rep.sum(axis=1), max_sum, atol=tolerance)
+    else:
+        assert (subspace.exp_rep.sum(axis=1) <= max_sum + tolerance).all()
+
+
+p_d1 = NumericalDiscreteParameter(name="d1", values=[0.0, 0.5, 1.0])
+p_d2 = NumericalDiscreteParameter(name="d2", values=[0.0, 0.5, 1.0])
+p_t1 = TaskParameter(name="t1", values=["A", "B"])
+p_t2 = TaskParameter(name="t2", values=["A", "B"])
+
+
+@pytest.mark.parametrize(
+    ("simplex_parameters", "product_parameters", "n_elements"),
+    [
+        param([p_d1, p_d2], [p_t1, p_t2], 6 * 4, id="both"),
+        param([p_d1, p_d2], [], 6, id="simplex-only"),
+        param([], [p_t1, p_t2], 4, id="task_only"),
+    ],
+)
+def test_discrete_space_creation_from_simplex_mixed(
+    simplex_parameters, product_parameters, n_elements
+):
+    """Additional non-simplex parameters enter in form of a Cartesian product."""
+    max_sum = 1.0
+    subspace = SubspaceDiscrete.from_simplex(
+        max_sum, simplex_parameters, product_parameters, boundary_only=False
+    )
+    assert len(subspace.exp_rep) == n_elements  # <-- (# simplex part) x (# task part)
+    assert not any(subspace.exp_rep.duplicated())

tests/hypothesis_strategies/parameters.py

@@ -1,5 +1,7 @@
 """Hypothesis strategies for parameters."""
 
+from typing import Optional
+
 import hypothesis.strategies as st
 import numpy as np
 from hypothesis.extra.pandas import columns, data_frames
@@ -76,14 +78,18 @@ def custom_descriptors(draw: st.DrawFn):
 @st.composite
 def numerical_discrete_parameter(
     draw: st.DrawFn,
+    min_value: Optional[float] = None,
+    max_value: Optional[float] = None,
 ):
     """Generate :class:`baybe.parameters.numerical.NumericalDiscreteParameter`."""
     name = draw(parameter_name)
     values = draw(
         st.lists(
-            st.one_of(
-                st.integers(),
-                st.floats(allow_infinity=False, allow_nan=False),
+            st.floats(
+                allow_infinity=False,
+                allow_nan=False,
+                min_value=min_value,
+                max_value=max_value,
             ),
             min_size=2,
             unique=True,

tests/serialization/test_searchspace_serialization.py

@@ -4,7 +4,10 @@
 
 import pytest
 
+from baybe.parameters.numerical import NumericalDiscreteParameter
 from baybe.searchspace import SearchSpace
+from baybe.searchspace.discrete import SubspaceDiscrete
+from baybe.serialization.core import converter
 
 
 @pytest.mark.parametrize(
@@ -40,3 +43,25 @@ def test_from_dataframe_deserialization(searchspace):
     )
     deserialized = SearchSpace.from_json(config)
     assert searchspace == deserialized, (searchspace, deserialized)
+
+
+def test_from_simplex_deserialization():
+    """Deserialization from simplex yields back the original object."""
+    parameters = [
+        NumericalDiscreteParameter("p1", [0, 0.5, 1]),
+        NumericalDiscreteParameter("p2", [0, 0.5, 1]),
+    ]
+    max_sum = 1.0
+    subspace = SubspaceDiscrete.from_simplex(max_sum, simplex_parameters=parameters)
+    parameters_string = json.dumps(converter.unstructure(parameters))
+    config = """
+    {
+        "constructor": "from_simplex",
+        "max_sum": __fillin__max_sum__,
+        "simplex_parameters": __fillin__parameters__
+    }
+    """.replace("__fillin__max_sum__", str(max_sum)).replace(
+        "__fillin__parameters__", parameters_string
+    )
+    deserialized = SubspaceDiscrete.from_json(config)
+    assert subspace == deserialized, (subspace, deserialized)