Off-grid sensors. Band-limited interpolation and tests. (#192)

* Add BLI tests
* Add example for off grid sensors.
* Add off-grid sensors.

Author: tomelse

.gitignore

@@ -1,36 +1,37 @@
-# Project
-*.ipynb_checkpoints
-*.pytest_cache
-*.vscode
-*.egg-info
-*__pycache__
-.venv
-site/*
-build/*
-
-# poetry
-poetry.lock
-
-# Testing
-.coverage
-coverage.xml
-
-# Pages to be generated
-docs/test_reports/test_accuracy.png
-docs/test_reports/test_data.txt
-docs/test_reports/test_report.md
-
-# Data
-*.mat
-docs/notebooks/others/ct.jpg
-
-#k-wave outputs
-docs/benchmarks/kwave/*.png
-
-
-# Misc
-*.DS_Store
-*.env.local
-*.env.development.local
-*.env.test.local
-*.env.production.local
+.idea
+# Project
+*.ipynb_checkpoints
+*.pytest_cache
+*.vscode
+*.egg-info
+*__pycache__
+.venv
+site/*
+build/*
+
+# poetry
+poetry.lock
+
+# Testing
+.coverage
+coverage.xml
+
+# Pages to be generated
+docs/test_reports/test_accuracy.png
+docs/test_reports/test_data.txt
+docs/test_reports/test_report.md
+
+# Data
+*.mat
+docs/notebooks/others/ct.jpg
+
+#k-wave outputs
+docs/benchmarks/kwave/*.png
+
+
+# Misc
+*.DS_Store
+*.env.local
+*.env.development.local
+*.env.test.local
+*.env.production.local

CHANGELOG.md

@@ -4,6 +4,8 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 ## [Unreleased]
+### Added
+- Added off grid sensors [@tomelse]
 
 ## [0.1.2] - 2023-06-22
 ### Changed
@@ -64,4 +66,3 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 [0.0.3]: https://github.com/ucl-bug/jwave/compare/0.0.2...0.0.3
 [0.0.2]: https://github.com/ucl-bug/jwave/compare/0.0.1...0.0.2
 [0.0.1]: https://github.com/ucl-bug/jwave/releases/tag/0.0.1
-

docs/notebooks/ivp/off_grid_sensors.ipynb

@@ -16,6 +16,7 @@
 import math
 from dataclasses import dataclass
 from typing import List, Tuple, Union
+from numpy.typing import ArrayLike
 
 import numpy as np
 from jax import numpy as jnp
@@ -35,7 +36,7 @@ class Medium:
 
     Attributes:
       domain (Domain): domain of the medium
-      sound_speed (jnp.darray): speed of sound map, can be a scalar
+      sound_speed (jnp.ndarray): speed of sound map, can be a scalar
       density (jnp.ndarray): density map, can be a scalar
       attenuation (jnp.ndarray): attenuation map, can be a scalar
       pml_size (int): size of the PML layer in grid-points
@@ -148,8 +149,8 @@ class MediumType(Medium):
 @type_of.dispatch
 def type_of(m: Medium):
     return MediumType[type(m.sound_speed),
-                      type(m.density),
-                      type(m.attenuation)]
+    type(m.density),
+    type(m.attenuation)]
 
 
 MediumAllScalars = MediumType[object, object, object]
@@ -187,11 +188,11 @@ def _unit_fibonacci_sphere(
         coordinates of the points on the sphere.
     """
     points = []
-    phi = math.pi * (3.0 - math.sqrt(5.0))    # golden angle in radians
+    phi = math.pi * (3.0 - math.sqrt(5.0))  # golden angle in radians
     for i in range(samples):
-        y = 1 - (i / float(samples - 1)) * 2    # y goes from 1 to -1
-        radius = math.sqrt(1 - y * y)    # radius at y
-        theta = phi * i    # golden angle increment
+        y = 1 - (i / float(samples - 1)) * 2  # y goes from 1 to -1
+        radius = math.sqrt(1 - y * y)  # radius at y
+        theta = phi * i  # golden angle increment
         x = math.cos(theta) * radius
         z = math.sin(theta) * radius
         points.append((x, y, z))
@@ -228,15 +229,15 @@ def _fibonacci_sphere(
 
 def _circ_mask(N, radius, centre):
     x, y = np.mgrid[0:N[0], 0:N[1]]
-    dist_from_centre = np.sqrt((x - centre[0])**2 + (y - centre[1])**2)
+    dist_from_centre = np.sqrt((x - centre[0]) ** 2 + (y - centre[1]) ** 2)
     mask = (dist_from_centre < radius).astype(int)
     return mask
 
 
 def _sphere_mask(N, radius, centre):
     x, y, z = np.mgrid[0:N[0], 0:N[1], 0:N[2]]
-    dist_from_centre = np.sqrt((x - centre[0])**2 + (y - centre[1])**2 +
-                               (z - centre[2])**2)
+    dist_from_centre = np.sqrt((x - centre[0]) ** 2 + (y - centre[1]) ** 2 +
+                               (z - centre[2]) ** 2)
     mask = (dist_from_centre < radius).astype(int)
     return mask
 
@@ -327,7 +328,7 @@ def __init__(self, mask, signal, dt, domain):
 
     def tree_flatten(self):
         children = (self.mask, self.signal, self.dt)
-        aux = (self.domain, )
+        aux = (self.domain,)
         return (children, aux)
 
     @classmethod
@@ -430,7 +431,7 @@ def __init__(self, positions):
 
     def tree_flatten(self):
         children = None
-        aux = (self.positions, )
+        aux = (self.positions,)
         return (children, aux)
 
     @classmethod
@@ -461,10 +462,115 @@ def __call__(self, p: Field, u: Field, rho: Field):
             return p.on_grid[self.positions[0]]
         elif len(self.positions) == 2:
             return p.on_grid[self.positions[0],
-                             self.positions[1]]    # type: ignore
+            self.positions[1]]  # type: ignore
         elif len(self.positions) == 3:
             return p.on_grid[self.positions[0], self.positions[1],
-                             self.positions[2]]    # type: ignore
+            self.positions[2]]  # type: ignore
+        else:
+            raise ValueError(
+                "Sensors positions must be 1, 2 or 3 dimensional. Not {}".
+                format(len(self.positions)))
+
+
+def _bli_function(x0: jnp.ndarray, x: jnp.ndarray, n: int, include_imag: bool = False) -> jnp.ndarray:
+    """
+    The function used to compute the band limited interpolation function.
+
+    Args:
+        x0 (jnp.ndarray): Position of the sensors along the axis.
+        x (jnp.ndarray): Grid positions.
+        n (int): Size of the grid
+        include_imag (bool): Include the imaginary component?
+
+    Returns:
+    jnp.ndarray: The values of the function at the grid positions.
+    """
+    dx = jnp.where((x - x0[:, None]) == 0, 1, x - x0[:, None])  # https://github.com/google/jax/issues/1052
+    dx_nonzero = (x - x0[:, None]) != 0
+
+    if n % 2 == 0:
+        y = jnp.sin(jnp.pi * dx) / \
+            jnp.tan(jnp.pi * dx / n) / n
+        y -= jnp.sin(jnp.pi * x0[:, None]) * jnp.sin(jnp.pi * x) / n
+        if include_imag:
+            y += 1j * jnp.cos(jnp.pi * x0[:, None]) * jnp.sin(jnp.pi * x) / n
+    else:
+        y = jnp.sin(jnp.pi * dx) / \
+            jnp.sin(jnp.pi * dx / n) / n
+
+    # Deal with case of precisely on grid.
+    y = y * jnp.all(dx_nonzero, axis=1)[:, None] + (1 - dx_nonzero) * (~jnp.all(dx_nonzero, axis=1)[:, None])
+    return y
+
+
+@register_pytree_node_class
+class BLISensors:
+    """ Band-limited interpolant (off-grid) sensors.
+
+    Args:
+        positions (Tuple of List of float): Sensor positions.
+        n (Tuple of int): Grid size.
+
+    Attributes:
+        positions (Tuple[jnp.ndarray]): Sensor positions
+        n (Tuple[int]): Grid size.
+    """
+
+    positions: Tuple[jnp.ndarray]
+    n: Tuple[int]
+
+    def __init__(self, positions: Tuple[jnp.ndarray], n: Tuple[int]):
+        self.positions = positions
+        self.n = n
+
+        # Calculate the band-limited interpolant weights if not provided.
+        x = jnp.arange(n[0])[None]
+        self.bx = jnp.expand_dims(_bli_function(positions[0], x, n[0]),
+                                  axis=range(2, 2 + len(n)))
+
+        if len(n) > 1:
+            y = jnp.arange(n[1])[None]
+            self.by = jnp.expand_dims(_bli_function(positions[1], y, n[1]),
+                                      axis=range(2, 2 + len(n) - 1))
+        else:
+            self.by = None
+
+        if len(n) > 2:
+            z = jnp.arange(n[2])[None]
+            self.bz = jnp.expand_dims(_bli_function(positions[2], z, n[2]),
+                                      axis=range(2, 2 + len(n) - 2))
+        else:
+            self.bz = None
+
+    def tree_flatten(self):
+        children = self.positions,
+        aux = self.n,
+        return children, aux
+
+    @classmethod
+    def tree_unflatten(cls, aux, children):
+        return cls(*children, *aux)
+
+    def __call__(self, p: Field, u, v):
+        r"""Returns the values of the field p at the sensors positions.
+        Args:
+          p (Field): The field to be sampled.
+        """
+        if len(self.positions) == 1:
+            # 1D
+            pw = jnp.sum(p.on_grid[None] * self.bx, axis=1)
+            return pw
+        elif len(self.positions) == 2:
+            # 2D
+            pw = jnp.sum(p.on_grid[None] * self.bx, axis=1)
+            pw = jnp.sum(pw * self.by, axis=1)
+            return pw
+        elif len(self.positions) == 3:
+            # 3D
+            pw = jnp.sum(p.on_grid[None] * self.bx, axis=1)
+            pw = jnp.sum(pw * self.by, axis=1)
+            pw = jnp.sum(pw * self.bz, axis=1)
+            return pw
         else:
             raise ValueError(
                 "Sensors positions must be 1, 2 or 3 dimensional. Not {}".
@@ -488,7 +594,7 @@ def __init__(self, dt, t_end):
         self.t_end = t_end
 
     def tree_flatten(self):
-        children = (None, )
+        children = (None,)
         aux = (self.dt, self.t_end)
         return (children, aux)
 
@@ -522,7 +628,7 @@ def from_medium(medium: Medium, cfl: float = 0.3, t_end=None):
             np.max)
         if t_end is None:
             t_end = np.sqrt(
-                sum((x[-1] - x[0])**2
+                sum((x[-1] - x[0]) ** 2
                     for x in medium.domain.spatial_axis)) / functional(
-                        medium.sound_speed)(np.min)
+                medium.sound_speed)(np.min)
         return TimeAxis(dt=float(dt), t_end=float(t_end))

jwave/geometry.py

@@ -74,6 +74,7 @@ nav:
       - Homogeneous wave propagation: notebooks/ivp/homogeneous_medium.ipynb
       - 3D simulations: notebooks/ivp/3d.ipynb
       - Sensors: notebooks/ivp/homogeneous_medium_sensors.ipynb
+      - Off-Grid Sensors: notebooks/ivp/off_grid_sensors.ipynb
       - Custom sensors: notebooks/ivp/custom_sensors.ipynb
       - Automatic differentiation: notebooks/ivp/homogeneous_medium_backprop.ipynb
       - Heterogeneous medium: notebooks/ivp/heterogeneous_medium.ipynb

mkdocs.yml

@@ -0,0 +1,82 @@
+from jwave.geometry import _bli_function, BLISensors, Domain, FourierSeries
+import numpy as np
+import pytest
+
+
+@pytest.mark.parametrize("n_grid", [100, 101])
+def test_bli_function(n_grid):
+    # Make a load of sensors on the grid. Check the bli function.
+    y = _bli_function(np.arange(n_grid), np.arange(n_grid), n_grid)
+    # Assert that the bli function is 1 at one place for each detector.
+    assert (np.all(np.sum(y != 0, axis=1) == 1))
+    # Assert that the bli function is 0 everywhere else
+    assert (np.all(np.sum(y == 0, axis=1) == n_grid - 1))
+    # Assert that the 1 is in the correct place.
+    assert (np.all(y[np.arange(n_grid), np.arange(n_grid)] == 1))
+
+    # Check off-grid points:
+    y = _bli_function(np.arange(0, n_grid - 1) + 0.25, np.arange(n_grid), n_grid)
+    # Check that the sensor is non-zero at more than one place.
+    assert (np.all(np.sum(y != 0, axis=1) > 1))
+    assert (np.all(np.isclose(np.sum(y, axis=1), 1)))
+
+
+@pytest.mark.parametrize("nx,ny,nz", [(100, 100, 100), (100, 101, 102)])
+def test_sensor(nx, ny, nz):
+    # Check that it is identical to on-grid detectors in that case.
+    n_detectors = min(nx, ny, nz)
+
+    xi = np.arange(n_detectors)
+    yi = np.arange(n_detectors)
+    zi = np.arange(n_detectors)
+    np.random.shuffle(xi)
+    np.random.shuffle(yi)
+    np.random.shuffle(zi)
+
+    x = xi.astype(float)
+    y = yi.astype(float)
+    z = zi.astype(float)
+
+    p = np.random.random((nx, ny, nz))
+
+    s1d = BLISensors((x,), (nx,))
+    s2d = BLISensors((x, y), (nx, ny))
+    s3d = BLISensors((x, y, z), (nx, ny, nz))
+
+    domain1d = Domain((nx,), (1,))
+    p1d = FourierSeries(p[:, 0, 0], domain1d)
+
+    domain2d = Domain((nx, ny), (1, 1))
+    p2d = FourierSeries(p[:, :, 0], domain2d)
+
+    domain3d = Domain((nx, ny, nz), (1, 1, 1))
+    p3d = FourierSeries(p, domain3d)
+
+    result = s1d(p1d, None, None)
+    assert (np.all(result[..., 0] == p[xi, 0, 0]))
+
+    result = s2d(p2d, None, None)
+    assert (np.all(result[..., 0] == p[xi, yi, 0]))
+
+    result = s3d(p3d, None, None)
+    assert (np.all(result[..., 0] == p[xi, yi, zi]))
+
+    # Check off-grid (perturb a bit):
+    s3d = BLISensors((x + 0.25, y + 0.3, z + 0.1), (nx, ny, nz))
+    domain3d = Domain((nx, ny, nz), (1, 1, 1))
+    # Check ones in ones out.
+    p3d = FourierSeries(np.ones((nx, ny, nz)), domain3d)
+    y = s3d(p3d, None, None)
+    assert (np.all(np.isclose(y, 1)))
+
+    # Check zeros in zeros out
+    p3d = FourierSeries(np.zeros((nx, ny, nz)), domain3d)
+    y = s3d(p3d, None, None)
+    assert (np.all(y == 0))
+
+
+if __name__ == "__main__":
+    test_bli_function(100)
+    test_bli_function(101)
+    test_sensor()
+    test_sensor(100, 101, 102)