ennbo 0.1.0__py3-none-any.whl → 0.1.7__py3-none-any.whl

enn/enn/enn_conditional.py

@@ -0,0 +1,325 @@
+from __future__ import annotations
+import numpy as np
+from .candidates import Candidates
+from .conditional_posterior_draw_internals import ConditionalPosteriorDrawInternals
+from .enn_like_protocol import ENNLike
+from .enn_params import ENNParams, PosteriorFlags
+from .neighbors import Neighbors
+
+_ENNLike = ENNLike
+_Candidates = Candidates
+_Neighbors = Neighbors
+
+
+def _pairwise_sq_l2(a: np.ndarray, b: np.ndarray) -> np.ndarray:
+    a = np.asarray(a, dtype=float)
+    b = np.asarray(b, dtype=float)
+    aa = np.sum(a * a, axis=1, keepdims=True)
+    bb = np.sum(b * b, axis=1, keepdims=True).T
+    dist2 = aa + bb - 2.0 * (a @ b.T)
+    return np.maximum(dist2, 0.0)
+
+
+def _validate_x(enn: ENNLike, x: np.ndarray) -> np.ndarray:
+    x = np.asarray(x, dtype=float)
+    if x.ndim != 2 or x.shape[1] != enn._num_dim:
+        raise ValueError(x.shape)
+    return x
+
+
+def _validate_whatif(
+    enn: ENNLike, x_whatif: np.ndarray, y_whatif: np.ndarray
+) -> tuple[np.ndarray, np.ndarray]:
+    x_whatif = np.asarray(x_whatif, dtype=float)
+    y_whatif = np.asarray(y_whatif, dtype=float)
+    if x_whatif.ndim != 2 or x_whatif.shape[1] != enn._num_dim:
+        raise ValueError(x_whatif.shape)
+    if y_whatif.ndim != 2 or y_whatif.shape[1] != enn._num_metrics:
+        raise ValueError(y_whatif.shape)
+    if x_whatif.shape[0] != y_whatif.shape[0]:
+        raise ValueError((x_whatif.shape, y_whatif.shape))
+    return x_whatif, y_whatif
+
+
+def _scale_x_if_needed(enn: ENNLike, x: np.ndarray) -> np.ndarray:
+    return x / enn._x_scale if enn._scale_x else x
+
+
+def _compute_total_n(enn: ENNLike, num_whatif: int, flags: PosteriorFlags) -> int:
+    total_n = len(enn) + int(num_whatif)
+    if flags.exclude_nearest and total_n <= 1:
+        raise ValueError(total_n)
+    return total_n
+
+
+def _compute_search_k(params: ENNParams, flags: PosteriorFlags, total_n: int) -> int:
+    return int(
+        min(params.k_num_neighbors + (1 if flags.exclude_nearest else 0), total_n)
+    )
+
+
+def _get_train_candidates(enn: ENNLike, x: np.ndarray, *, search_k: int) -> Candidates:
+    batch_size = x.shape[0]
+    if len(enn) == 0 or search_k == 0:
+        return Candidates(
+            dist2=np.zeros((batch_size, 0), dtype=float),
+            ids=np.zeros((batch_size, 0), dtype=int),
+            y=np.zeros((batch_size, 0, enn._num_metrics), dtype=float),
+            yvar=(
+                np.zeros((batch_size, 0, enn._num_metrics), dtype=float)
+                if enn._train_yvar is not None
+                else None
+            ),
+        )
+    train_search_k = int(min(search_k, len(enn)))
+    dist2_train, idx_train = enn._enn_index.search(
+        x, search_k=train_search_k, exclude_nearest=False
+    )
+    y_train = enn._train_y[idx_train]
+    yvar_train = enn._train_yvar[idx_train] if enn._train_yvar is not None else None
+    return Candidates(dist2=dist2_train, ids=idx_train, y=y_train, yvar=yvar_train)
+
+
+def _get_whatif_candidates(
+    enn: ENNLike,
+    x: np.ndarray,
+    x_whatif: np.ndarray,
+    y_whatif: np.ndarray,
+) -> tuple[np.ndarray, np.ndarray]:
+    x_scaled = _scale_x_if_needed(enn, x)
+    x_whatif_scaled = _scale_x_if_needed(enn, x_whatif)
+    dist2_whatif = _pairwise_sq_l2(x_scaled, x_whatif_scaled)
+    batch_size = x.shape[0]
+    y_whatif_batched = np.broadcast_to(
+        y_whatif[np.newaxis, :, :], (batch_size, y_whatif.shape[0], y_whatif.shape[1])
+    )
+    return dist2_whatif, y_whatif_batched
+
+
+_WhatifCandidateTuple = tuple[np.ndarray, np.ndarray, np.ndarray]
+
+
+def _merge_candidates(
+    enn: ENNLike,
+    *,
+    train: Candidates,
+    whatif: _WhatifCandidateTuple,
+) -> Candidates:
+    dist2_whatif, ids_whatif, y_whatif_batched = whatif
+    dist2_all = np.concatenate([train.dist2, dist2_whatif], axis=1)
+    ids_all = np.concatenate([train.ids, ids_whatif], axis=1)
+    y_all = np.concatenate([train.y, y_whatif_batched], axis=1)
+    if train.yvar is None:
+        return Candidates(dist2=dist2_all, ids=ids_all, y=y_all, yvar=None)
+    batch_size = dist2_all.shape[0]
+    num_whatif = dist2_whatif.shape[1]
+    yvar_whatif = np.zeros((batch_size, num_whatif, enn._num_metrics))
+    yvar_all = np.concatenate([train.yvar, yvar_whatif], axis=1)
+    return Candidates(dist2=dist2_all, ids=ids_all, y=y_all, yvar=yvar_all)
+
+
+def _select_sorted_candidates(dist2_all: np.ndarray, *, search_k: int) -> np.ndarray:
+    batch_size, num_candidates = dist2_all.shape
+    if search_k < num_candidates:
+        sel = np.argpartition(dist2_all, kth=search_k - 1, axis=1)[:, :search_k]
+    else:
+        sel = np.broadcast_to(np.arange(num_candidates), (batch_size, num_candidates))
+    sel_dist2 = np.take_along_axis(dist2_all, sel, axis=1)
+    sel_order = np.argsort(sel_dist2, axis=1)
+    return np.take_along_axis(sel, sel_order, axis=1)
+
+
+def _take_along_axis_3d(a: np.ndarray, idx_2d: np.ndarray) -> np.ndarray:
+    return np.take_along_axis(a, idx_2d[:, :, np.newaxis], axis=1)
+
+
+def _make_empty_normal(enn: ENNLike, batch_size: int):
+    from .enn_normal import ENNNormal
+
+    internals = enn._empty_posterior_internals(batch_size)
+    return ENNNormal(internals.mu, internals.se)
+
+
+def _build_candidates(
+    enn: ENNLike,
+    x: np.ndarray,
+    x_whatif: np.ndarray,
+    y_whatif: np.ndarray,
+    *,
+    search_k: int,
+) -> Candidates:
+    train_candidates = _get_train_candidates(enn, x, search_k=search_k)
+    dist2_whatif, y_whatif_batched = _get_whatif_candidates(enn, x, x_whatif, y_whatif)
+    n_train = int(len(enn))
+    ids_whatif = np.broadcast_to(
+        n_train + np.arange(x_whatif.shape[0], dtype=int), dist2_whatif.shape
+    )
+    return _merge_candidates(
+        enn,
+        train=train_candidates,
+        whatif=(dist2_whatif, ids_whatif, y_whatif_batched),
+    )
+
+
+def _select_effective_neighbors(
+    candidates: Candidates,
+    *,
+    search_k: int,
+    k: int,
+    exclude_nearest: bool,
+) -> Neighbors | None:
+    sel = _select_sorted_candidates(candidates.dist2, search_k=search_k)
+    if exclude_nearest:
+        sel = sel[:, 1:]
+    sel = sel[:, : int(min(k, sel.shape[1]))]
+    if sel.shape[1] == 0:
+        return None
+    dist2s = np.take_along_axis(candidates.dist2, sel, axis=1)
+    ids = np.take_along_axis(candidates.ids, sel, axis=1)
+    y_neighbors = _take_along_axis_3d(candidates.y, sel)
+    yvar_neighbors = (
+        _take_along_axis_3d(candidates.yvar, sel)
+        if candidates.yvar is not None
+        else None
+    )
+    return Neighbors(dist2=dist2s, ids=ids, y=y_neighbors, yvar=yvar_neighbors)
+
+
+def _compute_mu_se(
+    enn: ENNLike,
+    neighbors: Neighbors,
+    *,
+    params: ENNParams,
+    flags: PosteriorFlags,
+    y_scale: np.ndarray,
+) -> tuple[np.ndarray, np.ndarray]:
+    stats = enn._compute_weighted_stats(
+        neighbors.dist2,
+        neighbors.y,
+        yvar_neighbors=neighbors.yvar,
+        params=params,
+        observation_noise=flags.observation_noise,
+        y_scale=y_scale,
+    )
+    return stats.mu, stats.se
+
+
+def _compute_draw_internals(
+    enn: ENNLike,
+    neighbors: Neighbors,
+    *,
+    params: ENNParams,
+    flags: PosteriorFlags,
+    y_scale: np.ndarray,
+) -> ConditionalPosteriorDrawInternals:
+    stats = enn._compute_weighted_stats(
+        neighbors.dist2,
+        neighbors.y,
+        yvar_neighbors=neighbors.yvar,
+        params=params,
+        observation_noise=flags.observation_noise,
+        y_scale=y_scale,
+    )
+    return ConditionalPosteriorDrawInternals(
+        idx=neighbors.ids.astype(int, copy=False),
+        w_normalized=stats.w_normalized,
+        l2=stats.l2,
+        mu=stats.mu,
+        se=stats.se,
+    )
+
+
+def _conditional_neighbors_nonempty_whatif(
+    enn: ENNLike,
+    x_whatif: np.ndarray,
+    y_whatif: np.ndarray,
+    x: np.ndarray,
+    *,
+    params: ENNParams,
+    flags: PosteriorFlags,
+) -> tuple[int, int, Neighbors | None]:
+    batch_size = x.shape[0]
+    search_k = _compute_search_k(
+        params, flags, _compute_total_n(enn, x_whatif.shape[0], flags)
+    )
+    if search_k == 0:
+        return batch_size, search_k, None
+    candidates = _build_candidates(enn, x, x_whatif, y_whatif, search_k=search_k)
+    neighbors = _select_effective_neighbors(
+        candidates,
+        search_k=search_k,
+        k=params.k_num_neighbors,
+        exclude_nearest=flags.exclude_nearest,
+    )
+    return batch_size, search_k, neighbors
+
+
+def _compute_conditional_posterior_impl(
+    enn: ENNLike,
+    x_whatif: np.ndarray,
+    y_whatif: np.ndarray,
+    x: np.ndarray,
+    *,
+    params: ENNParams,
+    flags: PosteriorFlags,
+    y_scale: np.ndarray,
+):
+    from .enn_normal import ENNNormal
+
+    x = _validate_x(enn, x)
+    x_whatif, y_whatif = _validate_whatif(enn, x_whatif, y_whatif)
+    if x_whatif.shape[0] == 0:
+        return enn.posterior(x, params=params, flags=flags)
+    batch_size, search_k, neighbors = _conditional_neighbors_nonempty_whatif(
+        enn, x_whatif, y_whatif, x, params=params, flags=flags
+    )
+    if search_k == 0 or neighbors is None:
+        return _make_empty_normal(enn, batch_size)
+    mu, se = _compute_mu_se(enn, neighbors, params=params, flags=flags, y_scale=y_scale)
+    return ENNNormal(mu, se)
+
+
+def compute_conditional_posterior(
+    enn: ENNLike,
+    x_whatif: np.ndarray,
+    y_whatif: np.ndarray,
+    x: np.ndarray,
+    *,
+    params: ENNParams,
+    flags: PosteriorFlags,
+    y_scale: np.ndarray,
+):
+    return _compute_conditional_posterior_impl(
+        enn, x_whatif, y_whatif, x, params=params, flags=flags, y_scale=y_scale
+    )
+
+
+def compute_conditional_posterior_draw_internals(
+    enn: ENNLike,
+    x_whatif: np.ndarray,
+    y_whatif: np.ndarray,
+    x: np.ndarray,
+    *,
+    params: ENNParams,
+    flags: PosteriorFlags,
+    y_scale: np.ndarray,
+) -> ConditionalPosteriorDrawInternals:
+    x = _validate_x(enn, x)
+    x_whatif, y_whatif = _validate_whatif(enn, x_whatif, y_whatif)
+    if x_whatif.shape[0] == 0:
+        raise ValueError("x_whatif must be non-empty for conditional draw internals")
+    batch_size, search_k, neighbors = _conditional_neighbors_nonempty_whatif(
+        enn, x_whatif, y_whatif, x, params=params, flags=flags
+    )
+    if search_k == 0 or neighbors is None:
+        empty_internals = enn._empty_posterior_internals(batch_size)
+        return ConditionalPosteriorDrawInternals(
+            idx=empty_internals.idx,
+            w_normalized=empty_internals.w_normalized,
+            l2=empty_internals.l2,
+            mu=empty_internals.mu,
+            se=empty_internals.se,
+        )
+    return _compute_draw_internals(
+        enn, neighbors, params=params, flags=flags, y_scale=y_scale
+    )

enn/enn/enn_fit.py

@@ -1,15 +1,49 @@
 from __future__ import annotations
-
 from typing import TYPE_CHECKING, Any
 
 if TYPE_CHECKING:
     import numpy as np
     from numpy.random import Generator
-
-    from .enn import EpistemicNearestNeighbors
+    from .enn_class import EpistemicNearestNeighbors
     from .enn_params import ENNParams
 
-from .enn_util import standardize_y
+
+def _validate_subsample_inputs(
+    x: np.ndarray | Any, y: np.ndarray | Any, P: int, paramss: list
+) -> tuple[np.ndarray, np.ndarray]:
+    import numpy as np
+
+    x_array = np.asarray(x, dtype=float)
+    if x_array.ndim != 2:
+        raise ValueError(x_array.shape)
+    y_array = np.asarray(y, dtype=float)
+    if y_array.ndim == 1:
+        y_array = y_array.reshape(-1, 1)
+    if y_array.ndim != 2:
+        raise ValueError(y_array.shape)
+    if x_array.shape[0] != y_array.shape[0]:
+        raise ValueError((x_array.shape, y_array.shape))
+    if P <= 0:
+        raise ValueError(P)
+    if len(paramss) == 0:
+        raise ValueError("paramss must be non-empty")
+    return x_array, y_array
+
+
+def _compute_single_loglik(
+    y_scaled: np.ndarray, mu_i: np.ndarray, se_i: np.ndarray
+) -> float:
+    import numpy as np
+
+    if not np.isfinite(mu_i).all() or not np.isfinite(se_i).all():
+        return 0.0
+    if np.any(se_i <= 0.0):
+        return 0.0
+    var_scaled = se_i**2
+    loglik = -0.5 * np.sum(
+        np.log(2.0 * np.pi * var_scaled) + (y_scaled - mu_i) ** 2 / var_scaled
+    )
+    return float(loglik) if np.isfinite(loglik) else 0.0
 
 
 def subsample_loglik(
@@ -23,68 +57,37 @@ def subsample_loglik(
 ) -> list[float]:
     import numpy as np
 
-    if x.ndim != 2:
-        raise ValueError(x.shape)
-    if y.ndim != 1:
-        raise ValueError(y.shape)
-    if x.shape[0] != y.shape[0]:
-        raise ValueError((x.shape, y.shape))
-    if P <= 0:
-        raise ValueError(P)
-    if len(paramss) == 0:
-        raise ValueError("paramss must be non-empty")
-    n = x.shape[0]
-    if n == 0:
-        return [0.0] * len(paramss)
-    if len(model) <= 1:
+    x_array, y_array = _validate_subsample_inputs(x, y, P, paramss)
+    n = x_array.shape[0]
+    if n == 0 or len(model) <= 1:
         return [0.0] * len(paramss)
     P_actual = min(P, n)
-    if P_actual == n:
-        indices = np.arange(n, dtype=int)
-    else:
-        indices = rng.permutation(n)[:P_actual]
-    x_selected = x[indices]
-    y_selected = y[indices]
-    if not np.isfinite(y_selected).all():
+    indices = (
+        np.arange(n, dtype=int) if P_actual == n else rng.permutation(n)[:P_actual]
+    )
+    x_sel, y_sel = x_array[indices], y_array[indices]
+    if not np.isfinite(y_sel).all():
         return [0.0] * len(paramss)
-    post_batch = model.batch_posterior(
-        x_selected, paramss, exclude_nearest=True, observation_noise=True
+    from .enn_params import PosteriorFlags
+
+    post = model.batch_posterior(
+        x_sel,
+        paramss,
+        flags=PosteriorFlags(exclude_nearest=True, observation_noise=True),
     )
-    mu_batch = post_batch.mu
-    se_batch = post_batch.se
-    if mu_batch.shape[2] == 1:
-        mu_batch = mu_batch[:, :, 0]
-        se_batch = se_batch[:, :, 0]
-    num_params = len(paramss)
-    if mu_batch.shape != (num_params, P_actual) or se_batch.shape != (
-        num_params,
-        P_actual,
-    ):
-        raise ValueError((mu_batch.shape, se_batch.shape, (num_params, P_actual)))
-    _, y_std = standardize_y(y)
-    y_scaled = y_selected / y_std
-    mu_scaled = mu_batch / y_std
-    se_scaled = se_batch / y_std
-    result = []
-    for i in range(num_params):
-        mu_i = mu_scaled[i]
-        se_i = se_scaled[i]
-        if not np.isfinite(mu_i).all() or not np.isfinite(se_i).all():
-            result.append(0.0)
-            continue
-        if np.any(se_i <= 0.0):
-            result.append(0.0)
-            continue
-        diff = y_scaled - mu_i
-        var_scaled = se_i**2
-        log_term = np.log(2.0 * np.pi * var_scaled)
-        quad = diff**2 / var_scaled
-        loglik = -0.5 * np.sum(log_term + quad)
-        if not np.isfinite(loglik):
-            result.append(0.0)
-            continue
-        result.append(float(loglik))
-    return result
+    num_params, num_outputs = len(paramss), y_sel.shape[1]
+    expected_shape = (num_params, P_actual, num_outputs)
+    if post.mu.shape != expected_shape or post.se.shape != expected_shape:
+        raise ValueError((post.mu.shape, post.se.shape, expected_shape))
+    y_std = np.std(y_array, axis=0, keepdims=True).astype(float)
+    y_std = np.where(np.isfinite(y_std) & (y_std > 0.0), y_std, 1.0)
+    y_scaled = y_sel / y_std
+    mu_scaled = post.mu / y_std
+    se_scaled = post.se / y_std
+    return [
+        _compute_single_loglik(y_scaled, mu_scaled[i], se_scaled[i])
+        for i in range(num_params)
+    ]
 
 
 def enn_fit(
@@ -100,12 +103,6 @@ def enn_fit(
 
     train_x = model.train_x
     train_y = model.train_y
-    train_yvar = model.train_yvar
-    if train_y.shape[1] != 1:
-        raise ValueError(train_y.shape)
-    if train_yvar is not None and train_yvar.shape[1] != 1:
-        raise ValueError(train_yvar.shape)
-    y = train_y[:, 0]
     log_min = -3.0
     log_max = 3.0
     epi_var_scale_log_values = rng.uniform(log_min, log_max, size=num_fit_candidates)
@@ -116,26 +113,30 @@ def enn_fit(
     ale_homoscedastic_values = 10**ale_homoscedastic_log_values
     paramss = [
         ENNParams(
-            k=k,
-            epi_var_scale=float(epi_val),
-            ale_homoscedastic_scale=float(ale_val),
+            k_num_neighbors=k,
+            epistemic_variance_scale=float(epi_val),
+            aleatoric_variance_scale=float(ale_val),
         )
         for epi_val, ale_val in zip(epi_var_scale_values, ale_homoscedastic_values)
     ]
     if params_warm_start is not None:
         paramss.append(
             ENNParams(
-                k=k,
-                epi_var_scale=params_warm_start.epi_var_scale,
-                ale_homoscedastic_scale=params_warm_start.ale_homoscedastic_scale,
+                k_num_neighbors=k,
+                epistemic_variance_scale=params_warm_start.epistemic_variance_scale,
+                aleatoric_variance_scale=params_warm_start.aleatoric_variance_scale,
             )
         )
     if len(paramss) == 0:
-        return ENNParams(k=k, epi_var_scale=1.0, ale_homoscedastic_scale=0.0)
+        return ENNParams(
+            k_num_neighbors=k,
+            epistemic_variance_scale=1.0,
+            aleatoric_variance_scale=0.0,
+        )
     import numpy as np
 
     logliks = subsample_loglik(
-        model, train_x, y, paramss=paramss, P=num_fit_samples, rng=rng
+        model, train_x, train_y, paramss=paramss, P=num_fit_samples, rng=rng
     )
     if len(logliks) == 0:
         return paramss[0]

enn/enn/enn_hash.py

@@ -0,0 +1,79 @@
+from __future__ import annotations
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import numpy as np
+
+
+def normal_hash_batch_multi_seed(
+    function_seeds: np.ndarray, data_indices: np.ndarray, num_metrics: int
+) -> np.ndarray:
+    import numpy as np
+    from scipy.special import ndtri
+
+    num_seeds = len(function_seeds)
+    unique_indices, inverse = np.unique(data_indices, return_inverse=True)
+    num_unique = len(unique_indices)
+    seed_grid, idx_grid, metric_grid = np.meshgrid(
+        function_seeds.astype(np.uint64),
+        unique_indices.astype(np.uint64),
+        np.arange(num_metrics, dtype=np.uint64),
+        indexing="ij",
+    )
+    seed_flat = seed_grid.ravel()
+    idx_flat = idx_grid.ravel()
+    metric_flat = metric_grid.ravel()
+    combined_seeds = (seed_flat * np.uint64(1_000_003) + idx_flat) * np.uint64(
+        1_000_003
+    ) + metric_flat
+    uniform_vals = np.empty(len(combined_seeds), dtype=float)
+    for i, seed in enumerate(combined_seeds):
+        rng = np.random.Generator(np.random.Philox(int(seed)))
+        uniform_vals[i] = rng.random()
+    uniform_vals = np.clip(uniform_vals, 1e-10, 1.0 - 1e-10)
+    normal_vals = ndtri(uniform_vals).reshape(num_seeds, num_unique, num_metrics)
+    return normal_vals[:, inverse.ravel(), :].reshape(
+        num_seeds, *data_indices.shape, num_metrics
+    )
+
+
+def normal_hash_batch_multi_seed_fast(
+    function_seeds: np.ndarray, data_indices: np.ndarray, num_metrics: int
+) -> np.ndarray:
+    import numpy as np
+
+    function_seeds = np.asarray(function_seeds, dtype=np.int64)
+    data_indices = np.asarray(data_indices)
+    if num_metrics <= 0:
+        raise ValueError(num_metrics)
+    num_seeds = len(function_seeds)
+    unique_indices, inverse = np.unique(data_indices, return_inverse=True)
+
+    def _splitmix64(x: np.ndarray) -> np.ndarray:
+        with np.errstate(over="ignore"):
+            x = x + np.uint64(0x9E3779B97F4A7C15)
+            z = x
+            z = (z ^ (z >> np.uint64(30))) * np.uint64(0xBF58476D1CE4E5B9)
+            z = (z ^ (z >> np.uint64(27))) * np.uint64(0x94D049BB133111EB)
+            z = z ^ (z >> np.uint64(31))
+            return z
+
+    seeds_u64 = function_seeds.astype(np.uint64, copy=False)
+    unique_u64 = unique_indices.astype(np.uint64, copy=False)
+    metric_u64 = np.arange(num_metrics, dtype=np.uint64)
+    normal_vals = np.empty((num_seeds, unique_indices.size, num_metrics), dtype=float)
+    p = np.uint64(1_000_003)
+    inv_2p53 = 1.0 / 9007199254740992.0
+    for si, s in enumerate(seeds_u64):
+        with np.errstate(over="ignore"):
+            base = (s * p + unique_u64) * p
+        combined = base[:, None] + metric_u64[None, :]
+        r1 = _splitmix64(combined)
+        r2 = _splitmix64(combined ^ np.uint64(0xD2B74407B1CE6E93))
+        u1 = (r1 >> np.uint64(11)).astype(np.float64) * inv_2p53
+        u2 = (r2 >> np.uint64(11)).astype(np.float64) * inv_2p53
+        u1 = np.clip(u1, 1e-12, 1.0 - 1e-12)
+        normal_vals[si, :, :] = np.sqrt(-2.0 * np.log(u1)) * np.cos(2.0 * np.pi * u2)
+    return normal_vals[:, inverse.ravel(), :].reshape(
+        num_seeds, *data_indices.shape, num_metrics
+    )

enn/enn/enn_index.py

@@ -0,0 +1,92 @@
+from __future__ import annotations
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    import numpy as np
+
+
+class ENNIndex:
+    def __init__(
+        self,
+        train_x_scaled: np.ndarray,
+        num_dim: int,
+        x_scale: np.ndarray,
+        scale_x: bool,
+        driver: Any = None,
+    ) -> None:
+        from enn.turbo.config.enums import ENNIndexDriver
+
+        if driver is None:
+            driver = ENNIndexDriver.FLAT
+        self._train_x_scaled = train_x_scaled
+        self._num_dim = num_dim
+        self._x_scale = x_scale
+        self._scale_x = scale_x
+        self._driver = driver
+        self._index: Any | None = None
+        self._build_index()
+
+    def _build_index(self) -> None:
+        import faiss
+        import numpy as np
+        from enn.turbo.config.enums import ENNIndexDriver
+
+        if len(self._train_x_scaled) == 0:
+            return
+        x_f32 = self._train_x_scaled.astype(np.float32, copy=False)
+        if self._driver == ENNIndexDriver.FLAT:
+            index = faiss.IndexFlatL2(self._num_dim)
+        elif self._driver == ENNIndexDriver.HNSW:
+            # TODO: Make M configurable
+            index = faiss.IndexHNSWFlat(self._num_dim, 32)
+        else:
+            raise ValueError(f"Unknown driver: {self._driver}")
+        index.add(x_f32)
+        self._index = index
+
+    def add(self, x: np.ndarray) -> None:
+        import numpy as np
+        from enn.turbo.config.enums import ENNIndexDriver
+
+        x = np.asarray(x, dtype=float)
+        if x.ndim != 2 or x.shape[1] != self._num_dim:
+            raise ValueError(x.shape)
+        x_scaled = x / self._x_scale if self._scale_x else x
+        x_f32 = x_scaled.astype(np.float32, copy=False)
+        if self._index is None:
+            import faiss
+
+            if self._driver == ENNIndexDriver.FLAT:
+                self._index = faiss.IndexFlatL2(self._num_dim)
+            elif self._driver == ENNIndexDriver.HNSW:
+                self._index = faiss.IndexHNSWFlat(self._num_dim, 32)
+            else:
+                raise ValueError(f"Unknown driver: {self._driver}")
+        self._index.add(x_f32)
+
+    def search(
+        self,
+        x: np.ndarray,
+        *,
+        search_k: int,
+        exclude_nearest: bool,
+    ) -> tuple[np.ndarray, np.ndarray]:
+        import numpy as np
+
+        search_k = int(search_k)
+        if search_k <= 0:
+            raise ValueError(search_k)
+        x = np.asarray(x, dtype=float)
+        if x.ndim != 2 or x.shape[1] != self._num_dim:
+            raise ValueError(x.shape)
+        if self._index is None:
+            raise RuntimeError("index is not initialized")
+        x_scaled = x / self._x_scale if self._scale_x else x
+        x_f32 = x_scaled.astype(np.float32, copy=False)
+        dist2s_full, idx_full = self._index.search(x_f32, search_k)
+        dist2s_full = dist2s_full.astype(float)
+        idx_full = idx_full.astype(int)
+        if exclude_nearest:
+            dist2s_full = dist2s_full[:, 1:]
+            idx_full = idx_full[:, 1:]
+        return dist2s_full, idx_full