Support NaN-masked training and batch updates

examples/TimeSeriesWithNaNs.py

@@ -0,0 +1,163 @@
+# %%
+import matplotlib.pyplot as plt
+import numpy as np
+
+import python_som
+
+# %%
+# Generate two sets of time series data with errors modeled after a possion distributions
+# appropriate for count data. The first set of time series will be constant, while the second set
+# will have underlying red noise.
+rng = np.random.default_rng(42)
+n_steps = 30
+
+# Constant time series
+n_constant_timeseries = 1000
+mean_values = rng.uniform(low=0, high=10, size=n_constant_timeseries)
+sigmas = np.sqrt(mean_values)
+x_const = rng.normal(loc=mean_values, scale=sigmas, size=(30, n_constant_timeseries)).T
+x_const_err = np.sqrt(abs(x_const))
+
+i = rng.integers(low=0, high=n_constant_timeseries)
+plt.errorbar(np.arange(n_steps), x_const[i], yerr=x_const_err[i], fmt="o")
+plt.axhline(mean_values[i], color="k", ls=":")
+plt.savefig("constant_timeseries_example.pdf", dpi=300)
+plt.close()
+# %%
+# Non-constant time series with red noise
+n_non_constant_timeseries = 1000
+spectrum = np.zeros(n_steps, dtype=np.complex128)
+indices = np.arange(1, n_steps - 1)
+spectrum[indices] = 1 / indices**2
+all_spectra = spectrum[np.newaxis, :].repeat(n_non_constant_timeseries, axis=0)
+all_spectra[:, 0] = 1  # rng.uniform(low=1000, high=100000, size=n_constant_timeseries)
+phases = np.exp(1j * rng.uniform(0, 2 * np.pi, (n_non_constant_timeseries, n_steps)))
+n = phases * all_spectra
+x_true_rednoise = np.abs(np.fft.ifft(n, axis=1)) * 100
+x_true_rednoise_err = np.sqrt(abs(x_true_rednoise))
+x_rednoise = rng.normal(loc=x_true_rednoise, scale=x_true_rednoise_err)
+x_rednoise_err = np.sqrt(abs(x_rednoise))
+
+i = rng.integers(low=0, high=n_constant_timeseries)
+plt.plot(np.arange(n_steps), x_true_rednoise[i], color="k", ls=":")
+plt.errorbar(np.arange(n_steps), x_rednoise[i], yerr=x_rednoise_err[i], fmt="o")
+plt.savefig("rednoise_timeseries_example.pdf", dpi=300)
+plt.close()
+# %%
+# Combine the two sets of time series into one dataset
+data = np.concatenate([x_const, x_rednoise], axis=0)
+data_err = np.concatenate([x_const_err, x_rednoise_err], axis=0)
+# %%
+medians = np.median(data, axis=1)[:, np.newaxis]
+normed_data = data / medians
+normed_data_err = data_err / np.abs(medians)
+# %%
+# randomly drop epochs
+n_exp_missing = 10
+n_missing = rng.poisson(lam=n_exp_missing, size=data.shape[0])
+for i, inm in enumerate(n_missing):
+    if inm > 0:
+        missing_indices = rng.choice(np.arange(n_steps), size=inm, replace=False)
+        normed_data[i, missing_indices] = np.nan
+        normed_data_err[i, missing_indices] = np.nan
+# %%
+X = normed_data
+
+# Train a self-organizing map on the time series data
+somsize = (10, 10)
+state = np.random.RandomState(42)
+som = python_som.SOM(
+    x=somsize[0],
+    y=somsize[1],
+    input_len=n_steps,
+    learning_rate=0.5,
+    neighborhood_radius=1.0,
+    neighborhood_function="gaussian",
+    cyclic_x=True,
+    cyclic_y=True,
+    data=normed_data,
+    random_seed=42,
+)
+som.fit(X, verbose=True, mode="batch")
+
+win_map = np.array(np.unravel_index(som.predict(X), som.get_shape())).T
+
+fig, axs = plt.subplots(*somsize, figsize=(7, 7))
+for position in np.unique(win_map, axis=0):
+    mask = (win_map[:, 0] == position[0]) & (win_map[:, 1] == position[1])
+    if not any(mask):
+        continue
+    ax = axs[somsize[0] - 1 - position[0], position[1]] if somsize[1] > 1 else axs[position[0]]
+    ax.plot(np.nanmean(normed_data[mask], axis=0), c="k")
+    ax.fill_between(
+        np.arange(n_steps),
+        *np.nanquantile(normed_data[mask], [0.05, 0.95], axis=0),
+        color="gray",
+        alpha=0.5,
+    )
+    ax.xaxis.set_ticklabels([])
+    ax.yaxis.set_ticklabels([])
+fig.savefig("som_timeseries.pdf", dpi=300)
+plt.close()
+
+constants_counts = np.unique(
+    som.predict(normed_data[:n_constant_timeseries]), return_counts=True, axis=0
+)
+constants_map = np.zeros(somsize)
+for p, c in zip(
+    np.array(np.unravel_index(constants_counts[0], som.get_shape())).T,
+    constants_counts[1],
+    strict=False,
+):
+    constants_map[p[0], p[1]] = c
+
+rednoise_counts = np.unique(
+    som.predict(normed_data[n_constant_timeseries:]), return_counts=True, axis=0
+)
+rednoise_map = np.zeros(somsize)
+for p, c in zip(
+    np.array(np.unravel_index(rednoise_counts[0], som.get_shape())).T,
+    rednoise_counts[1],
+    strict=False,
+):
+    rednoise_map[p[0], p[1]] = c
+
+purity_map = rednoise_map / (constants_map + rednoise_map)
+recall_map = rednoise_map / rednoise_map.sum()
+
+fig, axs = plt.subplots(ncols=4, figsize=(20, 5))
+for cmap, pmap, ax in zip(
+    ["Reds", "Blues", "copper", "Reds"],
+    [rednoise_map, constants_map, purity_map, recall_map],
+    axs,
+    strict=False,
+):
+    mesh = ax.pcolormesh(pmap, cmap=cmap)  # plotting the distance map as background
+    fig.colorbar(mesh, ax=ax)
+fig.savefig("som_timeseries_purity.pdf", dpi=300)
+plt.close()
+
+rednoise = rednoise_map.flatten()
+constants = constants_map.flatten()
+probs = purity_map.flatten()
+
+recall = []
+precision = []
+xx = np.linspace(0, 1, 100)
+for i in xx:
+    m = probs >= i
+    precision.append(rednoise[m].sum() / (rednoise[m].sum() + constants[m].sum()))
+    recall.append(rednoise[m].sum() / rednoise.sum())
+
+fig, ax = plt.subplots()
+ax.plot(xx, precision, label="Precision")
+ax.plot(xx, recall, label="Recall")
+ax.set_xlabel("Precision")
+ax.set_ylabel("Score")
+ax.set_xlim(0, 1)
+ax.set_ylim(0, 1)
+ax.legend()
+fig.savefig("som_timeseries_f1.pdf", dpi=300)
+plt.close()
+
+# %

src/python_som/_core/_distance.py

@@ -27,7 +27,9 @@ def euclidean_distance(a: npt.ArrayLike, b: npt.ArrayLike) -> npt.NDArray[np.flo
     :param b: Array-like of values. Must not be a scalar.
     :return: Distances between ``a`` and ``b``.
     """
-    result: npt.NDArray[np.floating] = np.linalg.norm(np.subtract(a, b), ord=2, axis=-1)
+    result: npt.NDArray[np.floating] = np.linalg.norm(
+        np.nan_to_num(np.subtract(a, b)), ord=2, axis=-1
+    )
     return result
 
 

src/python_som/_core/_linalg.py

@@ -51,12 +51,12 @@ def pca(data: npt.NDArray[Any], n_components: int = _N_COMPONENTS) -> PrincipalC
     :return: The fitted mean, components and explained variance.
     """
     array = np.asarray(data, dtype=float)
-    mean = array.mean(axis=0)
+    mean = np.nanmean(array, axis=0)
     _, singular_values, right_vectors = np.linalg.svd(array - mean, full_matrices=False)
 
     # Orient each component on its largest-magnitude loading. The flip is applied to all of them
     # before truncation, which is the order scikit-learn does it in.
-    dominant = np.argmax(np.abs(right_vectors), axis=1)
+    dominant = np.nanargmax(np.abs(right_vectors), axis=1)
     signs = np.sign(right_vectors[np.arange(right_vectors.shape[0]), dominant])
     right_vectors = right_vectors * signs[:, None]
 
@@ -84,8 +84,8 @@ def standardize(data: npt.NDArray[Any]) -> npt.NDArray[np.floating]:
     :return: The standardized array.
     """
     array = np.asarray(data, dtype=float)
-    mean = array.mean(axis=0)
-    variance = array.var(axis=0)
+    mean = np.nanmean(array, axis=0)
+    variance = np.nanvar(array, axis=0)
 
     n_samples = array.shape[0]
     eps = np.finfo(np.float64).eps

src/python_som/_core/_match.py

@@ -107,11 +107,13 @@ def bmu_indices(
     if distance is not euclidean_distance:
         return np.array([np.asarray(distance(x, flat)).argmin() for x in data], dtype=np.intp)
 
+    if not np.isfinite(flat).all():
+        raise ValueError("weights must contain only finite values")
     shift = flat.mean(axis=0)
     centred = flat - shift
     squared = np.einsum("nf,nf->n", centred, centred)
 
-    if kernel is not None:  # pragma: no cover  reached only when numba is installed
+    if kernel is not None and not np.isnan(data).any():  # pragma: no cover
         return kernel(data - shift, centred, squared)
 
     n_nodes = len(flat)
@@ -120,11 +122,17 @@ def bmu_indices(
     out = np.empty(len(data), dtype=np.intp)
     for start in range(0, len(data), chunk):
         block = data[start : start + chunk]
-        # Into a preallocated buffer: allocating one per chunk was 1.5x slower and 8x heavier.
-        np.matmul(block - shift, centred.T, out=scores[: len(block)])
+        if np.isinf(block).any():
+            raise ValueError("data must contain only finite values or NaN")
+        mask = ~np.isnan(block)
+        if not mask.any(axis=1).all():
+            raise ValueError("each sample must contain at least one finite value")
+        centred_block = np.where(mask, block - shift, 0.0)
         block_scores = scores[: len(block)]
+        np.matmul(centred_block, centred.T, out=block_scores)
         block_scores *= -2.0
-        block_scores += squared
+        block_scores += np.matmul(mask, (centred**2).T)
+        block_scores += np.einsum("nf,nf->n", centred_block, centred_block)[:, None]
         out[start : start + len(block)] = block_scores.argmin(axis=1)
     return out
 
@@ -136,21 +144,24 @@ def accumulate(
     distance: DistanceFunction,
     kernel: BmuKernel | None = None,
 ) -> tuple[npt.NDArray[np.floating], npt.NDArray[np.floating]]:
-    """Sum the samples mapped to each node, and count them.
+    """Sum observed sample features mapped to each node, and count them.
 
-    These are the ``n_j`` and ``n_j * xbar_j`` of Kohonen (2013), Eq. (8): the count of samples
-    whose best match is node ``j``, and their sum.
+    These are the per-feature versions of ``n_j`` and ``n_j * xbar_j`` of Kohonen (2013), Eq. (8).
+    Missing features contribute to neither the sum nor the count.
 
     :param data: Dataset of shape ``(n_samples, n_features)``.
     :param weights: Models, of shape ``(x, y, n_features)``.
     :param shape: Shape of the grid.
     :param distance: Dissimilarity measure.
     :param kernel: Optional accelerated search; see :func:`bmu_indices`.
-    :return: Per-node sums of shape ``(x, y, n_features)`` and counts of shape ``(x, y)``.
+    :return: Per-node sums and feature counts, both of shape ``(x, y, n_features)``.
     """
     nodes = bmu_indices(data, weights, distance, kernel)
     n_nodes = shape[0] * shape[1]
-    sums = np.zeros((n_nodes, weights.shape[-1]))
-    np.add.at(sums, nodes, data)
-    counts = np.bincount(nodes, minlength=n_nodes).astype(float)
-    return sums.reshape(*shape, weights.shape[-1]), counts.reshape(shape)
+    n_features = weights.shape[-1]
+    valid = np.isfinite(data)
+    sums = np.zeros((n_nodes, n_features))
+    counts = np.zeros((n_nodes, n_features))
+    np.add.at(sums, nodes, np.where(valid, data, 0.0))
+    np.add.at(counts, nodes, valid)
+    return sums.reshape(*shape, n_features), counts.reshape(*shape, n_features)

src/python_som/_core/_update.py

@@ -71,13 +71,16 @@ def batch_update(
 
     :param weights: Current models, of shape ``(x, y, n_features)``.
     :param sums: Per-node sums of the samples mapped to each node.
-    :param counts: Per-node counts of the samples mapped to each node.
+    :param counts: Per-node, per-feature counts of observed samples, of shape
+        ``(x, y, n_features)``.
     :param hx: Per-axis neighborhood factor for the first axis, of shape ``(x, x)``.
     :param hy: Per-axis neighborhood factor for the second axis, of shape ``(y, y)``.
     :return: The updated models, as a new array.
     """
     numerator = np.einsum("ac,bd,cdf->abf", hx, hy, sums, optimize=True)
-    denominator = np.einsum("ac,bd,cd->ab", hx, hy, counts, optimize=True)
+    if counts.ndim == 2:
+        counts = counts[..., None]
+    denominator = np.einsum("ac,bd,cdf->abf", hx, hy, counts, optimize=True)
     updated = weights.copy()
-    np.divide(numerator, denominator[..., None], out=updated, where=denominator[..., None] > 0)
+    np.divide(numerator, denominator, out=updated, where=denominator > 0)
     return updated

tests/test_batch_update_equivalence.py

@@ -72,7 +72,7 @@ def _per_node_reference(
 
     :param weights: Current models.
     :param sums: Per-node sums.
-    :param counts: Per-node counts.
+    :param counts: Per-node, per-feature counts.
     :param shape: Grid shape.
     :param name: Neighborhood function name.
     :param sigma: Neighborhood radius.
@@ -84,17 +84,21 @@ def _per_node_reference(
     for node in np.ndindex(shape):
         node_2d = (int(node[0]), int(node[1]))
         h = evaluate(shape, node_2d, sigma, cyclic)
-        denominator = float(np.sum(h * counts))
-        if denominator > 0:
-            updated[node_2d] = np.einsum("xy,xyf->f", h, sums) / denominator
+        denominator = np.einsum("xy,xyf->f", h, counts)
+        np.divide(
+            np.einsum("xy,xyf->f", h, sums),
+            denominator,
+            out=updated[node_2d],
+            where=denominator > 0,
+        )
     return updated
 
 
 def _case(shape: tuple[int, int], n_features: int = 3) -> tuple[np.ndarray, ...]:
     """Build models, per-node sums and per-node counts for one grid.
 
-    Counts are drawn with zeros in them on purpose: a node with no data in reach is the case that
-    must keep its previous value, and it is the one a naive implementation destroys.
+    Per-feature counts are drawn with zeros in them on purpose: a node with no data in reach is the
+    case that must keep its previous value, and it is the one a naive implementation destroys.
 
     :param shape: Grid shape.
     :param n_features: Number of features.
@@ -104,7 +108,7 @@ def _case(shape: tuple[int, int], n_features: int = 3) -> tuple[np.ndarray, ...]
     return (
         rng.normal(size=(*shape, n_features)),
         rng.normal(size=(*shape, n_features)),
-        rng.integers(0, 3, size=shape).astype(float),
+        rng.integers(0, 3, size=(*shape, n_features)).astype(float),
     )
 
 
@@ -196,7 +200,7 @@ def test_the_update_is_concurrent_over_every_node() -> None:
     # node's new value, this would disagree.
     late = (shape[0] - 1, shape[1] - 1)
     h = gaussian(shape, late, sigma, (False, False))
-    expected = np.einsum("xy,xyf->f", h, sums) / float(np.sum(h * counts))
+    expected = np.einsum("xy,xyf->f", h, sums) / np.einsum("xy,xyf->f", h, counts)
     np.testing.assert_allclose(updated[late], expected, rtol=1e-12)
 
     assert not np.shares_memory(updated, weights), "the update must not alias its input"

tests/test_bmu_search.py

@@ -229,9 +229,12 @@ def test_accumulate_and_quantization_go_through_the_same_search() -> None:
     _, counts = accumulate(data, weights, shape, euclidean_distance)
     nodes = bmu_indices(data, weights, euclidean_distance)
 
-    expected_counts = np.bincount(nodes, minlength=shape[0] * shape[1]).reshape(shape)
+    expected_counts = np.broadcast_to(
+        np.bincount(nodes, minlength=shape[0] * shape[1]).reshape(shape)[..., None],
+        counts.shape,
+    )
     np.testing.assert_array_equal(counts, expected_counts.astype(float))
-    assert counts.sum() == len(data)
+    assert counts.sum() == len(data) * data.shape[1]
 
     flat = weights.reshape(-1, weights.shape[-1])
     errors = quantization(data, weights, euclidean_distance)