Add _ReducedSearchSpace for parameter-only reduced search spaces by kalama-ai · Pull Request #819 · emdgroup/baybe

kalama-ai · 2026-06-02T17:04:14Z

Introduces _ReducedSearchSpace (private SearchSpace subclass) that exposes only parameter information without building the expensive discrete Cartesian product. This is constructed via SearchSpace._without_parameters(names), which returns a reduced version with the specified parameters removed.

The class blocks access to anything beyond parameter-related properties.

Motivation

When building composite GP kernels, we need to call kernel factories on search spaces with certain parameters stripped away. A full SearchSpace.from_product() would rebuild the entire discrete candidate set. _ReducedSearchSpace avoids this by storing the parameter tuple directly and computing everything from it.

Changes to SearchSpace and kernel factories

The reduced search space exposes only name-based properties, no integer indices. To support this, we made two changes to existing code:

Replaced task_idx sentinel checks with n_tasks > 1 / n_tasks == 1 in kernel factories and fit criterion factories. (Equivalent because TaskParameter requires a minimum of 2 values).
Added SearchSpace._get_n_comp_rep_columns(selector) that returns the number of comp-rep columns for a parameter selection. This replaces the previous len(get_comp_rep_parameter_indices(...)) in _get_effective_dimensionality, avoiding the need to expose index-returning methods on the reduced space.

Design of _ReducedSearchSpace

Subclasses SearchSpace so it passes type checks and works with existing factory signatures
Stores a tuple[Parameter, ...] and derives parameters, parameter_names, comp_rep_columns from it
Overrides __getattribute__ with an allowlist — accessing transform, constraints, discrete, continuous, etc. raises NotImplementedError
Skips __attrs_post_init__

Allowed attributes

_ALLOWED_ATTRIBUTES: ClassVar[frozenset[str]] = frozenset({
    "parameters",
    "parameter_names",
    "comp_rep_columns",
    "_task_parameter",
    "n_tasks",
    "_get_n_comp_rep_columns",
    "_parameters",
    "_ALLOWED_ATTRIBUTES",
})

Outlook

This class is a building block for kernel override dispatching in the GP surrogate. The dispatching logic will:

Strip the task parameter from the search space
Call kernel factories on the reduced space to obtain a base kernel
Separately construct the task kernel
Multiply both at the gpytorch level after resolving indices against the full search space
Ensure that factories called on a reduced search space return BayBE Kernel objects (which use parameter names) rather than raw gpytorch kernels (which bake in integer indices)

Copilot

Pull request overview

This PR introduces a private, lightweight _ReducedSearchSpace to support GP kernel factory calls on “parameter-only” search spaces (e.g., with some parameters stripped) without rebuilding expensive discrete Cartesian products, via SearchSpace._without_parameters(names).

Changes:

Added SearchSpace._without_parameters(...) and _ReducedSearchSpace (parameter-only SearchSpace subclass) in baybe/searchspace/core.py.
Added a dedicated test suite validating exposed/blocked attributes and a kernel split/product equivalence check in tests/test_searchspace.py.
Added a TODO in the GP surrogate fit path about validating factory return types when operating on reduced search spaces.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 4 comments.

File	Description
`baybe/searchspace/core.py`	Adds `_without_parameters` and new `_ReducedSearchSpace` with allowlisted attribute access and parameter-derived properties.
`tests/test_searchspace.py`	Adds tests for reduced search space behavior and kernel factory interoperability (active dims consistency).
`baybe/surrogates/gaussian_process/core.py`	Adds a TODO note about validating kernel factory outputs when reduced search spaces are used.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

+    @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."
+        )


AdrianSosic · 2026-06-09T07:13:09Z

            names
        ) + self.continuous.get_parameters_by_name(names)

+    def _without_parameters(self, names: Collection[str], /) -> _ReducedSearchSpace:


I'm very open to discussing the method names once we have a truly public interface, but for now I suggest to rename it to _drop_parameters, since we already have it on other classes as well

Will rename.

you need to also update the PR description, which still contains the old name. I guess there are now also other things to be changed in the description, since your latest commits?

AVHopp

I am in general happy with the PR, mostly minor comments (and one where we might want to discuss a bit)

AVHopp · 2026-06-10T13:56:30Z

+        Returns:
+            A reduced search space containing only parameter information.
+        """
+        current_names = {p.name for p in self.parameters}


Couldn't you also precompute set_names = set(names) here to avoid the double calculation in both the if and the calculation of remaining?

AVHopp · 2026-06-10T14:33:32Z

+        allowed = object.__getattribute__(self, "_ALLOWED_ATTRIBUTES")
+        if name in allowed:
+            return object.__getattribute__(self, name)
+        raise NotImplementedError(


Shouldn't this rather raise an AttributeError?

AVHopp · 2026-06-10T14:36:25Z

        )

        ### Kernel
+        # TODO: When calling a factory on a `_ReducedSearchSpace`, validate that


Just to be sure, this is intended to be done in a follow-up PR I assume and not in this one?

Yes, this was only meant to highlight where the new _ReducedSearchSPaceclass would be used for the first time. Our plan is to pass it only to BayBE factories and block gpytorch kernels. This is why the current allow list in the class blocks everything index related.

AVHopp · 2026-06-10T14:40:38Z

+        reduced = searchspace._drop_parameters({"Task"})
+        assert reduced.n_tasks == 1
+
+    def test_get_n_comp_rep_columns(self, searchspace):


Wouldn't it make sense to also add a ParameterSelctor-based test and not only str-based ones?

AVHopp · 2026-06-10T14:42:38Z

+    def test_blocked_attributes(self, searchspace):
+        """Verify that accessing non-parameter attributes raises an error."""
+        reduced = searchspace._drop_parameters({"Task"})
+        with pytest.raises(NotImplementedError):


Is there any programmatic way of accessing all of these fields while also ensuring that we actually get all? If we'd now add a new field to SearchSpace, one would need to remember that this test here should also be changed. I would imagine some mechanism to go through all fields that are not part of the allowed list, verify that trying to access them raises an error and check whether any of the fields has did not raise an error.

Added here: e341f4a

AVHopp · 2026-06-10T14:44:59Z

        assert np.isclose(interpoint_constraint_result, 0.6, atol=1e-6)
+
+
+class TestReducedSearchSpace:


We usually don't use test classes, so I would prefer if this follows the same style that we use in other places as well (looping in @AdrianSosic since he might say "No, this is the better approach, we should keep it and not use the bad one that we use everywhere else just for consistency!")

so far we haven't used them. Perhaps we will in the future, perhaps not 🙃

But you know what: we need an easy way to remove all the corresponding tests later, without having to read each test body, so some grouping is definitely needed. So can you either:

go back to the class approach or

more the tests to a dedicated test file

AVHopp · 2026-06-10T14:49:06Z

+
+
+@define(slots=False)
+class _ReducedSearchSpace(SearchSpace):


Do we think that inheritance is the way to go here? This implementation clearly violates LSP, and we could probably also use Protocol instead. I am not saying that we should - in fact, my gut feeling is that using the inheritance here is the better way right now, but I want to at least open up the discussion and make sure that we all are happy with this design. Hence also looping in @Scienfitz and @AdrianSosic for their opinions :)

Fair point. But tbh I this is hopefully only a temporary solution and I think we can live with it.

Have also talked with @kalama-ai offline about a related issue, but my opinion: this is clearly a hack anyway and not a first-class new feature, so do whatever it takes to bring you to the goal as quickly as possible 😬

- 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

- 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`

Co-authored-by: AdrianSosic <adrian.sosic@merckgroup.com>

- Build real SubspaceDiscrete/SubspaceContinuous objects with zero rows instead of storing a flat _parameters tuple - Remove _parameters field, all property overrides, __attrs_post_init, override, and _get_n_comp_rep_columns override - Convert TestReducedSearchSpace class to standalone test functions - Add programmatic test that discovers all SearchSpace attributes and verifies non-allowed ones raise NotImplementedError

AdrianSosic · 2026-06-12T06:44:03Z

+            "_task_parameter",
+            "n_tasks",
+            "_get_n_comp_rep_columns",
+            "get_comp_rep_parameter_indices",


Now you have added this one here to the list, which shouldn't be there. I know, _get_n_comp_rep_columns needs it, but you can work around that in its implementation

kalama-ai commented Jun 2, 2026

View reviewed changes

Comment thread baybe/searchspace/core.py Outdated

kalama-ai commented Jun 3, 2026

View reviewed changes

Comment thread baybe/searchspace/core.py

kalama-ai commented Jun 3, 2026

View reviewed changes

Comment thread baybe/searchspace/core.py

kalama-ai marked this pull request as ready for review June 3, 2026 08:54

kalama-ai requested review from AdrianSosic and Scienfitz as code owners June 3, 2026 08:54

Copilot AI review requested due to automatic review settings June 3, 2026 08:54

kalama-ai requested a review from AVHopp as a code owner June 3, 2026 08:54

Copilot started reviewing on behalf of kalama-ai June 3, 2026 08:55 View session

Copilot AI reviewed Jun 3, 2026

View reviewed changes

Scienfitz assigned kalama-ai Jun 3, 2026

Scienfitz added this to the 0.16.0 milestone Jun 3, 2026

kalama-ai changed the title ~~Add _ReducedSearchSpace for parameter-only reduced search paces~~ Add _ReducedSearchSpace for parameter-only reduced search spaces Jun 8, 2026

AdrianSosic requested changes Jun 9, 2026

View reviewed changes

AVHopp reviewed Jun 10, 2026

View reviewed changes

kalama-ai and others added 6 commits June 12, 2026 08:01

Simplify counting of comp_rep_columns

c2c3f72

Co-authored-by: AdrianSosic <adrian.sosic@merckgroup.com>

Simplify if statement

2a4d7f5

Co-authored-by: AdrianSosic <adrian.sosic@merckgroup.com>

Rename _without_parameters() to _drop_parameters()

dae8b3c

kalama-ai force-pushed the feat/reduce-searchspaces branch from e341f4a to fa31cc4 Compare June 12, 2026 06:01

Add get_parameters_by_name to allow list for default kernel factory

99c6f31

AdrianSosic reviewed Jun 12, 2026

View reviewed changes

		assert np.isclose(interpoint_constraint_result, 0.6, atol=1e-6)


		class TestReducedSearchSpace:

Conversation

kalama-ai commented Jun 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

AVHopp left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

kalama-ai commented Jun 2, 2026 •

edited

Loading