feat: Re-introduce graphs

graph_playground.py

@@ -0,0 +1,47 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from graph import Grammar, mutations, parse, select, to_string
+
+
+# Leafs
+@dataclass
+class T:
+    s: str
+
+    # This is the `op()`
+    def __call__(self) -> str:
+        return self.s
+
+
+def join(*s: str) -> str:
+    return "[" + "".join(s) + "]"
+
+
+grammar_1 = Grammar.from_dict(
+    {
+        "s": (["a", "b", "p a", "p p"], join),
+        "p": ["a b", "s"],
+        "a": T("a"),
+        "b": T("b"),
+    }
+)
+
+root = parse(grammar_1, "s(p(s(a), a))")
+
+selections = list(select(root, how=("climb", range(1, 3))))
+mutants = mutations(
+    root=root,
+    grammar=grammar_1,
+    which=selections,
+    max_mutation_depth=3,
+)
+mutants = list(mutants)
+
+import rich
+
+rich.print("grammar", grammar_1)
+rich.print("root", f"{to_string(root)}")
+rich.print("selections", [to_string(s) for s in selections])
+rich.print("mutants", [to_string(m) for m in mutants])

neps/__init__.py

@@ -4,7 +4,7 @@
 from neps.optimizers.optimizer import SampledConfig
 from neps.plot.plot import plot
 from neps.plot.tensorboard_eval import tblogger
-from neps.space import Categorical, Constant, Float, Integer, SearchSpace
+from neps.space import Categorical, Constant, Float, Grammar, Integer, SearchSpace
 from neps.state import BudgetInfo, Trial
 from neps.status.status import status
 from neps.utils.files import load_and_merge_yamls as load_yamls
@@ -15,6 +15,7 @@
     "Categorical",
     "Constant",
     "Float",
+    "Grammar",
     "Integer",
     "SampledConfig",
     "SearchSpace",

neps/api.py

@@ -18,7 +18,7 @@
     from ConfigSpace import ConfigurationSpace
 
     from neps.optimizers.algorithms import CustomOptimizer
-    from neps.space import Parameter, SearchSpace
+    from neps.space import Constant, Grammar, Parameter, SearchSpace
     from neps.state import EvaluatePipelineReturn
 
 logger = logging.getLogger(__name__)
@@ -27,7 +27,7 @@
 def run(  # noqa: PLR0913
     evaluate_pipeline: Callable[..., EvaluatePipelineReturn] | str,
     pipeline_space: (
-        Mapping[str, dict | str | int | float | Parameter]
+        Mapping[str, dict | str | int | float | Parameter | Constant | Grammar]
         | SearchSpace
         | ConfigurationSpace
     ),

neps/optimizers/algorithms.py

@@ -82,7 +82,7 @@ def _bo(
             f" Got: {pipeline_space.fidelities}"
         )
 
-    parameters = pipeline_space.searchables
+    parameters = {**pipeline_space.numerical, **pipeline_space.categoricals}
 
     match initial_design_size:
         case "ndim":
@@ -126,9 +126,6 @@ def _bracket_optimizer(  # noqa: C901, PLR0912, PLR0915
     sampler: Literal["uniform", "prior", "priorband"] | PriorBandSampler | Sampler,
     bayesian_optimization_kick_in_point: int | float | None,
     sample_prior_first: bool | Literal["highest_fidelity"],
-    # NOTE: This is the only argument to get a default, since it
-    # is not required for hyperband style algorithms, only single bracket
-    # style ones.
     early_stopping_rate: int | None,
     device: torch.device | None,
 ) -> BracketOptimizer:
@@ -183,7 +180,7 @@ def _bracket_optimizer(  # noqa: C901, PLR0912, PLR0915
     """
     assert pipeline_space.fidelity is not None
     fidelity_name, fidelity = pipeline_space.fidelity
-    parameters = pipeline_space.searchables
+    parameters = {**pipeline_space.numerical, **pipeline_space.categoricals}
 
     if len(pipeline_space.fidelities) != 1:
         raise ValueError(
@@ -324,9 +321,8 @@ def _bracket_optimizer(  # noqa: C901, PLR0912, PLR0915
 
 
 def determine_optimizer_automatically(space: SearchSpace) -> str:
-    has_prior = any(
-        parameter.prior is not None for parameter in space.searchables.values()
-    )
+    parameters = {**space.numerical, **space.categoricals}
+    has_prior = any(parameter.prior is not None for parameter in parameters.values())
     has_fidelity = len(space.fidelities) > 0
 
     match (has_prior, has_fidelity):
@@ -360,14 +356,18 @@ def random_search(
             In this case, the max fidelity is always used.
     """
     if ignore_fidelity:
-        parameters = pipeline_space.searchables
+        parameters = {**pipeline_space.numerical, **pipeline_space.categoricals}
     else:
-        parameters = {**pipeline_space.searchables, **pipeline_space.fidelities}
+        parameters = {
+            **pipeline_space.numerical,
+            **pipeline_space.categoricals,
+            **pipeline_space.fidelities,
+        }
 
     return RandomSearch(
         space=pipeline_space,
         encoder=ConfigEncoder.from_parameters(parameters),
-        sampler=(
+        numerical_sampler=(
             Prior.from_parameters(parameters)
             if use_priors
             else Uniform(ndim=len(parameters))
@@ -384,6 +384,9 @@ def grid_search(pipeline_space: SearchSpace) -> GridSearch:
     """
     from neps.optimizers.utils.grid import make_grid
 
+    if pipeline_space.grammar is not None:
+        raise NotImplementedError("Grammars not supported for `grid_search` yet.")
+
     return GridSearch(configs_list=make_grid(pipeline_space))
 
 
@@ -445,7 +448,7 @@ def ifbo(
     space, fid_bins = _adjust_space_to_match_stepsize(pipeline_space, step_size)
     assert space.fidelity is not None
     fidelity_name, fidelity = space.fidelity
-    parameters = space.searchables
+    parameters = {**pipeline_space.numerical, **pipeline_space.categoricals}
 
     match initial_design_size:
         case "ndim":

neps/optimizers/bayesian_optimization.py

@@ -86,7 +86,11 @@ def __call__(
         n: int | None = None,
     ) -> SampledConfig | list[SampledConfig]:
         assert self.space.fidelity is None, "Fidelity not supported yet."
-        parameters = self.space.searchables
+        parameters = {
+            **self.space.numerical,
+            **self.space.categoricals,
+            **self.space.grammars,
+        }
 
         n_to_sample = 1 if n is None else n
         n_sampled = len(trials)

neps/optimizers/bracket_optimizer.py

@@ -249,6 +249,12 @@ class BracketOptimizer:
     fid_name: str
     """The name of the fidelity in the space."""
 
+    def __post_init__(self) -> None:
+        if self.space.grammar is not None:
+            raise NotImplementedError(
+                "Grammars not supported for `BracketOptimizer` yet."
+            )
+
     def __call__(  # noqa: C901, PLR0912
         self,
         trials: Mapping[str, Trial],
@@ -257,7 +263,7 @@ def __call__(  # noqa: C901, PLR0912
     ) -> SampledConfig | list[SampledConfig]:
         assert n is None, "TODO"
         space = self.space
-        parameters = space.searchables
+        parameters = {**self.space.numerical, **self.space.categoricals}
 
         # If we have no trials, we either go with the prior or just a sampled config
         if len(trials) == 0:

neps/optimizers/ifbo.py

@@ -129,6 +129,10 @@ class IFBO:
     Each one will be treated as an individual fidelity level.
     """
 
+    def __post_init__(self) -> None:
+        if self.space.grammar is not None:
+            raise NotImplementedError("Grammars not supported for `IFBO` yet.")
+
     def __call__(
         self,
         trials: Mapping[str, Trial],
@@ -137,7 +141,7 @@ def __call__(
     ) -> SampledConfig | list[SampledConfig]:
         assert self.space.fidelity is not None
         fidelity_name, fidelity = self.space.fidelity
-        parameters = self.space.searchables
+        parameters = {**self.space.numerical, **self.space.categoricals}
 
         assert n is None, "TODO"
         ids = [int(config_id.split("_", maxsplit=1)[0]) for config_id in trials]

neps/optimizers/models/graphs/context_managers.py

@@ -0,0 +1,65 @@
+from __future__ import annotations
+
+from collections.abc import Iterator
+from contextlib import contextmanager
+from typing import TYPE_CHECKING
+
+from botorch.models import SingleTaskGP
+
+from neps.optimizers.models.graphs.kernels import BoTorchWLKernel, compute_kernel
+
+if TYPE_CHECKING:
+    import networkx as nx
+    from botorch.models.gp_regression_mixed import Kernel
+
+
+@contextmanager
+def set_graph_lookup(
+    kernel_or_gp: Kernel | SingleTaskGP,
+    new_graphs: list[nx.Graph],
+    *,
+    append: bool = True,
+) -> Iterator[None]:
+    """Context manager to temporarily set the graph lookup for a kernel or GP model.
+
+    Args:
+        kernel_or_gp (Kernel | SingleTaskGP): The kernel or GP model whose graph lookup is
+            to be set.
+        new_graphs (list[nx.Graph]): The new graphs to set in the graph lookup.
+        append (bool, optional): Whether to append the new graphs to the existing graph
+            lookup. Defaults to True.
+    """
+    kernel_prev_graphs: list[tuple[Kernel, list[nx.Graph]]] = []
+
+    # Determine the modules to update based on the input type
+    if isinstance(kernel_or_gp, SingleTaskGP):
+        modules = [
+            k
+            for k in kernel_or_gp.covar_module.sub_kernels()
+            if isinstance(k, BoTorchWLKernel)
+        ]
+    elif isinstance(kernel_or_gp, BoTorchWLKernel):
+        modules = [kernel_or_gp]
+    else:
+        assert hasattr(kernel_or_gp, "sub_kernels"), (
+            "Kernel module must have sub_kernels method."
+        )
+        modules = [
+            k for k in kernel_or_gp.sub_kernels() if isinstance(k, BoTorchWLKernel)
+        ]
+
+    # Save the current graph lookup and set the new graph lookup
+    for kern in modules:
+        compute_kernel.cache_clear()
+
+        kernel_prev_graphs.append((kern, kern.graph_lookup))
+        if append:
+            kern.set_graph_lookup([*kern.graph_lookup, *new_graphs])
+        else:
+            kern.set_graph_lookup(new_graphs)
+
+    yield
+
+    # Restore the original graph lookup after the context manager exits
+    for kern, prev_graphs in kernel_prev_graphs:
+        kern.set_graph_lookup(prev_graphs)