feat: add pre-computed affinity matrices

README.md

@@ -22,7 +22,7 @@
 ```python
 import numpy as np
 
-from snf2 import fuse, make_affinity
+from snf2 import affinity_matrix, fuse, make_affinity
 
 modality_a = np.array(
     [[0.0, 1.0], [0.2, 0.8], [1.0, 0.1], [0.9, 0.2]],
@@ -66,58 +66,26 @@ Metric-specific data requirements follow SciPy. SNF2 raises an error if a
 metric produces non-finite or negative pairwise distances for the supplied
 data.
 
-SNF2 currently provides the two core algorithm stages: constructing an
-affinity matrix from one feature matrix and fusing affinity matrices across
-modalities.
-
-## Usage
+For precomputed distances, use `affinity_matrix` directly:
 
 ```python
-import numpy as np
-
-from snf2 import fuse, make_affinity
-
-modality_a = np.array(
-    [[0.0, 1.0], [0.2, 0.8], [1.0, 0.1], [0.9, 0.2]],
+distances = np.array(
+    [
+        [0.0, 0.3, 1.2, 1.0],
+        [0.3, 0.0, 1.0, 0.8],
+        [1.2, 1.0, 0.0, 0.2],
+        [1.0, 0.8, 0.2, 0.0],
+    ],
 )
-modality_b = np.array(
-    [[1.0, 0.0], [0.8, 0.1], [0.1, 1.0], [0.2, 0.9]],
-)
-
-affinities = [
-    make_affinity(modality_a, n_neighbors=2),
-    make_affinity(modality_b, n_neighbors=2),
-]
-fused_network = fuse(affinities, n_neighbors=2)
-```
-
-Rows are samples and columns are features. SNF2 does not standardize or align
-inputs: callers must preprocess each modality and ensure identical sample
-ordering before constructing affinities.
-
-Affinity construction defaults to squared Euclidean distance and accepts every
-named metric supported by
-[`scipy.spatial.distance.pdist`](https://docs.scipy.org/doc/scipy/reference/generated/scipy.spatial.distance.pdist.html).
-Use `metric_kwargs` for metric-specific arguments:
-
-```python
-correlation_network = make_affinity(
-    modality_a,
-    metric="correlation",
-    n_neighbors=2,
-)
-
-minkowski_network = make_affinity(
-    modality_a,
-    metric="minkowski",
-    metric_kwargs={"p": 3.5},
+precomputed_network = affinity_matrix(
+    distances,
     n_neighbors=2,
 )
 ```
 
-Metric-specific data requirements follow SciPy. SNF2 raises an error if a
-metric produces non-finite or negative pairwise distances for the supplied
-data.
+`affinity_matrix` accepts distances, not similarities. Convert a similarity
+matrix with a transformation appropriate to that measure before calling it;
+for a similarity bounded to `[0, 1]`, that may be `1 - similarity`.
 
 ## Development setup
 

docs/devnotes.md

@@ -5,11 +5,16 @@ that are important to future SNF2 contributors.
 
 ## Design decisions
 
-No entries yet.
+- `make_affinity` computes pairwise distances from sample-by-feature data and
+  delegates the shared SNF kernel to `affinity_matrix`.
+- `affinity_matrix` accepts precomputed distances rather than similarities, so
+  callers remain responsible for choosing a scientifically appropriate
+  similarity-to-distance transformation.
 
 ## Compatibility notes
 
-No entries yet.
+- Missing-value handling belongs in preprocessing or in the distance
+  calculation. SNF2 does not define a generic `nan_policy` for SciPy metrics.
 
 ## Open questions
 

docs/index.md

@@ -7,7 +7,7 @@ The initial API separates affinity construction from network fusion:
 ```python
 import numpy as np
 
-from snf2 import fuse, make_affinity
+from snf2 import affinity_matrix, fuse, make_affinity
 
 rna = np.array([[0.0, 1.0], [0.2, 0.8], [1.0, 0.1], [0.9, 0.2]])
 protein = np.array([[1.0, 0.0], [0.8, 0.1], [0.1, 1.0], [0.2, 0.9]])
@@ -33,5 +33,24 @@ at least two finite, nonnegative, symmetric affinity matrices with the same
 shape. Metric-specific data requirements follow SciPy; SNF2 rejects
 non-finite or negative pairwise distances before constructing affinities.
 
+Use `affinity_matrix` when distances have already been computed:
+
+```python
+distances = np.array(
+    [
+        [0.0, 0.3, 1.2, 1.0],
+        [0.3, 0.0, 1.0, 0.8],
+        [1.2, 1.0, 0.0, 0.2],
+        [1.0, 0.8, 0.2, 0.0],
+    ],
+)
+precomputed_network = affinity_matrix(distances, n_neighbors=2)
+```
+
+The input to `affinity_matrix` is a distance matrix, not a similarity matrix.
+Convert similarities with a transformation appropriate to the similarity
+measure first; for a similarity bounded to `[0, 1]`, that may be
+`1 - similarity`.
+
 For setup and development commands, see the
 [project README](https://github.com/bhklab/snf2#readme).

src/snf2/__init__.py

@@ -1,6 +1,6 @@
 """SNF2: a modern Python implementation of Similarity Network Fusion."""
 
-from snf2.affinity import make_affinity
+from snf2.affinity import affinity_matrix, make_affinity
 from snf2.fusion import fuse
 
-__all__ = ["fuse", "make_affinity"]
+__all__ = ["affinity_matrix", "fuse", "make_affinity"]

src/snf2/_validation.py

@@ -37,6 +37,47 @@ def as_feature_matrix(data: ArrayLike) -> NDArray[np.float64]:
     return matrix
 
 
+def as_distance_matrix(distances: ArrayLike) -> NDArray[np.float64]:
+    """Return a validated pairwise-distance matrix as an owned float64 array."""
+    try:
+        raw = np.asarray(distances)
+    except (TypeError, ValueError) as error:
+        raise TypeError("distances must be a rectangular real-valued array") from error
+
+    if raw.ndim != 2 or raw.shape[0] != raw.shape[1]:
+        raise ValueError("distances must be a square matrix")
+    if raw.shape[0] < 2:
+        raise ValueError("distances must contain at least two samples")
+    if not np.issubdtype(raw.dtype, np.number) or np.issubdtype(
+        raw.dtype, np.complexfloating
+    ):
+        raise TypeError("distances must contain real numeric values")
+
+    matrix = np.array(raw, dtype=np.float64, copy=True)
+    if not np.all(np.isfinite(matrix)):
+        raise ValueError("distances must contain only finite values")
+    if np.any(matrix < 0):
+        raise ValueError("distances must be nonnegative")
+    if not np.allclose(
+        matrix,
+        matrix.T,
+        rtol=SYMMETRY_RTOL,
+        atol=SYMMETRY_ATOL,
+    ):
+        raise ValueError("distances must be symmetric")
+    if not np.allclose(
+        np.diag(matrix),
+        0,
+        rtol=0,
+        atol=SYMMETRY_ATOL,
+    ):
+        raise ValueError("distances must have a zero diagonal")
+
+    matrix = np.asarray((matrix + matrix.T) / 2, dtype=np.float64)
+    np.fill_diagonal(matrix, 0)
+    return matrix
+
+
 def validate_n_neighbors(n_neighbors: int, n_samples: int) -> int:
     """Validate and return the requested neighborhood size."""
     if isinstance(n_neighbors, bool) or not isinstance(n_neighbors, Integral):

src/snf2/affinity.py

@@ -6,8 +6,10 @@
 import numpy as np
 from numpy.typing import ArrayLike, NDArray
 from scipy.spatial.distance import pdist, squareform
+from scipy.stats import norm
 
 from snf2._validation import (
+    as_distance_matrix,
     as_feature_matrix,
     validate_n_neighbors,
     validate_positive_float,
@@ -78,9 +80,6 @@ def make_affinity(
     distances = squareform(
         pdist_with_named_metric(matrix, metric=metric, **distance_kwargs)
     )
-    distances = np.asarray((distances + distances.T) / 2, dtype=np.float64)
-    np.fill_diagonal(distances, 0)
-
     if not np.all(np.isfinite(distances)):
         raise ValueError(
             f"metric {metric!r} produced non-finite pairwise distances; "
@@ -89,20 +88,79 @@ def make_affinity(
     if np.any(distances < 0):
         raise ValueError(f"metric {metric!r} produced negative pairwise distances")
 
-    epsilon = np.finfo(np.float64).eps
-    sorted_distances = np.sort(distances, axis=1)
-    neighborhood_means = sorted_distances[:, 1 : neighbors + 1].mean(axis=1) + epsilon
-    local_widths = (
-        neighborhood_means[:, None] + neighborhood_means[None, :] + distances
-    ) / 3 + epsilon
-    local_widths = np.maximum(local_widths, epsilon)
+    return affinity_matrix(
+        distances,
+        n_neighbors=neighbors,
+        scale=kernel_scale,
+    )
+
+
+def affinity_matrix(
+    distances: ArrayLike,
+    *,
+    n_neighbors: int = 20,
+    scale: float = 0.5,
+) -> NDArray[np.float64]:
+    """Construct an SNF affinity matrix from pairwise distances.
+
+    Parameters
+    ----------
+    distances
+        Square, finite, nonnegative, symmetric pairwise-distance matrix with a
+        zero diagonal. Similarities must first be converted to distances using
+        a transformation appropriate to the similarity measure.
+    n_neighbors
+        Number of nearest neighbors used to estimate each local scale.
+    scale
+        Positive multiplier applied to the locally estimated kernel width.
+
+    Returns
+    -------
+    numpy.ndarray
+        A symmetric ``float64`` sample-by-sample affinity matrix.
 
-    kernel_widths = kernel_scale * local_widths
-    affinities = np.exp(-(distances**2) / (2 * kernel_widths**2))
-    affinities /= np.sqrt(2 * np.pi) * kernel_widths
-    affinities = np.asarray((affinities + affinities.T) / 2, dtype=np.float64)
+    Raises
+    ------
+    TypeError
+        If the distances or parameters have incompatible types.
+    ValueError
+        If the distances or parameters have invalid values.
+    """
+    diff = as_distance_matrix(distances)
+    n = diff.shape[0]
+    k = validate_n_neighbors(n_neighbors, n)
+    kernel_scale = validate_positive_float(scale, name="scale")
+    eps = np.finfo(np.float64).eps
+
+    # Symmetrize and remove self-distances.
+    diff = np.asarray((diff + diff.T) / 2, dtype=np.float64)
+    np.fill_diagonal(diff, 0.0)
+
+    # For each row, sort distances and take the first k non-self values.
+    sorted_rows = np.sort(diff, axis=1)
+    nearest = sorted_rows[:, 1 : k + 1]
+
+    # R's mean(x[is.finite(x)]), applied row-wise.
+    finite = np.isfinite(nearest)
+    counts = finite.sum(axis=1)
+    means = np.where(finite, nearest, 0.0).sum(axis=1) / counts
+    means += eps
+
+    # Equivalent to:
+    # outer(means, means, avg) / 3 * 2 + Diff / 3 + eps
+    sig = (2.0 / 3.0) * ((means[:, None] + means[None, :]) / 2)
+    sig += diff / 3.0
+    sig += eps
+    sig = np.maximum(sig, eps)
+
+    # Gaussian density: dnorm(Diff, mean=0, sd=scale * Sig)
+    densities = np.asarray(
+        norm.pdf(diff, loc=0.0, scale=kernel_scale * sig),
+        dtype=np.float64,
+    )
 
+    # Ensure the resulting affinity matrix is symmetric.
+    affinities = np.asarray((densities + densities.T) / 2, dtype=np.float64)
     if not np.all(np.isfinite(affinities)):
         raise ValueError("affinity computation produced non-finite values")
-
     return affinities