WIP Add option to change jacobian dtype on EQL gridders

harmonica/equivalent_layer/harmonic.py

@@ -71,6 +71,10 @@ class EQLHarmonic(vdb.BaseGridder):
  this constant *relative_depth*. Use positive numbers (negative numbers
  would mean point sources are above the data points). Ignored if
  *points* is specified.
+ jacobian_dtype : str, type or numpy dtype
+ Variable type used for defining Jacobian elements. Use
+ ``"float32"`` for reducing memory consumption, but you might expect
+ some accuracy on fitted coefficients. Default ``"float64"``.
 
  Attributes
  ----------
@@ -93,11 +97,12 @@ class EQLHarmonic(vdb.BaseGridder):
  extra_coords_name = "upward"
 
  def __init__(
- self, damping=None, points=None, relative_depth=500,
+ self, damping=None, points=None, relative_depth=500, jacobian_dtype="float64"
  ):
  self.damping = damping
  self.points = points
  self.relative_depth = relative_depth
+ self.jacobian_dtype = jacobian_dtype
  # Define Green's function for Cartesian coordinates
  self.greens_function = greens_func_cartesian
 
@@ -176,9 +181,7 @@ def predict(self, coordinates):
  )
  return data.reshape(shape)
 
- def jacobian(
- self, coordinates, points, dtype="float64"
- ): # pylint: disable=no-self-use
+ def jacobian(self, coordinates, points): # pylint: disable=no-self-use
  """
  Make the Jacobian matrix for the equivalent layer.
 
@@ -196,8 +199,6 @@ def jacobian(
  Tuple of arrays containing the coordinates of the point sources
  used as equivalent layer in the following order:
  (``easting``, ``northing``, ``upward``).
- dtype : str or numpy dtype
- The type of the Jacobian array.
 
  Returns
  -------
@@ -207,7 +208,7 @@ def jacobian(
  # Compute Jacobian matrix
  n_data = coordinates[0].size
  n_points = points[0].size
- jac = np.zeros((n_data, n_points), dtype=dtype)
+ jac = np.zeros((n_data, n_points), dtype=self.jacobian_dtype)
  jacobian_numba(coordinates, points, jac, self.greens_function)
  return jac
 
@@ -272,6 +273,10 @@ class EQLHarmonicSpherical(EQLHarmonic):
  this constant *relative_depth*. Use positive numbers (negative numbers
  would mean point sources are above the data points). Ignored if
  *points* is specified.
+ jacobian_dtype : str, type or numpy dtype
+ Variable type used for defining Jacobian elements. Use
+ ``"float32"`` for reducing memory consumption, but you might expect
+ some accuracy on fitted coefficients. Default ``"float64"``.
 
  Attributes
  ----------
@@ -294,9 +299,14 @@ class EQLHarmonicSpherical(EQLHarmonic):
  extra_coords_name = "radius"
 
  def __init__(
- self, damping=None, points=None, relative_depth=500,
+ self, damping=None, points=None, relative_depth=500, jacobian_dtype="float64"
  ):
- super().__init__(damping=damping, points=points, relative_depth=relative_depth)
+ super().__init__(
+ damping=damping,
+ points=points,
+ relative_depth=relative_depth,
+ jacobian_dtype=jacobian_dtype,
+ )
  # Define Green's function for spherical coordinates
  self.greens_function = greens_func_spherical
 
@@ -355,9 +365,7 @@ def predict(self, coordinates):
  data = super().predict(coordinates)
  return data
 
- def jacobian(
- self, coordinates, points, dtype="float64"
- ): # pylint: disable=no-self-use
+ def jacobian(self, coordinates, points): # pylint: disable=no-self-use
  """
  Make the Jacobian matrix for the equivalent layer.
 
@@ -384,7 +392,7 @@ def jacobian(
  The (n_data, n_points) Jacobian matrix.
  """
  # Overwrite method just to change the docstring
- jacobian = super().jacobian(coordinates, points, dtype=dtype)
+ jacobian = super().jacobian(coordinates, points)
  return jacobian
 
 

harmonica/tests/test_eql_harmonic.py

@@ -134,6 +134,56 @@ def test_eql_harmonic_jacobian_cartesian():
  npt.assert_allclose(jacobian[nearest_neighbours][0], jacobian[nearest_neighbours])
 
 
+@pytest.mark.use_numba
+def test_eql_harmonic_jacobian_dtype():
+ """
+ Check if jacobian dtype can be properly changed
+ Use EQLHarmonic
+ """
+ size = 50
+ region = (-100, 100, -100, 100)
+ coordinates = vd.scatter_points(region, size=size, extra_coords=100)
+ points = (coordinates[0], coordinates[1], coordinates[2] - 200)
+ # Create EQL with default jacobian_dtype
+ eql = EQLHarmonic()
+ jacobian = eql.jacobian(coordinates, points)
+ assert jacobian.nbytes == 8 * size ** 2
+ # Create EQL with jacobian_dtype=float32
+ eql = EQLHarmonic(jacobian_dtype="float32")
+ jacobian = eql.jacobian(coordinates, points)
+ assert jacobian.nbytes == 4 * size ** 2
+
+
+@require_numba
+def test_eql_harmonic_jacobian_dtype_accuracy():
+ """
+ Check if setting jacobian elements to float32 generate accurate predictions
+ Use EQLHarmonic.
+ """
+ region = (-3e3, -1e3, 5e3, 7e3)
+ # Build synthetic point masses
+ points = vd.grid_coordinates(region=region, shape=(6, 6), extra_coords=-1e3)
+ masses = vd.datasets.CheckerBoard(amplitude=1e13, region=region).predict(points)
+ # Define a set of observation points
+ coordinates = vd.scatter_points(region=region, size=600, extra_coords=0)
+ # Get synthetic data
+ data = point_mass_gravity(coordinates, points, masses, field="g_z")
+
+ # Fit EQL gridder using float32 jacobian elements
+ eql_32 = EQLHarmonic(jacobian_dtype="float32")
+ eql_32.fit(coordinates, data)
+
+ # The interpolation should be perfect on the observation points
+ npt.assert_allclose(data, eql_32.predict(coordinates), rtol=1e-2)
+
+ # Compare predicted results obtained by an EQL gridder that uses the
+ # default dtype for jacobian elements
+ eql = EQLHarmonic()
+ eql.fit(coordinates, data)
+ grid = vd.grid_coordinates(region=region, shape=(40, 40), extra_coords=0)
+ npt.assert_allclose(eql.predict(grid), eql_32.predict(grid), rtol=1e-2)
+
+
 @require_numba
 def test_eql_harmonic_spherical():
  """
@@ -252,3 +302,23 @@ def test_eql_harmonic_custom_points_spherical():
  grid, points, masses, field="g_z", coordinate_system="spherical"
  )
  npt.assert_allclose(true, eql.predict(grid), rtol=0.05)
+
+
+@pytest.mark.use_numba
+def test_eql_harmonic_spherical_jacobian_dtype():
+ """
+ Check if jacobian dtype can be properly changed
+ Use EQLHarmonicSpherical
+ """
+ size = 50
+ region = (-10, 10, -10, 10)
+ coordinates = vd.scatter_points(region, size=size, extra_coords=6400e3)
+ points = (coordinates[0], coordinates[1], coordinates[2] - 2000)
+ # Create EQL with default jacobian_dtype
+ eql = EQLHarmonicSpherical()
+ jacobian = eql.jacobian(coordinates, points)
+ assert jacobian.nbytes == 8 * size ** 2
+ # Create EQL with jacobian_dtype=float32
+ eql = EQLHarmonicSpherical(jacobian_dtype="float32")
+ jacobian = eql.jacobian(coordinates, points)
+ assert jacobian.nbytes == 4 * size ** 2