change the creation of sparse matrices to support regional maps (#27)

healpix_convolution/kernels/common.py

@@ -2,18 +2,25 @@
 import sparse
 
 
-def create_sparse(cell_ids, neighbours, weights, shape):
+def create_sparse(cell_ids, neighbours, weights):
     neighbours_ = np.reshape(neighbours, (-1,))
-    mask = neighbours_ != -1
 
-    n_neighbours = neighbours.shape[-1]
-    cell_ids_ = np.reshape(
-        np.repeat(cell_ids[..., None], repeats=n_neighbours, axis=-1), (-1,)
+    all_cell_ids = np.unique(neighbours_)
+    if all_cell_ids[0] == -1:
+        all_cell_ids = all_cell_ids[1:]
+
+    row_indices = np.reshape(
+        np.broadcast_to(np.arange(cell_ids.size)[:, None], shape=neighbours.shape),
+        (-1,),
     )
+    column_indices = np.searchsorted(all_cell_ids, neighbours_, side="left")
 
-    coords = np.reshape(np.stack([cell_ids_, neighbours_], axis=0), (2, -1))
+    coords = np.stack([row_indices, column_indices], axis=0)
+
+    mask = neighbours_ != -1
 
     weights_ = np.reshape(weights, (-1,))[mask]
     coords_ = coords[..., mask]
 
+    shape = (cell_ids.size, all_cell_ids.size)
     return sparse.COO(coords=coords_, data=weights_, shape=shape, fill_value=0)

healpix_convolution/kernels/gaussian.py

@@ -65,7 +65,4 @@ def gaussian_kernel(
     masked = np.where(nb == -1, 0, phi_x)
     normalized = masked / np.sum(masked, axis=1, keepdims=True)
 
-    # TODO (keewis): figure out a way to translate global healpix indices to local ones
-    # The kernel should still work for a subset of the full map.
-    shape = (12 * 4**resolution, 12 * 4**resolution)
-    return create_sparse(cell_ids, nb, normalized, shape=shape)
+    return create_sparse(cell_ids, nb, normalized)

healpix_convolution/tests/test_kernels.py

@@ -24,16 +24,21 @@
     ),
 )
 def test_create_sparse(cell_ids, neighbours, weights):
-    actual = kernels.common.create_sparse(cell_ids, neighbours, weights, shape=(48, 48))
+    input_cell_ids = np.unique(neighbours)
+    if input_cell_ids[0] == -1:
+        input_cell_ids = input_cell_ids[1:]
+
+    actual = kernels.common.create_sparse(cell_ids, neighbours, weights)
 
     nnz = np.sum(neighbours != -1, axis=1)
     value = nnz * weights[0]
 
+    expected_shape = (cell_ids.size, input_cell_ids.size)
     assert hasattr(actual, "nnz"), "not a sparse matrix"
     assert np.allclose(
-        np.sum(actual[cell_ids, :], axis=1).todense(), value
+        np.sum(actual, axis=1).todense(), value
     ), "rows have unexpected values"
-    assert actual.size == 48**2
+    assert actual.shape == expected_shape
 
 
 class TestGaussian:
@@ -78,7 +83,7 @@ def test_gaussian_kernel(self, cell_ids, kwargs):
         kernel_sum = np.sum(actual, axis=1)
 
         assert np.sum(np.isnan(actual)) == 0
-        np.testing.assert_allclose(kernel_sum[cell_ids].todense(), 1)
+        np.testing.assert_allclose(kernel_sum.todense(), 1)
 
         # try determining the sigma from the values for better tests