Extract 1D (#301)

CHANGELOG.rst

@@ -9,6 +9,10 @@ Changelog
 latest
 ------
 
+- Model instances have a new attribute ``exctract_1d``, which returns a layered
+  (1D) model, extracted from the 3D model according the provided parameters;
+  see :attr:`emg3d.models.Model.extract_1d`.
+
 - CLI takes new the boolean ``add_noise`` in the section ``[noise_opts]``
   (default is True).
 

emg3d/maps.py

@@ -812,10 +812,14 @@ def ellipse_indices(coo, p0, p1, radius, factor=1., minor=1., check_foci=True):
 
     Parameters
     ----------
-    coo : tuple of ndarray
-        Tuple of two arrays defining the points in x and y.
-
-    p0, p1 : ndarray
+    coo : tuple of two ndarrays
+        Tuple of two arrays defining the points in x and y:
+        - If two vectors are given (of same or different size), they are taken
+          as the x- and y-values of a regular grid.
+        - If two 2D-arrays are given of the same shape, they are taken as the
+          (regular or irregular) x- and y-values.
+
+    p0, p1 : array_like
         (x, y)-coordinates of two points.
 
     radius : float
@@ -847,10 +851,12 @@ def ellipse_indices(coo, p0, p1, radius, factor=1., minor=1., check_foci=True):
 
     """
     # Center coordinates
-    cx, cy = (p0 + p1) / 2
+    cx = (p0[0] + p1[0]) / 2
+    cy = (p0[1] + p1[1]) / 2
 
     # Adjacent and opposite sides
-    dx, dy = (p1 - p0) / 2
+    dx = (p1[0] - p0[0]) / 2
+    dy = (p1[1] - p0[1]) / 2
 
     # c: linear eccentricity
     dxy = np.linalg.norm([dx, dy])
@@ -874,4 +880,7 @@ def ellipse_indices(coo, p0, p1, radius, factor=1., minor=1., check_foci=True):
     A = (cos/major)**2 + (sin/minor)**2
     B = 2*cos*sin*(major**-2 - minor**-2)
     C = (sin/major)**2 + (cos/minor)**2
-    return A*X**2 + B*X*Y + C*Y**2 <= 1.0
+    if X.ndim == 1:
+        return A*X[:, None]**2 + B*np.outer(X, Y) + C*Y[None, :]**2 <= 1.0
+    else:
+        return A*X**2 + B*X*Y + C*Y**2 <= 1.0

emg3d/meshes.py

@@ -79,7 +79,9 @@ def __init__(self, h, origin, **kwargs):
         self.origin = np.array(origin)
 
         # Width of cells, cast to arrays.
-        self.h = [np.array(h[0]), np.array(h[1]), np.array(h[2])]
+        self.h = [np.array(h[0], dtype=float),
+                  np.array(h[1], dtype=float),
+                  np.array(h[2], dtype=float)]
 
         # Node related properties.
         shape_nodes = (self.h[0].size+1, self.h[1].size+1, self.h[2].size+1)

emg3d/models.py

@@ -365,6 +365,184 @@ def interpolate_to_grid(self, grid, **interpolate_opts):
         # Assemble new model.
         return Model(grid, mapping=self.map.name, **model_inp)
 
+    def extract_1d(self, method, p0, p1=None, ellipse=None, merge=False,
+                   return_imat=False):
+        """Return a layered (1D) model.
+
+        The returned model has shape (1, 1, nz), where nz is the original
+        number of cells in vertical direction (except if ``merge=True``). How
+        this 1D model is obtained depends on the input parameters.
+
+        .. note::
+
+            If ``method`` is either ``'cylinder'`` or ``'prism'``, an
+            ``ellipse``-dict has to be provided with at least the key-value
+            pair for ``'radius'``. E.g., ``ellipse={'radius': 1000}``.
+
+
+        Parameters
+        ----------
+        method : str
+            Method how to obtain the layered model. Implemented are:
+
+            - ``'midpoint'``: Returns the model at the midpoint between ``p0``
+              and ``p1``. If ``p1`` is not provided, the model at ``p0`` is
+              returned. The x- and y-width of the returned model corresponds to
+              the selected cell.
+            - ``'cylinder'``: Volume-averages the values of each layer within a
+              right elliptic cylinder, using the function
+              :func:`emg3d.maps.ellipse_indices` to obtain the ellipse. This
+              method requires the ``ellipse``-parameter. The x- and y-width of
+              the returned model corresponds to the x- and y-width of the
+              ellipse.
+            - ``'prism'``: Same as ``'cylinder'``, but an enveloping right
+              rectangular prism is fit around the cylinder.
+
+        p0, p1 : array_like
+            (x, y)-coordinates of two points. If ``p1`` is not provided, it is
+            set to ``p0``.
+
+        ellipse : dict, default: None
+            Options-dict passed through to :func:`emg3d.maps.ellipse_indices`.
+            This is only used in the methods ``'cylinder'`` and ``'prism'``,
+            for which it is a required parameter.
+
+        merge : bool, default: False
+            If True, adjacent layers of identical properties are combined. That
+            implies that the returned model might have less (but larger) cells
+            in z-direction than the original model.
+
+        return_imat : bool, default: False
+            If True, the interpolation matrix is returned in addition to the
+            model. This matrix can be used, for instance, to inspect the chosen
+            selection.
+
+        Returns
+        -------
+        layered : Model
+            The layered Model; a :class:`emg3d.models.Model` instance.
+
+        imat : ndarray, returned if return_imat=True
+            Interpolation matrix of shape (nx, ny) of the original model.
+
+        """
+        # Get ellipse dict.
+        ellipse = {} if ellipse is None else ellipse
+
+        # Check if input is consistent with method.
+        methods = ['midpoint', 'cylinder', 'prism']
+        if method not in methods:
+            raise ValueError(
+                f"Unknown method '{method}'; implemented: {methods}."
+            )
+        if method != 'midpoint' and 'radius' not in ellipse.keys():
+            raise TypeError(
+                f"Method '{method}' requires the dict 'ellipse' "
+                "containing at least the parameter 'radius'."
+            )
+
+        # Flag if midpoint is used or another method.
+        midpoint = method == 'midpoint'
+
+        # If p1 is not given, we set it to p0.
+        if p1 is None:
+            p1 = p0
+
+        # Get relevant indices and the used-cells matrix.
+        if not midpoint:
+
+            # Indices of used cells.
+            coo = (self.grid.cell_centers_x, self.grid.cell_centers_y)
+            use = maps.ellipse_indices(coo=coo, p0=p0, p1=p1, **ellipse)
+
+            # Start and end indices.
+            ix, iy = use.nonzero()
+            if not ix.size:  # Switch to 'midpoint' if selection is empty.
+                midpoint = True
+            else:
+                six, eix = ix.min(), ix.max()
+                siy, eiy = iy.min(), iy.max()
+
+        if midpoint:
+
+            # Find corresponding indices, limit to grid.
+            def index(nodes, coo):
+                """Return index for interval btw nodes in which coo resides."""
+                x = np.asarray(coo < np.r_[nodes, np.infty]).nonzero()[0][0]-1
+                return np.clip(x, 0, nodes.size-2)
+
+            # Start and end indices.
+            six = eix = index(self.grid.nodes_x, (p0[0] + p1[0])/2)
+            siy = eiy = index(self.grid.nodes_y, (p0[1] + p1[1])/2)
+
+        # Create interpolation matrix.
+        imat = np.zeros(self.shape[:2])
+        if not midpoint:
+            pp = np.outer(self.grid.h[0][six:eix+1], self.grid.h[1][siy:eiy+1])
+            if method == "cylinder":
+                pp *= use[six:eix+1, siy:eiy+1]
+            pp /= pp.sum()
+        else:
+            pp = 1.0
+        imat[six:eix+1, siy:eiy+1] = pp
+
+        # Interpolate property_{x;y;z}; mu_r; and epsilon_r.
+        props = {}
+        for prop in self._def_properties:
+            values = getattr(self, prop)
+
+            if not midpoint:
+                if not self.map.name.startswith('L'):
+                    values = np.log10(values)
+                val = np.einsum('ij,ijk->k', imat, values)
+                if not self.map.name.startswith('L'):
+                    val = 10**val
+            else:
+                val = values[six, siy, :]
+
+            props[prop] = val
+
+        # If merge, combine adjacent layers of identical properties.
+        if merge:
+
+            # Get indices of non-merge sequences.
+            diff = np.zeros(self.shape[2])
+            for k, v in props.items():
+                diff += abs(np.diff(np.r_[-1, v]))
+            ind = diff.nonzero()[0]
+
+            # Merge.
+            props = {k: v[ind] for k, v in props.items()}
+
+            # Get cell sizes in z.
+            hz = np.diff(np.r_[self.grid.nodes_z[ind], self.grid.nodes_z[-1]])
+
+        else:
+            hz = self.grid.h[2]
+
+        # Create (1, 1, nz) grid.
+        grid_out = meshes.TensorMesh(
+            h=(
+                [self.grid.nodes_x[eix+1] - self.grid.nodes_x[six], ],
+                [self.grid.nodes_y[eiy+1] - self.grid.nodes_y[siy], ],
+                hz
+            ),
+            origin=(
+                self.grid.nodes_x[six],
+                self.grid.nodes_y[siy],
+                self.grid.origin[2]
+            ),
+        )
+
+        # Get layered model on grid_out.
+        layered = Model(grid=grid_out, **props, mapping=self.map)
+
+        # Return "1D" model (only 1 cell in x/y).
+        if return_imat:
+            return layered, imat
+        else:
+            return layered
+
     # INTERNAL UTILITIES
     def _init_parameter(self, values, name):
         """Initiate parameter by casting and broadcasting."""

emg3d/surveys.py

@@ -66,7 +66,7 @@ class Survey:
 
     Parameters
     ----------
-    sources, receivers : {Tx*, Rx*, list, dict)
+    sources, receivers : {Tx*, Rx*, list, dict}
         Any of the available sources or receivers, e.g.,
         :class:`emg3d.electrodes.TxElectricDipole`, or a list or dict of
         Tx*/Rx* instances. If it is a dict, it is used as is, including the

tests/test_maps.py

@@ -709,7 +709,7 @@ def test_ellipse_indices():
     # Circle
     p0, p1 = np.array([0, 0]), np.array([0, 0])
     x = np.arange(5)-2
-    X, Y = np.meshgrid(x, x)
+    X, Y = np.meshgrid(x, x, indexing='ij')
     out = maps.ellipse_indices((X, Y), p0, p1, radius=2)
     res = np.array([[False, False,  True, False, False],
                     [False,  True,  True,  True, False],
@@ -720,10 +720,9 @@ def test_ellipse_indices():
     # plt.pcolormesh(out); plt.axis('equal')
 
     # Ellipse
-    p0, p1 = np.array([-1, 0]), np.array([2, 2])
+    p0, p1 = [-1, 0], (2, 2)
     x = np.arange(9)-4
-    X, Y = np.meshgrid(x, x)
-    out = maps.ellipse_indices((X, Y), p0, p1, radius=1, minor=0.5)
+    out = maps.ellipse_indices((x, x), p0, p1, radius=1, minor=0.5)
     res = np.array([
         [False, False, False, False, False, False, False, False, False],
         [False, False, False, False, False, False, False, False, False],
@@ -734,7 +733,7 @@ def test_ellipse_indices():
         [False, False, False,  True,  True,  True,  True,  True, False],
         [False, False, False, False,  True,  True,  True, False, False],
         [False, False, False, False, False, False, False, False, False]
-    ])
+    ]).T
     assert_allclose(out, res)
     # plt.pcolormesh(out); plt.axis('equal')
 

tests/test_models.py

@@ -417,6 +417,90 @@ def test_dict(self):
             models.Model.from_dict(mdict)
 
 
+class TestExtract1D:
+
+    grid = emg3d.TensorMesh(
+            [[1, 2, 3, 4], [3, 2, 2, 3, 4], [1, 1, 1, 1, 1]], (0, 0, 0))
+    num = np.arange(1, 21).reshape((4, 5), order='F')
+    property_x = np.zeros((4, 5, 5), order='F')
+    property_x[:, :, 0] = num+100
+    property_x[:, :, 1] = num+100
+    property_x[:, :, 2] = num
+    property_x[:, :, 3] = num+200
+    property_x[:, :, 4] = num+200
+    property_y = property_x/2.0
+    property_z = property_x*1.4
+    mu_r = property_x*1.11
+    epsilon_r = property_x*2.22
+
+    mlin = models.Model(
+        grid, property_x, property_y, property_z, mu_r, epsilon_r
+    )
+    mlog = models.Model(
+        grid, np.log10(property_x), np.log10(property_y), np.log10(property_z),
+        np.log10(mu_r), np.log10(epsilon_r), mapping='LgResistivity'
+    )
+
+    def test_errors(self):
+        # Unknown method
+        with pytest.raises(ValueError, match="Unknown method 'unknown'"):
+            self.mlin.extract_1d(method='unknown', p0=[0, 0])
+
+        # Missing ellipse
+        with pytest.raises(TypeError, match="requires the dict 'ellipse'"):
+            self.mlin.extract_1d(method='prism', p0=[0, 0])
+
+        # Ellipse but missing radius
+        with pytest.raises(TypeError, match="the parameter 'radius'"):
+            self.mlin.extract_1d(
+                method='cylinder', p0=[0, 0], ellipse={'factor': 1.2}
+            )
+
+    def test_midpoint(self):
+
+        # One point, outside
+        layered, imat = self.mlin.extract_1d(
+                'midpoint', [-3, -3], return_imat=True)
+        assert_allclose(layered.property_x[0, 0, :],
+                        [101.0, 101.0, 1.0, 201.0, 201.0])
+        assert_allclose(layered.property_y, layered.property_x/2.0)
+        assert_allclose(layered.property_z, layered.property_x*1.4)
+        assert_allclose(layered.mu_r, layered.property_x*1.11)
+        assert_allclose(layered.epsilon_r, layered.property_x*2.22)
+        assert imat[0, 0] == 1.0
+        assert imat.sum() == 1.0
+
+        # Two points, merge
+        layered = self.mlog.extract_1d('midpoint', [2, 4], [7, 8], merge=True)
+        assert_allclose(10**layered.property_x[0, 0, :], [111.0, 11.0, 211.0])
+
+    def test_averaging(self):
+
+        # Empty selection => same as midpoint
+        mpt = self.mlin.extract_1d('midpoint', [2.1, 7.1])
+        cyl = self.mlin.extract_1d(
+                'cylinder', [2.1, 7.1], ellipse={'radius': 0.001})
+        assert mpt == cyl
+
+        cylin, cmat = self.mlin.extract_1d(
+                'cylinder', [2, 4], [7, 8], ellipse={'radius': 1.5},
+                return_imat=True)
+        cylog, _ = self.mlog.extract_1d(
+                'cylinder', [2, 4], [7, 8], ellipse={'radius': 1.5},
+                return_imat=True)
+
+        vals = [110.4, 110.4, 9.4, 210.5, 210.5]
+        assert_allclose(cylin.property_x[0, 0, :], vals, atol=0.1)
+        assert_allclose(10**cylog.property_x[0, 0, :], vals, atol=0.1)
+
+        cylin, pmat = self.mlin.extract_1d(
+                'prism', [2, 4], [7, 8], ellipse={'radius': 1.5},
+                return_imat=True)
+
+        assert_allclose(pmat[:, :-1] > 0, 1)
+        assert_allclose(pmat[:, -1], 0.0)
+
+
 def test_expand_grid_model():
     grid = emg3d.TensorMesh([[4, 2, 2, 4], [2, 2, 2, 2], [1, 1]], (0, 0, 0))
     model = emg3d.Model(grid, 1, np.ones(grid.shape_cells)*2, mu_r=3,