edmkit-search 0.0.1a1__tar.gz

edmkit_search-0.0.1a1/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 temma
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

edmkit_search-0.0.1a1/PKG-INFO

@@ -0,0 +1,70 @@
+Metadata-Version: 2.4
+Name: edmkit-search
+Version: 0.0.1a1
+Summary: Trajectory construction on an energy landscape for edmkit
+Keywords: edm,empirical-dynamic-modeling,feature-selection,search,time-series
+Author: FUJISHIGE TEMMA
+Author-email: FUJISHIGE TEMMA <tenma.x0@gmail.com>
+License-Expression: MIT
+License-File: LICENSE
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Science/Research
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Topic :: Scientific/Engineering
+Classifier: Typing :: Typed
+Requires-Dist: edmkit>=0.0.7
+Requires-Dist: numpy>=2.4.4
+Requires-Python: >=3.13
+Project-URL: Homepage, https://github.com/FujishigeTemma/edmkit-search
+Project-URL: Repository, https://github.com/FujishigeTemma/edmkit-search
+Project-URL: Issues, https://github.com/FujishigeTemma/edmkit-search/issues
+Description-Content-Type: text/markdown
+
+# edmkit-search
+
+Trajectory construction on an energy landscape, packaged as a library on top
+of [`edmkit`](https://github.com/FujishigeTemma/edmkit).
+
+## Install
+
+```bash
+pip install edmkit-search
+# or
+uv add edmkit-search
+```
+
+## Overview
+
+The search loop builds a trajectory step by step:
+
+1. From the current frontier (batch of `(state, context)`),
+2. expand neighbors via a `Neighborhood`,
+3. score children with an `Energy` to obtain `(energies, contexts)`,
+4. and pick the next frontier with a `Strategy`.
+
+See `CODING.md` for the design rules. The subpackages mirror the four
+abstractions: `energy/`, `neighborhood/`, `state/`, `strategy/`, plus
+`dataset/` for the input containers.
+
+## Quick start
+
+```bash
+uv sync
+PYTHON_GIL=0 uv run python e2e/synthetic.py
+```
+
+`e2e/synthetic.py` generates a Lorenz-96 trajectory mixed with noise columns
+and uses greedy forward selection to recover the informative subset.
+
+## Tests
+
+```bash
+uv run pytest
+```
+
+## Experiments
+
+Real-data experiments (fly, fish) and the analysis pipeline live in a
+separate repository: [`edmkit-search-experiments`](../edmkit-search-experiments).

edmkit_search-0.0.1a1/README.md

@@ -0,0 +1,46 @@
+# edmkit-search
+
+Trajectory construction on an energy landscape, packaged as a library on top
+of [`edmkit`](https://github.com/FujishigeTemma/edmkit).
+
+## Install
+
+```bash
+pip install edmkit-search
+# or
+uv add edmkit-search
+```
+
+## Overview
+
+The search loop builds a trajectory step by step:
+
+1. From the current frontier (batch of `(state, context)`),
+2. expand neighbors via a `Neighborhood`,
+3. score children with an `Energy` to obtain `(energies, contexts)`,
+4. and pick the next frontier with a `Strategy`.
+
+See `CODING.md` for the design rules. The subpackages mirror the four
+abstractions: `energy/`, `neighborhood/`, `state/`, `strategy/`, plus
+`dataset/` for the input containers.
+
+## Quick start
+
+```bash
+uv sync
+PYTHON_GIL=0 uv run python e2e/synthetic.py
+```
+
+`e2e/synthetic.py` generates a Lorenz-96 trajectory mixed with noise columns
+and uses greedy forward selection to recover the informative subset.
+
+## Tests
+
+```bash
+uv run pytest
+```
+
+## Experiments
+
+Real-data experiments (fly, fish) and the analysis pipeline live in a
+separate repository: [`edmkit-search-experiments`](../edmkit-search-experiments).

edmkit_search-0.0.1a1/pyproject.toml

@@ -0,0 +1,48 @@
+[project]
+name = "edmkit-search"
+version = "0.0.1a1"
+description = "Trajectory construction on an energy landscape for edmkit"
+authors = [{ name = "FUJISHIGE TEMMA", email = "tenma.x0@gmail.com" }]
+license = "MIT"
+license-files = ["LICENSE"]
+readme = "README.md"
+requires-python = ">=3.13"
+keywords = ["edm", "empirical-dynamic-modeling", "feature-selection", "search", "time-series"]
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "Intended Audience :: Science/Research",
+    "Operating System :: OS Independent",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.13",
+    "Topic :: Scientific/Engineering",
+    "Typing :: Typed",
+]
+dependencies = [
+    "edmkit>=0.0.7",
+    "numpy>=2.4.4",
+]
+
+[project.urls]
+Homepage = "https://github.com/FujishigeTemma/edmkit-search"
+Repository = "https://github.com/FujishigeTemma/edmkit-search"
+Issues = "https://github.com/FujishigeTemma/edmkit-search/issues"
+
+[dependency-groups]
+dev = ["hypothesis>=6.152.5", "pytest>=9.0.3", "ruff>=0.15.12", "ty>=0.0.35"]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+addopts = "-x --tb=short"
+
+[build-system]
+requires = ["uv_build>=0.10.9,<0.11.0"]
+build-backend = "uv_build"
+
+[tool.uv.build-backend]
+module-name = "edmkit.search"
+
+[tool.ruff]
+exclude = [".claude"]
+
+[tool.ty.src]
+exclude = [".claude"]

edmkit_search-0.0.1a1/src/edmkit/search/__init__.py

@@ -0,0 +1,3 @@
+from . import dataset, energy, neighborhood, state, strategy
+
+__all__ = ["dataset", "energy", "neighborhood", "state", "strategy"]

edmkit_search-0.0.1a1/src/edmkit/search/dataset/__init__.py

@@ -0,0 +1,5 @@
+# ruff: noqa: F401
+"""Dataset subpackage: containers and transforms."""
+
+from .containers import Dataset, Subset
+from .transforms import Transform, compose, gaussian_noise, zscore_normalize

edmkit_search-0.0.1a1/src/edmkit/search/dataset/containers.py

@@ -0,0 +1,90 @@
+from functools import cached_property
+
+import numpy as np
+
+from .transforms import Transform
+
+
+class Dataset:
+    """Time series dataset for X -> Y mapping.
+
+    Parameters
+    ----------
+    `X` : `np.ndarray` of shape `(T, D_x)`
+        Input time series.
+    `Y` : `np.ndarray` of shape `(T, D_y)` or `(T,)`
+        Output time series. 1D is auto-promoted to `(T, 1)`.
+    `transform` : :type: `Transform` or `None`, default `None`
+        `(x, y)` -> `(x', y')` closure for preprocessing / augmentation.
+
+    Raises
+    ------
+    `ValueError`
+        If `X` is not 2-dimensional, `Y` is not 1D or 2D,
+        or if `T` dimensions mismatch.
+    """
+
+    def __init__(
+        self,
+        X: np.ndarray,
+        Y: np.ndarray,
+        *,
+        transform: Transform | None = None,
+    ):
+        if X.ndim != 2:
+            raise ValueError(f"X must be 2-dimensional (T, D_x), got shape {X.shape}")
+        if Y.ndim == 1:
+            Y = Y[:, np.newaxis]
+        if Y.ndim != 2:
+            raise ValueError(f"Y must be 1D or 2D, got shape {Y.shape}")
+        if X.shape[0] != Y.shape[0]:
+            raise ValueError(
+                f"T mismatch: X has {X.shape[0]} timesteps, Y has {Y.shape[0]}"
+            )
+
+        self.X = X.astype(np.float32)
+        self.Y = Y.astype(np.float32)
+        self.transform = transform
+
+    def __len__(self) -> int:
+        return self.X.shape[0]
+
+    def __getitem__(self, idx: int) -> tuple[np.ndarray, np.ndarray]:
+        x = self.X[idx].copy()
+        y = self.Y[idx].copy()
+        if self.transform is not None:
+            x, y = self.transform(x, y)
+        return x, y
+
+
+class Subset(Dataset):
+    """A view into a `Dataset` selected by index, without copying data.
+
+    This is a `Dataset` subtype, so row views can be passed anywhere a
+    dataset is expected.
+
+    Parameters
+    ----------
+    `dataset` : :type: `Dataset`
+        The underlying dataset.
+    `indices` : `np.ndarray`
+        Indices into `dataset` to expose.
+    """
+
+    def __init__(self, dataset: Dataset, indices: np.ndarray):
+        self.dataset = dataset
+        self.indices = indices
+
+    @cached_property
+    def X(self) -> np.ndarray:
+        return self.dataset.X[self.indices]
+
+    @cached_property
+    def Y(self) -> np.ndarray:
+        return self.dataset.Y[self.indices]
+
+    def __len__(self) -> int:
+        return len(self.indices)
+
+    def __getitem__(self, idx: int) -> tuple[np.ndarray, np.ndarray]:
+        return self.dataset[self.indices[idx]]

edmkit_search-0.0.1a1/src/edmkit/search/dataset/transforms.py

@@ -0,0 +1,110 @@
+from collections.abc import Callable
+from functools import reduce
+from typing import TypeAlias
+
+import numpy as np
+
+Transform: TypeAlias = Callable[[np.ndarray, np.ndarray], tuple[np.ndarray, np.ndarray]]
+"""Type alias for a function that takes `(x, y)` and returns `(x', y')`.
+Typically a closure returned by a higher-order preprocessing or data-augmentation function.
+"""
+
+
+def zscore_normalize(
+    data: np.ndarray,
+    *,
+    target: str,
+) -> Transform:
+    """Return a `Transform` that applies z-score normalization using statistics computed from `data`.
+
+    Parameters
+    ----------
+    `data` : `np.ndarray` of shape `(T, D)` or `(N, T, D)`
+        Data from which mean and std are computed.
+    `target` : `str`, default `"x"`
+        Which element of the `(x, y)` pair to normalize.
+        `"x"` normalizes input, `"y"` normalizes output, `"both"` normalizes both
+        (using the same statistics — only valid when `D_x == D_y`).
+
+    Returns
+    -------
+    :type: `Transform`
+        `(x, y)` -> `(x', y')` where the selected target(s) are normalized.
+
+    Examples
+    --------
+    ```python
+    zscore_x = zscore_normalize(X_train, target="x")        # normalize input
+    zscore_y = zscore_normalize(Y_train, target="y")        # normalize output
+    tf = compose(zscore_x, zscore_y)                        # both, independent stats
+    ```
+    """
+    if target not in ("x", "y", "both"):
+        raise ValueError(f"target must be 'x', 'y', or 'both', got {target!r}")
+
+    flat = data.reshape(-1, data.shape[-1])  # (N*T, D) or (T, D)
+    mean = flat.mean(axis=0).astype(np.float32)  # (D,)
+    std = np.clip(flat.std(axis=0).astype(np.float32), 1e-8, None)  # (D,)
+
+    def normalize(a: np.ndarray) -> np.ndarray:
+        return (a - mean) / std
+
+    if target == "x":
+
+        def transform(x: np.ndarray, y: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
+            return (normalize(x), y)
+    elif target == "y":
+
+        def transform(x: np.ndarray, y: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
+            return (x, normalize(y))
+    else:
+
+        def transform(x: np.ndarray, y: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
+            return (normalize(x), normalize(y))
+
+    return transform
+
+
+def gaussian_noise(
+    sigma: float = 0.1,
+    rng: np.random.Generator | None = None,
+) -> Transform:
+    """Return a `Transform` that adds Gaussian noise to the input (for data augmentation).
+
+    Parameters
+    ----------
+    `sigma` : `float`, default `0.1`
+        Standard deviation of the noise.
+    `rng` : `np.random.Generator` or `None`, default `None`
+        Random number generator for reproducibility.
+        If `None`, a new unseeded generator is created.
+
+    Returns
+    -------
+    :type: `Transform`
+        `(x, y)` -> `(x + noise, y)`
+    """
+    rng = np.random.default_rng(rng)
+
+    def transform(x: np.ndarray, y: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
+        noise = rng.standard_normal(x.shape).astype(x.dtype) * sigma
+        return (x + noise, y)
+
+    return transform
+
+
+def compose(*transforms: Transform) -> Transform:
+    """Return a `Transform` that applies multiple transforms in left-to-right order.
+
+    Examples
+    --------
+    ```python
+    transform = compose(zscore_normalize(X_train), gaussian_noise(0.05))
+    # zscore is applied first, then noise
+    ```
+    """
+
+    def transform(x: np.ndarray, y: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
+        return reduce(lambda xy, fn: fn(*xy), transforms, (x, y))
+
+    return transform

edmkit_search-0.0.1a1/src/edmkit/search/energy/__init__.py

@@ -0,0 +1,16 @@
+from edmkit.search.energy.energy import Contexts, Energies, Energy
+from edmkit.search.energy.folds import folds
+from edmkit.search.energy.holdout import holdout
+from edmkit.search.energy.loo import loo
+from edmkit.search.energy.weight import WeightFunc, softmax
+
+__all__ = [
+    "Contexts",
+    "Energies",
+    "Energy",
+    "WeightFunc",
+    "folds",
+    "holdout",
+    "loo",
+    "softmax",
+]

edmkit_search-0.0.1a1/src/edmkit/search/energy/energy.py

@@ -0,0 +1,28 @@
+from collections.abc import Callable
+from dataclasses import dataclass
+
+import numpy as np
+import numpy.typing as npt
+
+from edmkit.search.state import States
+
+type Energies = npt.NDArray[np.float64]
+"""
+Energies is a 1D array of energy values, one for each state.
+"""
+type Contexts = npt.NDArray[np.float64]
+"""
+Contexts is an opaque ndarray that can be used to store any additional information needed for the next energy computation.
+"""
+
+
+@dataclass(frozen=True)
+class Energy:
+    initial: Callable[[], npt.NDArray[np.float64]]
+    """
+    initial returns the initial Contexts entry to feed into the first step call.
+    """
+    step: Callable[
+        [States, Contexts],
+        tuple[Energies, Contexts],
+    ]

edmkit_search-0.0.1a1/src/edmkit/search/energy/folds.py

@@ -0,0 +1,64 @@
+from collections.abc import Sequence
+
+import numpy as np
+import numpy.typing as npt
+from edmkit.metrics import MetricFunc
+from edmkit.splits import Fold
+from edmkit.types import PredictFunc
+
+from edmkit.search import dataset
+from edmkit.search.state import States
+
+from .energy import Contexts, Energies, Energy
+from .weight import WeightFunc
+
+BATCH_SIZE = 10000
+
+
+def folds(
+    *,
+    data: dataset.Dataset,
+    folds: Sequence[Fold],
+    predict: PredictFunc,
+    metric: MetricFunc,
+    weight: WeightFunc,
+) -> Energy:
+    if len(folds) == 0:
+        raise ValueError("folds must be non-empty")
+
+    n_folds = len(folds)
+    subsets = [
+        (dataset.Subset(data, f.train), dataset.Subset(data, f.validation))
+        for f in folds
+    ]
+
+    def initial() -> npt.NDArray[np.float64]:
+        return np.zeros((1, n_folds), dtype=np.float64)
+
+    def step(
+        states: States,
+        contexts: Contexts,
+    ) -> tuple[Energies, Contexts]:
+        n = states.shape[0]
+        metrics = np.empty((n, n_folds), dtype=np.float64)
+        for start in range(0, n, BATCH_SIZE):
+            end = min(start + BATCH_SIZE, n)
+            size = end - start
+            idx = states[start:end].astype(np.intp)
+            for i, (train, validation) in enumerate(subsets):
+                X = np.ascontiguousarray(train.X[:, idx].transpose(1, 0, 2))
+                Y = np.broadcast_to(train.Y, (size, *train.Y.shape))
+                Q = np.ascontiguousarray(validation.X[:, idx].transpose(1, 0, 2))
+                metrics[start:end, i] = metric(
+                    predict(X, Y, Q),
+                    np.broadcast_to(validation.Y, (size, *validation.Y.shape)),
+                )
+
+        energies = np.array(
+            [weight(contexts[i]) @ (metrics[i] - contexts[i]) for i in range(n)],
+            dtype=np.float64,
+        )
+
+        return energies, metrics
+
+    return Energy(initial=initial, step=step)

edmkit_search-0.0.1a1/src/edmkit/search/energy/holdout.py

@@ -0,0 +1,47 @@
+import numpy as np
+import numpy.typing as npt
+from edmkit.metrics import MetricFunc
+from edmkit.splits import Fold
+from edmkit.types import PredictFunc
+
+from edmkit.search import dataset
+from edmkit.search.state import States
+
+from .energy import Contexts, Energies, Energy
+
+BATCH_SIZE = 10000
+
+
+def holdout(
+    *,
+    data: dataset.Dataset,
+    fold: Fold,
+    predict: PredictFunc,
+    metric: MetricFunc,
+) -> Energy:
+    train = dataset.Subset(data, fold.train)
+    validation = dataset.Subset(data, fold.validation)
+
+    def initial() -> npt.NDArray[np.float64]:
+        return np.empty((1, 0), dtype=np.float64)
+
+    def step(
+        states: States,
+        contexts: Contexts,  # not used
+    ) -> tuple[Energies, Contexts]:
+        n = states.shape[0]
+        energies = np.empty(n, dtype=np.float64)
+        for start in range(0, n, BATCH_SIZE):
+            end = min(start + BATCH_SIZE, n)
+            size = end - start
+            idx = states[start:end].astype(np.intp)
+            X = np.ascontiguousarray(train.X[:, idx].transpose(1, 0, 2))
+            Y = np.broadcast_to(train.Y, (size, *train.Y.shape))
+            Q = np.ascontiguousarray(validation.X[:, idx].transpose(1, 0, 2))
+            energies[start:end] = metric(
+                predict(X, Y, Q),
+                np.broadcast_to(validation.Y, (size, *validation.Y.shape)),
+            )
+        return energies, np.empty((n, 0), dtype=np.float64)
+
+    return Energy(initial=initial, step=step)

edmkit_search-0.0.1a1/src/edmkit/search/energy/loo.py

@@ -0,0 +1,46 @@
+import numpy as np
+import numpy.typing as npt
+from edmkit.metrics import MetricFunc
+
+from edmkit import simplex_projection
+from edmkit.search import dataset
+from edmkit.search.state import States
+
+from .energy import Contexts, Energies, Energy
+
+BATCH_SIZE = 10000
+
+
+def loo(
+    *,
+    data: dataset.Dataset,
+    metric: MetricFunc,
+    theiler_window: int = 0,
+    max_batch_size: int | None = None,
+) -> Energy:
+    if theiler_window < 0:
+        raise ValueError("theiler_window must be non-negative")
+
+    def initial() -> npt.NDArray[np.float64]:
+        return np.empty((1, 0), dtype=np.float64)
+
+    def step(
+        states: States,
+        contexts: Contexts,
+    ) -> tuple[Energies, Contexts]:
+        del contexts
+        n = states.shape[0]
+        energies = np.empty(n, dtype=np.float64)
+        for start in range(0, n, BATCH_SIZE):
+            end = min(start + BATCH_SIZE, n)
+            size = end - start
+            idx = states[start:end].astype(np.intp)
+            X = np.ascontiguousarray(data.X[:, idx].transpose(1, 0, 2))
+            Y = np.broadcast_to(data.Y, (size, *data.Y.shape))
+            energies[start:end] = metric(
+                simplex_projection.loo(X, Y, theiler_window=theiler_window),
+                Y,
+            )
+        return energies, np.empty((n, 0), dtype=np.float64)
+
+    return Energy(initial=initial, step=step)

edmkit_search-0.0.1a1/src/edmkit/search/energy/weight.py

@@ -0,0 +1,30 @@
+from collections.abc import Callable
+
+import numpy as np
+import numpy.typing as npt
+
+type WeightFunc = Callable[[npt.NDArray[np.float64]], npt.NDArray[np.float64]]
+
+
+def normalized(values: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
+    total = values.sum()
+    if not np.isfinite(total) or total <= 0:
+        raise ValueError("weights must have a positive finite sum")
+    return values / total
+
+
+def softmax(temperature: float = 1.0) -> WeightFunc:
+    if temperature <= 0:
+        raise ValueError(f"temperature must be positive, got {temperature}")
+
+    def weight(values: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
+        state = np.asarray(values, dtype=np.float64)
+        if state.ndim != 1:
+            raise ValueError(f"weight state must be 1D, got {state.ndim}D")
+        if state.size == 0:
+            return np.empty(0, dtype=np.float64)
+        logits = state / temperature
+        logits = logits - np.max(logits)
+        return normalized(np.exp(logits))
+
+    return weight

edmkit_search-0.0.1a1/src/edmkit/search/neighborhood/__init__.py

@@ -0,0 +1,4 @@
+from edmkit.search.neighborhood.forward import forward
+from edmkit.search.neighborhood.neighborhood import Neighborhood
+
+__all__ = ["Neighborhood", "forward"]

edmkit_search-0.0.1a1/src/edmkit/search/neighborhood/forward.py

@@ -0,0 +1,45 @@
+import numpy as np
+import numpy.typing as npt
+
+from edmkit.search.state import States
+
+from .neighborhood import Neighborhood
+
+
+def forward(n: int) -> Neighborhood:
+    """Forward search over `n` candidate indices.
+
+    Each parent state of length `d` (assumed to hold unique indices in `[0, n)`)
+    expands into exactly `n - d` children — one per index not yet selected. Within
+    a parent, children are emitted in a per-row random order; across parents, the
+    parent order is preserved.
+    """
+    if n < 0:
+        raise ValueError(f"n must be non-negative, got {n}")
+
+    def expand(
+        states: States, rng: np.random.Generator
+    ) -> tuple[States, npt.NDArray[np.int64]]:
+        N, d = states.shape
+        per_parent = n - d
+        if N == 0 or per_parent == 0:
+            return (
+                np.empty((0, d + 1), dtype=np.int64),
+                np.empty(0, dtype=np.int64),
+            )
+
+        # Per-row random permutation of [0, n).
+        perms = rng.permuted(np.tile(np.arange(n, dtype=np.int64), (N, 1)), axis=1)
+        # mask[i, j] = True iff j is not already in states[i].
+        mask = np.ones((N, n), dtype=bool)
+        mask[np.arange(N)[:, None], states] = False
+        # Each row of `perms` has exactly `per_parent` survivors (states are unique).
+        survivors = np.take_along_axis(mask, perms, axis=1)
+        chosen = perms[survivors].reshape(N, per_parent)
+
+        parents = np.repeat(np.arange(N, dtype=np.int64), per_parent)
+        prefix = np.repeat(states, per_parent, axis=0)
+        children = np.concatenate([prefix, chosen.reshape(-1, 1)], axis=1)
+        return children, parents
+
+    return expand

edmkit_search-0.0.1a1/src/edmkit/search/neighborhood/neighborhood.py

@@ -0,0 +1,23 @@
+from collections.abc import Callable
+
+import numpy as np
+import numpy.typing as npt
+
+from edmkit.search.state import States
+
+type Neighborhood = Callable[
+    [States, np.random.Generator],
+    tuple[States, npt.NDArray[np.int64]],
+]
+"""
+A Neighborhood expands a batch of N parent states into M children.
+
+Given parents of shape (N, d), it returns:
+  * `children`: shape (M, d') — the next-step states
+  * `parents`:  shape (M,)    — `parents[i] in [0, N)` points to the parent row that
+                                  produced `children[i]`
+
+`parents` lets callers (e.g. beam search) replicate parent-side data such as the
+energy context in lockstep with the children, without the neighborhood needing to
+know about that data.
+"""

edmkit_search-0.0.1a1/src/edmkit/search/state/__init__.py

@@ -0,0 +1,3 @@
+from edmkit.search.state.states import States, initial
+
+__all__ = ["States", "initial"]

edmkit_search-0.0.1a1/src/edmkit/search/state/states.py

@@ -0,0 +1,13 @@
+import numpy as np
+import numpy.typing as npt
+
+type States = npt.NDArray[np.int64]
+"""
+States is a 2D array of shape (N, d): N states, each holding d indices into the original dataset.
+A single step's batch always contains states of equal length d.
+"""
+
+
+def initial() -> States:
+    """Initial states batch: a single empty state, shape (1, 0)."""
+    return np.empty((1, 0), dtype=np.int64)

edmkit_search-0.0.1a1/src/edmkit/search/strategy/__init__.py

@@ -0,0 +1,6 @@
+from edmkit.search.strategy.beam import beam
+from edmkit.search.strategy.frontier import Frontier, Step
+from edmkit.search.strategy.greedy import greedy
+from edmkit.search.strategy.run import run
+
+__all__ = ["Frontier", "Step", "beam", "greedy", "run"]

edmkit_search-0.0.1a1/src/edmkit/search/strategy/beam.py

@@ -0,0 +1,37 @@
+import numpy as np
+
+from edmkit.search.energy import Energy
+from edmkit.search.neighborhood import Neighborhood
+
+from .frontier import Frontier, Step
+
+
+def beam(
+    E: Energy,
+    N: Neighborhood,
+    *,
+    width: int,
+    cutoff: float = float("inf"),
+) -> Step:
+    if width < 1:
+        raise ValueError(f"width must be >= 1, got {width}")
+
+    def step(frontier: Frontier, rng: np.random.Generator) -> Frontier:
+        children, parents = N(frontier.states, rng)
+        if children.shape[0] == 0:
+            return Frontier(
+                states=children,
+                contexts=frontier.contexts[:0],
+                energies=np.empty(0, dtype=np.float64),
+            )
+
+        energies, contexts = E.step(children, frontier.contexts[parents])
+        kept = np.flatnonzero(energies <= cutoff)
+        order = kept[np.argsort(energies[kept], kind="stable")[:width]]
+        return Frontier(
+            states=children[order],
+            contexts=contexts[order],
+            energies=energies[order],
+        )
+
+    return step

edmkit_search-0.0.1a1/src/edmkit/search/strategy/frontier.py

@@ -0,0 +1,23 @@
+from collections.abc import Callable
+from dataclasses import dataclass
+
+import numpy as np
+
+from edmkit.search.energy import Contexts, Energies
+from edmkit.search.state import States
+
+
+@dataclass(frozen=True)
+class Frontier:
+    states: States
+    contexts: Contexts
+    energies: Energies
+
+    def __len__(self) -> int:
+        return self.states.shape[0]
+
+
+type Step = Callable[
+    [Frontier, np.random.Generator],
+    Frontier,
+]

edmkit_search-0.0.1a1/src/edmkit/search/strategy/greedy.py

@@ -0,0 +1,13 @@
+from edmkit.search.energy import Energy
+from edmkit.search.neighborhood import Neighborhood
+from edmkit.search.strategy.beam import beam
+from edmkit.search.strategy.frontier import Step
+
+
+def greedy(
+    E: Energy,
+    N: Neighborhood,
+    *,
+    cutoff: float = float("inf"),
+) -> Step:
+    return beam(E, N, width=1, cutoff=cutoff)

edmkit_search-0.0.1a1/src/edmkit/search/strategy/run.py

@@ -0,0 +1,28 @@
+from collections.abc import Iterator
+
+import numpy as np
+
+from .frontier import Frontier, Step
+
+
+def run(
+    initial: Frontier,
+    step: Step,
+    *,
+    max_steps: int,
+    rng: np.random.Generator,
+) -> Iterator[Frontier]:
+    if max_steps < 0:
+        raise ValueError(f"max_steps must be non-negative, got {max_steps}")
+
+    frontier = initial
+    for _ in range(max_steps):
+        frontier = step(frontier, rng)
+        if not frontier:
+            return
+        i = int(np.argmin(frontier.energies))
+        yield Frontier(
+            states=frontier.states[i : i + 1],
+            contexts=frontier.contexts[i : i + 1],
+            energies=frontier.energies[i : i + 1],
+        )