edmkit 0.0.14__tar.gz → 0.0.15__tar.gz

{edmkit-0.0.14 → edmkit-0.0.15}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: edmkit
-Version: 0.0.14
+Version: 0.0.15
 Summary: Simple EDM (Empirical Dynamic Modeling) library
 Author: FUJISHIGE TEMMA
 Author-email: FUJISHIGE TEMMA <tenma.x0@gmail.com>

{edmkit-0.0.14 → edmkit-0.0.15}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "edmkit"
-version = "0.0.14"
+version = "0.0.15"
 description = "Simple EDM (Empirical Dynamic Modeling) library"
 readme = "README.md"
 requires-python = ">= 3.13"

{edmkit-0.0.14 → edmkit-0.0.15}/pyproject.toml.orig

@@ -1,6 +1,6 @@
 [project]
 name = "edmkit"
-version = "0.0.14"
+version = "0.0.15"
 description = "Simple EDM (Empirical Dynamic Modeling) library"
 authors = [{ name = "FUJISHIGE TEMMA", email = "tenma.x0@gmail.com" }]
 readme = "README.md"

edmkit-0.0.15/src/edmkit/simplex_projection/__init__.py

@@ -0,0 +1,7 @@
+from .knn import knn as knn
+from .linear_operator import linear_operator
+from .simplex_projection import simplex_projection
+from .soft import soft_simplex_projection
+from .soft_linear_operator import soft_linear_operator
+
+__all__ = ["linear_operator", "simplex_projection", "soft_linear_operator", "soft_simplex_projection"]

edmkit-0.0.15/src/edmkit/simplex_projection/linear_operator.py

@@ -0,0 +1,206 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, overload
+
+import numpy as np
+
+from edmkit.simplex_projection.knn import knn
+from edmkit.util import pairwise_distance
+
+__all__ = ["linear_operator"]
+
+if TYPE_CHECKING:
+    from tinygrad import Tensor
+
+
+@overload
+def linear_operator(
+    X: np.ndarray,
+    Q: np.ndarray,
+    *,
+    k: int | None = None,
+    mask: np.ndarray | None = None,
+) -> np.ndarray: ...
+
+
+@overload
+def linear_operator(
+    X: Tensor,
+    Q: Tensor,
+    *,
+    k: int | None = None,
+    mask: Tensor | None = None,
+) -> Tensor: ...
+
+
+def linear_operator(
+    X,
+    Q,
+    *,
+    k=None,
+    mask=None,
+):
+    """
+    Build the linear operator `S` of simplex projection, such that ``S @ Y`` equals ``simplex_projection(X, Y, Q)`` for any target `Y`.
+
+    Parameters
+    ----------
+    X : np.ndarray or Tensor
+        The input data of shape (N,) or (N, E) or (B, N, E)
+    Q : np.ndarray or Tensor
+        The query points of shape (M,) or (M, E) or (B, M, E) for which to find the nearest neighbors in `X`.
+    k : int or None, default None
+        The number of nearest neighbors to use. If None, uses E + 1, where E is the dimension of `X`.
+    mask : np.ndarray or Tensor or None
+        Boolean mask of shape (M, N) or (B, M, N) indicating, for each query in `Q`, which library points to include when finding nearest neighbors.
+
+    Returns
+    -------
+    S : np.ndarray or Tensor
+        The operator of shape (M, N) or (B, M, N).
+
+    Raises
+    ------
+    ValueError
+        - If `k` is not positive.
+        - If the batch size or dimension of `X` and `Q` do not match.
+        - If `mask` does not have shape (M, N) or (B, M, N).
+        - If fewer than `k` library points are unmasked for some query.
+
+    Examples
+    --------
+    ```python
+    import numpy as np
+
+    from edmkit.simplex_projection import linear_operator, simplex_projection
+
+    rng = np.random.default_rng(0)
+    X = rng.normal(size=(200, 3))
+    Q = rng.normal(size=(50, 3))
+    Y = rng.normal(size=(200, 4))
+
+    assert np.allclose(linear_operator(X, Q) @ Y, simplex_projection(X, Y, Q))
+    ```
+    """
+    if isinstance(X, np.ndarray):
+        return _numpy(X, Q, k=k, mask=mask)
+
+    return _tensor(X, Q, k=k, mask=mask)
+
+
+def _numpy(
+    X: np.ndarray,
+    Q: np.ndarray,
+    *,
+    k: int | None = None,
+    mask: np.ndarray | None = None,
+) -> np.ndarray:
+    # ensure at least 2D
+    if X.ndim == 1:
+        X = X[:, None]
+    if Q.ndim == 1:
+        Q = Q[:, None]
+
+    if not (X.ndim == Q.ndim and X.ndim in (2, 3)):
+        raise ValueError(f"X and Q must both be 2D or both be 3D arrays, got X.ndim={X.ndim}, Q.ndim={Q.ndim}")
+    if Q.shape[:-2] != X.shape[:-2] or Q.shape[-1] != X.shape[-1]:
+        raise ValueError(f"batch size and dimension of X and Q must match, got X.shape={X.shape} and Q.shape={Q.shape}")
+    if mask is not None and mask.shape != Q.shape[:-1] + X.shape[-2:-1]:
+        raise ValueError(f"mask shape must be {Q.shape[:-1] + X.shape[-2:-1]}, got mask.shape={mask.shape}")
+
+    return_ndim = X.ndim
+
+    # treat the unbatched case as a single batch
+    if X.ndim == 2:
+        X, Q = X[None], Q[None]
+        if mask is not None:
+            mask = mask[None]
+
+    B, N, E = X.shape
+    M = Q.shape[1]
+    k = E + 1 if k is None else k
+    if k <= 0:
+        raise ValueError(f"k must be positive, got k={k}")
+
+    distances = np.empty((B, M, k))
+    indices = np.empty((B, M, k), dtype=np.intp)
+    for b in range(B):
+        if mask is None:
+            distances[b], indices[b] = knn(X[b], Q[b], k)
+        else:  # over-fetch, then drop the masked-out candidates
+            n_exclude = int((~mask[b]).sum(axis=-1).max())
+            if N - n_exclude < k:
+                raise ValueError(
+                    f"Not enough unmasked points in X to find {k} neighbors, got N={N} with up to {n_exclude} points masked out per query"
+                )
+            _distances, _indices = knn(X[b], Q[b], k + n_exclude)
+            _distances = np.where(np.take_along_axis(mask[b], _indices, axis=-1), _distances, np.inf)
+            top_k = np.argsort(_distances, axis=-1)[..., :k]
+            distances[b] = np.take_along_axis(_distances, top_k, axis=-1)
+            indices[b] = np.take_along_axis(_indices, top_k, axis=-1)
+
+    # clamp to avoid division by zero
+    d_min = np.fmax(distances.min(axis=-1, keepdims=True), 1e-6)  # (B, M, 1)
+    weights = np.exp(-distances / d_min)  # (B, M, k)
+    weights /= weights.sum(axis=-1, keepdims=True)
+
+    S = np.zeros((B, M, N))
+    np.put_along_axis(S, indices, weights, axis=-1)
+
+    if return_ndim == 2:
+        return S[0]  # (M, N)
+    return S  # (B, M, N)
+
+
+def _tensor(
+    X: Tensor,
+    Q: Tensor,
+    *,
+    k: int | None = None,
+    mask: Tensor | None = None,
+) -> Tensor:
+    # ensure at least 2D
+    if X.ndim == 1:
+        X = X[:, None]
+    if Q.ndim == 1:
+        Q = Q[:, None]
+
+    if not (X.ndim == Q.ndim and X.ndim in (2, 3)):
+        raise ValueError(f"X and Q must both be 2D or both be 3D tensors, got X.ndim={X.ndim}, Q.ndim={Q.ndim}")
+    if Q.shape[:-2] != X.shape[:-2] or Q.shape[-1] != X.shape[-1]:
+        raise ValueError(f"batch size and dimension of X and Q must match, got X.shape={X.shape} and Q.shape={Q.shape}")
+    if mask is not None and mask.shape != Q.shape[:-1] + X.shape[-2:-1]:
+        raise ValueError(f"mask shape must be {Q.shape[:-1] + X.shape[-2:-1]}, got mask.shape={mask.shape}")
+
+    return_ndim = X.ndim
+
+    # treat the unbatched case as a single batch
+    if X.ndim == 2:
+        X, Q = X[None], Q[None]
+        if mask is not None:
+            mask = mask[None]
+
+    B, N, E = (int(X.shape[0]), int(X.shape[1]), int(X.shape[2]))
+    M = int(Q.shape[1])
+    k = E + 1 if k is None else k
+    if k <= 0:
+        raise ValueError(f"k must be positive, got k={k}")
+
+    # clamp to avoid NaN gradient
+    D = pairwise_distance(Q, X).clamp(min_=1e-12).sqrt()  # (B, M, N)
+
+    if mask is not None:
+        D = mask.where(D, float("inf"))
+
+    distances, indices = D.topk(k, dim=-1, largest=False, sorted_=True)  # (B, M, k)
+
+    # neighbors are sorted, so the first is the nearest; clamp to avoid division by zero
+    d_min = distances[..., :1].clamp(min_=1e-6)  # (B, M, 1)
+    weights = (-distances / d_min).exp()  # (B, M, k)
+    weights = weights / weights.sum(axis=-1, keepdim=True)
+
+    S = weights.zeros(B, M, N, dtype=weights.dtype, device=weights.device).scatter(-1, indices, weights)  # (B, M, N)
+
+    if return_ndim == 2:
+        return S[0]  # (M, N)
+    return S  # (B, M, N)

{edmkit-0.0.14 → edmkit-0.0.15}/src/edmkit/simplex_projection/simplex_projection.py

@@ -63,12 +63,14 @@ def simplex_projection(
     -------
     predictions : np.ndarray or Tensor
         The predicted values based on the weighted mean of the nearest neighbors in `Y`.
+        Shape (M,) if `Y` is 1D, (M, E') if `Y` is 2D, or (B, M, E') if `Y` is 3D.
 
     Raises
     ------
     ValueError
         - If `k` is not positive.
         - If the input arrays `X` and `Y` do not have the same number of points.
+        - If the batch size or dimension of `X` and `Q` do not match.
         - If `mask` does not have shape (M, N) or (B, M, N).
         - If fewer than `k` library points are unmasked for some query.
 
@@ -121,6 +123,8 @@ def _numpy(
     k: int | None = None,
     mask: np.ndarray | None = None,
 ):
+    return_ndim = Y.ndim
+
     # ensure at least 2D
     if X.ndim == 1:
         X = X[:, None]
@@ -139,8 +143,7 @@ def _numpy(
         raise ValueError(f"mask shape must be {Q.shape[:-1] + X.shape[-2:-1]}, got mask.shape={mask.shape}")
 
     # treat the unbatched case as a single batch
-    batched = X.ndim == 3
-    if not batched:
+    if X.ndim == 2:
         X, Y, Q = X[None], Y[None], Q[None]
         if mask is not None:
             mask = mask[None]
@@ -177,9 +180,11 @@ def _numpy(
 
     predictions = np.matmul(weights[..., None, :], Y_neighbors).squeeze(-2) / weights.sum(axis=-1, keepdims=True)  # (B, M, E')
 
-    if not batched:
-        return predictions.squeeze()  # (M,) or (M, E')
-    return predictions
+    if return_ndim == 1:
+        return predictions[0, :, 0]  # (M,)
+    if return_ndim == 2:
+        return predictions[0]  # (M, E')
+    return predictions  # (B, M, E')
 
 
 def _tensor(
@@ -194,6 +199,8 @@ def _tensor(
     # Loading it only when needed keeps the numpy path free of that scheduler pressure.
     from tinygrad import Tensor, dtypes
 
+    return_ndim = Y.ndim
+
     # ensure at least 2D
     if X.ndim == 1:
         X = X[:, None]
@@ -212,8 +219,7 @@ def _tensor(
         raise ValueError(f"mask shape must be {Q.shape[:-1] + X.shape[-2:-1]}, got mask.shape={mask.shape}")
 
     # treat the unbatched case as a single batch
-    batched = X.ndim == 3
-    if not batched:
+    if X.ndim == 2:
         X, Y, Q = X[None], Y[None], Q[None]
         if mask is not None:
             mask = mask[None]
@@ -242,9 +248,11 @@ def _tensor(
 
     predictions: Tensor = (weights.unsqueeze(-1) * Y_neighbors).sum(axis=-2) / weights.sum(axis=-1, keepdim=True)  # (B, M, E')
 
-    if not batched:
-        return predictions.squeeze()  # (M,) or (M, E')
-    return predictions
+    if return_ndim == 1:
+        return predictions[0, :, 0]  # (M,)
+    if return_ndim == 2:
+        return predictions[0]  # (M, E')
+    return predictions  # (B, M, E')
 
 
 if TYPE_CHECKING:

{edmkit-0.0.14 → edmkit-0.0.15}/src/edmkit/simplex_projection/soft.py

@@ -70,6 +70,7 @@ def soft_simplex_projection(
     -------
     predictions : np.ndarray or Tensor
         The predicted values based on the weighted mean of the nearest neighbors in `Y`.
+        Shape (M,) if `Y` is 1D, (M, E') if `Y` is 2D, or (B, M, E') if `Y` is 3D.
 
     Raises
     ------
@@ -77,8 +78,9 @@ def soft_simplex_projection(
         - If `k` is not positive.
         - If `softness` is not positive.
         - If the input arrays `X` and `Y` do not have the same number of points.
-        - If `X` does not contain at least `k + 1` points.
+        - If the batch size or dimension of `X` and `Q` do not match.
         - If `mask` does not have shape (M, N) or (B, M, N).
+        - If fewer than `k + 1` library points are unmasked for some query.
 
     Examples
     --------
@@ -134,6 +136,8 @@ def _numpy(
     if softness <= 0:
         raise ValueError(f"softness must be positive, got softness={softness}")
 
+    return_ndim = Y.ndim
+
     # ensure at least 2D
     if X.ndim == 1:
         X = X[:, None]
@@ -152,8 +156,7 @@ def _numpy(
         raise ValueError(f"mask shape must be {Q.shape[:-1] + X.shape[-2:-1]}, got mask.shape={mask.shape}")
 
     # treat the unbatched case as a single batch
-    batched = X.ndim == 3
-    if not batched:
+    if X.ndim == 2:
         X, Y, Q = X[None], Y[None], Q[None]
         if mask is not None:
             mask = mask[None]
@@ -162,8 +165,9 @@ def _numpy(
     k = E + 1 if k is None else k
     if k <= 0:
         raise ValueError(f"k must be positive, got k={k}")
-    if N < k + 1:
-        raise ValueError(f"Not enough points in X to find {k + 1} neighbors, got N={N}")
+    n_exclude = 0 if mask is None else int((~mask).sum(axis=-1).max())
+    if N - n_exclude < k + 1:
+        raise ValueError(f"Not enough unmasked points in X to find {k + 1} neighbors, got N={N} with up to {n_exclude} points masked out per query")
 
     # clamp to align with _tensor()
     D = np.sqrt(np.maximum(pairwise_distance_np(Q, X), 1e-12))  # (B, M, N)
@@ -186,9 +190,11 @@ def _numpy(
 
     predictions = np.matmul(weights, Y) / weights.sum(axis=-1, keepdims=True)  # (B, M, E')
 
-    if not batched:
-        return predictions.squeeze()  # (M,) or (M, E')
-    return predictions
+    if return_ndim == 1:
+        return predictions[0, :, 0]  # (M,)
+    if return_ndim == 2:
+        return predictions[0]  # (M, E')
+    return predictions  # (B, M, E')
 
 
 def _tensor(
@@ -203,6 +209,8 @@ def _tensor(
     if softness <= 0:
         raise ValueError(f"softness must be positive, got softness={softness}")
 
+    return_ndim = Y.ndim
+
     # ensure at least 2D
     if X.ndim == 1:
         X = X[:, None]
@@ -221,8 +229,7 @@ def _tensor(
         raise ValueError(f"mask shape must be {Q.shape[:-1] + X.shape[-2:-1]}, got mask.shape={mask.shape}")
 
     # treat the unbatched case as a single batch
-    batched = X.ndim == 3
-    if not batched:
+    if X.ndim == 2:
         X, Y, Q = X[None], Y[None], Q[None]
         if mask is not None:
             mask = mask[None]
@@ -257,6 +264,8 @@ def _tensor(
 
     predictions: Tensor = weights.matmul(Y) / weights.sum(axis=-1, keepdim=True)  # (B, M, E')
 
-    if not batched:
-        return predictions.squeeze()  # (M,) or (M, E')
-    return predictions
+    if return_ndim == 1:
+        return predictions[0, :, 0]  # (M,)
+    if return_ndim == 2:
+        return predictions[0]  # (M, E')
+    return predictions  # (B, M, E')

edmkit-0.0.15/src/edmkit/simplex_projection/soft_linear_operator.py

@@ -0,0 +1,203 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, overload
+
+import numpy as np
+from scipy.special import expit
+
+from edmkit.util import pairwise_distance, pairwise_distance_np
+
+__all__ = ["soft_linear_operator"]
+
+if TYPE_CHECKING:
+    from tinygrad import Tensor
+
+
+@overload
+def soft_linear_operator(
+    X: np.ndarray,
+    Q: np.ndarray,
+    *,
+    k: int | None = None,
+    mask: np.ndarray | None = None,
+    softness: float = 0.02,
+) -> np.ndarray: ...
+
+
+@overload
+def soft_linear_operator(
+    X: Tensor,
+    Q: Tensor,
+    *,
+    k: int | None = None,
+    mask: Tensor | None = None,
+    softness: float = 0.02,
+) -> Tensor: ...
+
+
+def soft_linear_operator(
+    X,
+    Q,
+    *,
+    k=None,
+    mask=None,
+    softness=0.02,
+):
+    """
+    Build the linear operator `S` of soft simplex projection, such that ``S @ Y`` equals ``soft_simplex_projection(X, Y, Q)`` for any target `Y`.
+
+    Parameters
+    ----------
+    X : np.ndarray or Tensor
+        The input data of shape (N,) or (N, E) or (B, N, E)
+    Q : np.ndarray or Tensor
+        The query points of shape (M,) or (M, E) or (B, M, E) for which to find the nearest neighbors in `X`.
+    k : int or None, default None
+        The number of nearest neighbors to use. If None, uses E + 1, where E is the dimension of `X`.
+    mask : np.ndarray or Tensor or None
+        Boolean mask of shape (M, N) or (B, M, N) indicating, for each query in `Q`, which library points to include when finding nearest neighbors.
+    softness : float, default 0.02
+        Width of the boundary between neighbors and non-neighbors, as a fraction of the neighborhood radius.
+
+    Returns
+    -------
+    S : np.ndarray or Tensor
+        The operator of shape (M, N) or (B, M, N).
+
+    Raises
+    ------
+    ValueError
+        - If `k` is not positive.
+        - If `softness` is not positive.
+        - If the batch size or dimension of `X` and `Q` do not match.
+        - If `mask` does not have shape (M, N) or (B, M, N).
+        - If fewer than `k + 1` library points are unmasked for some query.
+
+    Examples
+    --------
+    ```python
+    import numpy as np
+
+    from edmkit.simplex_projection import soft_linear_operator, soft_simplex_projection
+
+    rng = np.random.default_rng(0)
+    X = rng.normal(size=(200, 3))
+    Q = rng.normal(size=(50, 3))
+    Y = rng.normal(size=(200, 4))
+
+    assert np.allclose(soft_linear_operator(X, Q) @ Y, soft_simplex_projection(X, Y, Q))
+    ```
+    """
+    if isinstance(X, np.ndarray):
+        return _numpy(X, Q, k=k, mask=mask, softness=softness)
+
+    return _tensor(X, Q, k=k, mask=mask, softness=softness)
+
+
+def _numpy(
+    X: np.ndarray,
+    Q: np.ndarray,
+    *,
+    k: int | None = None,
+    mask: np.ndarray | None = None,
+    softness: float = 0.02,
+) -> np.ndarray:
+    if softness <= 0:
+        raise ValueError(f"softness must be positive, got softness={softness}")
+
+    # ensure at least 2D
+    if X.ndim == 1:
+        X = X[:, None]
+    if Q.ndim == 1:
+        Q = Q[:, None]
+
+    if not (X.ndim == Q.ndim and X.ndim in (2, 3)):
+        raise ValueError(f"X and Q must both be 2D or both be 3D arrays, got X.ndim={X.ndim}, Q.ndim={Q.ndim}")
+    if Q.shape[:-2] != X.shape[:-2] or Q.shape[-1] != X.shape[-1]:
+        raise ValueError(f"batch size and dimension of X and Q must match, got X.shape={X.shape} and Q.shape={Q.shape}")
+    if mask is not None and mask.shape != Q.shape[:-1] + X.shape[-2:-1]:
+        raise ValueError(f"mask shape must be {Q.shape[:-1] + X.shape[-2:-1]}, got mask.shape={mask.shape}")
+
+    N, E = X.shape[-2:]
+    k = E + 1 if k is None else k
+    if k <= 0:
+        raise ValueError(f"k must be positive, got k={k}")
+    n_exclude = 0 if mask is None else int((~mask).sum(axis=-1).max())
+    if N - n_exclude < k + 1:
+        raise ValueError(f"Not enough unmasked points in X to find {k + 1} neighbors, got N={N} with up to {n_exclude} points masked out per query")
+
+    # clamp to align with _tensor()
+    D = np.sqrt(np.maximum(pairwise_distance_np(Q, X), 1e-12))  # (M, N) or (B, M, N)
+
+    if mask is not None:
+        # a finite sentinel instead of inf to avoid NaN
+        D = np.where(mask, D, D.max(axis=-1, keepdims=True) * 2 + 1)
+
+    neighbors = np.sort(np.partition(D, k, axis=-1)[..., : k + 1], axis=-1)  # (M, k + 1) or (B, M, k + 1)
+
+    # clamp to avoid division by zero
+    d_min = np.maximum(neighbors[..., :1], 1e-6)  # (M, 1) or (B, M, 1)
+    # radius is the midpoint between the k-th and (k+1)-th nearest neighbors
+    radius = (neighbors[..., k - 1 : k] + neighbors[..., k : k + 1]) / 2  # (M, 1) or (B, M, 1)
+    radius = np.maximum(radius, 1e-6)
+
+    weights = np.exp(-D / d_min) * expit((radius - D) / (softness * radius))  # (M, N) or (B, M, N)
+    if mask is not None:
+        weights = np.where(mask, weights, 0.0)
+
+    return weights / weights.sum(axis=-1, keepdims=True)
+
+
+def _tensor(
+    X: Tensor,
+    Q: Tensor,
+    *,
+    k: int | None = None,
+    mask: Tensor | None = None,
+    softness: float = 0.02,
+) -> Tensor:
+    if softness <= 0:
+        raise ValueError(f"softness must be positive, got softness={softness}")
+
+    # ensure at least 2D
+    if X.ndim == 1:
+        X = X[:, None]
+    if Q.ndim == 1:
+        Q = Q[:, None]
+
+    if not (X.ndim == Q.ndim and X.ndim in (2, 3)):
+        raise ValueError(f"X and Q must both be 2D or both be 3D tensors, got X.ndim={X.ndim}, Q.ndim={Q.ndim}")
+    if Q.shape[:-2] != X.shape[:-2] or Q.shape[-1] != X.shape[-1]:
+        raise ValueError(f"batch size and dimension of X and Q must match, got X.shape={X.shape} and Q.shape={Q.shape}")
+    if mask is not None and mask.shape != Q.shape[:-1] + X.shape[-2:-1]:
+        raise ValueError(f"mask shape must be {Q.shape[:-1] + X.shape[-2:-1]}, got mask.shape={mask.shape}")
+
+    N = int(X.shape[-2])
+    E = int(X.shape[-1])
+    k = E + 1 if k is None else k
+    if k <= 0:
+        raise ValueError(f"k must be positive, got k={k}")
+    if N < k + 1:
+        raise ValueError(f"Not enough points in X to find {k + 1} neighbors, got N={N}")
+
+    # clamp to avoid NaN gradient
+    D = pairwise_distance(Q, X).clamp(min_=1e-12).sqrt()  # (M, N) or (B, M, N)
+
+    if mask is not None:
+        # a finite sentinel instead of inf to avoid NaN
+        D = mask.where(D, D.max(axis=-1, keepdim=True) * 2 + 1)
+
+    neighbors = D.topk(k + 1, dim=-1, largest=False, sorted_=True)[0]  # (M, k + 1) or (B, M, k + 1)
+
+    # clamp to avoid division by zero
+    d_min = neighbors[..., :1].clamp(min_=1e-6)  # (M, 1) or (B, M, 1)
+    # radius is the midpoint between the k-th and (k+1)-th nearest neighbors
+    radius = (neighbors[..., k - 1 : k] + neighbors[..., k : k + 1]) / 2  # (M, 1) or (B, M, 1)
+    radius = radius.clamp(min_=1e-6)
+
+    gate = ((radius - D) / (softness * radius)).sigmoid()  # (M, N) or (B, M, N)
+    weights = (-D / d_min).exp() * gate  # (M, N) or (B, M, N)
+    if mask is not None:
+        weights = mask.where(weights, 0)
+
+    return weights / weights.sum(axis=-1, keepdim=True)

edmkit-0.0.15/src/edmkit/smap/__init__.py

@@ -0,0 +1,5 @@
+from .coefficients import coefficients
+from .linear_operator import linear_operator
+from .smap import smap
+
+__all__ = ["coefficients", "linear_operator", "smap"]