From 3e2bdef9fc60f7125f04b98024793a6f7601a349 Mon Sep 17 00:00:00 2001 From: kalama-ai Date: Tue, 2 Jun 2026 18:51:03 +0200 Subject: [PATCH 1/7] Add _ReducedSearchSpace for parameter-only reduced search paces - Introduce _ReducedSearchSpace subclass with subset of accessible attribute - Add SearchSpace._without_parameters() factory method - Override parameters, parameter_names, comp_rep_columns, task_idx - Add TODO in GaussianProcessSurrogate._fit for future dispatching - Add tests including kernel product integration test --- baybe/searchspace/core.py | 107 +++++++++++++++++++- baybe/surrogates/gaussian_process/core.py | 2 + tests/test_searchspace.py | 114 ++++++++++++++++++++++ 3 files changed, 221 insertions(+), 2 deletions(-) diff --git a/baybe/searchspace/core.py b/baybe/searchspace/core.py index 8d75cae927..e953436d21 100644 --- a/baybe/searchspace/core.py +++ b/baybe/searchspace/core.py @@ -3,10 +3,10 @@ from __future__ import annotations import gc -from collections.abc import Iterable, Iterator, Sequence +from collections.abc import Collection, Iterable, Iterator, Sequence from enum import Enum from itertools import product -from typing import TYPE_CHECKING, cast +from typing import TYPE_CHECKING, ClassVar, cast import numpy as np import numpy.typing as npt @@ -515,6 +515,109 @@ def get_parameters_by_name(self, names: Sequence[str]) -> tuple[Parameter, ...]: names ) + self.continuous.get_parameters_by_name(names) + def _without_parameters(self, names: Collection[str], /) -> _ReducedSearchSpace: + """Return a reduced search space without the named parameters. + + The returned object exposes only parameter information and blocks + access to constraints, subspaces, and transformation. + + Args: + names: The names of the parameters to remove. + + Raises: + ValueError: If any name does not match a parameter in the space. + + Returns: + A reduced search space containing only parameter information. + """ + current_names = {p.name for p in self.parameters} + unknown = set(names) - current_names + if unknown: + raise ValueError( + f"Parameter name(s) {unknown} not found in the search space. " + f"Available: {current_names}." + ) + remaining = tuple(p for p in self.parameters if p.name not in set(names)) + return _ReducedSearchSpace(parameters=remaining) + + +@define(slots=False) +class _ReducedSearchSpace(SearchSpace): + """A lightweight search space exposing only parameter information. + + Provides access to parameter-related properties needed by kernel factory + calls. Blocks access to constraints, subspaces, transformation, and other + functionality requiring the full search space internals. + + This class is not intended for direct construction. Use + :meth:`SearchSpace._without_parameters` instead. + """ + + _parameters: tuple[Parameter, ...] = field(kw_only=True, alias="parameters") + """The parameters of the reduced search space.""" + + _ALLOWED_ATTRIBUTES: ClassVar[frozenset[str]] = frozenset( + { + "parameters", + "parameter_names", + "comp_rep_columns", + "_task_parameter", + "task_idx", + "n_tasks", + "get_comp_rep_parameter_indices", + "_parameters", + "_ALLOWED_ATTRIBUTES", + } + ) + """Attributes accessible on this reduced search space.""" + + @override + def __attrs_post_init__(self): + """Skip parent validation.""" + + @override + def __getattribute__(self, name: str): + """Guard attribute access, allowing only parameter-related attributes.""" + if name.startswith("__"): + return object.__getattribute__(self, name) + allowed = object.__getattribute__(self, "_ALLOWED_ATTRIBUTES") + if name in allowed: + return object.__getattribute__(self, name) + raise NotImplementedError( + f"'{object.__getattribute__(self, '__class__').__name__}' does not " + f"support attribute '{name}'. Only parameter information is available." + ) + + @override + @property + def comp_rep_columns(self) -> tuple[str, ...]: + """The columns spanning the computational representation.""" + return tuple( + col + for p in object.__getattribute__(self, "_parameters") + for col in p.comp_rep_columns + ) + + @override + @property + def parameter_names(self) -> tuple[str, ...]: + """Return tuple of parameter names.""" + return tuple(p.name for p in object.__getattribute__(self, "_parameters")) + + @override + @property + def parameters(self) -> tuple[Parameter, ...]: + """Return the parameters of the reduced search space.""" + return object.__getattribute__(self, "_parameters") + + @override + @property + def task_idx(self) -> int | None: + """The column index of the task parameter in computational representation.""" + if (task_param := self._task_parameter) is None: + return None + return self.comp_rep_columns.index(task_param.name) + def to_searchspace( x: Parameter | SubspaceDiscrete | SubspaceContinuous | SearchSpace, / diff --git a/baybe/surrogates/gaussian_process/core.py b/baybe/surrogates/gaussian_process/core.py index cec34a32a5..4c1eeb93f0 100644 --- a/baybe/surrogates/gaussian_process/core.py +++ b/baybe/surrogates/gaussian_process/core.py @@ -311,6 +311,8 @@ def _fit(self, train_x: Tensor, train_y: Tensor) -> None: ) ### Kernel + # TODO: When calling a factory on a `_ReducedSearchSpace`, validate that + # it returns a BayBE Kernel (not a raw gpytorch kernel). kernel = self.kernel_factory( context.searchspace, context.objective, context.measurements ) diff --git a/tests/test_searchspace.py b/tests/test_searchspace.py index 0579a82d44..0be4beb405 100644 --- a/tests/test_searchspace.py +++ b/tests/test_searchspace.py @@ -533,3 +533,117 @@ def test_sample_from_polytope_mixed_constraints_with_interpoint(): # Verify interpoint constraint is satisfied across the batch interpoint_constraint_result = samples["Conti_finite1"].sum() assert np.isclose(interpoint_constraint_result, 0.6, atol=1e-6) + + +class TestReducedSearchSpace: + """Tests for _ReducedSearchSpace.""" + + @pytest.fixture + def searchspace(self): + """Create a hybrid search space with a task parameter.""" + return SearchSpace.from_product( + [ + CategoricalParameter("Color", ["red", "blue"]), + NumericalContinuousParameter("x", (0.0, 1.0)), + TaskParameter("Task", ["A", "B"]), + ] + ) + + def test_parameters(self, searchspace): + """Verify that the reduced space exposes only the remaining parameters.""" + reduced = searchspace._without_parameters({"Task"}) + names = {p.name for p in reduced.parameters} + assert names == {"Color", "x"} + + def test_parameter_names(self, searchspace): + """Verify that parameter_names derives from the remaining parameters.""" + reduced = searchspace._without_parameters({"Task"}) + assert set(reduced.parameter_names) == {"Color", "x"} + + def test_comp_rep_columns(self, searchspace): + """Verify that comp_rep_columns matches the full space minus the removed one.""" + reduced = searchspace._without_parameters({"Task"}) + task_param = next(p for p in searchspace.parameters if p.name == "Task") + expected = tuple( + c + for c in searchspace.comp_rep_columns + if c not in task_param.comp_rep_columns + ) + assert reduced.comp_rep_columns == expected + + def test_task_idx_is_none(self, searchspace): + """Verify that task_idx is None when the task parameter is removed.""" + reduced = searchspace._without_parameters({"Task"}) + assert reduced.task_idx is None + + def test_n_tasks_is_one(self, searchspace): + """Verify that n_tasks is 1 when no task parameter is present.""" + reduced = searchspace._without_parameters({"Task"}) + assert reduced.n_tasks == 1 + + def test_get_comp_rep_parameter_indices(self, searchspace): + """Verify that index resolution works for remaining parameters.""" + reduced = searchspace._without_parameters({"Task"}) + indices = reduced.get_comp_rep_parameter_indices("x") + assert len(indices) > 0 + + def test_blocked_attributes(self, searchspace): + """Verify that accessing non-parameter attributes raises an error.""" + reduced = searchspace._without_parameters({"Task"}) + with pytest.raises(NotImplementedError): + reduced.transform + with pytest.raises(NotImplementedError): + reduced.constraints + with pytest.raises(NotImplementedError): + reduced.discrete + with pytest.raises(NotImplementedError): + reduced.continuous + + def test_unknown_parameter_raises(self, searchspace): + """Verify that removing a nonexistent parameter raises an error.""" + with pytest.raises(ValueError, match="not found"): + searchspace._without_parameters({"nonexistent"}) + + def test_multiple_parameters_removed(self, searchspace): + """Verify that multiple parameters can be removed at once.""" + reduced = searchspace._without_parameters({"Task", "Color"}) + assert reduced.parameter_names == ("x",) + + def test_kernel_product_matches_default_factory(self): + """Verify that manual kernel split matches BayBE's default factory.""" + import torch + + from baybe.objectives.single import SingleTargetObjective + from baybe.surrogates.gaussian_process.presets.baybe import ( + BayBEKernelFactory, + _BayBENumericalKernelFactory, + _BayBETaskKernelFactory, + ) + from baybe.targets import NumericalTarget + + task_param = TaskParameter("Task", ["A", "B", "C"]) + cat_param = CategoricalParameter("Color", ["red", "blue", "green"]) + cont_param = NumericalContinuousParameter("x", (0.0, 1.0)) + searchspace = SearchSpace.from_product([cat_param, cont_param, task_param]) + + objective = SingleTargetObjective(NumericalTarget("y")) + measurements = pd.DataFrame( + {"Color": ["red"], "x": [0.5], "Task": ["A"], "y": [1.0]} + ) + + kernel_default = BayBEKernelFactory()(searchspace, objective, measurements) + + reduced_ss = searchspace._without_parameters({"Task"}) + base_baybe = _BayBENumericalKernelFactory()(reduced_ss, objective, measurements) + base_gpytorch = base_baybe.to_gpytorch(searchspace) + + task_only_ss = SearchSpace.from_product([task_param]) + task_baybe = _BayBETaskKernelFactory()(task_only_ss, objective, measurements) + task_gpytorch = task_baybe.to_gpytorch(searchspace) + + kernel_manual = base_gpytorch * task_gpytorch + + assert type(kernel_default) is type(kernel_manual) + for k_def, k_man in zip(kernel_default.kernels, kernel_manual.kernels): + assert type(k_def) is type(k_man) + assert torch.equal(k_def.active_dims, k_man.active_dims) From c0289f0d5315f024c61e7395e5b93f5a25a2710d Mon Sep 17 00:00:00 2001 From: kalama-ai Date: Wed, 3 Jun 2026 16:08:11 +0200 Subject: [PATCH 2/7] Remove index properties from `_ReducedSearchSpace` - Replace `task_idx` sentinel checks with `n_tasks == 1 / > 1` across kernel factories and fit criterion (equivalent because `TaskParameter` enforces min 2 values) - Add `SearchSpace._get_n_comp_rep_columns(selector)` for counting number of indices in comp representation - Remove `task_idx` and `get_comp_rep_parameter_indices` from the reduced searchspace allowlist since these return integer indices that are only meaningful relative to the reduced space, not the full training tensor - Override `_get_n_comp_rep_columns` on `_ReducedSearchSpace` to avoid calling the blocked `get_comp_rep_parameter_indices` --- baybe/searchspace/core.py | 42 +++++++++++++++---- .../components/fit_criterion.py | 2 +- .../gaussian_process/components/kernel.py | 10 ++--- .../gaussian_process/presets/baybe.py | 2 +- tests/test_searchspace.py | 17 ++++---- 5 files changed, 48 insertions(+), 25 deletions(-) diff --git a/baybe/searchspace/core.py b/baybe/searchspace/core.py index e953436d21..8801c0465e 100644 --- a/baybe/searchspace/core.py +++ b/baybe/searchspace/core.py @@ -433,6 +433,23 @@ def get_comp_rep_parameter_indices( if col in p.comp_rep_columns ) + def _get_n_comp_rep_columns( + self, + name_or_selector: str | ParameterSelectorProtocol, + /, + ) -> int: + """Get the number of comp-rep columns for a parameter selection. + + Args: + name_or_selector: Either the name of a single parameter or a selector + that filters parameters to be included. + + Returns: + The number of columns in the computational representation associated + with the selected parameter(s). + """ + return len(self.get_comp_rep_parameter_indices(name_or_selector)) + @staticmethod def estimate_product_space_size(parameters: Iterable[Parameter]) -> MemorySize: """Estimate an upper bound for the memory size of a product space. @@ -562,9 +579,8 @@ class _ReducedSearchSpace(SearchSpace): "parameter_names", "comp_rep_columns", "_task_parameter", - "task_idx", "n_tasks", - "get_comp_rep_parameter_indices", + "_get_n_comp_rep_columns", "_parameters", "_ALLOWED_ATTRIBUTES", } @@ -611,12 +627,22 @@ def parameters(self) -> tuple[Parameter, ...]: return object.__getattribute__(self, "_parameters") @override - @property - def task_idx(self) -> int | None: - """The column index of the task parameter in computational representation.""" - if (task_param := self._task_parameter) is None: - return None - return self.comp_rep_columns.index(task_param.name) + def _get_n_comp_rep_columns( + self, + name_or_selector: str | ParameterSelectorProtocol, + /, + ) -> int: + """Get the number of comp-rep columns for a parameter selection.""" + if isinstance(name_or_selector, str): + params = [p for p in self.parameters if p.name == name_or_selector] + else: + params = [p for p in self.parameters if name_or_selector(p)] + return sum( + 1 + for p in params + for col in p.comp_rep_columns + if col in self.comp_rep_columns + ) def to_searchspace( diff --git a/baybe/surrogates/gaussian_process/components/fit_criterion.py b/baybe/surrogates/gaussian_process/components/fit_criterion.py index 40bad24188..3910864b86 100644 --- a/baybe/surrogates/gaussian_process/components/fit_criterion.py +++ b/baybe/surrogates/gaussian_process/components/fit_criterion.py @@ -66,7 +66,7 @@ class _MLLForNonTLFitCriterionFactory(FitCriterionFactoryProtocol): def __call__( self, searchspace: SearchSpace, objective: Objective, measurements: pd.DataFrame ) -> FitCriterion: - if searchspace.task_idx is None: + if searchspace.n_tasks == 1: return FitCriterion.MARGINAL_LOG_LIKELIHOOD from baybe.surrogates.gaussian_process.presets.baybe import ( diff --git a/baybe/surrogates/gaussian_process/components/kernel.py b/baybe/surrogates/gaussian_process/components/kernel.py index 1ea78aae58..e21b7bce67 100644 --- a/baybe/surrogates/gaussian_process/components/kernel.py +++ b/baybe/surrogates/gaussian_process/components/kernel.py @@ -83,10 +83,8 @@ def get_parameter_names(self, searchspace: SearchSpace) -> tuple[str, ...]: def _get_effective_dimensionality(self, searchspace: SearchSpace) -> int: """Get the number of computational columns for the selected parameters.""" - return len( - searchspace.get_comp_rep_parameter_indices( - self.parameter_selector or (lambda _: True) - ) + return searchspace._get_n_comp_rep_columns( + self.parameter_selector or (lambda _: True) ) def _validate_parameter_kinds(self, parameters: Iterable[Parameter]) -> None: @@ -213,7 +211,7 @@ def __call__( target_cls._supported_parameter_kinds = broadened_kinds self.parameter_selector = original_selector - if searchspace.task_idx is not None: + if searchspace.n_tasks > 1: icm = ICMKernelFactory(base_kernel_or_factory=base_kernel) return icm(searchspace, objective, measurements) return base_kernel @@ -301,7 +299,7 @@ def _validate_task_kernel_factory(self, _, factory: KernelFactoryProtocol): def __call__( self, searchspace: SearchSpace, objective: Objective, measurements: pd.DataFrame ) -> Kernel | GPyTorchKernel: - if searchspace.task_idx is None: + if searchspace.n_tasks == 1: raise IncompatibleSearchSpaceError( f"'{type(self).__name__}' can only be used with a searchspace that " f"contains a '{TaskParameter.__name__}'." diff --git a/baybe/surrogates/gaussian_process/presets/baybe.py b/baybe/surrogates/gaussian_process/presets/baybe.py index 889aafe87d..f68717f6dd 100644 --- a/baybe/surrogates/gaussian_process/presets/baybe.py +++ b/baybe/surrogates/gaussian_process/presets/baybe.py @@ -276,7 +276,7 @@ def __call__( ) -> FitCriterion: return ( FitCriterion.MARGINAL_LOG_LIKELIHOOD - if searchspace.task_idx is None + if searchspace.n_tasks == 1 else FitCriterion.LEAVE_ONE_OUT_PSEUDOLIKELIHOOD ) diff --git a/tests/test_searchspace.py b/tests/test_searchspace.py index 0be4beb405..9c4fbca151 100644 --- a/tests/test_searchspace.py +++ b/tests/test_searchspace.py @@ -571,21 +571,16 @@ def test_comp_rep_columns(self, searchspace): ) assert reduced.comp_rep_columns == expected - def test_task_idx_is_none(self, searchspace): - """Verify that task_idx is None when the task parameter is removed.""" - reduced = searchspace._without_parameters({"Task"}) - assert reduced.task_idx is None - def test_n_tasks_is_one(self, searchspace): """Verify that n_tasks is 1 when no task parameter is present.""" reduced = searchspace._without_parameters({"Task"}) assert reduced.n_tasks == 1 - def test_get_comp_rep_parameter_indices(self, searchspace): - """Verify that index resolution works for remaining parameters.""" + def test_get_n_comp_rep_columns(self, searchspace): + """Verify that _get_n_comp_rep_columns works on the reduced space.""" reduced = searchspace._without_parameters({"Task"}) - indices = reduced.get_comp_rep_parameter_indices("x") - assert len(indices) > 0 + assert reduced._get_n_comp_rep_columns("x") == 1 + assert reduced._get_n_comp_rep_columns("Color") == 2 def test_blocked_attributes(self, searchspace): """Verify that accessing non-parameter attributes raises an error.""" @@ -598,6 +593,10 @@ def test_blocked_attributes(self, searchspace): reduced.discrete with pytest.raises(NotImplementedError): reduced.continuous + with pytest.raises(NotImplementedError): + reduced.task_idx + with pytest.raises(NotImplementedError): + reduced.get_comp_rep_parameter_indices("x") def test_unknown_parameter_raises(self, searchspace): """Verify that removing a nonexistent parameter raises an error.""" From c2c3f7231d6371212f4980df64ada80cf8520ac5 Mon Sep 17 00:00:00 2001 From: kalama-ai Date: Tue, 9 Jun 2026 13:15:09 +0200 Subject: [PATCH 3/7] Simplify counting of comp_rep_columns Co-authored-by: AdrianSosic --- baybe/searchspace/core.py | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) diff --git a/baybe/searchspace/core.py b/baybe/searchspace/core.py index 8801c0465e..5fdc134a5a 100644 --- a/baybe/searchspace/core.py +++ b/baybe/searchspace/core.py @@ -637,12 +637,7 @@ def _get_n_comp_rep_columns( params = [p for p in self.parameters if p.name == name_or_selector] else: params = [p for p in self.parameters if name_or_selector(p)] - return sum( - 1 - for p in params - for col in p.comp_rep_columns - if col in self.comp_rep_columns - ) + return sum(len(p.comp_rep_columns) for p in params) def to_searchspace( From 2a4d7f539298dd4d8a2f8ce38573421984cdedc0 Mon Sep 17 00:00:00 2001 From: kalama-ai Date: Tue, 9 Jun 2026 13:15:36 +0200 Subject: [PATCH 4/7] Simplify if statement Co-authored-by: AdrianSosic --- baybe/searchspace/core.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/baybe/searchspace/core.py b/baybe/searchspace/core.py index 5fdc134a5a..0a2a1e60f8 100644 --- a/baybe/searchspace/core.py +++ b/baybe/searchspace/core.py @@ -548,8 +548,7 @@ def _without_parameters(self, names: Collection[str], /) -> _ReducedSearchSpace: A reduced search space containing only parameter information. """ current_names = {p.name for p in self.parameters} - unknown = set(names) - current_names - if unknown: + if unknown := set(names) - current_names: raise ValueError( f"Parameter name(s) {unknown} not found in the search space. " f"Available: {current_names}." From dae8b3c7430962a1fe0926c01bb43194afc57c70 Mon Sep 17 00:00:00 2001 From: kalama-ai Date: Tue, 9 Jun 2026 13:26:35 +0200 Subject: [PATCH 5/7] Rename _without_parameters() to _drop_parameters() --- baybe/searchspace/core.py | 6 +++--- tests/test_searchspace.py | 18 +++++++++--------- 2 files changed, 12 insertions(+), 12 deletions(-) diff --git a/baybe/searchspace/core.py b/baybe/searchspace/core.py index 0a2a1e60f8..d9c9738f93 100644 --- a/baybe/searchspace/core.py +++ b/baybe/searchspace/core.py @@ -532,7 +532,7 @@ def get_parameters_by_name(self, names: Sequence[str]) -> tuple[Parameter, ...]: names ) + self.continuous.get_parameters_by_name(names) - def _without_parameters(self, names: Collection[str], /) -> _ReducedSearchSpace: + def _drop_parameters(self, names: Collection[str], /) -> _ReducedSearchSpace: """Return a reduced search space without the named parameters. The returned object exposes only parameter information and blocks @@ -548,7 +548,7 @@ def _without_parameters(self, names: Collection[str], /) -> _ReducedSearchSpace: A reduced search space containing only parameter information. """ current_names = {p.name for p in self.parameters} - if unknown := set(names) - current_names: + if unknown := set(names) - current_names: raise ValueError( f"Parameter name(s) {unknown} not found in the search space. " f"Available: {current_names}." @@ -566,7 +566,7 @@ class _ReducedSearchSpace(SearchSpace): functionality requiring the full search space internals. This class is not intended for direct construction. Use - :meth:`SearchSpace._without_parameters` instead. + :meth:`SearchSpace._drop_parameters` instead. """ _parameters: tuple[Parameter, ...] = field(kw_only=True, alias="parameters") diff --git a/tests/test_searchspace.py b/tests/test_searchspace.py index 9c4fbca151..facff1b1ec 100644 --- a/tests/test_searchspace.py +++ b/tests/test_searchspace.py @@ -551,18 +551,18 @@ def searchspace(self): def test_parameters(self, searchspace): """Verify that the reduced space exposes only the remaining parameters.""" - reduced = searchspace._without_parameters({"Task"}) + reduced = searchspace._drop_parameters({"Task"}) names = {p.name for p in reduced.parameters} assert names == {"Color", "x"} def test_parameter_names(self, searchspace): """Verify that parameter_names derives from the remaining parameters.""" - reduced = searchspace._without_parameters({"Task"}) + reduced = searchspace._drop_parameters({"Task"}) assert set(reduced.parameter_names) == {"Color", "x"} def test_comp_rep_columns(self, searchspace): """Verify that comp_rep_columns matches the full space minus the removed one.""" - reduced = searchspace._without_parameters({"Task"}) + reduced = searchspace._drop_parameters({"Task"}) task_param = next(p for p in searchspace.parameters if p.name == "Task") expected = tuple( c @@ -573,18 +573,18 @@ def test_comp_rep_columns(self, searchspace): def test_n_tasks_is_one(self, searchspace): """Verify that n_tasks is 1 when no task parameter is present.""" - reduced = searchspace._without_parameters({"Task"}) + reduced = searchspace._drop_parameters({"Task"}) assert reduced.n_tasks == 1 def test_get_n_comp_rep_columns(self, searchspace): """Verify that _get_n_comp_rep_columns works on the reduced space.""" - reduced = searchspace._without_parameters({"Task"}) + reduced = searchspace._drop_parameters({"Task"}) assert reduced._get_n_comp_rep_columns("x") == 1 assert reduced._get_n_comp_rep_columns("Color") == 2 def test_blocked_attributes(self, searchspace): """Verify that accessing non-parameter attributes raises an error.""" - reduced = searchspace._without_parameters({"Task"}) + reduced = searchspace._drop_parameters({"Task"}) with pytest.raises(NotImplementedError): reduced.transform with pytest.raises(NotImplementedError): @@ -601,11 +601,11 @@ def test_blocked_attributes(self, searchspace): def test_unknown_parameter_raises(self, searchspace): """Verify that removing a nonexistent parameter raises an error.""" with pytest.raises(ValueError, match="not found"): - searchspace._without_parameters({"nonexistent"}) + searchspace._drop_parameters({"nonexistent"}) def test_multiple_parameters_removed(self, searchspace): """Verify that multiple parameters can be removed at once.""" - reduced = searchspace._without_parameters({"Task", "Color"}) + reduced = searchspace._drop_parameters({"Task", "Color"}) assert reduced.parameter_names == ("x",) def test_kernel_product_matches_default_factory(self): @@ -632,7 +632,7 @@ def test_kernel_product_matches_default_factory(self): kernel_default = BayBEKernelFactory()(searchspace, objective, measurements) - reduced_ss = searchspace._without_parameters({"Task"}) + reduced_ss = searchspace._drop_parameters({"Task"}) base_baybe = _BayBENumericalKernelFactory()(reduced_ss, objective, measurements) base_gpytorch = base_baybe.to_gpytorch(searchspace) From fa31cc48da5cc8db39518f6a087ef49989d23519 Mon Sep 17 00:00:00 2001 From: kalama-ai Date: Thu, 11 Jun 2026 14:26:22 +0200 Subject: [PATCH 6/7] =?UTF-8?q?=C2=A0=20Address=20reviewer=20comments:?= =?UTF-8?q?=C2=A0=20=C2=A0=20=C2=A0=20=C2=A0=20=C2=A0=20=C2=A0=20=C2=A0=20?= =?UTF-8?q?=C2=A0=20=C2=A0=20=C2=A0=20=C2=A0=20=C2=A0=20=C2=A0=20=C2=A0=20?= =?UTF-8?q?=C2=A0=20=C2=A0=20=C2=A0=20=C2=A0=20=C2=A0=20=C2=A0=20=C2=A0=20?= =?UTF-8?q?=C2=A0=20=C2=A0=20=C2=A0=20=C2=A0=20=C2=A0=20=C2=A0=20=C2=A0=20?= =?UTF-8?q?=C2=A0=20=C2=A0=20=C2=A0=20=C2=A0=20=C2=A0=20=C2=A0=20=C2=A0=20?= =?UTF-8?q?=C2=A0=20=C2=A0=20-=20Build=20real=20SubspaceDiscrete/SubspaceC?= =?UTF-8?q?ontinuous=20objects=20with=20zero=20rows=20instead=20of=20stori?= =?UTF-8?q?ng=20a=20flat=20=5Fparameters=20tuple=20=C2=A0=20-=20Remove=20?= =?UTF-8?q?=5Fparameters=20field,=20all=20property=20overrides,=20=5F=5Fat?= =?UTF-8?q?trs=5Fpost=5Finit,=20override,=20and=20=5Fget=5Fn=5Fcomp=5Frep?= =?UTF-8?q?=5Fcolumns=20override=C2=A0=20=C2=A0=20=C2=A0=20=C2=A0=20=C2=A0?= =?UTF-8?q?=20=C2=A0=20=C2=A0=20=C2=A0=20=C2=A0=20=C2=A0=20=C2=A0=20=C2=A0?= =?UTF-8?q?=20=C2=A0=20=C2=A0=20=C2=A0=20=C2=A0=20=C2=A0=20=C2=A0=20=C2=A0?= =?UTF-8?q?=20=C2=A0=20=C2=A0=20=C2=A0=20=C2=A0=20=C2=A0=20=C2=A0=20=C2=A0?= =?UTF-8?q?=20=C2=A0=20=C2=A0=20=C2=A0=20=C2=A0=20=C2=A0=20=C2=A0=20=C2=A0?= =?UTF-8?q?=20=C2=A0=20=C2=A0=20=C2=A0=20=C2=A0=20-=20Convert=20TestReduce?= =?UTF-8?q?dSearchSpace=20class=20to=20standalone=20test=20functions=20?= =?UTF-8?q?=C2=A0=20-=20Add=20programmatic=20test=20that=20discovers=20all?= =?UTF-8?q?=20SearchSpace=20attributes=20and=20verifies=20non-allowed=20on?= =?UTF-8?q?es=20raise=20NotImplementedError?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- baybe/searchspace/core.py | 81 ++++++------ tests/test_searchspace.py | 251 ++++++++++++++++++++++---------------- 2 files changed, 182 insertions(+), 150 deletions(-) diff --git a/baybe/searchspace/core.py b/baybe/searchspace/core.py index d9c9738f93..0b614b83ea 100644 --- a/baybe/searchspace/core.py +++ b/baybe/searchspace/core.py @@ -553,8 +553,33 @@ def _drop_parameters(self, names: Collection[str], /) -> _ReducedSearchSpace: f"Parameter name(s) {unknown} not found in the search space. " f"Available: {current_names}." ) - remaining = tuple(p for p in self.parameters if p.name not in set(names)) - return _ReducedSearchSpace(parameters=remaining) + remaining = [p for p in self.parameters if p.name not in set(names)] + + disc_params = [p for p in remaining if p.is_discrete] + cont_params = [p for p in remaining if p.is_continuous] + + # Explicit comp_rep needed because transform() drops columns for empty inputs. + discrete = ( + SubspaceDiscrete( + parameters=disc_params, + exp_rep=pd.DataFrame(columns=[p.name for p in disc_params]), + comp_rep=pd.DataFrame( + columns=[c for p in disc_params for c in p.comp_rep_columns] + ), + ) + if disc_params + else SubspaceDiscrete.empty() + ) + + continuous = ( + SubspaceContinuous( + parameters=cont_params, + ) + if cont_params + else SubspaceContinuous.empty() + ) + + return _ReducedSearchSpace(discrete=discrete, continuous=continuous) @define(slots=False) @@ -562,34 +587,31 @@ class _ReducedSearchSpace(SearchSpace): """A lightweight search space exposing only parameter information. Provides access to parameter-related properties needed by kernel factory - calls. Blocks access to constraints, subspaces, transformation, and other - functionality requiring the full search space internals. + calls. Blocks access to transformation, index-based lookups, and other + functionality requiring actual candidate data. This class is not intended for direct construction. Use :meth:`SearchSpace._drop_parameters` instead. """ - _parameters: tuple[Parameter, ...] = field(kw_only=True, alias="parameters") - """The parameters of the reduced search space.""" - _ALLOWED_ATTRIBUTES: ClassVar[frozenset[str]] = frozenset( { + "discrete", + "continuous", "parameters", "parameter_names", "comp_rep_columns", + "constraints", + "type", "_task_parameter", "n_tasks", "_get_n_comp_rep_columns", - "_parameters", + "get_comp_rep_parameter_indices", "_ALLOWED_ATTRIBUTES", } ) """Attributes accessible on this reduced search space.""" - @override - def __attrs_post_init__(self): - """Skip parent validation.""" - @override def __getattribute__(self, name: str): """Guard attribute access, allowing only parameter-related attributes.""" @@ -603,41 +625,6 @@ def __getattribute__(self, name: str): f"support attribute '{name}'. Only parameter information is available." ) - @override - @property - def comp_rep_columns(self) -> tuple[str, ...]: - """The columns spanning the computational representation.""" - return tuple( - col - for p in object.__getattribute__(self, "_parameters") - for col in p.comp_rep_columns - ) - - @override - @property - def parameter_names(self) -> tuple[str, ...]: - """Return tuple of parameter names.""" - return tuple(p.name for p in object.__getattribute__(self, "_parameters")) - - @override - @property - def parameters(self) -> tuple[Parameter, ...]: - """Return the parameters of the reduced search space.""" - return object.__getattribute__(self, "_parameters") - - @override - def _get_n_comp_rep_columns( - self, - name_or_selector: str | ParameterSelectorProtocol, - /, - ) -> int: - """Get the number of comp-rep columns for a parameter selection.""" - if isinstance(name_or_selector, str): - params = [p for p in self.parameters if p.name == name_or_selector] - else: - params = [p for p in self.parameters if name_or_selector(p)] - return sum(len(p.comp_rep_columns) for p in params) - def to_searchspace( x: Parameter | SubspaceDiscrete | SubspaceContinuous | SearchSpace, / diff --git a/tests/test_searchspace.py b/tests/test_searchspace.py index facff1b1ec..f4111e04c1 100644 --- a/tests/test_searchspace.py +++ b/tests/test_searchspace.py @@ -535,114 +535,159 @@ def test_sample_from_polytope_mixed_constraints_with_interpoint(): assert np.isclose(interpoint_constraint_result, 0.6, atol=1e-6) -class TestReducedSearchSpace: - """Tests for _ReducedSearchSpace.""" - - @pytest.fixture - def searchspace(self): - """Create a hybrid search space with a task parameter.""" - return SearchSpace.from_product( - [ - CategoricalParameter("Color", ["red", "blue"]), - NumericalContinuousParameter("x", (0.0, 1.0)), - TaskParameter("Task", ["A", "B"]), - ] - ) +@pytest.fixture(name="reduced_searchspace") +def fixture_reduced_searchspace(): + """A reduced search space with the task parameter removed.""" + ss = SearchSpace.from_product( + [ + CategoricalParameter("Color", ["red", "blue"]), + NumericalContinuousParameter("x", (0.0, 1.0)), + TaskParameter("Task", ["A", "B"]), + ] + ) + return ss._drop_parameters({"Task"}) - def test_parameters(self, searchspace): - """Verify that the reduced space exposes only the remaining parameters.""" - reduced = searchspace._drop_parameters({"Task"}) - names = {p.name for p in reduced.parameters} - assert names == {"Color", "x"} - - def test_parameter_names(self, searchspace): - """Verify that parameter_names derives from the remaining parameters.""" - reduced = searchspace._drop_parameters({"Task"}) - assert set(reduced.parameter_names) == {"Color", "x"} - - def test_comp_rep_columns(self, searchspace): - """Verify that comp_rep_columns matches the full space minus the removed one.""" - reduced = searchspace._drop_parameters({"Task"}) - task_param = next(p for p in searchspace.parameters if p.name == "Task") - expected = tuple( - c - for c in searchspace.comp_rep_columns - if c not in task_param.comp_rep_columns - ) - assert reduced.comp_rep_columns == expected - - def test_n_tasks_is_one(self, searchspace): - """Verify that n_tasks is 1 when no task parameter is present.""" - reduced = searchspace._drop_parameters({"Task"}) - assert reduced.n_tasks == 1 - - def test_get_n_comp_rep_columns(self, searchspace): - """Verify that _get_n_comp_rep_columns works on the reduced space.""" - reduced = searchspace._drop_parameters({"Task"}) - assert reduced._get_n_comp_rep_columns("x") == 1 - assert reduced._get_n_comp_rep_columns("Color") == 2 - - def test_blocked_attributes(self, searchspace): - """Verify that accessing non-parameter attributes raises an error.""" - reduced = searchspace._drop_parameters({"Task"}) - with pytest.raises(NotImplementedError): - reduced.transform - with pytest.raises(NotImplementedError): - reduced.constraints - with pytest.raises(NotImplementedError): - reduced.discrete - with pytest.raises(NotImplementedError): - reduced.continuous - with pytest.raises(NotImplementedError): - reduced.task_idx - with pytest.raises(NotImplementedError): - reduced.get_comp_rep_parameter_indices("x") - - def test_unknown_parameter_raises(self, searchspace): - """Verify that removing a nonexistent parameter raises an error.""" - with pytest.raises(ValueError, match="not found"): - searchspace._drop_parameters({"nonexistent"}) - - def test_multiple_parameters_removed(self, searchspace): - """Verify that multiple parameters can be removed at once.""" - reduced = searchspace._drop_parameters({"Task", "Color"}) - assert reduced.parameter_names == ("x",) - - def test_kernel_product_matches_default_factory(self): - """Verify that manual kernel split matches BayBE's default factory.""" - import torch - - from baybe.objectives.single import SingleTargetObjective - from baybe.surrogates.gaussian_process.presets.baybe import ( - BayBEKernelFactory, - _BayBENumericalKernelFactory, - _BayBETaskKernelFactory, - ) - from baybe.targets import NumericalTarget - task_param = TaskParameter("Task", ["A", "B", "C"]) - cat_param = CategoricalParameter("Color", ["red", "blue", "green"]) - cont_param = NumericalContinuousParameter("x", (0.0, 1.0)) - searchspace = SearchSpace.from_product([cat_param, cont_param, task_param]) +def test_reduced_parameters(reduced_searchspace): + """Verify that the reduced space exposes only the remaining parameters.""" + names = {p.name for p in reduced_searchspace.parameters} + assert names == {"Color", "x"} - objective = SingleTargetObjective(NumericalTarget("y")) - measurements = pd.DataFrame( - {"Color": ["red"], "x": [0.5], "Task": ["A"], "y": [1.0]} - ) - kernel_default = BayBEKernelFactory()(searchspace, objective, measurements) +def test_reduced_parameter_names(reduced_searchspace): + """Verify that parameter_names derives from the remaining parameters.""" + assert set(reduced_searchspace.parameter_names) == {"Color", "x"} + + +def test_reduced_comp_rep_columns(reduced_searchspace): + """Verify that comp_rep_columns matches the remaining parameters' columns.""" + expected = set() + for p in reduced_searchspace.parameters: + expected.update(p.comp_rep_columns) + assert set(reduced_searchspace.comp_rep_columns) == expected + + +def test_reduced_n_tasks(reduced_searchspace): + """Verify that n_tasks is 1 when no task parameter is present.""" + assert reduced_searchspace.n_tasks == 1 + + +def test_reduced_get_n_comp_rep_columns(reduced_searchspace): + """Verify that _get_n_comp_rep_columns works on the reduced space.""" + assert reduced_searchspace._get_n_comp_rep_columns("x") == 1 + assert reduced_searchspace._get_n_comp_rep_columns("Color") == 2 + + +def test_reduced_get_comp_rep_parameter_indices(reduced_searchspace): + """Verify that get_comp_rep_parameter_indices works on the reduced space.""" + indices = reduced_searchspace.get_comp_rep_parameter_indices("x") + assert len(indices) == 1 + indices = reduced_searchspace.get_comp_rep_parameter_indices("Color") + assert len(indices) == 2 + + +def test_reduced_blocked_attributes(reduced_searchspace): + """Verify that all non-allowed attributes raise NotImplementedError.""" + from baybe.searchspace.core import _ReducedSearchSpace + + allowed = _ReducedSearchSpace._ALLOWED_ATTRIBUTES + + # All public non-dunder attributes of SearchSpace + all_attrs = {name for name in dir(SearchSpace) if not name.startswith("_")} + + # Attributes that should be blocked (not in allowlist) + blocked = all_attrs - allowed + + for attr in sorted(blocked): + with pytest.raises(NotImplementedError, match=attr): + getattr(reduced_searchspace, attr) + + +def test_reduced_repr(reduced_searchspace): + """Verify that repr does not crash on a reduced search space.""" + result = repr(reduced_searchspace) + assert "_ReducedSearchSpace" in result + + +def test_reduced_str(reduced_searchspace): + """Verify that str does not crash on a reduced search space.""" + result = str(reduced_searchspace) + assert isinstance(result, str) + + +def test_reduced_eq(reduced_searchspace): + """Verify that equality comparison works on reduced search spaces.""" + ss = SearchSpace.from_product( + [ + CategoricalParameter("Color", ["red", "blue"]), + NumericalContinuousParameter("x", (0.0, 1.0)), + TaskParameter("Task", ["A", "B"]), + ] + ) + other = ss._drop_parameters({"Task"}) + assert reduced_searchspace == other + + +def test_reduced_unknown_parameter_raises(): + """Verify that removing a nonexistent parameter raises an error.""" + ss = SearchSpace.from_product( + [ + CategoricalParameter("Color", ["red", "blue"]), + NumericalContinuousParameter("x", (0.0, 1.0)), + ] + ) + with pytest.raises(ValueError, match="not found"): + ss._drop_parameters({"nonexistent"}) + + +def test_reduced_multiple_parameters_removed(): + """Verify that multiple parameters can be removed at once.""" + ss = SearchSpace.from_product( + [ + CategoricalParameter("Color", ["red", "blue"]), + NumericalContinuousParameter("x", (0.0, 1.0)), + TaskParameter("Task", ["A", "B"]), + ] + ) + reduced = ss._drop_parameters({"Task", "Color"}) + assert reduced.parameter_names == ("x",) + + +def test_reduced_kernel_product_matches_default_factory(): + """Verify that manual kernel split matches BayBE's default factory.""" + import torch + + from baybe.objectives.single import SingleTargetObjective + from baybe.surrogates.gaussian_process.presets.baybe import ( + BayBEKernelFactory, + _BayBENumericalKernelFactory, + _BayBETaskKernelFactory, + ) + from baybe.targets import NumericalTarget + + task_param = TaskParameter("Task", ["A", "B", "C"]) + cat_param = CategoricalParameter("Color", ["red", "blue", "green"]) + cont_param = NumericalContinuousParameter("x", (0.0, 1.0)) + searchspace = SearchSpace.from_product([cat_param, cont_param, task_param]) + + objective = SingleTargetObjective(NumericalTarget("y")) + measurements = pd.DataFrame( + {"Color": ["red"], "x": [0.5], "Task": ["A"], "y": [1.0]} + ) + + kernel_default = BayBEKernelFactory()(searchspace, objective, measurements) - reduced_ss = searchspace._drop_parameters({"Task"}) - base_baybe = _BayBENumericalKernelFactory()(reduced_ss, objective, measurements) - base_gpytorch = base_baybe.to_gpytorch(searchspace) + reduced_ss = searchspace._drop_parameters({"Task"}) + base_baybe = _BayBENumericalKernelFactory()(reduced_ss, objective, measurements) + base_gpytorch = base_baybe.to_gpytorch(searchspace) - task_only_ss = SearchSpace.from_product([task_param]) - task_baybe = _BayBETaskKernelFactory()(task_only_ss, objective, measurements) - task_gpytorch = task_baybe.to_gpytorch(searchspace) + task_only_ss = SearchSpace.from_product([task_param]) + task_baybe = _BayBETaskKernelFactory()(task_only_ss, objective, measurements) + task_gpytorch = task_baybe.to_gpytorch(searchspace) - kernel_manual = base_gpytorch * task_gpytorch + kernel_manual = base_gpytorch * task_gpytorch - assert type(kernel_default) is type(kernel_manual) - for k_def, k_man in zip(kernel_default.kernels, kernel_manual.kernels): - assert type(k_def) is type(k_man) - assert torch.equal(k_def.active_dims, k_man.active_dims) + assert type(kernel_default) is type(kernel_manual) + for k_def, k_man in zip(kernel_default.kernels, kernel_manual.kernels): + assert type(k_def) is type(k_man) + assert torch.equal(k_def.active_dims, k_man.active_dims) From 99c6f3137e0ad0cf663dd6c4588ed2c25a8ea1ae Mon Sep 17 00:00:00 2001 From: kalama-ai Date: Fri, 12 Jun 2026 08:04:13 +0200 Subject: [PATCH 7/7] Add get_parameters_by_name to allow list for default kernel factory --- baybe/searchspace/core.py | 1 + 1 file changed, 1 insertion(+) diff --git a/baybe/searchspace/core.py b/baybe/searchspace/core.py index 0b614b83ea..d35fefaf7a 100644 --- a/baybe/searchspace/core.py +++ b/baybe/searchspace/core.py @@ -607,6 +607,7 @@ class _ReducedSearchSpace(SearchSpace): "n_tasks", "_get_n_comp_rep_columns", "get_comp_rep_parameter_indices", + "get_parameters_by_name", "_ALLOWED_ATTRIBUTES", } )