diff --git a/doc/api/index.rst b/doc/api/index.rst index 0bef91208..cf87e025b 100644 --- a/doc/api/index.rst +++ b/doc/api/index.rst @@ -81,6 +81,8 @@ Magnetic fields: dipole_magnetic dipole_magnetic_component + prism_magnetic + prism_magnetic_component Layers and meshes: diff --git a/doc/user_guide/forward_modelling/prism.rst b/doc/user_guide/forward_modelling/prism.rst index 9093b146c..ff464b792 100644 --- a/doc/user_guide/forward_modelling/prism.rst +++ b/doc/user_guide/forward_modelling/prism.rst @@ -49,7 +49,7 @@ computation point: Gravitational fields -^^^^^^^^^^^^^^^^^^^^ +-------------------- The :func:`harmonica.prism_gravity` is able to compute the gravitational potential (``"potential"``), the acceleration components (``"g_e"``, ``"g_n"``, @@ -139,8 +139,9 @@ Plot the results: plt.colorbar(tmp, ax=ax, label="Eotvos", orientation="horizontal", pad=0.08) plt.show() + Passing multiple prisms -^^^^^^^^^^^^^^^^^^^^^^^ +~~~~~~~~~~~~~~~~~~~~~~~ We can compute the gravitational field of a set of prisms by passing a list of them, where each prism is defined as mentioned before, and then making @@ -200,10 +201,123 @@ Lets plot this gravitational field: fig.show() +Magnetic fields +--------------- + +The :func:`harmonica.prism_magnetic` and +:func:`harmonica.prism_magnetic_component` functions allows us to calculate the +magnetic fields generated by rectangular prisms on a set of observation points. +Each rectangular prism is defined in the same way we did for the gravity +forward modelling, although we now need to define a magnetization vector for +each one of them. + +For example: + +.. jupyter-execute:: + + import numpy as np + + prisms = [ + [-5e3, -3e3, -5e3, -2e3, -10e3, -1e3], + [3e3, 4e3, 4e3, 5e3, -9e3, -1e3], + ] + + magnetization = np.array([ + [0.5, 0.5, -0.78989], + [-0.4, 0.3, 0.2], + ]) + +Each row of the ``magnetization`` 2D vector corresponds to the easting, +northing and upward components of the magnetization vector of each prism in +``prisms``, provided in :math:`Am^{-1}`. + +With the :func:`harmonica.prism_magnetic` function we can compute the three +components of the magnetic field the prisms generates on any set of observation +points. + +.. jupyter-execute:: + + region = (-10e3, 10e3, -10e3, 10e3) + shape = (51, 51) + height = 10 + coordinates = vd.grid_coordinates(region, shape=shape, extra_coords=height) + +.. jupyter-execute:: + + b_e, b_n, b_u = hm.prism_magnetic(coordinates, prisms, magnetization) + +.. jupyter-execute:: + + fig, axes = plt.subplots(nrows=1, ncols=3, sharey=True, figsize=(12, 6)) + + for ax, mag_component, title in zip(axes, (b_e, b_n, b_u), ("Be", "Bn", "Bu")): + maxabs = vd.maxabs(mag_component) + tmp = ax.pcolormesh( + coordinates[0], + coordinates[1], + mag_component, + vmin=-maxabs, + vmax=maxabs, + cmap="RdBu_r", + ) + ax.contour( + coordinates[0], + coordinates[1], + mag_component, + colors="k", + linewidths=0.5, + ) + ax.set_title(title) + ax.set_aspect("equal") + plt.colorbar( + tmp, + ax=ax, + orientation="horizontal", + label="nT", + pad=0.08, + aspect=42, + shrink=0.8, + ) + + plt.show() + + +Alternatively, we can use :func:`harmonica.prism_magnetic_component` to +calculate only a single component of the magnetic field. + +.. important:: + + Using :func:`harmonica.prism_magnetic_component` to compute several magnetic + components is less efficient that using :func:`harmonica.prism_magnetic`. + Use the former only when a single component is needed. + +For example, we can calculate only the upward component of the magnetic field +generated by these two prisms: + +.. jupyter-execute:: + + b_u = hm.prism_magnetic_component( + coordinates, prisms, magnetization, component="upward" + ) + +.. jupyter-execute:: + + maxabs = vd.maxabs(b_u) + + tmp = plt.pcolormesh( + coordinates[0], coordinates[1], b_u, vmin=-maxabs, vmax=maxabs, cmap="RdBu_r" + ) + plt.contour(coordinates[0], coordinates[1], b_u, colors="k", linewidths=0.5) + plt.title("Bu") + plt.gca().set_aspect("equal") + plt.colorbar(tmp, label="nT", pad=0.03, aspect=42, shrink=0.8) + plt.show() + + .. _prism_layer: Prism layer -^^^^^^^^^^^ +----------- One of the most common usage of prisms is to model geologic structures. Harmonica offers the possibility to define a layer of prisms through the @@ -241,8 +355,6 @@ sample a trigonometric function for this simple example: .. jupyter-execute:: - import numpy as np - wavelength = 24 * spacing surface = np.abs(np.sin(easting * 2 * np.pi / wavelength)) diff --git a/harmonica/__init__.py b/harmonica/__init__.py index 33dc1b2ef..62f6901be 100644 --- a/harmonica/__init__.py +++ b/harmonica/__init__.py @@ -11,8 +11,9 @@ from ._equivalent_sources.spherical import EquivalentSourcesSph from ._forward.dipole import dipole_magnetic, dipole_magnetic_component from ._forward.point import point_gravity -from ._forward.prism import prism_gravity +from ._forward.prism_gravity import prism_gravity from ._forward.prism_layer import DatasetAccessorPrismLayer, prism_layer +from ._forward.prism_magnetic import prism_magnetic, prism_magnetic_component from ._forward.tesseroid import tesseroid_gravity from ._forward.tesseroid_layer import DatasetAccessorTesseroidLayer, tesseroid_layer from ._gravity_corrections import bouguer_correction diff --git a/harmonica/_forward/_tesseroid_variable_density.py b/harmonica/_forward/_tesseroid_variable_density.py index d2be1c8fe..f5ce59e06 100644 --- a/harmonica/_forward/_tesseroid_variable_density.py +++ b/harmonica/_forward/_tesseroid_variable_density.py @@ -165,6 +165,7 @@ def _density_based_discretization(tesseroid, density): tesseroids : list List containing the boundaries of discretized tesseroids. """ + # Define normalized density def normalized_density(radius): return (density(radius) - density_min) / (density_max - density_min) diff --git a/harmonica/_forward/prism.py b/harmonica/_forward/prism_gravity.py similarity index 100% rename from harmonica/_forward/prism.py rename to harmonica/_forward/prism_gravity.py diff --git a/harmonica/_forward/prism_layer.py b/harmonica/_forward/prism_layer.py index 52461277f..c0f65f777 100644 --- a/harmonica/_forward/prism_layer.py +++ b/harmonica/_forward/prism_layer.py @@ -14,7 +14,7 @@ import xarray as xr from ..visualization import prism_to_pyvista -from .prism import prism_gravity +from .prism_gravity import prism_gravity def prism_layer( diff --git a/harmonica/_forward/prism_magnetic.py b/harmonica/_forward/prism_magnetic.py new file mode 100644 index 000000000..4fa447e30 --- /dev/null +++ b/harmonica/_forward/prism_magnetic.py @@ -0,0 +1,410 @@ +# Copyright (c) 2018 The Harmonica Developers. +# Distributed under the terms of the BSD 3-Clause License. +# SPDX-License-Identifier: BSD-3-Clause +# +# This code is part of the Fatiando a Terra project (https://www.fatiando.org) +# +""" +Compute magnetic field generated by rectangular prisms +""" +import numpy as np +from choclo.prism import magnetic_e, magnetic_field, magnetic_n, magnetic_u +from numba import jit, prange + +from .prism_gravity import _check_prisms +from .utils import initialize_progressbar + + +def prism_magnetic( + coordinates, + prisms, + magnetization, + parallel=True, + dtype=np.float64, + progressbar=False, + disable_checks=False, +): + """ + Magnetic field of right-rectangular prisms in Cartesian coordinates + + Parameters + ---------- + coordinates : list of arrays + List of arrays containing the ``easting``, ``northing`` and ``upward`` + coordinates of the computation points defined on a Cartesian coordinate + system. All coordinates should be in meters. + prisms : list, 1d-array, or 2d-array + List or array containing the coordinates of the prism(s) in the + following order: + west, east, south, north, bottom, top in a Cartesian coordinate system. + All coordinates should be in meters. Coordinates for more than one + prism can be provided. In this case, *prisms* should be a list of lists + or 2d-array (with one prism per row). + magnetization : list or array + List or array containing the magnetization vector of each prism in + :math:`Am^{-1}`. Each vector should be an array with three elements + in the following order: ``magnetization_e``, ``magnetization_n``, + ``magnetization_u``. + parallel : bool (optional) + If True the computations will run in parallel using Numba built-in + parallelization. If False, the forward model will run on a single core. + Might be useful to disable parallelization if the forward model is run + by an already parallelized workflow. Default to True. + dtype : data-type (optional) + Data type assigned to the resulting gravitational field. Default to + ``np.float64``. + progressbar : bool (optional) + If True, a progress bar of the computation will be printed to standard + error (stderr). Requires :mod:`numba_progress` to be installed. + Default to ``False``. + disable_checks : bool (optional) + Flag that controls whether to perform a sanity check on the model. + Should be set to ``True`` only when it is certain that the input model + is valid and it does not need to be checked. + Default to ``False``. + + Returns + ------- + magnetic_field : tuple of array + Tuple containing each component of the magnetic field generated by the + prisms as arrays. The three components are returned in the following + order: ``b_e``, ``b_n``, ``b_u``. + """ + # Figure out the shape and size of the output array(s) + cast = np.broadcast(*coordinates[:3]) + # Convert coordinates, prisms and magnetization to arrays with proper shape + coordinates = tuple(np.atleast_1d(i).ravel() for i in coordinates[:3]) + prisms = np.atleast_2d(prisms) + magnetization = np.atleast_2d(magnetization) + # Sanity checks + if not disable_checks: + _run_sanity_checks(prisms, magnetization) + # Discard null prisms (zero volume or null magnetization) + prisms, magnetization = _discard_null_prisms(prisms, magnetization) + # Run computations + b_e, b_n, b_u = tuple(np.zeros(cast.size, dtype=dtype) for _ in range(3)) + with initialize_progressbar(coordinates[0].size, progressbar) as progress_proxy: + if parallel: + _jit_prism_magnetic_field_parallel( + coordinates, prisms, magnetization, b_e, b_n, b_u, progress_proxy + ) + else: + _jit_prism_magnetic_field_serial( + coordinates, prisms, magnetization, b_e, b_n, b_u, progress_proxy + ) + # Convert to nT + b_e *= 1e9 + b_n *= 1e9 + b_u *= 1e9 + return b_e.reshape(cast.shape), b_n.reshape(cast.shape), b_u.reshape(cast.shape) + + +def prism_magnetic_component( + coordinates, + prisms, + magnetization, + component, + parallel=True, + dtype=np.float64, + progressbar=False, + disable_checks=False, +): + """ + Compute single component of the magnetic field of right-rectangular prisms + + .. important:: + + Use this function only if you need to compute a single component of the + magnetic field. Use :func:`harmonica.prism_magnetic` to compute the + three components more efficiently. + + Parameters + ---------- + coordinates : list or 1d-array + List or array containing ``easting``, ``northing`` and ``upward`` of + the computation points defined on a Cartesian coordinate system. + All coordinates should be in meters. + prisms : list, 1d-array, or 2d-array + List or array containing the coordinates of the prism(s) in the + following order: + west, east, south, north, bottom, top in a Cartesian coordinate system. + All coordinates should be in meters. Coordinates for more than one + prism can be provided. In this case, *prisms* should be a list of lists + or 2d-array (with one prism per line). + magnetization : list or array + List or array containing the magnetization vector of each prism in + :math:`Am^{-1}`. Each vector should be an array with three elements + in the following order: ``magnetization_e``, ``magnetization_n``, + ``magnetization_u``. + component : str + Computed that will be computed. Available options are: ``"easting"``, + ``"northing"`` or ``"upward"``. + parallel : bool (optional) + If True the computations will run in parallel using Numba built-in + parallelization. If False, the forward model will run on a single core. + Might be useful to disable parallelization if the forward model is run + by an already parallelized workflow. Default to True. + dtype : data-type (optional) + Data type assigned to the resulting gravitational field. Default to + ``np.float64``. + progressbar : bool (optional) + If True, a progress bar of the computation will be printed to standard + error (stderr). Requires :mod:`numba_progress` to be installed. + Default to ``False``. + disable_checks : bool (optional) + Flag that controls whether to perform a sanity check on the model. + Should be set to ``True`` only when it is certain that the input model + is valid and it does not need to be checked. + Default to ``False``. + + Returns + ------- + b_component : array + Array with the component of the magnetic field generated by the + prisms on every observation point. + """ + # Figure out the shape and size of the output array(s) + cast = np.broadcast(*coordinates[:3]) + # Convert coordinates, prisms and magnetization to arrays with proper shape + coordinates = tuple(np.atleast_1d(i).ravel() for i in coordinates[:3]) + prisms = np.atleast_2d(prisms) + magnetization = np.atleast_2d(magnetization) + # Choose forward modelling function based on the chosen component + forward_function = _get_magnetic_forward_function(component) + # Sanity checks + if not disable_checks: + _run_sanity_checks(prisms, magnetization) + # Discard null prisms (zero volume or null magnetization) + prisms, magnetization = _discard_null_prisms(prisms, magnetization) + # Run computations + result = np.zeros(cast.size, dtype=dtype) + with initialize_progressbar(coordinates[0].size, progressbar) as progress_proxy: + if parallel: + _jit_prism_magnetic_component_parallel( + coordinates, + prisms, + magnetization, + result, + forward_function, + progress_proxy, + ) + else: + _jit_prism_magnetic_component_serial( + coordinates, + prisms, + magnetization, + result, + forward_function, + progress_proxy, + ) + # Convert to nT + result *= 1e9 + return result.reshape(cast.shape) + + +def _jit_prism_magnetic_field( + coordinates, prisms, magnetization, b_e, b_n, b_u, progress_proxy=None +): + """ + Compute magnetic fields of prisms on computation points + + Parameters + ---------- + coordinates : tuple + Tuple containing ``easting``, ``northing`` and ``upward`` of the + computation points as arrays, all defined on a Cartesian coordinate + system and in meters. + prisms : 2d-array + Two dimensional array containing the coordinates of the prism(s) in the + following order: west, east, south, north, bottom, top in a Cartesian + coordinate system. + All coordinates should be in meters. + magnetization : 2d-array + Array containing the magnetization vector of each prism in + :math:`Am^{-1}`. Each vector will be a row in the 2d-array. + b_e : 1d-array + Array where the resulting values of the easting component of the + magnetic field will be stored. + b_n : 1d-array + Array where the resulting values of the northing component of the + magnetic field will be stored. + b_u : 1d-array + Array where the resulting values of the upward component of the + magnetic field will be stored. + progress_proxy : :class:`numba_progress.ProgressBar` or None + Instance of :class:`numba_progress.ProgressBar` that gets updated after + each iteration on the observation points. Use None if no progress bar + is should be used. + """ + # Check if we need to update the progressbar on each iteration + update_progressbar = progress_proxy is not None + # Iterate over computation points and prisms + for l in prange(coordinates[0].size): + for m in range(prisms.shape[0]): + easting_comp, northing_comp, upward_comp = magnetic_field( + coordinates[0][l], + coordinates[1][l], + coordinates[2][l], + prisms[m, 0], + prisms[m, 1], + prisms[m, 2], + prisms[m, 3], + prisms[m, 4], + prisms[m, 5], + magnetization[m, 0], + magnetization[m, 1], + magnetization[m, 2], + ) + b_e[l] += easting_comp + b_n[l] += northing_comp + b_u[l] += upward_comp + # Update progress bar if called + if update_progressbar: + progress_proxy.update(1) + + +def _jit_prism_magnetic_component( + coordinates, prisms, magnetization, result, forward_function, progress_proxy=None +): + """ + Compute a single component of the magnetic field of prisms + + Parameters + ---------- + coordinates : tuple + Tuple containing ``easting``, ``northing`` and ``upward`` of the + computation points as arrays, all defined on a Cartesian coordinate + system and in meters. + prisms : 2d-array + Two dimensional array containing the coordinates of the prism(s) in the + following order: west, east, south, north, bottom, top in a Cartesian + coordinate system. + All coordinates should be in meters. + magnetization : 2d-array + Array containing the magnetization vector of each prism in + :math:`Am^{-1}`. Each vector will be a row in the 2d-array. + result : 1d-array + Array where the resulting values of the desired component of the + magnetic field will be stored. + forward_function : callable + Forward function to be used to compute the desired component of the + magnetic field. Choose one of :func:`choclo.prism.magnetic_easting`, + :func:`choclo.prism.magnetic_northing` or + :func:`choclo.prism.magnetic_upward`. + progress_proxy : :class:`numba_progress.ProgressBar` or None + Instance of :class:`numba_progress.ProgressBar` that gets updated after + each iteration on the observation points. Use None if no progress bar + is should be used. + """ + # Check if we need to update the progressbar on each iteration + update_progressbar = progress_proxy is not None + # Iterate over computation points and prisms + for l in prange(coordinates[0].size): + for m in range(prisms.shape[0]): + result[l] += forward_function( + coordinates[0][l], + coordinates[1][l], + coordinates[2][l], + prisms[m, 0], + prisms[m, 1], + prisms[m, 2], + prisms[m, 3], + prisms[m, 4], + prisms[m, 5], + magnetization[m, 0], + magnetization[m, 1], + magnetization[m, 2], + ) + # Update progress bar if called + if update_progressbar: + progress_proxy.update(1) + + +def _discard_null_prisms(prisms, magnetization): + """ + Discard prisms with zero volume or null magnetization + + Parameters + ---------- + prisms : 2d-array + Array containing the boundaries of the prisms in the following order: + ``w``, ``e``, ``s``, ``n``, ``bottom``, ``top``. + The array must have the following shape: (``n_prisms``, 6), where + ``n_prisms`` is the total number of prisms. + This array of prisms must have valid boundaries. + Run ``_check_prisms`` before. + magnetization : 2d-array + Array containing the magnetization vector of each prism in + :math:`Am^{-1}`. Each vector will be a row in the 2d-array. + + Returns + ------- + prisms : 2d-array + A copy of the ``prisms`` array that doesn't include the null prisms + (prisms with zero volume or zero density). + magnetization : 2d-array + A copy of the ``magnetization`` array that doesn't include the + magnetization vectors for null prisms (prisms with zero volume or + null magnetization). + """ + west, east, south, north, bottom, top = tuple(prisms[:, i] for i in range(6)) + # Mark prisms with zero volume as null prisms + null_prisms = (west == east) | (south == north) | (bottom == top) + # Mark prisms with null magnetization as null prisms + null_prisms[(magnetization == 0).all(axis=1)] = True + # Keep only non null prisms + prisms = prisms[np.logical_not(null_prisms), :] + magnetization = magnetization[np.logical_not(null_prisms)] + return prisms, magnetization + + +def _run_sanity_checks(prisms, magnetization): + """ + Run sanity checks on prisms and their magnetization + """ + if magnetization.shape[0] != prisms.shape[0]: + raise ValueError( + f"Number of magnetization vectors ({magnetization.shape[0]}) " + + f"mismatch the number of prisms ({prisms.shape[0]})" + ) + if magnetization.shape[1] != 3: + raise ValueError( + f"Found magnetization vectors with '{magnetization.shape[1]}' " + + "elements. Magnetization vectors should have only 3 elements." + ) + _check_prisms(prisms) + + +def _get_magnetic_forward_function(component): + """ + Returns the Choclo magnetic forward modelling function for the desired + component + + Parameters + ---------- + component : str + Magnetic field component. + + Returns + ------- + forward_function : callable + Forward modelling function for the desired component. + """ + if component not in ("easting", "northing", "upward"): + raise ValueError( + f"Invalid component '{component}'. " + "It must be either 'easting', 'northing' or 'upward'." + ) + functions = {"easting": magnetic_e, "northing": magnetic_n, "upward": magnetic_u} + return functions[component] + + +# Define jitted versions of the forward modelling function +_jit_prism_magnetic_field_serial = jit(nopython=True)(_jit_prism_magnetic_field) +_jit_prism_magnetic_field_parallel = jit(nopython=True, parallel=True)( + _jit_prism_magnetic_field +) +_jit_prism_magnetic_component_serial = jit(nopython=True)(_jit_prism_magnetic_component) +_jit_prism_magnetic_component_parallel = jit(nopython=True, parallel=True)( + _jit_prism_magnetic_component +) diff --git a/harmonica/tests/test_point_gravity.py b/harmonica/tests/test_point_gravity.py index 88a31cb25..560b61d83 100644 --- a/harmonica/tests/test_point_gravity.py +++ b/harmonica/tests/test_point_gravity.py @@ -359,7 +359,6 @@ def test_tensor_non_diagonal_components(field, flipped_field): class TestTensorSymmetryCartesian: - # Define sample point source and its mass point = [1.1, 1.2, 1.3] mass = [2670] diff --git a/harmonica/tests/test_prism.py b/harmonica/tests/test_prism.py index 0500d8433..a8d946620 100644 --- a/harmonica/tests/test_prism.py +++ b/harmonica/tests/test_prism.py @@ -32,7 +32,7 @@ ProgressBar = None from .. import bouguer_correction -from .._forward.prism import _check_prisms, _discard_null_prisms, prism_gravity +from .._forward.prism_gravity import _check_prisms, _discard_null_prisms, prism_gravity from .utils import run_only_with_numba diff --git a/harmonica/tests/test_prism_magnetic.py b/harmonica/tests/test_prism_magnetic.py new file mode 100644 index 000000000..101511e5d --- /dev/null +++ b/harmonica/tests/test_prism_magnetic.py @@ -0,0 +1,388 @@ +# Copyright (c) 2018 The Harmonica Developers. +# Distributed under the terms of the BSD 3-Clause License. +# SPDX-License-Identifier: BSD-3-Clause +# +# This code is part of the Fatiando a Terra project (https://www.fatiando.org) +# +""" +Test forward functions for magnetic field of prisms +""" +import choclo +import numpy as np +import numpy.testing as npt +import pytest +import verde as vd +from choclo.prism import magnetic_field + +try: + from numba_progress import ProgressBar +except ImportError: + ProgressBar = None + +from .. import prism_magnetic, prism_magnetic_component +from .utils import run_only_with_numba + + +def test_invalid_component(): + "Check if passing an invalid component raises an error" + prism = [-100, 100, -100, 100, -200, -100] + magnetization = [1000, 1, 2] + coordinates = [0, 0, 0] + with pytest.raises(ValueError, match="Invalid component"): + prism_magnetic_component( + coordinates, prism, magnetization, component="Not a valid field" + ) + + +@pytest.mark.use_numba +@pytest.mark.skipif(ProgressBar is None, reason="requires numba_progress") +@pytest.mark.parametrize("component", (None, "easting", "northing", "upward")) +def test_progress_bar(component): + """ + Check if forward modelling results with and without progress bar match + """ + prisms = [ + [-100, 0, -100, 0, -10, 0], + [0, 100, -100, 0, -10, 0], + ] + magnetizations = [ + [1.0, 1.0, 1.0], + [1.0, -1.0, 5.0], + ] + coordinates = vd.grid_coordinates( + region=(-100, 100, -100, 100), spacing=20, extra_coords=10 + ) + if component is None: + result_progress_true = prism_magnetic( + coordinates, prisms, magnetizations, progressbar=True + ) + result_progress_false = prism_magnetic( + coordinates, prisms, magnetizations, progressbar=False + ) + else: + result_progress_true = prism_magnetic_component( + coordinates, prisms, magnetizations, component, progressbar=True + ) + result_progress_false = prism_magnetic_component( + coordinates, prisms, magnetizations, component, progressbar=False + ) + npt.assert_allclose(result_progress_true, result_progress_false) + + +class TestSerialVsParallel: + """ + Test serial vs parallel + """ + + @pytest.mark.parametrize("component", (None, "easting", "northing", "upward")) + def test_prisms_parallel_vs_serial_no_numba(self, component): + """ + Check results of parallelized and serials runs + + Run a small problem with Numba disable to count for test coverage. + """ + prisms = [ + [-100, 0, -100, 0, -10, 0], + [0, 100, -100, 0, -10, 0], + ] + magnetizations = [ + [1.0, 1.0, 1.0], + [1.0, -1.0, 5.0], + ] + coordinates = ([0, 10], [0, 10], [0, 10]) + if component is None: + parallel = prism_magnetic( + coordinates, prisms, magnetizations, parallel=True + ) + serial = prism_magnetic(coordinates, prisms, magnetizations, parallel=False) + else: + parallel = prism_magnetic_component( + coordinates, prisms, magnetizations, component, parallel=True + ) + serial = prism_magnetic_component( + coordinates, prisms, magnetizations, component, parallel=False + ) + npt.assert_allclose(parallel, serial) + + @run_only_with_numba + @pytest.mark.parametrize("component", (None, "easting", "northing", "upward")) + def test_prisms_parallel_vs_serial(self, component): + """ + Check results of parallelized and serials runs + + Run a large problem only with Numba enabled. + """ + prisms = [ + [-100, 0, -100, 0, -10, 0], + [0, 100, -100, 0, -10, 0], + [-100, 0, 0, 100, -10, 0], + [0, 100, 0, 100, -10, 0], + ] + magnetizations = [ + [1.0, 1.0, 1.0], + [1.0, -1.0, 5.0], + [-2.0, 1.0, 3.0], + [5.0, 4.0, 1.0], + ] + coordinates = vd.grid_coordinates( + region=(-100, 100, -100, 100), spacing=20, extra_coords=10 + ) + if component is None: + parallel = prism_magnetic( + coordinates, prisms, magnetizations, parallel=True + ) + serial = prism_magnetic(coordinates, prisms, magnetizations, parallel=False) + else: + parallel = prism_magnetic_component( + coordinates, prisms, magnetizations, component, parallel=True + ) + serial = prism_magnetic_component( + coordinates, prisms, magnetizations, component, parallel=False + ) + npt.assert_allclose(parallel, serial) + + +class TestInvalidPrisms: + """ + Test forward modelling functions against invalid prisms + """ + + boundaries = [ + [100, -100, -100, 100, -200, -100], # w > e + [-100, 100, 100, -100, -200, -100], # s > n + [-100, 100, -100, -00, -100, -200], # bottom > top + ] + + @pytest.fixture() + def sample_coordinates(self): + """Return sample coordinates""" + return [0, 0, 0] + + @pytest.fixture() + def sample_magnetization(self): + """Return sample magnetization""" + return [1.0, 1.0, 1.0] + + @pytest.fixture(params=boundaries) + def invalid_prism(self, request): + """Return sample valid prism (with zero and non-zero volume)""" + return np.atleast_2d(request.param) + + @pytest.mark.parametrize("component", (None, "easting", "northing", "upward")) + def test_invalid_prisms( + self, sample_coordinates, invalid_prism, sample_magnetization, component + ): + """ + Test forward modelling functions with invalid prisms + + It should raise any error + """ + msg = "boundary can't be greater than the" + if component is None: + with pytest.raises(ValueError, match=msg): + prism_magnetic(sample_coordinates, invalid_prism, sample_magnetization) + else: + with pytest.raises(ValueError, match=msg): + prism_magnetic_component( + sample_coordinates, invalid_prism, sample_magnetization, component + ) + + @pytest.mark.parametrize("component", (None, "easting", "northing", "upward")) + def test_disable_checks(self, invalid_prism, component): + """Test if disabling checks doesn't raise errors on invalid prisms""" + magnetization = [100, 10, -10] + coordinates = [0, 0, 0] + if component is None: + prism_magnetic( + coordinates, invalid_prism, magnetization, disable_checks=True + ) + else: + prism_magnetic_component( + coordinates, + invalid_prism, + magnetization, + component, + disable_checks=True, + ) + + +class TestInvalidMagnetization: + """ + Test errors after invalid magnetization arrays + """ + + @pytest.fixture() + def sample_coordinates(self): + """Return sample coordinates""" + return [0, 0, 0] + + @pytest.fixture() + def sample_prisms(self): + """Return sample prisms""" + prisms = [ + [-100, 0, -100, 0, -200, -100], + [-100, 0, 0, 100, -200, -100], + [0, 100, -100, 0, -200, -100], + [0, 100, 0, 100, -200, -100], + ] + return prisms + + @pytest.mark.parametrize("component", (None, "easting", "northing", "upward")) + def test_invalid_number_of_vectors( + self, sample_coordinates, sample_prisms, component + ): + """Check error when magnetization has invalid number of vectors""" + # Generate an array with only two magnetization vectors + magnetizations = [ + [1.0, 1.0, 1.0], + [-2.0, 3.0, -5.0], + ] + msg = "Number of magnetization vectors" + if component is None: + with pytest.raises(ValueError, match=msg): + prism_magnetic(sample_coordinates, sample_prisms, magnetizations) + else: + with pytest.raises(ValueError, match=msg): + prism_magnetic_component( + sample_coordinates, sample_prisms, magnetizations, component + ) + + @pytest.mark.parametrize("component", (None, "easting", "northing", "upward")) + def test_invalid_number_of_elements( + self, sample_coordinates, sample_prisms, component + ): + """Check error when magnetization has invalid number of elements""" + # Generate an array with only two magnetization vectors + magnetizations = [ + [1.0, 1.0, 1.0, 3.4], + [-2.0, 3.0, -5.0, 3.4], + [-2.0, 3.0, -5.0, 3.4], + [-2.0, 3.0, -5.0, 3.4], + ] + msg = "Found magnetization vectors with" + if component is None: + with pytest.raises(ValueError, match=msg): + prism_magnetic(sample_coordinates, sample_prisms, magnetizations) + else: + with pytest.raises(ValueError, match=msg): + prism_magnetic_component( + sample_coordinates, sample_prisms, magnetizations, component + ) + + +class TestAgainstChoclo: + """ + Test forward modelling functions against dumb Choclo runs + """ + + @pytest.fixture() + def sample_prisms(self): + """ + Return four sample prisms + """ + prisms = np.array( + [ + [-10, 10, -10, 0, -10, 0], + [-10, 0, -10, 10, -20, -10], + [5, 15, 5, 15, -15, -5], + [-5, 5, -5, 5, -35, -20], + ], + dtype=float, + ) + return prisms + + @pytest.fixture() + def sample_magnetizations(self): + """ + Return sample magnetization vectors for the prisms + """ + magnetizations = np.array( + [ + [1.0, 1.0, 1], + [-1.0, -0.5, 2.0], + [3, -1.0, -3.0], + [-1.5, 2.0, -2.0], + ], + dtype=np.float64, + ) + return magnetizations + + @pytest.fixture() + def sample_coordinates(self): + """ + Return four sample observation points + """ + easting = np.array([-5, 10, 0, 15], dtype=float) + northing = np.array([14, -4, 11, 0], dtype=float) + upward = np.array([9, 6, 6, 12], dtype=float) + return (easting, northing, upward) + + def test_against_choclo( + self, sample_coordinates, sample_prisms, sample_magnetizations + ): + """ + Test prism_magnetic against raw Choclo runs + """ + easting, northing, upward = sample_coordinates + # Compute expected results with dumb Choclo runs + n_coords = easting.size + n_prisms = sample_prisms.shape[0] + expected_magnetic_e = np.zeros(n_coords, dtype=np.float64) + expected_magnetic_n = np.zeros(n_coords, dtype=np.float64) + expected_magnetic_u = np.zeros(n_coords, dtype=np.float64) + for i in range(n_coords): + for j in range(n_prisms): + b_e, b_n, b_u = magnetic_field( + easting[i], + northing[i], + upward[i], + *sample_prisms[j, :], + *sample_magnetizations[j, :], + ) + expected_magnetic_e[i] += b_e + expected_magnetic_n[i] += b_n + expected_magnetic_u[i] += b_u + # Convert to nT + expected_magnetic_e *= 1e9 + expected_magnetic_n *= 1e9 + expected_magnetic_u *= 1e9 + # Compare with harmonica results + b_e, b_n, b_u = prism_magnetic( + sample_coordinates, sample_prisms, sample_magnetizations + ) + npt.assert_allclose(b_e, expected_magnetic_e) + npt.assert_allclose(b_n, expected_magnetic_n) + npt.assert_allclose(b_u, expected_magnetic_u) + + @pytest.mark.parametrize("component", ("easting", "northing", "upward")) + def test_component_against_choclo( + self, sample_coordinates, sample_prisms, sample_magnetizations, component + ): + """ + Test prism_magnetic_component against raw Choclo runs + """ + easting, northing, upward = sample_coordinates + # Compute expected results with dumb Choclo runs + n_coords = easting.size + n_prisms = sample_prisms.shape[0] + expected_result = np.zeros(n_coords, dtype=np.float64) + forward_func = getattr(choclo.prism, f"magnetic_{component[0]}") + for i in range(n_coords): + for j in range(n_prisms): + expected_result[i] += forward_func( + easting[i], + northing[i], + upward[i], + *sample_prisms[j, :], + *sample_magnetizations[j, :], + ) + # Convert to nT + expected_result *= 1e9 + # Compare with harmonica results + result = prism_magnetic_component( + sample_coordinates, + sample_prisms, + sample_magnetizations, + component, + ) + npt.assert_allclose(result, expected_result)