Drop null prisms when converting a prism layer to pyvista

harmonica/_forward/prism_layer.py

@@ -470,21 +470,37 @@ def get_prism(self, indices):
  top = self._obj.top.values[indices]
  return west, east, south, north, bottom, top
 
- def to_pyvista(self):
+ def to_pyvista(self, drop_null_prisms=False):
  """
  Return a pyvista UnstructuredGrid to plot the PrismLayer
 
+ Parameters
+ ----------
+ drop_null_prisms : bool (optional)
+ If True, prisms with zero volume are not going to be included in
+ the :class:`pyvista.UnstructuredGrid`.
+ If False, every prism in the layer will be included.
+ Default False.
+
  Returns
  -------
  pv_grid : :class:`pyvista.UnstructuredGrid`
  :class:`pyvista.UnstructuredGrid` containing each prism of the
  layer as a hexahedron along with their properties.
  """
  prisms = self._to_prisms()
+ null_prisms = np.zeros_like(prisms[:, 0], dtype=bool)
+ if drop_null_prisms:
+ bottom, top = prisms[:, -2], prisms[:, -1]
+ null_prisms = top == bottom
+ prisms = prisms[np.logical_not(null_prisms)]
+ # Define properties
  properties = None
  if self._obj.data_vars:
  properties = {
- data_var: np.asarray(self._obj[data_var]).ravel()
+ data_var: np.asarray(self._obj[data_var]).ravel()[
+ np.logical_not(null_prisms)
+ ]
  for data_var in self._obj.data_vars
  }
  return prism_to_pyvista(prisms, properties=properties)

harmonica/tests/test_prism_layer.py

@@ -419,7 +419,7 @@ def test_to_pyvista(dummy_layer, properties):
  assert pv_grid.n_points == 20 * 8
  # Check coordinates of prisms
  for i, prism in enumerate(layer.prism_layer._to_prisms()):
- npt.assert_allclose(prism, pv_grid.cell_bounds(i))
+ npt.assert_allclose(prism, pv_grid.cell[i].bounds)
  # Check properties of the prisms
  if properties is None:
  assert pv_grid.n_arrays == 0
@@ -431,6 +431,31 @@ def test_to_pyvista(dummy_layer, properties):
  npt.assert_allclose(pv_grid.get_array("density"), layer.density.values.ravel())
 
 
+@pytest.mark.skipif(pyvista is None, reason="requires pyvista")
+def test_to_pyvista_zero_volume(dummy_layer):
+ """
+ Test the conversion of the prism layer to pyvista.UnstructuredGrid when
+ some prisms have zero volume
+ """
+ (easting, northing), surface, reference, density = dummy_layer
+ # Build the layer with or without properties
+ properties = {"density": density}
+ layer = prism_layer((easting, northing), surface, reference, properties=properties)
+ # Assign zero volume to some prisms
+ layer.top.values[0, 0] = layer.bottom.values[0, 0]
+ layer.top.values[2, 1] = layer.bottom.values[2, 1]
+ # Convert the layer to pyvista UnstructuredGrid
+ pv_grid = layer.prism_layer.to_pyvista(drop_null_prisms=True)
+ # Check properties of the pyvista grid
+ expected_n_prisms = 20 - 2
+ assert pv_grid.n_cells == expected_n_prisms
+ assert pv_grid.n_points == expected_n_prisms * 8
+ assert pv_grid.n_arrays == 1
+ assert pv_grid.array_names == ["density"]
+ assert pv_grid.get_array("density").ndim == 1
+ assert pv_grid.get_array("density").size == expected_n_prisms
+
+
 @pytest.mark.skipif(ProgressBar is None, reason="requires numba_progress")
 @pytest.mark.use_numba
 def test_progress_bar(dummy_layer):