Candidates Engine

[AdrianSosic]

SubspaceDiscrete Refactor — Dev Branch

This PR tracks an incremental refactor of SubspaceDiscrete toward a more general, lazy, and composable design. Individual feature PRs target this branch and merge here in sequence; the branch merges to main only when the refactor is complete.

Goals

1. Backend-agnostic dataframe operations — support pandas, Polars, PyArrow at API boundaries.
2. Lazy candidate construction — replace eager materialization of the full search space with Polars-backed lazy evaluation internally.
3. Single entry point for candidate access — get_candidates(...) with exclusion support and lazy CandidateResult.
4. Pluggable candidate selection for large/infinite spaces — CandidatePolicy as a two-level class hierarchy supporting both top-down filtering and bottom-up generative strategies, owned by the recommender, injected into get_candidates.

Problems with the current design

• Unenforced invariants. parameters, constraints, exp_rep, comp_rep are four independent attributes that must remain mutually consistent (exp_rep = product filtered by constraints; comp_rep = encoding of exp_rep). Consistency is maintained by convention only — nothing in the type prevents drift via evolve, manual construction, or future refactoring.
• Backend lock-in. Pandas is hardcoded; Polars support is partial and produces different row orders.
• Eager materialization breaks on large/infinite spaces. Modest formulation/mixture spaces produce millions of rows in RAM even when only a tiny subset is ever scored. There is no mechanism for a recommender to request a tractable subset — the entire space must fit in memory at construction time. Countably infinite spaces (e.g., string-valued parameters defined by an alphabet) cannot even be represented.
• Static encoded-columns hack. The transform method uses comp_rep[self.comp_rep.columns] to anchor the column set. Investigation confirms this is purely defensive (no subspace-level column modification occurs), but it's fragile and unnecessary once encoding is purely parameter-level.

Note: the existing allow_recommending_* flags and toggle_discrete_candidates on Campaign are not problems. They provide legitimate trajectory/history-based control and stay exactly as they are. Their implementation simply moves from FilteredSubspaceDiscrete mask manipulation to building an exclude frame passed to get_candidates().

Target design

@define(frozen=True)
class SubspaceDiscrete:
    parameters: tuple[Parameter, ...]   # schema (columns, types, encodings)
    candidates: Candidates              # extent (which rows are in the space)

• Candidates is a protocol with to_lazy(parameters) -> nw.LazyFrame and an is_enumerable property. Concrete types: ProductCandidates(constraints) (workhorse), TableCandidates(frame) (user-supplied tables), GeneratedCandidates(generator) (escape hatch for non-enumerable spaces). Non-enumerable candidates raise InfiniteSpaceError from to_lazy().
• Constraints are a flat tuple on ProductCandidates. Each DiscreteConstraint subclass implements to_narwhals_expr(parameters) -> nw.Expr. No predicate AST — just a conjunctive list. subspace.constraints is a simple accessor via getattr(self.candidates, "constraints", ()).
• Encoding lives on each Parameter via a new to_computational_expr(col) method that returns Narwhals expression(s). A free encode(frame, parameters) composes them. No separate Encoder class.
• get_candidates(n_candidates, *, exclude, policy, rng) is the single entry point. Two orthogonal axes:
  • policy — selection/construction strategy for candidate sets (owned by the recommender)
  • exclude — trajectory-based row removal (owned by Campaign), applied after the policy
• Returns a CandidateResult — lazy, cached, backend-aware.
• Backend: Narwhals for boundary translation; Polars-backed lazy frames internally always (true laziness regardless of user's input format).

Key design choices

• Two attributes, no redundancy. parameters is schema, candidates is extent. comp_rep and constraints become derived views — there's no cached representation that can drift.
• No predicate AST. Constraints are a flat conjunctive list on ProductCandidates with to_narwhals_expr() as the extension point. A Predicate protocol with Or/Not can be added non-breakingly in release 2 by widening the constraint tuple type.
• exclude: DataFrame instead of predicate objects. Simplest possible interface for "remove these rows." Campaign builds the combined exclusion frame from its own state. No exclude_measured/exclude_pending/exclude_recommended kwargs.
• CandidateResult instead of tuple[DataFrame, DataFrame]. Lazy (nothing computed until asked), cached (no recomputation), backend-aware (returns same type user passed in; falls back to baybe.Settings.default_backend).
• CandidatePolicy lives on the recommender, not on SubspaceDiscrete or Campaign. Three responsibilities, three owners: subspace defines the space (accepts policy at call time), Campaign owns exclusion (trajectory control), the recommender owns selection strategy. Different recommenders need different policies for the same space.
• Policy receives (Candidates, parameters), not a LazyFrame. This gives policies full access to the space definition, enabling both top-down and bottom-up strategies. A FilteringCandidatePolicy base class uses the template method pattern to shield concrete implementations from Candidates — they only implement _filter(lazy_frame, ...). A GenerativeCandidatePolicy base class provides full access for bottom-up candidate construction from infinite/non-enumerable spaces.
• Two base classes (FilteringCandidatePolicy, GenerativeCandidatePolicy). These represent genuinely different operations: filtering is a frame transformation (LazyFrame → LazyFrame), generation is candidate construction (SpaceDefinition → LazyFrame). The template method on FilteringCandidatePolicy handles Candidates interaction (enumeration, enumerability validation); the concrete subclass only sees a LazyFrame. This eliminates coupling. GenerativeCandidatePolicy implements select() directly with full space access. The hierarchy makes the infinite-space contract statically checkable via isinstance.
• Self-enforcing validation. FilteringCandidatePolicy.select() raises InfiniteSpaceError on non-enumerable candidates. get_candidates() only checks the degenerate "no policy + non-enumerable" case. No redundant validation needed.
• ChainedPolicy with construction-time ordering validation. Only the first policy may be generative; all subsequent must be FilteringCandidatePolicy. The chain calls _filter directly on positions 2+. Validated at construction time — TypeError when building the chain beats a runtime crash during recommendation.
• Generative policies are not restricted to infinite spaces. RandomSamplingPolicy on a large-but-finite 10⁹-row space is legitimate. The restriction is one-directional: filtering requires enumerable (enforced), generation works on anything.
• Parameter.sample(n, rng) as the extension point for generative policies. Each parameter knows its domain — sample is a natural extension. Finite parameters draw from enumerated values; infinite parameters (e.g., SubstanceParameter with alphabet + length) generate random valid values.
• Constraints are subclass-specific, not protocol-level. Only ProductCandidates carries constraints. Ad-hoc filtering across all implementations is handled by get_candidates(exclude=...).
• Encoding on Parameter. Each subclass already owns its encoding rule. Statefulness (decorrelation, constant-column-dropping) is encapsulated in the parameter's comp_df cached property. The subspace never modifies encoded columns.
• Polars internally, Narwhals at boundaries. Narwhals cannot provide true lazy evaluation over pandas. Internal computation always uses Polars for genuine laziness.
• allow_* flags stay on Campaign. They are trajectory-based controls, not a design smell. Implementation moves from mask-based FilteredSubspaceDiscrete to building an exclude frame.

---

PRs

PR 0 — Groundwork

Audit, dependency setup, and performance validation. No behavior change.

• Add narwhals and polars as hard dependencies.
• Audit every external access to subspace.exp_rep/comp_rep → MIGRATION_AUDIT.md (drives PR 4).
• Audit each Parameter subclass's current encoding method → PARAMETER_ENCODING_AUDIT.md (drives PR 6).
• Proof-of-concept benchmark: Polars cross-join + filter vs. current eager construction on real parameter configs.
• Add benchmarks/searchspace/ directory.

PR 1 — Inert new abstractions (including CandidatePolicy hierarchy)

Add the new types without wiring them in.

• New module baybe/searchspace/candidates.py: Candidates protocol with to_lazy(parameters) and is_enumerable property, ProductCandidates(constraints), TableCandidates(frame), GeneratedCandidates(generator).
• Add to_narwhals_expr(parameters) -> nw.Expr to every existing DiscreteConstraint subclass.
• New module baybe/searchspace/result.py: CandidateResult with .exp_rep(backend=) and .comp_rep(backend=).
• New module baybe/searchspace/policy.py:
  • CandidatePolicy protocol with select(candidates: Candidates, parameters, n_candidates, rng) -> nw.LazyFrame.
  • FilteringCandidatePolicy ABC — template method: select() validates enumerability + calls to_lazy(), delegates to abstract _filter(lazy_frame, ...). Concrete implementations decoupled from Candidates.
  • GenerativeCandidatePolicy ABC — select() receives full space definition, constructs candidates bottom-up.
  • Filtering policies: TakeFirstPolicy, RandomSubsamplePolicy, HashSubsamplePolicy, FarthestPointPolicy.
  • Generative policies: RandomSamplingPolicy.
  • ChainedPolicy(policies) — sequential composition with construction-time validation (all but first must be FilteringCandidatePolicy).
• Add sample(n, rng) method to each Parameter subclass.
• Hypothesis strategies and unit tests in isolation. No SubspaceDiscrete changes, no recommender changes.

PR 2 — Dual-storage in SubspaceDiscrete

Add candidates attribute alongside existing fields; both representations populated.

• Add candidates: Candidates | None attribute.
• All classmethod constructors populate it (constraints passed directly to ProductCandidates).
• __attrs_post_init__ reverse-engineers ProductCandidates for raw constructor calls (temporary; removed in PR 8).
• New private _eager_materialize wraps existing construction logic.
• Tests verify bit-identical exp_rep/comp_rep to pre-PR-2.

PR 3 — get_candidates() with policy support

Establish the single public entry point with the full two-axis signature. Implementation still eager.

• Implement get_candidates(n_candidates, *, exclude, policy, rng) returning CandidateResult.
• Policy receives (self.candidates, self.parameters) — full space definition.
• Exclusion via anti-join semantics (reuses fuzzy_row_match), applied after policy.
• Non-enumerable space without policy → InfiniteSpaceError.
• FilteringCandidatePolicy on non-enumerable space → InfiniteSpaceError (self-enforced by the policy).
• GenerativeCandidatePolicy on non-enumerable space → constructs valid candidates.
• Document in user guide; show Campaign-side exclusion pattern, recommender-side policy pattern, and infinite-space generative pattern.

PR 4 — Wire policy into recommenders; migrate callers

Add candidate_policy to recommenders. Migrate internal callers from exp_rep/comp_rep to get_candidates(). Split into sub-PRs.

• 4a Add candidate_policy: CandidatePolicy | None = None to recommender base class. Pass through to get_candidates(policy=self.candidate_policy) in _recommend_with_discrete_parts. Default None = no behavioral change.
• 4b Migrate non-predictive recommenders (RandomRecommender, FPSRecommender).
• 4c Migrate BoTorch-based recommenders.
• 4d Migrate Campaign — consolidates allow_recommending_* flags + toggle_discrete_candidates into a single exclude frame per recommend() call. Passes recommender's candidate_policy through.
• 4e Constraint validation/checking code.
• 4f Remaining call sites from MIGRATION_AUDIT.md.

By end of PR 4: zero internal references to exp_rep/comp_rep; candidate_policy configurable on recommenders; FilteredSubspaceDiscrete no longer used internally; no public API break.

PR 5 — Deprecate old attributes

Mark exp_rep/comp_rep deprecated; update everything user-facing.

• exp_rep/comp_rep become properties emitting DeprecationWarning.
• constraints becomes a read-only property delegating to self.candidates.constraints (not deprecated).
• All examples/ and docs/userguide/ updated to new patterns (including candidate_policy in large-space and infinite-space examples).
• Internal tree passes pytest -W error::DeprecationWarning.

PR 6 — Lazy backend (the big one)

Switch internals to Narwhals + Polars lazy evaluation. Can split into 6a (lazy candidates) and 6b (per-parameter encoding migration).

• Strand A Lazy ProductCandidates.to_lazy (with is_enumerable check) / TableCandidates.to_lazy / GeneratedCandidates.to_lazy; port exclusion logic to Narwhals (anti-join); make every DiscreteConstraint.to_narwhals_expr Narwhals-native; drop BAYBE_DEACTIVATE_POLARS.
• Strand B Add to_computational_expr(col) to every Parameter subclass; implement free encode(frame, parameters); keep old pandas methods as deprecated shims.
• Strand C _exp_rep/_comp_rep become cached_property; get_candidates operates on the lazy frame end-to-end, returns lazy CandidateResult. Filtering policies now operate on truly lazy frames. Generative policies construct candidates without attempting enumeration.
• Benchmark target informed by PR 0 results; memory drops to O(n_candidates). RandomSubsamplePolicy on 10⁹-row product space completes in bounded memory. RandomSamplingPolicy on non-enumerable space completes in O(n) time/memory.

PR 7 — Deprecate old attributes

If deprecation was already completed in PR 5, this is a no-op. Allows flexibility in release cadence — deprecation can happen before or after the lazy backend switch.

PR 8 — Major version cleanup (baybe 2.0)

Remove deprecations; arrive at the final two-attribute design.

• Remove exp_rep/comp_rep properties; remove constraints from constructor.
• Remove __attrs_post_init__ reverse-engineering hack.
• Remove deprecated pandas-based encoding shims on Parameter subclasses.
• Remove transform method's comp_rep[self.comp_rep.columns] hack.
• Remove FilteredDiscreteSubspace (subsumed by get_candidates(exclude=...)).
• Final shape: SubspaceDiscrete(parameters, candidates).
• Migration guide as top-level doc.
• Version → 2.0.0.

---

Cross-cutting

• Backward compatibility maintained throughout PRs 0–7; breaking changes only in PR 8 (the 2.0 release).
• candidate_policy=None (the default) preserves existing behavior exactly — no policy means head(n) fallback.
• Performance benchmarks run on every PR touching searchspace/. Includes non-enumerable space benchmarks for generative policies.
• Serialization tested round-trip for every new class via existing SerialMixin pattern. CandidatePolicy subclasses are attrs-defined; ChainedPolicy serializes as a list. Parameter comp_df column sets serialized to prevent decorrelation drift.
• Property tests for constraint expression equivalence, candidate/parameter consistency, encoding equivalence (old vs new), policy composition, FilteringCandidatePolicy raises on non-enumerable, ChainedPolicy rejects non-filtering policies in non-first positions at construction time.

To keep in mind

• FilteredDiscreteSubspace becomes obsolete (removed in PR 8).
• The static encoded-columns hack (comp_rep[self.comp_rep.columns]) disappears — encoding is purely parameter-level; no subspace-level column anchoring needed.
• The allow_* flags remain on Campaign — they are trajectory controls, not a design smell.
• Non-enumerable spaces require a GenerativeCandidatePolicy. This is enforced by get_candidates() (raises without policy) and FilteringCandidatePolicy.select() (raises on non-enumerable candidates). Error messages guide users to the correct fix.
• Parameter.sample(n, rng) is the extension point for generative policies — each parameter must know how to produce random valid values from its domain.
• Specialized Candidates implementations (e.g., SimplexCandidates) should be added when benchmarks show the generic ProductCandidates cross-join path is insufficient for heavily-constrained spaces.

Future improvements

• Predicate AST with Or/Not: Widen ProductCandidates.constraints to accept a Predicate protocol; add Or, Not, RawExprPredicate. Non-breaking extension.
• AcquisitionPrescreenPolicy: A materializing filtering policy using surrogate inference to pre-filter by expected acquisition value. Requires surrogate access — separate design discussion.
• Functionality for adjusting existing spaces (shrink/expand/filter).
• SimplexCandidates: Preserve smart incremental enumeration from from_simplex.
• UnionCandidates: Disjoint union for transfer learning / task parameter spaces.
• Advanced generative policies: LatinHypercubePolicy, EvolutionaryPolicy, GrammarGuidedPolicy — more sophisticated bottom-up construction strategies. The GenerativeCandidatePolicy base class and Parameter.sample() provide the extension points.

[Copilot]

Copilot wasn't able to review any files in this pull request.

---

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

[Scienfitz]

@AdrianSosic please can you use the instrument of Issues and Subissues for this? Its kind of what that was made for and allows for a more fine-grained optional discussion compared to a PR that will not be reviewed anyway etc