Repository: fatiando/harmonica
Pull request #243 (2021-09-03 19:43:44)

Allow prism_layer to take xarray objects as arguments
Extend the support of `prism_layer` to accept `surface`, `reference`, and
properties as `xarray.DataArray`. Improve the test functions by creating
a `pytest` fixture that parametrizes the data types of the output arrays:
either Numpy arrays or `xarray.DataArray`s. Generate a new `pytest` fixture
that returns a layer of prisms with some holes (missing prisms), ready to be
used as the expected results of some test functions. Parametrize the `field` on
each test function that computes gravitational fields of prism layers. Simplify
the `prism_layer_topo` example so the surface and the density of the layer are
passed as xarray objects.


**Reminders**:

- [ ] Run `make format` and `make check` to make sure the code follows the style guide.
- [ ] Add tests for new features or tests that would have caught the bug that you're fixing.
- [ ] Add new public functions/methods/classes to `doc/api/index.rst` and the base `__init__.py` file for the package.
- [ ] Write detailed docstrings for all functions/classes/methods. It often helps to design better code if you write the docstrings first.
- [ ] If adding new functionality, add an example to the docstring, gallery, and/or tutorials.
- [ ] Add your full name, affiliation, and [ORCID](https://orcid.org) (optional) to the `AUTHORS.md` file (if you haven't already) in case you'd like to be listed as an author on the [Zenodo](https://zenodo.org/communities/fatiando) archive of the next release.



Patch:
diff --git a/examples/forward/prisms_topo_gravity.py b/examples/forward/prisms_topo_gravity.py
index be25914..de9b8bf 100644
--- a/examples/forward/prisms_topo_gravity.py
+++ b/examples/forward/prisms_topo_gravity.py
@@ -45,7 +45,7 @@ density = density.where(south_africa_topo >= 0, 1000 - 2900)
 # Create layer of prisms
 prisms = hm.prism_layer(
     (south_africa_topo.easting, south_africa_topo.northing),
-    surface=south_africa_topo.values,
+    surface=south_africa_topo,
     reference=0,
     properties={"density": density},
 )
diff --git a/harmonica/forward/prism_layer.py b/harmonica/forward/prism_layer.py
index 15756c5..fdc0eac 100644
--- a/harmonica/forward/prism_layer.py
+++ b/harmonica/forward/prism_layer.py
@@ -118,7 +118,7 @@ def prism_layer(
     # and values, respectively
     if properties:
         data_names = tuple(p for p in properties.keys())
-        data = tuple(p for p in properties.values())
+        data = tuple(np.asarray(p) for p in properties.values())
     # Create xr.Dataset for prisms
     prisms = vd.make_xarray_grid(
         coordinates, data=data, data_names=data_names, dims=dims
@@ -279,13 +279,14 @@ class DatasetAccessorPrismLayer:
             prisms layer. It can be either a plane or an irregular surface
             passed as 2d array. Height(s) must be in meters.
         """
+        surface, reference = np.asarray(surface), np.asarray(reference)
         if surface.shape != self.shape:
             raise ValueError(
                 f"Invalid surface array with shape '{surface.shape}'. "
                 + "Its shape should be compatible with the coordinates "
                 + "of the layer of prisms."
             )
-        if isinstance(reference, np.ndarray):
+        if reference.ndim != 0:
             if reference.shape != self.shape:
                 raise ValueError(
                     f"Invalid reference array with shape '{reference.shape}'. "
@@ -293,7 +294,7 @@ class DatasetAccessorPrismLayer:
                     + "of the layer of prisms."
                 )
         else:
-            reference *= np.ones(self.shape)
+            reference = reference * np.ones(self.shape)
         top = surface.copy()
         bottom = reference.copy()
         reverse = surface < reference


Test patch:
diff --git a/harmonica/tests/test_prism_layer.py b/harmonica/tests/test_prism_layer.py
index 05965e5..ecf6395 100644
--- a/harmonica/tests/test_prism_layer.py
+++ b/harmonica/tests/test_prism_layer.py
@@ -13,18 +13,62 @@ import pytest
 import numpy as np
 import numpy.testing as npt
 import verde as vd
+import xarray as xr
 
 from .. import prism_layer, prism_gravity
 
 
-def test_prism_layer():
+@pytest.fixture(params=("numpy", "xarray"))
+def dummy_layer(request):
     """
-    Check if a layer of prisms is property constructed
+    Generate dummy arrays for defining prism layers
     """
     easting = np.linspace(-1, 3, 5)
     northing = np.linspace(7, 13, 4)
+    shape = (northing.size, easting.size)
     reference = 0
-    surface = np.arange(20, dtype=float).reshape(4, 5)
+    surface = np.arange(20, dtype=float).reshape(*shape)
+    density = 2670 * np.ones(shape)
+    if request.param == "xarray":
+        easting = xr.DataArray(easting, dims=("easting",))
+        northing = xr.DataArray(northing, dims=("northing",))
+        reference, surface = xr.DataArray(reference), xr.DataArray(surface)
+        density = xr.DataArray(density)
+    return (easting, northing), surface, reference, density
+
+
+@pytest.fixture
+def prism_layer_with_holes(dummy_layer):  # pylint: disable=redefined-outer-name
+    """
+    Return a set of prisms with some missing elements
+
+    The prisms are returned as a tuple of boundaries, ready to be passed to
+    ``hm.prism_gravity``. They would represent the same prisms that the
+    ``dummy_layer`` generated, but with two missing prisms: the ``(3, 3)`` and
+    the ``(2, 1)``.
+    """
+    (easting, northing), surface, reference, density = dummy_layer
+    layer = prism_layer(
+        (easting, northing), surface, reference, properties={"density": density}
+    )
+    indices = [(3, 3), (2, 1)]
+    prisms = list(
+        layer.prism_layer.get_prism((i, j))
+        for i in range(4)
+        for j in range(5)
+        if (i, j) not in indices
+    )
+    density = list(
+        density[i, j] for i in range(4) for j in range(5) if (i, j) not in indices
+    )
+    return prisms, density
+
+
+def test_prism_layer(dummy_layer):  # pylint: disable=redefined-outer-name
+    """
+    Check if a layer of prisms is property constructed
+    """
+    (easting, northing), surface, reference, _ = dummy_layer
     layer = prism_layer((easting, northing), surface, reference)
     assert "easting" in layer.coords
     assert "northing" in layer.coords
@@ -37,7 +81,7 @@ def test_prism_layer():
     # Surface below reference on a single point
     surface[1, 1] = -1
     expected_top = surface.copy()
-    expected_bottom = reference * np.ones_like(surface)
+    expected_bottom = np.zeros_like(surface)
     expected_top[1, 1], expected_bottom[1, 1] = reference, surface[1, 1]
     layer = prism_layer((easting, northing), surface, reference)
     assert "easting" in layer.coords
@@ -50,35 +94,32 @@ def test_prism_layer():
     npt.assert_allclose(layer.bottom, expected_bottom)
 
 
-def test_prism_layer_invalid_surface_reference():
+def test_prism_layer_invalid_surface_reference(
+    dummy_layer,
+):  # pylint: disable=redefined-outer-name
     """
     Check if invalid surface and/or reference are caught
     """
-    coordinates = np.linspace(-1, 3, 5), np.linspace(7, 10, 4)
+    coordinates, surface, reference, _ = dummy_layer
     # Surface with wrong shape
-    reference = 0
-    surface = np.arange(20, dtype=float)
+    surface_invalid = np.arange(20, dtype=float)
     with pytest.raises(ValueError):
-        prism_layer(coordinates, surface, reference)
+        prism_layer(coordinates, surface_invalid, reference)
     # Reference with wrong shape
-    reference = np.zeros(20)
+    reference_invalid = np.zeros(20)
     surface = np.arange(20, dtype=float).reshape(4, 5)
     with pytest.raises(ValueError):
-        prism_layer(coordinates, surface, reference)
+        prism_layer(coordinates, surface, reference_invalid)
 
 
-def test_prism_layer_properties():
+def test_prism_layer_properties(dummy_layer):  # pylint: disable=redefined-outer-name
     """
     Check passing physical properties to the prisms layer
     """
-    easting = np.linspace(-1, 3, 5)
-    northing = np.linspace(7, 10, 4)
-    reference = 0
-    surface = np.arange(20, dtype=float).reshape(4, 5)
-    density = 2670 * np.ones_like(surface)
-    suceptibility = np.arange(20, dtype=float).reshape(4, 5)
+    coordinates, surface, reference, density = dummy_layer
+    suceptibility = 0 * density + 1e-3
     layer = prism_layer(
-        (easting, northing),
+        coordinates,
         surface,
         reference,
         properties={"density": density, "suceptibility": suceptibility},
@@ -87,29 +128,29 @@ def test_prism_layer_properties():
     npt.assert_allclose(layer.suceptibility, suceptibility)
 
 
-def test_prism_layer_no_regular_grid():
+def test_prism_layer_no_regular_grid(
+    dummy_layer,
+):  # pylint: disable=redefined-outer-name
     """
     Check if error is raised if easting and northing are not regular
     """
-    reference = 0
-    surface = np.arange(20, dtype=float).reshape(4, 5)
+    (easting, northing), surface, reference, _ = dummy_layer
     # Easting as non evenly spaced set of coordinates
-    easting = np.linspace(-1, 3, 5)
-    northing = np.linspace(7, 13, 4)
-    easting[3] = -22
+    easting_invalid = easting.copy()
+    easting_invalid[3] = -22
     with pytest.raises(ValueError):
         prism_layer(
-            (easting, northing),
+            (easting_invalid, northing),
             surface,
             reference,
         )
     # Northing as non evenly spaced set of coordinates
-    easting = np.linspace(-1, 3, 5)
-    northing = np.linspace(7, 13, 4)
+    northing_invalid = northing.copy()
+    northing_invalid[3] = -22
     northing[3] = 12.98
     with pytest.raises(ValueError):
         prism_layer(
-            (easting, northing),
+            (easting, northing_invalid),
             surface,
             reference,
         )
@@ -172,15 +213,12 @@ def test_prism_layer_get_prism_by_index():
             )
 
 
-def test_nonans_prisms_mask():
+def test_nonans_prisms_mask(dummy_layer):  # pylint: disable=redefined-outer-name
     """
     Check if the mask for nonans prism is correctly created
     """
-    easting = np.linspace(1, 3, 5)
-    northing = np.linspace(7, 10, 4)
+    (easting, northing), surface, reference, _ = dummy_layer
     shape = (northing.size, easting.size)
-    reference = 0
-    surface = np.arange(20, dtype=float).reshape(shape)
 
     # No nan in top nor bottom
     # ------------------------
@@ -223,16 +261,14 @@ def test_nonans_prisms_mask():
     npt.assert_allclose(mask, expected_mask)
 
 
-def test_nonans_prisms_mask_property():
+def test_nonans_prisms_mask_property(
+    dummy_layer,
+):  # pylint: disable=redefined-outer-name
     """
     Check if the method masks the property and raises a warning
     """
-    easting = np.linspace(1, 3, 5)
-    northing = np.linspace(7, 10, 4)
+    (easting, northing), surface, reference, density = dummy_layer
     shape = (northing.size, easting.size)
-    reference = 0
-    surface = np.arange(20, dtype=float).reshape(shape)
-    density = 2670 * np.ones_like(surface)
 
     # Nans in top and property (on the same prisms)
     # ---------------------------------------------
@@ -279,96 +315,79 @@ def test_nonans_prisms_mask_property():
 
 
 @pytest.mark.use_numba
-def test_prism_layer_gravity():
+@pytest.mark.parametrize("field", ["potential", "g_z"])
+def test_prism_layer_gravity(
+    field, dummy_layer
+):  # pylint: disable=redefined-outer-name
     """
     Check if gravity method works as expected
     """
     coordinates = vd.grid_coordinates((1, 3, 7, 10), spacing=1, extra_coords=30.0)
-    easting = np.linspace(1, 3, 5)
-    northing = np.linspace(7, 10, 4)
-    shape = (northing.size, easting.size)
-    reference = 0
-    surface = np.arange(20, dtype=float).reshape(shape)
-    density = np.ones_like(surface, dtype=float)
+    (easting, northing), surface, reference, density = dummy_layer
     layer = prism_layer(
         (easting, northing), surface, reference, properties={"density": density}
     )
-    for field in ("potential", "g_z"):
-        expected_result = prism_gravity(
-            coordinates,
-            prisms=layer.prism_layer._to_prisms(),
-            density=density,
-            field=field,
-        )
-        npt.assert_allclose(
-            expected_result, layer.prism_layer.gravity(coordinates, field=field)
-        )
+    expected_result = prism_gravity(
+        coordinates,
+        prisms=layer.prism_layer._to_prisms(),
+        density=density,
+        field=field,
+    )
+    npt.assert_allclose(
+        expected_result, layer.prism_layer.gravity(coordinates, field=field)
+    )
 
 
 @pytest.mark.use_numba
-def test_prism_layer_gravity_with_nans():
+@pytest.mark.parametrize("field", ["potential", "g_z"])
+def test_prism_layer_gravity_surface_nans(
+    field, dummy_layer, prism_layer_with_holes
+):  # pylint: disable=redefined-outer-name
     """
-    Check if gravity method works as expected when one of the prisms has nans
+    Check if gravity method works as expected when surface has nans
     """
     coordinates = vd.grid_coordinates((1, 3, 7, 10), spacing=1, extra_coords=30.0)
-    easting = np.linspace(1, 3, 5)
-    northing = np.linspace(7, 10, 4)
-    shape = (northing.size, easting.size)
-    reference = 0
+    (easting, northing), surface, reference, density = dummy_layer
     # Create one layer that has nans on the surface array
-    surface = np.arange(20, dtype=float).reshape(shape)
+    surface_w_nans = surface.copy()
     indices = [(3, 3), (2, 1)]
     for index in indices:
-        surface[index] = np.nan
-    density = np.ones_like(surface, dtype=float)
-    layer_nans = prism_layer(
-        (easting, northing), surface, reference, properties={"density": density}
+        surface_w_nans[index] = np.nan
+    layer = prism_layer(
+        (easting, northing), surface_w_nans, reference, properties={"density": density}
     )
-    # Create one layer that has zero density but no nans
-    surface = np.arange(20, dtype=float).reshape(shape)
-    density = np.ones_like(surface, dtype=float)
-    for index in indices:
-        density[index] = 0
-    layer_nonans = prism_layer(
-        (easting, northing), surface, reference, properties={"density": density}
+    # Check if it generates the expected gravity field
+    prisms, rho = prism_layer_with_holes
+    npt.assert_allclose(
+        layer.prism_layer.gravity(coordinates, field=field),
+        prism_gravity(coordinates, prisms, rho, field=field),
     )
-    # Check if the two layers generate the same gravity field
-    for field in ("potential", "g_z"):
-        npt.assert_allclose(
-            layer_nans.prism_layer.gravity(coordinates, field=field),
-            layer_nonans.prism_layer.gravity(coordinates, field=field),
-        )
 
 
-def test_prism_layer_gravity_density_nans():
+@pytest.mark.use_numba
+@pytest.mark.parametrize("field", ["potential", "g_z"])
+def test_prism_layer_gravity_density_nans(
+    field, dummy_layer, prism_layer_with_holes
+):  # pylint: disable=redefined-outer-name
     """
     Check if prisms is ignored after a nan is found in density array
     """
     coordinates = vd.grid_coordinates((1, 3, 7, 10), spacing=1, extra_coords=30.0)
-    easting = np.linspace(1, 3, 5)
-    northing = np.linspace(7, 10, 4)
-    shape = (northing.size, easting.size)
-    reference = 0
+    prisms_coords, surface, reference, density = dummy_layer
     # Create one layer that has nans on the density array
-    surface = np.arange(20, dtype=float).reshape(shape)
     indices = [(3, 3), (2, 1)]
     for index in indices:
-        surface[index] = np.nan
-    density = np.ones_like(surface, dtype=float)
-    layer_nans = prism_layer(
-        (easting, northing), surface, reference, properties={"density": density}
+        density[index] = np.nan
+    layer = prism_layer(
+        prisms_coords, surface, reference, properties={"density": density}
     )
-    # Create one layer that has zero density but no nans
-    surface = np.arange(20, dtype=float).reshape(shape)
-    density = np.ones_like(surface, dtype=float)
-    for index in indices:
-        density[index] = 0
-    layer_nonans = prism_layer(
-        (easting, northing), surface, reference, properties={"density": density}
+    # Check if warning is raised after passing density with nans
+    with warnings.catch_warnings(record=True) as warn:
+        result = layer.prism_layer.gravity(coordinates, field=field)
+        assert len(warn) == 1
+    # Check if it generates the expected gravity field
+    prisms, rho = prism_layer_with_holes
+    npt.assert_allclose(
+        result,
+        prism_gravity(coordinates, prisms, rho, field=field),
     )
-    # Check if the two layers generate the same gravity field
-    for field in ("potential", "g_z"):
-        npt.assert_allclose(
-            layer_nans.prism_layer.gravity(coordinates, field=field),
-            layer_nonans.prism_layer.gravity(coordinates, field=field),
-        )