diff --git a/doc/gallery_src/transformations/reduction_to_pole.py b/doc/gallery_src/transformations/reduction_to_pole.py index 8a785a36b..5a8f74165 100644 --- a/doc/gallery_src/transformations/reduction_to_pole.py +++ b/doc/gallery_src/transformations/reduction_to_pole.py @@ -13,7 +13,6 @@ import pygmt import verde as vd import xarray as xr -import xrft import harmonica as hm @@ -22,15 +21,6 @@ fname = ensaio.fetch_lightning_creek_magnetic(version=1) magnetic_grid = xr.load_dataarray(fname) -# Pad the grid to increase accuracy of the FFT filter -pad_width = { - "easting": magnetic_grid.easting.size // 3, - "northing": magnetic_grid.northing.size // 3, -} -# drop the extra height coordinate -magnetic_grid_no_height = magnetic_grid.drop_vars("height") -magnetic_grid_padded = xrft.pad(magnetic_grid_no_height, pad_width) - # Define the inclination and declination of the region by the time of the data # acquisition (1990). inclination, declination = -52.98, 6.51 @@ -39,16 +29,11 @@ # that the sources share the same inclination and declination as the # geomagnetic field. rtp_grid = hm.reduction_to_pole( - magnetic_grid_padded, inclination=inclination, declination=declination + magnetic_grid, inclination=inclination, declination=declination ) - -# Unpad the reduced to the pole grid -rtp_grid = xrft.unpad(rtp_grid, pad_width) - # Show the reduced to the pole grid print("\nReduced to the pole magnetic grid:\n", rtp_grid) - # Plot original magnetic anomaly and the reduced to the pole fig = pygmt.Figure() with fig.subplot(nrows=1, ncols=2, figsize=("28c", "15c"), sharey="l"): diff --git a/doc/gallery_src/transformations/tga.py b/doc/gallery_src/transformations/tga.py index 21c710e5a..24d3b1bbe 100644 --- a/doc/gallery_src/transformations/tga.py +++ b/doc/gallery_src/transformations/tga.py @@ -13,7 +13,6 @@ import pygmt import verde as vd import xarray as xr -import xrft import harmonica as hm @@ -22,21 +21,8 @@ fname = ensaio.fetch_lightning_creek_magnetic(version=1) magnetic_grid = xr.load_dataarray(fname) -# Pad the grid to increase accuracy of the FFT filter -pad_width = { - "easting": magnetic_grid.easting.size // 3, - "northing": magnetic_grid.northing.size // 3, -} -# drop the extra height coordinate -magnetic_grid_no_height = magnetic_grid.drop_vars("height") -magnetic_grid_padded = xrft.pad(magnetic_grid_no_height, pad_width) - # Compute the total gradient amplitude of the grid -tga = hm.total_gradient_amplitude(magnetic_grid_padded) - -# Unpad the total gradient amplitude grid -tga = xrft.unpad(tga, pad_width) - +tga = hm.total_gradient_amplitude(magnetic_grid) # Show the total gradient amplitude print("\nTotal Gradient Amplitude:\n", tga) diff --git a/doc/gallery_src/transformations/tilt.py b/doc/gallery_src/transformations/tilt.py index 9f9c1b579..56331eccb 100644 --- a/doc/gallery_src/transformations/tilt.py +++ b/doc/gallery_src/transformations/tilt.py @@ -13,7 +13,6 @@ import pygmt import verde as vd import xarray as xr -import xrft import harmonica as hm @@ -22,21 +21,8 @@ fname = ensaio.fetch_lightning_creek_magnetic(version=1) magnetic_grid = xr.load_dataarray(fname) -# Pad the grid to increase accuracy of the FFT filter -pad_width = { - "easting": magnetic_grid.easting.size // 3, - "northing": magnetic_grid.northing.size // 3, -} -# drop the extra height coordinate -magnetic_grid_no_height = magnetic_grid.drop_vars("height") -magnetic_grid_padded = xrft.pad(magnetic_grid_no_height, pad_width) - # Compute the tilt of the grid -tilt_grid = hm.tilt_angle(magnetic_grid_padded) - -# Unpad the tilt grid -tilt_grid = xrft.unpad(tilt_grid, pad_width) - +tilt_grid = hm.tilt_angle(magnetic_grid) # Show the tilt print("\nTilt:\n", tilt_grid) @@ -47,19 +33,12 @@ # Apply a reduction to the pole over the magnetic anomaly grid. We will assume # that the sources share the same inclination and declination as the # geomagnetic field. -rtp_grid_padded = hm.reduction_to_pole( - magnetic_grid_padded, inclination=inclination, declination=declination +rtp_grid = hm.reduction_to_pole( + magnetic_grid, inclination=inclination, declination=declination ) -# Unpad the reduced to the pole grid -rtp_grid = xrft.unpad(rtp_grid_padded, pad_width) - # Compute the tilt of the padded rtp grid -tilt_rtp_grid = hm.tilt_angle(rtp_grid_padded) - -# Unpad the tilt grid -tilt_rtp_grid = xrft.unpad(tilt_rtp_grid, pad_width) - +tilt_rtp_grid = hm.tilt_angle(rtp_grid) # Show the tilt from RTP print("\nTilt from RTP:\n", tilt_rtp_grid) diff --git a/doc/gallery_src/transformations/upward_continuation.py b/doc/gallery_src/transformations/upward_continuation.py index db6ce9829..182b29261 100644 --- a/doc/gallery_src/transformations/upward_continuation.py +++ b/doc/gallery_src/transformations/upward_continuation.py @@ -12,7 +12,6 @@ import ensaio import pygmt import xarray as xr -import xrft import harmonica as hm @@ -21,22 +20,9 @@ fname = ensaio.fetch_lightning_creek_magnetic(version=1) magnetic_grid = xr.load_dataarray(fname) -# Pad the grid to increase accuracy of the FFT filter -pad_width = { - "easting": magnetic_grid.easting.size // 3, - "northing": magnetic_grid.northing.size // 3, -} -# drop the extra height coordinate -magnetic_grid_no_height = magnetic_grid.drop_vars("height") -magnetic_grid_padded = xrft.pad(magnetic_grid_no_height, pad_width) - # Upward continue the magnetic grid, from 500 m to 1000 m # (a height displacement of 500m) -upward_continued = hm.upward_continuation(magnetic_grid_padded, height_displacement=500) - -# Unpad the upward continued grid -upward_continued = xrft.unpad(upward_continued, pad_width) - +upward_continued = hm.upward_continuation(magnetic_grid, height_displacement=500) # Show the upward continued grid print("\nUpward continued magnetic grid:\n", upward_continued) diff --git a/doc/gallery_src/transformations/upward_derivative.py b/doc/gallery_src/transformations/upward_derivative.py index 2d48b89db..b432a0b57 100644 --- a/doc/gallery_src/transformations/upward_derivative.py +++ b/doc/gallery_src/transformations/upward_derivative.py @@ -13,7 +13,6 @@ import pygmt import verde as vd import xarray as xr -import xrft import harmonica as hm @@ -22,21 +21,8 @@ fname = ensaio.fetch_lightning_creek_magnetic(version=1) magnetic_grid = xr.load_dataarray(fname) -# Pad the grid to increase accuracy of the FFT filter -pad_width = { - "easting": magnetic_grid.easting.size // 3, - "northing": magnetic_grid.northing.size // 3, -} -# drop the extra height coordinate -magnetic_grid_no_height = magnetic_grid.drop_vars("height") -magnetic_grid_padded = xrft.pad(magnetic_grid_no_height, pad_width) - # Compute the upward derivative of the grid -deriv_upward = hm.derivative_upward(magnetic_grid_padded) - -# Unpad the derivative grid -deriv_upward = xrft.unpad(deriv_upward, pad_width) - +deriv_upward = hm.derivative_upward(magnetic_grid) # Show the upward derivative print("\nUpward derivative:\n", deriv_upward) diff --git a/doc/user_guide/transformations.rst b/doc/user_guide/transformations.rst index ce9ac9734..98ac637ad 100644 --- a/doc/user_guide/transformations.rst +++ b/doc/user_guide/transformations.rst @@ -15,8 +15,11 @@ We can load the data file using :mod:`xarray`: .. jupyter-execute:: - import ensaio import xarray as xr + import matplotlib.pyplot as plt + import verde as vd + import harmonica as hm + import ensaio fname = ensaio.fetch_lightning_creek_magnetic(version=1) magnetic_grid = xr.load_dataarray(fname) @@ -26,8 +29,6 @@ And plot it: .. jupyter-execute:: - import matplotlib.pyplot as plt - tmp = magnetic_grid.plot(cmap="seismic", center=0, add_colorbar=False) plt.gca().set_aspect("equal") plt.title("Magnetic anomaly grid") @@ -42,41 +43,6 @@ And plot it: :func:`verde.project_grid` function and a map projection like the ones available in :mod:`pyproj`. -Since all the grid transformations we are going to apply are based on FFT -methods, we usually want to pad them in order their increase the accuracy. -We can easily do it through the :func:`xrft.pad` function. -First we need to define how much padding we want to add along each direction. -We will add one third of the width and height of the grid to each side: - -.. jupyter-execute:: - - pad_width = { - "easting": magnetic_grid.easting.size // 3, - "northing": magnetic_grid.northing.size // 3, - } - -And then we can pad it, but dropping the ``height`` coordinate first (this is -needed by the :func:`xrft.pad` function): - -.. jupyter-execute:: - - import xrft - - magnetic_grid_no_height = magnetic_grid.drop_vars("height") - magnetic_grid_padded = xrft.pad(magnetic_grid_no_height, pad_width) - magnetic_grid_padded - -.. jupyter-execute:: - - tmp = magnetic_grid_padded.plot(cmap="seismic", center=0, add_colorbar=False) - plt.gca().set_aspect("equal") - plt.title("Padded magnetic anomaly grid") - plt.gca().ticklabel_format(style="sci", scilimits=(0, 0)) - plt.colorbar(tmp, label="nT") - plt.show() - -Now that we have the padded grid, we can apply any grid transformation. - Upward derivative ----------------- @@ -86,17 +52,7 @@ magnetic anomaly grid using the :func:`harmonica.derivative_upward` function: .. jupyter-execute:: - import harmonica as hm - - deriv_upward = hm.derivative_upward(magnetic_grid_padded) - deriv_upward - -This grid includes all the padding we added to the original magnetic grid, so -we better unpad it using :func:`xrft.unpad`: - -.. jupyter-execute:: - - deriv_upward = xrft.unpad(deriv_upward, pad_width) + deriv_upward = hm.derivative_upward(magnetic_grid) deriv_upward And plot it: @@ -160,14 +116,12 @@ frequency domain: .. jupyter-execute:: - deriv_easting = hm.derivative_easting(magnetic_grid_padded, method="fft") - deriv_easting = xrft.unpad(deriv_easting, pad_width) + deriv_easting = hm.derivative_easting(magnetic_grid, method="fft") deriv_easting .. jupyter-execute:: - deriv_northing = hm.derivative_northing(magnetic_grid_padded, method="fft") - deriv_northing = xrft.unpad(deriv_northing, pad_width) + deriv_northing = hm.derivative_northing(magnetic_grid, method="fft") deriv_northing .. jupyter-execute:: @@ -198,11 +152,10 @@ frequency domain: and have less artifacts than their FFT-based counterpart. - Upward continuation ------------------- -We can also upward continue the original magnetic grid. +We can also upward continue the original magnetic grid using :func:`harmonica.upward_continuation`. This is, estimating the magnetic field generated by the same sources at a higher altitude. The original magnetic anomaly grid is located at 500 m above the ellipsoid, as @@ -212,19 +165,26 @@ to upward continue it a height displacement of 500m: .. jupyter-execute:: + change_in_height = 500 # meters upward_continued = hm.upward_continuation( - magnetic_grid_padded, height_displacement=500 + magnetic_grid, height_displacement=change_in_height ) + upward_continued -This grid includes all the padding we added to the original magnetic grid, so -we better unpad it using :func:`xrft.unpad`: +Did you notice that the ``height`` coordinate is gone from the +upward-continued grid? We drop any non-dimensional coordinates when +doing upward continuation because we don't know the name of the +vertical coordinate and any values there would be wrong after continuation. +If we want it back, we need to assign an updated version of it: .. jupyter-execute:: - upward_continued = xrft.unpad(upward_continued, pad_width) + upward_continued = upward_continued.assign_coords( + {"height": magnetic_grid.height + change_in_height} + ) upward_continued -And plot it: +Now we can plot it: .. jupyter-execute:: @@ -269,11 +229,8 @@ remanence), then we can apply the reduction to the pole passing only the .. jupyter-execute:: rtp_grid = hm.reduction_to_pole( - magnetic_grid_padded, inclination=inclination, declination=declination + magnetic_grid, inclination=inclination, declination=declination ) - - # Unpad the reduced to the pole grid - rtp_grid = xrft.unpad(rtp_grid, pad_width) rtp_grid And plot it: @@ -296,15 +253,12 @@ magnetization vector of the sources, we can specify the mag_inclination, mag_declination = -25, 21 tmp = rtp_grid = hm.reduction_to_pole( - magnetic_grid_padded, + magnetic_grid, inclination=inclination, declination=declination, magnetization_inclination=mag_inclination, magnetization_declination=mag_declination, ) - - # Unpad the reduced to the pole grid - rtp_grid = xrft.unpad(rtp_grid, pad_width) rtp_grid .. jupyter-execute:: @@ -335,59 +289,49 @@ Let's define a cutoff wavelength of 5 kilometers: .. jupyter-execute:: - cutoff_wavelength = 5e3 + cutoff_wavelength = 5e3 # meters -Then apply the two filters to our padded magnetic grid: +Then apply the two filters to our magnetic grid: .. jupyter-execute:: - magnetic_low_freqs = hm.gaussian_lowpass( - magnetic_grid_padded, wavelength=cutoff_wavelength - ) - magnetic_high_freqs = hm.gaussian_highpass( - magnetic_grid_padded, wavelength=cutoff_wavelength + magnetic_low = hm.gaussian_lowpass( + magnetic_grid, wavelength=cutoff_wavelength ) - -And unpad them: + magnetic_low .. jupyter-execute:: - magnetic_low_freqs = xrft.unpad(magnetic_low_freqs, pad_width) - magnetic_high_freqs = xrft.unpad(magnetic_high_freqs, pad_width) - -.. jupyter-execute:: - - magnetic_low_freqs - -.. jupyter-execute:: - - magnetic_high_freqs + magnetic_high = hm.gaussian_highpass( + magnetic_grid, wavelength=cutoff_wavelength + ) + magnetic_high Let's plot the results side by side: .. jupyter-execute:: - import verde as vd - - fig, (ax1, ax2) = plt.subplots( - nrows=1, ncols=2, sharey=True, figsize=(12, 8) + fig, (ax1, ax2, ax3) = plt.subplots( + nrows=1, ncols=3, sharey=True, figsize=(14, 8) ) - maxabs = vd.maxabs(magnetic_low_freqs, magnetic_high_freqs) + maxabs = vd.maxabs(magnetic_grid, magnetic_low, magnetic_high) kwargs = dict(cmap="seismic", vmin=-maxabs, vmax=maxabs, add_colorbar=False) - tmp = magnetic_low_freqs.plot(ax=ax1, **kwargs) - tmp = magnetic_high_freqs.plot(ax=ax2, **kwargs) + tmp = magnetic_grid.plot(ax=ax1, **kwargs) + tmp = magnetic_low.plot(ax=ax2, **kwargs) + tmp = magnetic_high.plot(ax=ax3, **kwargs) - ax1.set_title("Magnetic anomaly after low-pass filter") - ax2.set_title("Magnetic anomaly after high-pass filter") - for ax in (ax1, ax2): + ax1.set_title("Original") + ax2.set_title("After low-pass filter") + ax3.set_title("After high-pass filter") + for ax in (ax1, ax2, ax3): ax.set_aspect("equal") ax.ticklabel_format(style="sci", scilimits=(0, 0)) plt.colorbar( tmp, - ax=[ax1, ax2], + ax=[ax1, ax2, ax3], label="nT", orientation="horizontal", aspect=42, @@ -422,19 +366,14 @@ We can apply it through the :func:`harmonica.total_gradient_amplitude` function. .. jupyter-execute:: tga_grid = hm.total_gradient_amplitude( - magnetic_grid_padded + magnetic_grid ) - - # Unpad the total gradient amplitude grid - tga_grid = xrft.unpad(tga_grid, pad_width) tga_grid And plot it: .. jupyter-execute:: - import verde as vd - tmp = tga_grid.plot(cmap="viridis", add_colorbar=False) plt.gca().set_aspect("equal") plt.title("Total gradient amplitude of the magnetic anomaly") diff --git a/harmonica/_transformations.py b/harmonica/_transformations.py index d505d3645..300570fb7 100644 --- a/harmonica/_transformations.py +++ b/harmonica/_transformations.py @@ -22,7 +22,7 @@ from .filters._utils import apply_filter, grid_sanity_checks -def derivative_upward(grid, order=1): +def derivative_upward(grid, order=1, pad=True, pad_kwargs=None): """ Calculate the derivative of a potential field grid in the upward direction. @@ -38,6 +38,16 @@ def derivative_upward(grid, order=1): same units. order : int The order of the derivative. Default to 1. + pad : bool + If True, will add padding to the grid before taking the Fourier Transform + and applying the filter and remove it after the inverse Fourier Transform. + Adding padding usually helps reduce edge effects from signal truncation. + Default is True. + pad_kwargs : dict or None + Any additional keyword arguments that should be passed to the + :meth:`xarray.DataArray.pad` function in the form of a dictionary. If none + are given, the default padding of 25% the dimensions of the grid will be + added using the "edge" method. Returns ------- @@ -54,10 +64,16 @@ def derivative_upward(grid, order=1): -------- harmonica.filters.derivative_upward_kernel """ - return apply_filter(grid, derivative_upward_kernel, order=order) + return apply_filter( + grid, + derivative_upward_kernel, + filter_kwargs={"order": order}, + pad=pad, + pad_kwargs=pad_kwargs, + ) -def derivative_easting(grid, order=1, method="finite-diff"): +def derivative_easting(grid, order=1, method="finite-diff", pad=True, pad_kwargs=None): """ Calculate the derivative of a regular grid in the easting direction. @@ -84,8 +100,17 @@ def derivative_easting(grid, order=1, method="finite-diff"): Method that will be used for computing the easting derivative. It can be either ``"finite-diff"``, for computing using :func:`xarray.differentiate`, or ``"fft"``, for using FFT-based - filters. - Default ``"finite-diff"``. + filters. Default ``"finite-diff"``. + pad : bool + If True, will add padding to the grid before taking the Fourier Transform + and applying the filter and remove it after the inverse Fourier Transform. + Adding padding usually helps reduce edge effects from signal truncation. + Default is True. + pad_kwargs : dict or None + Any additional keyword arguments that should be passed to the + :meth:`xarray.DataArray.pad` function in the form of a dictionary. If none + are given, the default padding of 25% the dimensions of the grid will be + added using the "edge" method. Returns ------- @@ -107,7 +132,13 @@ def derivative_easting(grid, order=1, method="finite-diff"): for _ in range(order): grid = grid.differentiate(coord=coordinate) elif method == "fft": - grid = apply_filter(grid, derivative_easting_kernel, order=order) + grid = apply_filter( + grid, + derivative_easting_kernel, + filter_kwargs={"order": order}, + pad=pad, + pad_kwargs=pad_kwargs, + ) else: msg = ( f"Invalid method '{method}'. Please select one from 'finite-diff' or 'fft'." @@ -116,7 +147,7 @@ def derivative_easting(grid, order=1, method="finite-diff"): return grid -def derivative_northing(grid, order=1, method="finite-diff"): +def derivative_northing(grid, order=1, method="finite-diff", pad=True, pad_kwargs=None): """ Calculate the derivative of a regular grid in the northing direction. @@ -143,8 +174,17 @@ def derivative_northing(grid, order=1, method="finite-diff"): Method that will be used for computing the easting derivative. It can be either ``"finite-diff"``, for computing using :func:`xarray.differentiate`, or ``"fft"``, for using FFT-based - filters. - Default ``"finite-diff"``. + filters. Default ``"finite-diff"``. + pad : bool + If True, will add padding to the grid before taking the Fourier Transform + and applying the filter and remove it after the inverse Fourier Transform. + Adding padding usually helps reduce edge effects from signal truncation. + Default is True. + pad_kwargs : dict or None + Any additional keyword arguments that should be passed to the + :meth:`xarray.DataArray.pad` function in the form of a dictionary. If none + are given, the default padding of 25% the dimensions of the grid will be + added using the "edge" method. Returns ------- @@ -166,7 +206,13 @@ def derivative_northing(grid, order=1, method="finite-diff"): for _ in range(order): grid = grid.differentiate(coord=coordinate) elif method == "fft": - return apply_filter(grid, derivative_northing_kernel, order=order) + return apply_filter( + grid, + derivative_northing_kernel, + filter_kwargs={"order": order}, + pad=pad, + pad_kwargs=pad_kwargs, + ) else: msg = ( f"Invalid method '{method}'. Please select one from 'finite-diff' or 'fft'." @@ -175,13 +221,18 @@ def derivative_northing(grid, order=1, method="finite-diff"): return grid -def upward_continuation(grid, height_displacement): +def upward_continuation(grid, height_displacement, pad=True, pad_kwargs=None): """ Calculate the upward continuation of a potential field grid. Compute the upward continuation of regular gridded data using frequency domain calculations through Fast Fourier Transform. + .. note:: + + Any non-dimensional coordinates of the grid will be dropped since + upward continuation may have made them no longer correct. + Parameters ---------- grid : :class:`xarray.DataArray` @@ -193,6 +244,16 @@ def upward_continuation(grid, height_displacement): The height displacement of upward continuation. For upward continuation, the height displacement should be positive. Its units are the same units of the ``grid`` coordinates. + pad : bool + If True, will add padding to the grid before taking the Fourier Transform + and applying the filter and remove it after the inverse Fourier Transform. + Adding padding usually helps reduce edge effects from signal truncation. + Default is True. + pad_kwargs : dict or None + Any additional keyword arguments that should be passed to the + :meth:`xarray.DataArray.pad` function in the form of a dictionary. If none + are given, the default padding of 25% the dimensions of the grid will be + added using the "edge" method. Returns ------- @@ -209,11 +270,16 @@ def upward_continuation(grid, height_displacement): harmonica.filters.upward_continuation_kernel """ return apply_filter( - grid, upward_continuation_kernel, height_displacement=height_displacement + grid, + upward_continuation_kernel, + filter_kwargs={"height_displacement": height_displacement}, + pad=pad, + pad_kwargs=pad_kwargs, + drop_coords=True, ) -def gaussian_lowpass(grid, wavelength): +def gaussian_lowpass(grid, wavelength, pad=True, pad_kwargs=None): """ Calculate the Gaussian low-pass of a potential field grid. @@ -230,6 +296,16 @@ def gaussian_lowpass(grid, wavelength): wavelength : float The cutoff wavelength in low-pass filter. Its units are the same units of the ``grid`` coordinates. + pad : bool + If True, will add padding to the grid before taking the Fourier Transform + and applying the filter and remove it after the inverse Fourier Transform. + Adding padding usually helps reduce edge effects from signal truncation. + Default is True. + pad_kwargs : dict or None + Any additional keyword arguments that should be passed to the + :meth:`xarray.DataArray.pad` function in the form of a dictionary. If none + are given, the default padding of 25% the dimensions of the grid will be + added using the "edge" method. Returns ------- @@ -245,10 +321,16 @@ def gaussian_lowpass(grid, wavelength): -------- harmonica.filters.gaussian_lowpass_kernel """ - return apply_filter(grid, gaussian_lowpass_kernel, wavelength=wavelength) + return apply_filter( + grid, + gaussian_lowpass_kernel, + pad=pad, + pad_kwargs=pad_kwargs, + filter_kwargs={"wavelength": wavelength}, + ) -def gaussian_highpass(grid, wavelength): +def gaussian_highpass(grid, wavelength, pad=True, pad_kwargs=None): """ Calculate the Gaussian high-pass of a potential field grid. @@ -265,6 +347,16 @@ def gaussian_highpass(grid, wavelength): wavelength : float The cutoff wavelength in high-pass filter. Its units are the same units of the ``grid`` coordinates. + pad : bool + If True, will add padding to the grid before taking the Fourier Transform + and applying the filter and remove it after the inverse Fourier Transform. + Adding padding usually helps reduce edge effects from signal truncation. + Default is True. + pad_kwargs : dict or None + Any additional keyword arguments that should be passed to the + :meth:`xarray.DataArray.pad` function in the form of a dictionary. If none + are given, the default padding of 25% the dimensions of the grid will be + added using the "edge" method. Returns ------- @@ -280,7 +372,13 @@ def gaussian_highpass(grid, wavelength): -------- harmonica.filters.gaussian_highpass_kernel """ - return apply_filter(grid, gaussian_highpass_kernel, wavelength=wavelength) + return apply_filter( + grid, + gaussian_highpass_kernel, + pad=pad, + pad_kwargs=pad_kwargs, + filter_kwargs={"wavelength": wavelength}, + ) def reduction_to_pole( @@ -289,6 +387,8 @@ def reduction_to_pole( declination, magnetization_inclination=None, magnetization_declination=None, + pad=True, + pad_kwargs=None, ): """ Calculate the reduction to the pole of a magnetic field grid. @@ -317,6 +417,16 @@ def reduction_to_pole( None, the ``magnetization_declination`` will be set equal to the ``declination``, neglecting remanent magnetization and self demagnetization. Default None. + pad : bool + If True, will add padding to the grid before taking the Fourier Transform + and applying the filter and remove it after the inverse Fourier Transform. + Adding padding usually helps reduce edge effects from signal truncation. + Default is True. + pad_kwargs : dict or None + Any additional keyword arguments that should be passed to the + :meth:`xarray.DataArray.pad` function in the form of a dictionary. If none + are given, the default padding of 25% the dimensions of the grid will be + added using the "edge" method. Returns ------- @@ -335,14 +445,18 @@ def reduction_to_pole( return apply_filter( grid, reduction_to_pole_kernel, - inclination=inclination, - declination=declination, - magnetization_inclination=magnetization_inclination, - magnetization_declination=magnetization_declination, + filter_kwargs={ + "inclination": inclination, + "declination": declination, + "magnetization_inclination": magnetization_inclination, + "magnetization_declination": magnetization_declination, + }, + pad=pad, + pad_kwargs=pad_kwargs, ) -def total_gradient_amplitude(grid): +def total_gradient_amplitude(grid, pad=True, pad_kwargs=None): r""" Calculate the total gradient amplitude of a potential field grid. @@ -357,12 +471,22 @@ def total_gradient_amplitude(grid): evenly spaced (regular grid). Its dimensions should be in the following order: *northing*, *easting*. Its coordinates should be defined in the same units. + pad : bool + If True, will add padding to the grid before taking the Fourier Transform + and applying the filter and remove it after the inverse Fourier Transform. + Adding padding usually helps reduce edge effects from signal truncation. + Default is True. + pad_kwargs : dict or None + Any additional keyword arguments that should be passed to the + :meth:`xarray.DataArray.pad` function in the form of a dictionary. If none + are given, the default padding of 25% the dimensions of the grid will be + added using the "edge" method. Returns ------- total_gradient_amplitude_grid : :class:`xarray.DataArray` - A :class:`xarray.DataArray` after calculating the - total gradient amplitude of the passed ``grid``. + A :class:`xarray.DataArray` after calculating the total gradient + amplitude of the passed ``grid``. Notes ----- @@ -386,22 +510,21 @@ def total_gradient_amplitude(grid): grid_sanity_checks(grid) # Calculate the gradients of the grid gradient = ( - derivative_easting(grid, order=1), - derivative_northing(grid, order=1), - derivative_upward(grid, order=1), + derivative_easting(grid, order=1, pad=pad, pad_kwargs=pad_kwargs), + derivative_northing(grid, order=1, pad=pad, pad_kwargs=pad_kwargs), + derivative_upward(grid, order=1, pad=pad, pad_kwargs=pad_kwargs), ) # return the total gradient amplitude return np.sqrt(gradient[0] ** 2 + gradient[1] ** 2 + gradient[2] ** 2) -def tilt_angle(grid): +def tilt_angle(grid, pad=True, pad_kwargs=None): r""" Calculate the tilt angle of a potential field grid. - Compute the tilt of a regular gridded potential field - :math:`M`. The horizontal derivatives are calculated - through finite-differences while the upward derivative - is calculated using FFT. + Compute the tilt of a regular gridded potential field :math:`M`. The + horizontal derivatives are calculated through finite-differences while the + upward derivative is calculated using FFT. Parameters ---------- @@ -410,6 +533,16 @@ def tilt_angle(grid): evenly spaced (regular grid). Its dimensions should be in the following order: *northing*, *easting*. Its coordinates should be defined in the same units. + pad : bool + If True, will add padding to the grid before taking the Fourier Transform + and applying the filter and remove it after the inverse Fourier Transform. + Adding padding usually helps reduce edge effects from signal truncation. + Default is True. + pad_kwargs : dict or None + Any additional keyword arguments that should be passed to the + :meth:`xarray.DataArray.pad` function in the form of a dictionary. If none + are given, the default padding of 25% the dimensions of the grid will be + added using the "edge" method. Returns ------- @@ -442,17 +575,12 @@ def tilt_angle(grid): [Blakely1995]_ [MillerSingh1994]_ """ - # Run sanity checks on the grid grid_sanity_checks(grid) - # Calculate the gradients of the grid - gradient = ( - derivative_easting(grid, order=1), - derivative_northing(grid, order=1), - derivative_upward(grid, order=1), - ) - # Calculate and return the tilt - horiz_deriv = np.sqrt(gradient[0] ** 2 + gradient[1] ** 2) - tilt = np.arctan2(gradient[2], horiz_deriv) + deriv_east = derivative_easting(grid, order=1, pad=pad, pad_kwargs=pad_kwargs) + deriv_north = derivative_northing(grid, order=1, pad=pad, pad_kwargs=pad_kwargs) + deriv_up = derivative_upward(grid, order=1, pad=pad, pad_kwargs=pad_kwargs) + horiz_deriv = np.hypot(deriv_east, deriv_north) + tilt = np.arctan2(deriv_up, horiz_deriv) return tilt diff --git a/harmonica/filters/_fft.py b/harmonica/filters/_fft.py index 74bbd46f6..d65d2de32 100644 --- a/harmonica/filters/_fft.py +++ b/harmonica/filters/_fft.py @@ -8,11 +8,10 @@ Wrap xrft functions to compute FFTs and inverse FFTs. """ -from xrft.xrft import fft as _fft -from xrft.xrft import ifft as _ifft +import xrft -def fft(grid, true_phase=True, true_amplitude=True, drop_bad_coords=True, **kwargs): +def fft(grid, true_phase=True, true_amplitude=True, **kwargs): """ Compute Fast Fourier Transform of a 2D regular grid. @@ -32,21 +31,17 @@ def fft(grid, true_phase=True, true_amplitude=True, drop_bad_coords=True, **kwar If True, the FFT is multiplied by the spacing of the transformed variables to match theoretical FT amplitude. Defaults to True. - drop_bad_coords : bool (optional) - If True, only the indexes of the array will be kept before passing it - to :func:`xrft.fft`. Any extra coordinate should be drooped, otherwise - :func:`xrft.fft` raises an error after finding *bad coordinates*. - Defaults to True. + **kwargs + Any extra keyword arguments will be passed the :func:`xrft.fft` function. Returns ------- fourier_transform : :class:`xarray.DataArray` Array with the Fourier transform of the original grid. """ - if drop_bad_coords: - bad_coords = tuple(c for c in grid.coords if c not in grid.indexes) - grid = grid.drop(bad_coords) - return _fft(grid, true_phase=true_phase, true_amplitude=true_amplitude, **kwargs) + return xrft.fft( + grid, true_phase=true_phase, true_amplitude=true_amplitude, **kwargs + ) def ifft(fourier_transform, true_phase=True, true_amplitude=True, **kwargs): @@ -69,15 +64,18 @@ def ifft(fourier_transform, true_phase=True, true_amplitude=True, **kwargs): If True, output is divided by the spacing of the transformed variables to match theoretical IFT amplitude. Defaults to True. + **kwargs + Any extra keyword arguments will be passed the :func:`xrft.ifft` function. Returns ------- grid : :class:`xarray.DataArray` Array with the inverse Fourier transform of the passed grid. """ - return _ifft( + return xrft.ifft( fourier_transform, true_phase=true_phase, true_amplitude=true_amplitude, + lag=(None, None), # Mutes an annoying FutureWarning from xrft **kwargs, ) diff --git a/harmonica/filters/_utils.py b/harmonica/filters/_utils.py index d34617d73..ac9bb0238 100644 --- a/harmonica/filters/_utils.py +++ b/harmonica/filters/_utils.py @@ -9,17 +9,27 @@ """ import numpy as np +import xrft from ._fft import fft, ifft -def apply_filter(grid, fft_filter, **kwargs): +def apply_filter( + grid, fft_filter, filter_kwargs=None, pad=True, pad_kwargs=None, drop_coords=False +): """ Apply a filter to a grid and return the transformed grid in spatial domain. Computes the Fourier transform of the given grid, builds the filter, applies it and returns the inverse Fourier transform of the filtered grid. + .. note:: + + Any non-dimensional coordinates in the original grid will be dropped + from the filtered grid. This is because we can't know if the filter + invalidates the coordinate values (for example, upward continuation + would invalidate any height coordinates). So it's safer to drop them. + Parameters ---------- grid : :class:`xarray.DataArray` @@ -29,9 +39,23 @@ def apply_filter(grid, fft_filter, **kwargs): same units. fft_filter : func Callable that builds the filter in the frequency domain. - kwargs : + filter_kwargs : dict Any additional keyword argument that should be passed to the - ``fft_filter``. + ``fft_filter`` in the form of a dictionary. + pad : bool + If True, will add padding to the grid before taking the Fourier Transform + and applying the filter and remove it after the inverse Fourier Transform. + Adding padding usually helps reduce edge effects from signal truncation. + Default is True. + pad_kwargs : dict + Any additional keyword arguments that should be passed to the + :meth:`xarray.DataArray.pad` function. If none are given, the default + padding of 25% the dimensions of the grid will be added using the + "edge" method. + drop_coords : bool + If True, non-dimensional coordinates of the grid will be dropped after + filtering. This is useful if the filter could move the grid, like in upward + continuation, which could make these coordinates incorrect. Returns ------- @@ -39,24 +63,52 @@ def apply_filter(grid, fft_filter, **kwargs): A :class:`xarray.DataArray` with the filtered version of the passed ``grid``. Defined are in the spatial domain. """ - # Run sanity checks on the grid + if filter_kwargs is None: + filter_kwargs = {} + if pad_kwargs is None: + pad_kwargs = {} grid_sanity_checks(grid) - # Catch the dims of the grid dims = grid.dims - # Compute Fourier Transform of the grid - fft_grid = fft(grid) - # Build the filter - da_filter = fft_filter(fft_grid, **kwargs) - # Apply the filter - filtered_fft_grid = fft_grid * da_filter - # Compute inverse FFT + # Need to remove non-dimensional coordinates before padding and FFT because + # xrft doesn't know what to do with them. + non_dim_coords = {c: grid[c] for c in grid.coords if c not in grid.indexes} + grid = grid.drop_vars(non_dim_coords.keys()) + if pad: + # By default, use a padding width of 25% of the largest grid dimension. + # Fedi et al. (2012; doi:10.1111/j.1365-246X.2011.05259.x) suggest + # a padding of 100% but that seems exaggerated. + if "pad_width" not in pad_kwargs: + width = int(0.25 * max(grid[d].size for d in dims)) + pad_kwargs["pad_width"] = {d: width for d in dims} + if "mode" not in pad_kwargs: + pad_kwargs["mode"] = "edge" + if "constant_values" not in pad_kwargs: + # Has to be included explicitly as None since xrft always passes + # it to xarray.DataArray.pad. + pad_kwargs["constant_values"] = None + fft_grid = fft(xrft.pad(grid, **pad_kwargs)) + else: + fft_grid = fft(grid) + # The filter convolution in the frequency domain is a multiplication + filtered_fft_grid = fft_grid * fft_filter(fft_grid, **filter_kwargs) + # Keep only the real part since the inverse transform returns complex + # number by default filtered_grid = ifft(filtered_fft_grid).real - - # use original coordinates on the filtered grid + if pad: + filtered_grid = xrft.unpad(filtered_grid, pad_kwargs["pad_width"]) + # Restore the original coordinates to the grid because the inverse + # transform calculates coordinates from the frequencies, which can lead to + # rounding errors and coordinates that are slightly off. This causes errors + # when doing operations with the transformed grids. Restoring the original + # coordinates avoids these issues. filtered_grid = filtered_grid.assign_coords( - {dims[1]: grid[dims[1]].values, dims[0]: grid[dims[0]].values} + {dims[1]: grid[dims[1]], dims[0]: grid[dims[0]]} ) - + # Restore the non-dimensional coordinates if desired + if not drop_coords: + filtered_grid = filtered_grid.assign_coords( + {name: non_dim_coords[name] for name in non_dim_coords} + ) return filtered_grid diff --git a/harmonica/tests/data/filter.nc b/harmonica/tests/data/filter.nc deleted file mode 100644 index ea70fc37b..000000000 Binary files a/harmonica/tests/data/filter.nc and /dev/null differ diff --git a/harmonica/tests/test_euler_deconvolution.py b/harmonica/tests/test_euler_deconvolution.py index 1338a23a2..ab7f51544 100644 --- a/harmonica/tests/test_euler_deconvolution.py +++ b/harmonica/tests/test_euler_deconvolution.py @@ -19,6 +19,7 @@ def test_euler_with_numeric_derivatives(): + "Check the Euler solution against a synthetic using numerical derivatives" # Add dipole source dipole_coordinates = (10e3, 15e3, -10e3) dipole_moments = magnetic_angles_to_vec(1.0e14, 0, 0) @@ -48,15 +49,16 @@ def test_euler_with_numeric_derivatives(): coordinates = (grid_table.easting, grid_table.northing, grid_table.upward) euler.fit( - (grid_table.easting, grid_table.northing, grid_table.upward), + coordinates, (grid_table.tfa, grid_table.d_east, grid_table.d_north, grid_table.d_up), ) - npt.assert_allclose(euler.location_, dipole_coordinates, atol=1.0e-3, rtol=1.0e-3) - npt.assert_allclose(euler.base_level_, true_base_level, atol=1.0e-3, rtol=1.0e-3) + npt.assert_allclose(euler.location_, dipole_coordinates, rtol=1.0e-2) + npt.assert_allclose(euler.base_level_, true_base_level, rtol=1.0e-3) def test_euler_with_analytic_derivatives(): + "Check the Euler solution against a synthetic using analytical derivatives" # Add dipole source masses_coordinates = (10e3, 15e3, -10e3) masses = 1.0e12 diff --git a/harmonica/tests/test_filters.py b/harmonica/tests/test_filters.py index 78576ed78..60d71cda1 100644 --- a/harmonica/tests/test_filters.py +++ b/harmonica/tests/test_filters.py @@ -29,6 +29,10 @@ ) from ..filters._utils import apply_filter +# ------------------------------- +# Fixtures +# ------------------------------- + @pytest.fixture(name="region") def fixture_region(): @@ -48,21 +52,6 @@ def fixture_sample_grid(region): return make_xarray_grid((easting, northing), data, data_names="sample").sample -@pytest.fixture(name="sample_grid_upward") -def fixture_sample_grid_upward(region): - """ - Return a sample grid as an :class:`xarray.DataArray` with upward coord. - """ - easting, northing, upward = grid_coordinates(region, spacing=500, extra_coords=100) - data = np.sin(easting / 1e3) + np.cos(northing / 1e3) - return make_xarray_grid( - (easting, northing, upward), - data, - data_names="sample", - extra_coords_names="upward", - ).sample - - @pytest.fixture(name="sample_grid_multiple_coords") def fixture_sample_grid_multiple_coords(region): """ @@ -86,50 +75,76 @@ def fixture_sample_grid_multiple_coords(region): ) -def test_fft_round_trip(sample_grid): +@pytest.fixture(name="invalid_grid_single_dim") +def fixture_invalid_grid_single_dim(): """ - Test if the wrapped fft and ifft functions satisfy a round trip. + Return a sample grid with a single dimension. + + This fixture is meant to test if apply_filter raises an error on a grid + with a single dimension. """ - xrt.assert_allclose(sample_grid, ifft(fft(sample_grid))) + x = np.linspace(0, 10, 11) + y = x**2 + grid = xr.DataArray(y, coords={"x": x}, dims=("x",)) + return grid -def test_fft_round_trip_upward(sample_grid_upward): +@pytest.fixture(name="invalid_grid_3_dims") +def fixture_invalid_grid_3_dims(): """ - Test if the wrapped fft and ifft functions satisfy a round trip. + Return a sample grid with 3 dimensions. + + This fixture is meant to test if apply_filter raises an error on a grid + with 3 dimensions. """ - round_trip = ifft(fft(sample_grid_upward)) - assert "upward" not in round_trip - # Assert if both arrays are close enough, but dropping upward first - xrt.assert_allclose( - sample_grid_upward.drop("upward"), ifft(fft(sample_grid_upward)) + x = np.linspace(0, 10, 11) + y = np.linspace(-4, 4, 9) + z = np.linspace(20, 30, 5) + xx, yy, zz = np.meshgrid(x, y, z) + data = xx + yy + zz + grid = xr.DataArray(data, coords={"x": x, "y": y, "z": z}, dims=("y", "x", "z")) + return grid + + +@pytest.fixture(name="sample_fft_grid") +def fixture_sample_fft_grid(): + """ + Return a sample fft_grid to be used in test functions. + """ + domain = (-9e-4, 9e-4, -8e-4, 8e-4) + freq_easting, freq_northing = grid_coordinates(region=domain, spacing=8e-4) + dummy_fft = np.ones_like(freq_easting) + fft_grid = make_xarray_grid( + (freq_easting, freq_northing), + dummy_fft, + data_names=["sample_fft"], + dims=("freq_northing", "freq_easting"), ) + return fft_grid.sample_fft -def test_fft_no_drop_bad_coords(sample_grid_upward): +@pytest.fixture(name="invalid_grid_with_nans") +def fixture_invalid_grid_with_nans(sample_grid): """ - Check if no dropping bad coordinates raises ValueError on upward coord. + Return a sample grid with nans. - ``xrft`` complains when *bad coordinates* are present in the input array. - This test, along with the ``drop_bad_coordinates`` argument, should be - removed if ``xrft`` changes this behaviour + This fixture is meant to test if apply_filter raises an error on a grid + with a nans. """ - msg = "Please drop these coordinates" - with pytest.raises(ValueError, match=msg): - fft(sample_grid_upward, drop_bad_coords=False) + sample_grid[2, 4] = np.nan + return sample_grid -def test_fft_no_drop_bad_coords_multi(sample_grid_multiple_coords): - """ - Check if no dropping bad coordinates raises ValueError on multiple coords. +# ------------------------------- +# Tests for FFT functions +# ------------------------------- - This test should fail because ``xrft`` complains when *bad coordinates* are - present in the input array. - This test, along with the ``drop_bad_coordinates`` argument, should be - removed if ``xrft`` changes this behaviour + +def test_fft_round_trip(sample_grid): """ - msg = "Please drop these coordinates" - with pytest.raises(ValueError, match=msg): - fft(sample_grid_multiple_coords, drop_bad_coords=False) + Test if the wrapped fft and ifft functions satisfy a round trip. + """ + xrt.assert_allclose(sample_grid, ifft(fft(sample_grid))) # ------------------------------- @@ -148,7 +163,7 @@ def dummy_filter(fourier_transform): def test_apply_filter(sample_grid): """ - Test apply_filter function using the dummy_filter. + Test apply_filter function when the grid has no extra coordinates. """ # Apply the dummy filter filtered_grid = apply_filter(sample_grid, dummy_filter) @@ -157,35 +172,87 @@ def test_apply_filter(sample_grid): xrt.assert_allclose(filtered_grid, expected) -@pytest.fixture(name="invalid_grid_single_dim") -def fixture_invalid_grid_single_dim(): +def test_apply_filter_keep_coords(sample_grid_multiple_coords): """ - Return a sample grid with a single dimension. + Test apply_filter function on a grid with extra coordinates. + """ + # Apply the dummy filter + filtered_grid = apply_filter(sample_grid_multiple_coords, dummy_filter) + # Compare output with expected results + expected = sample_grid_multiple_coords * 0 + # Non-dimensional coordinates should have been dropped by apply_filter. + assert set(sample_grid_multiple_coords.coords) == set(filtered_grid.coords) + npt.assert_allclose(filtered_grid.values, expected.values) - This fixture is meant to test if apply_filter raises an error on a grid - with a single dimension. + +def test_apply_filter_drop_coords(sample_grid_multiple_coords): """ - x = np.linspace(0, 10, 11) - y = x**2 - grid = xr.DataArray(y, coords={"x": x}, dims=("x",)) - return grid + Test apply_filter function on a grid with extra coordinates. + """ + # Apply the dummy filter + filtered_grid = apply_filter( + sample_grid_multiple_coords, dummy_filter, drop_coords=True + ) + # Compare output with expected results + expected = sample_grid_multiple_coords * 0 + # Non-dimensional coordinates should have been dropped by apply_filter. + assert set(sample_grid_multiple_coords.coords) != set(filtered_grid.coords) + assert set(sample_grid_multiple_coords.dims) == set(filtered_grid.coords) + npt.assert_allclose(filtered_grid.values, expected.values) -@pytest.fixture(name="invalid_grid_3_dims") -def fixture_invalid_grid_3_dims(): +def test_apply_filter_no_padding(sample_grid): """ - Return a sample grid with 3 dimensions. + Test apply_filter function with padding off. + """ + # Apply the dummy filter + filtered_grid = apply_filter(sample_grid, dummy_filter, pad=False) + # Compare output with expected results + expected = sample_grid * 0 + xrt.assert_allclose(filtered_grid, expected) - This fixture is meant to test if apply_filter raises an error on a grid - with 3 dimensions. + +def test_apply_filter_pad_width(sample_grid): """ - x = np.linspace(0, 10, 11) - y = np.linspace(-4, 4, 9) - z = np.linspace(20, 30, 5) - xx, yy, zz = np.meshgrid(x, y, z) - data = xx + yy + zz - grid = xr.DataArray(data, coords={"x": x, "y": y, "z": z}, dims=("y", "x", "z")) - return grid + Test apply_filter function with custom padding width. + """ + # Apply the dummy filter + filtered_grid = apply_filter( + sample_grid, + dummy_filter, + pad_kwargs={"pad_width": {"easting": 20, "northing": 20}}, + ) + # Compare output with expected results + expected = sample_grid * 0 + xrt.assert_allclose(filtered_grid, expected) + + +def test_apply_filter_pad_mode(sample_grid): + """ + Test apply_filter function with custom padding method. + """ + # Apply the dummy filter + filtered_grid = apply_filter( + sample_grid, dummy_filter, pad_kwargs={"mode": "median"} + ) + # Compare output with expected results + expected = sample_grid * 0 + xrt.assert_allclose(filtered_grid, expected) + + +def test_apply_filter_pad_with_constant(sample_grid): + """ + Test apply_filter function with the constant padding method. + """ + # Apply the dummy filter + filtered_grid = apply_filter( + sample_grid, + dummy_filter, + pad_kwargs={"mode": "constant", "constant_values": 120}, + ) + # Compare output with expected results + expected = sample_grid * 0 + xrt.assert_allclose(filtered_grid, expected) def test_apply_filter_grid_single_dimension(invalid_grid_single_dim): @@ -206,18 +273,6 @@ def test_apply_filter_grid_three_dimensions(invalid_grid_3_dims): apply_filter(invalid_grid_3_dims, dummy_filter) -@pytest.fixture(name="invalid_grid_with_nans") -def fixture_invalid_grid_with_nans(sample_grid): - """ - Return a sample grid with nans. - - This fixture is meant to test if apply_filter raises an error on a grid - with a nans. - """ - sample_grid[2, 4] = np.nan - return sample_grid - - def test_apply_filter_grid_with_nans(invalid_grid_with_nans): """ Check if apply_filter raises error on grid with single dimension. @@ -239,27 +294,10 @@ def test_coordinate_rounding_fix(sample_grid): # ----------------------------- -# Test upward derivative filter +# Test filter functions # ----------------------------- -@pytest.fixture(name="sample_fft_grid") -def fixture_sample_fft_grid(): - """ - Return a sample fft_grid to be used in test functions. - """ - domain = (-9e-4, 9e-4, -8e-4, 8e-4) - freq_easting, freq_northing = grid_coordinates(region=domain, spacing=8e-4) - dummy_fft = np.ones_like(freq_easting) - fft_grid = make_xarray_grid( - (freq_easting, freq_northing), - dummy_fft, - data_names=["sample_fft"], - dims=("freq_northing", "freq_easting"), - ) - return fft_grid.sample_fft - - @pytest.mark.parametrize("order", [1, 2, 3]) def test_derivative_upward_kernel(sample_fft_grid, order): """ diff --git a/harmonica/tests/test_transformations.py b/harmonica/tests/test_transformations.py index 55d3ceaec..cab87b756 100644 --- a/harmonica/tests/test_transformations.py +++ b/harmonica/tests/test_transformations.py @@ -12,13 +12,18 @@ from pathlib import Path import numpy as np +import numpy.testing as npt import pytest import verde as vd import xarray as xr import xarray.testing as xrt -import xrft -from .. import point_gravity +from .. import ( + dipole_magnetic, + magnetic_angles_to_vec, + point_gravity, + total_field_anomaly, +) from .._transformations import ( _get_dataarray_coordinate, derivative_easting, @@ -58,7 +63,7 @@ def fixture_sample_grid_coords(): Define sample grid coordinates. """ grid_coords = vd.grid_coordinates( - region=(-150e3, 150e3, -150e3, 150e3), shape=(41, 41), extra_coords=0 + region=(-300e3, 300e3, -300e3, 300e3), spacing=5000, extra_coords=0 ) return grid_coords @@ -69,7 +74,7 @@ def fixture_upward_grid_coords(): Define upward grid coordinates. """ grid_coords = vd.grid_coordinates( - region=(-150e3, 150e3, -150e3, 150e3), shape=(41, 41), extra_coords=10e3 + region=(-300e3, 300e3, -300e3, 300e3), spacing=5000, extra_coords=10e3 ) return grid_coords @@ -202,6 +207,38 @@ def fixture_sample_g_ee(sample_grid_coords, sample_sources): return g_ee.g_ee +@pytest.fixture(name="sample_g_ez") +def fixture_sample_g_ez(sample_grid_coords, sample_sources): + """ + Return g_ez field of sample points on sample grid coords. + """ + points, masses = sample_sources + g_ez = point_gravity(sample_grid_coords, points, masses, field="g_ez") + g_ez = vd.make_xarray_grid( + sample_grid_coords, + g_ez, + data_names="g_ez", + extra_coords_names="upward", + ) + return g_ez.g_ez + + +@pytest.fixture(name="sample_g_nz") +def fixture_sample_g_nz(sample_grid_coords, sample_sources): + """ + Return g_nz field of sample points on sample grid coords. + """ + points, masses = sample_sources + g_nz = point_gravity(sample_grid_coords, points, masses, field="g_nz") + g_nz = vd.make_xarray_grid( + sample_grid_coords, + g_nz, + data_names="g_nz", + extra_coords_names="upward", + ) + return g_nz.g_nz + + @pytest.mark.parametrize( ("index", "expected_dimension"), [(1, "easting"), (0, "northing")] ) @@ -253,19 +290,7 @@ def test_derivative_upward(sample_potential, sample_g_z): """ Test derivative_upward function against the synthetic model. """ - # Pad the potential field grid to improve accuracy - pad_width = { - "easting": sample_potential.easting.size // 3, - "northing": sample_potential.northing.size // 3, - } - # need to drop upward coordinate (bug in xrft) - potential_padded = xrft.pad( - sample_potential.drop_vars("upward"), - pad_width=pad_width, - ) - # Calculate upward derivative and unpad it - derivative = derivative_upward(potential_padded) - derivative = xrft.unpad(derivative, pad_width) + derivative = derivative_upward(sample_potential) # Compare against g_up (trim the borders to ignore boundary effects) trim = 6 derivative = derivative[trim:-trim, trim:-trim] @@ -281,19 +306,8 @@ def test_derivative_upward_order2(sample_potential, sample_g_zz): Note: We omit the minus sign here because the second derivative is positive for both downward (negative) and upward (positive) derivatives. """ - # Pad the potential field grid to improve accuracy - pad_width = { - "easting": sample_potential.easting.size // 3, - "northing": sample_potential.northing.size // 3, - } - # need to drop upward coordinate (bug in xrft) - potential_padded = xrft.pad( - sample_potential.drop_vars("upward"), - pad_width=pad_width, - ) # Calculate second upward derivative and unpad it - second_deriv = derivative_upward(potential_padded, order=2) - second_deriv = xrft.unpad(second_deriv, pad_width) + second_deriv = derivative_upward(sample_potential, order=2) # Compare against g_zz (trim the borders to ignore boundary effects) trim = 6 second_deriv = second_deriv[trim:-trim, trim:-trim] @@ -347,19 +361,7 @@ def test_derivative_easting_fft(sample_potential, sample_g_e): """ Test derivative_easting function against the synthetic model using FFTs. """ - # Pad the potential field grid to improve accuracy - pad_width = { - "easting": sample_potential.easting.size // 3, - "northing": sample_potential.northing.size // 3, - } - # need to drop upward coordinate (bug in xrft) - potential_padded = xrft.pad( - sample_potential.drop_vars("upward"), - pad_width=pad_width, - ) - # Calculate easting derivative and unpad it - derivative = derivative_easting(potential_padded) - derivative = xrft.unpad(derivative, pad_width) + derivative = derivative_easting(sample_potential) # Compare against g_e (trim the borders to ignore boundary effects) trim = 6 derivative = derivative[trim:-trim, trim:-trim] @@ -372,19 +374,7 @@ def test_derivative_easting_order2(sample_potential, sample_g_ee): """ Test higher order of derivative_easting function against the sample grid. """ - # Pad the potential field grid to improve accuracy - pad_width = { - "easting": sample_potential.easting.size // 3, - "northing": sample_potential.northing.size // 3, - } - # need to drop upward coordinate (bug in xrft) - potential_padded = xrft.pad( - sample_potential.drop_vars("upward"), - pad_width=pad_width, - ) - # Calculate second easting derivative and unpad it - second_deriv = derivative_easting(potential_padded, order=2) - second_deriv = xrft.unpad(second_deriv, pad_width) + second_deriv = derivative_easting(sample_potential, order=2) # Compare against g_ee (trim the borders to ignore boundary effects) trim = 6 second_deriv = second_deriv[trim:-trim, trim:-trim] @@ -423,19 +413,7 @@ def test_derivative_northing(sample_potential, sample_g_n): """ Test derivative_northing function against the synthetic model. """ - # Pad the potential field grid to improve accuracy - pad_width = { - "easting": sample_potential.easting.size // 3, - "northing": sample_potential.northing.size // 3, - } - # need to drop upward coordinate (bug in xrft) - potential_padded = xrft.pad( - sample_potential.drop_vars("upward"), - pad_width=pad_width, - ) - # Calculate northing derivative and unpad it - derivative = derivative_northing(potential_padded) - derivative = xrft.unpad(derivative, pad_width) + derivative = derivative_northing(sample_potential) # Compare against g_n (trim the borders to ignore boundary effects) trim = 6 derivative = derivative[trim:-trim, trim:-trim] @@ -448,19 +426,7 @@ def test_derivative_northing_order2(sample_potential, sample_g_nn): """ Test higher order of derivative_northing function against the sample grid. """ - # Pad the potential field grid to improve accuracy - pad_width = { - "easting": sample_potential.easting.size // 3, - "northing": sample_potential.northing.size // 3, - } - # need to drop upward coordinate (bug in xrft) - potential_padded = xrft.pad( - sample_potential.drop_vars("upward"), - pad_width=pad_width, - ) - # Calculate second northing derivative and unpad it - second_deriv = derivative_northing(potential_padded, order=2) - second_deriv = xrft.unpad(second_deriv, pad_width) + second_deriv = derivative_northing(sample_potential, order=2) # Compare against g_nn (trim the borders to ignore boundary effects) trim = 6 second_deriv = second_deriv[trim:-trim, trim:-trim] @@ -475,24 +441,10 @@ def test_laplace_fft(sample_potential): We will use FFT computations only. """ - # Pad the potential field grid to improve accuracy - pad_width = { - "easting": sample_potential.easting.size // 3, - "northing": sample_potential.northing.size // 3, - } - # need to drop upward coordinate (bug in xrft) - potential_padded = xrft.pad( - sample_potential.drop_vars("upward"), - pad_width=pad_width, - ) - # Calculate second northing derivative and unpad it method = "fft" - second_deriv_ee = derivative_easting(potential_padded, order=2, method=method) - second_deriv_nn = derivative_northing(potential_padded, order=2, method=method) - second_deriv_zz = derivative_upward(potential_padded, order=2) - second_deriv_ee = xrft.unpad(second_deriv_ee, pad_width) - second_deriv_nn = xrft.unpad(second_deriv_nn, pad_width) - second_deriv_zz = xrft.unpad(second_deriv_zz, pad_width) + second_deriv_ee = derivative_easting(sample_potential, order=2, method=method) + second_deriv_nn = derivative_northing(sample_potential, order=2, method=method) + second_deriv_zz = derivative_upward(sample_potential, order=2) # Compare g_nn + g_ee against -g_zz (trim the borders to ignore boundary # effects) trim = 6 @@ -506,36 +458,97 @@ def test_upward_continuation(sample_g_z, sample_g_z_upward): """ Test upward_continuation function against the synthetic model. """ - # Pad the potential field grid to improve accuracy - pad_width = { - "easting": sample_g_z.easting.size // 3, - "northing": sample_g_z.northing.size // 3, - } - # need to drop upward coordinate (bug in xrft) - gravity_padded = xrft.pad( - sample_g_z.drop_vars("upward"), - pad_width=pad_width, - ) - # Calculate upward continuation and unpad it - continuation = upward_continuation(gravity_padded, 10e3) - continuation = xrft.unpad(continuation, pad_width) + continuation = upward_continuation(sample_g_z, 10e3) # Compare against g_z_upward (trim the borders to ignore boundary effects) trim = 6 continuation = continuation[trim:-trim, trim:-trim] g_z_upward = sample_g_z_upward[trim:-trim, trim:-trim] # Drop upward for comparison - g_z_upward = g_z_upward.drop("upward") + g_z_upward = g_z_upward.drop_vars("upward") xrt.assert_allclose(continuation, g_z_upward, atol=1e-8) -def test_reduction_to_pole(sample_potential): +def test_reduction_to_pole_remanent(): + """ + Test reduction_to_pole against an analytical solution with remanent magnetization. + """ + coordinates = vd.grid_coordinates( + (-70e3, 20e3, -20e3, 60e3), spacing=0.5e3, extra_coords=500 + ) + finc, fdec = -45, 13 + minc, mdec = -14, -24 + dipole = [-25e3, 20e3, -5000] + moment = 1e12 + magnetic_field_pole = dipole_magnetic( + coordinates, + dipoles=dipole, + magnetic_moments=magnetic_angles_to_vec(moment, 90, 0), + field="b", + ) + anomaly_pole = total_field_anomaly(magnetic_field_pole, 90, 0) + magnetic_field = dipole_magnetic( + coordinates, + dipoles=dipole, + magnetic_moments=magnetic_angles_to_vec(moment, minc, mdec), + field="b", + ) + anomaly = total_field_anomaly(magnetic_field, finc, fdec) + grid = vd.make_xarray_grid(coordinates[:2], anomaly, data_names="anomaly") + anomaly_reduced = reduction_to_pole(grid.anomaly, finc, fdec, minc, mdec) + # Relative tol doesn't work because the anomaly at the pole is zero in + # a ring around the source and the rtol blows up at those points. + np.testing.assert_allclose( + anomaly_reduced.values, + anomaly_pole, + rtol=0, + atol=0.01 * np.abs(anomaly_pole).max(), + ) + + +def test_reduction_to_pole_induced(): + """ + Test reduction_to_pole against an analytical solution with induced magnetization. + """ + coordinates = vd.grid_coordinates( + (-70e3, 20e3, -20e3, 60e3), spacing=0.5e3, extra_coords=500 + ) + finc, fdec = -45, 13 + dipole = [-25e3, 20e3, -5000] + moment = 1e12 + magnetic_field_pole = dipole_magnetic( + coordinates, + dipoles=dipole, + magnetic_moments=magnetic_angles_to_vec(moment, 90, 0), + field="b", + ) + anomaly_pole = total_field_anomaly(magnetic_field_pole, 90, 0) + magnetic_field = dipole_magnetic( + coordinates, + dipoles=dipole, + magnetic_moments=magnetic_angles_to_vec(moment, finc, fdec), + field="b", + ) + anomaly = total_field_anomaly(magnetic_field, finc, fdec) + grid = vd.make_xarray_grid(coordinates[:2], anomaly, data_names="anomaly") + anomaly_reduced = reduction_to_pole(grid.anomaly, finc, fdec) + # Relative tol doesn't work because the anomaly at the pole is zero in + # a ring around the source and the rtol blows up at those points. + np.testing.assert_allclose( + anomaly_reduced.values, + anomaly_pole, + rtol=0, + atol=0.01 * np.abs(anomaly_pole).max(), + ) + + +def test_reduction_to_pole_dim_names(sample_potential): """ Test reduction_to_pole function with non-typical dim names. """ renamed_dims_grid = sample_potential.rename( {"easting": "name_one", "northing": "name_two"} ) - reduction_to_pole(renamed_dims_grid, 60, 45) + reduction_to_pole(renamed_dims_grid, 60, 45, 60, 45) class TestTotalGradientAmplitude: @@ -549,19 +562,7 @@ def test_against_synthetic( """ Test total_gradient_amplitude function against the synthetic model. """ - # Pad the potential field grid to improve accuracy - pad_width = { - "easting": sample_potential.easting.size // 3, - "northing": sample_potential.northing.size // 3, - } - # need to drop upward coordinate (bug in xrft) - potential_padded = xrft.pad( - sample_potential.drop_vars("upward"), - pad_width=pad_width, - ) - # Calculate total gradient amplitude and unpad it - tga = total_gradient_amplitude(potential_padded) - tga = xrft.unpad(tga, pad_width) + tga = total_gradient_amplitude(sample_potential) # Compare against g_tga (trim the borders to ignore boundary effects) trim = 6 tga = tga[trim:-trim, trim:-trim] @@ -613,37 +614,19 @@ class TestTilt: Test tilt function. """ - def test_against_synthetic( - self, sample_potential, sample_g_n, sample_g_e, sample_g_z - ): + def test_against_synthetic(self, sample_g_z, sample_g_nz, sample_g_ez, sample_g_zz): """ Test tilt function against the synthetic model. """ - # Pad the potential field grid to improve accuracy - pad_width = { - "easting": sample_potential.easting.size // 3, - "northing": sample_potential.northing.size // 3, - } - # need to drop upward coordinate (bug in xrft) - potential_padded = xrft.pad( - sample_potential.drop_vars("upward"), - pad_width=pad_width, - ) - # Calculate the tilt and unpad it - tilt_grid = tilt_angle(potential_padded) - tilt_grid = xrft.unpad(tilt_grid, pad_width) - # Compare against g_tilt (trim the borders to ignore boundary effects) - trim = 6 - tilt_grid = tilt_grid[trim:-trim, trim:-trim] - g_e = sample_g_e[trim:-trim, trim:-trim] - g_n = sample_g_n[trim:-trim, trim:-trim] - g_z = sample_g_z[trim:-trim, trim:-trim] - g_horiz_deriv = np.sqrt(g_e**2 + g_n**2) - g_tilt = np.arctan2( - -g_z, g_horiz_deriv - ) # use -g_z to use the _upward_ derivative - rms = root_mean_square_error(tilt_grid, g_tilt) - assert rms / np.abs(tilt_grid).max() < 0.1 + # Use the gz because the potential is too smooth and messes up at the edges + numerical = tilt_angle(sample_g_z) + # Use minus to use the upward derivative (z is downward) + analytical = np.arctan2(-sample_g_zz, np.sqrt(sample_g_ez**2 + sample_g_nz**2)) + # Compare only an inner region because the tilt is terrible at the edges + crop = [-125e3, 125e3, -125e3, 125e3] + numerical = numerical.sel(easting=slice(*crop[:2]), northing=slice(*crop[2:])) + analytical = analytical.sel(easting=slice(*crop[:2]), northing=slice(*crop[2:])) + np.testing.assert_allclose(numerical, analytical, atol=0.15) def test_invalid_grid_single_dimension(self): """ @@ -681,35 +664,69 @@ def test_invalid_grid_with_nans(self, sample_potential): tilt_angle(sample_potential) -class Testfilter: +class TestGaussianFilters: """ - Test filter result against the output from oasis montaj. + Test the Gaussian low and high pass filters against a synthetic. """ - expected_grid = xr.open_dataset(TEST_DATA_DIR / "filter.nc") - - def test_gaussian_lowpass_grid(self): - """ - Test gaussian_lowpass function against the output from oasis montaj. - """ - low_pass = gaussian_lowpass(self.expected_grid.filter_data, 10) - xrt.assert_allclose(self.expected_grid.filter_lp10, low_pass, atol=1e-6) - - def test_gaussian_highpass_grid(self): + def test_gaussian_lowpass(self): """ - Test gaussian_highpass function against the output from oasis montaj. + Test gaussian_lowpass against a synthetic. """ - high_pass = gaussian_highpass(self.expected_grid.filter_data, 10) - xrt.assert_allclose(self.expected_grid.filter_hp10, high_pass, atol=1e-6) + # Make a synthetic with only 2 known wavelengths + coordinates = vd.grid_coordinates((0, 10, 0, 10), spacing=0.05) + wavelength_high = 0.5 + wavelength_low = 10 + component_low = ( + 5 + * np.sin(2 * np.pi / wavelength_low * coordinates[0]) + * np.cos(2 * np.pi / wavelength_low * coordinates[1]) + ) + component_high = np.cos(2 * np.pi / wavelength_high * coordinates[0]) * np.sin( + 2 * np.pi / wavelength_high * coordinates[1] + ) + grid_low = xr.DataArray( + component_low, + coords={"x": coordinates[0][0, :], "y": coordinates[1][:, 0]}, + dims=("y", "x"), + ) + grid_high = xr.DataArray( + component_high, + coords={"x": coordinates[0][0, :], "y": coordinates[1][:, 0]}, + dims=("y", "x"), + ) + grid = grid_low + grid_high + # Try to isolate the long wavelength component + low = gaussian_lowpass(grid, wavelength=1) + # This is about a 10% error tolerance + npt.assert_allclose(low, grid_low, atol=0.4) - def test_reduction_to_pole_grid(self): + def test_gaussian_highpass(self): """ - Test reduction_to_pole function against the output from oasis montaj. + Test gaussian_highpass against a synthetic. """ - rtp = reduction_to_pole(self.expected_grid.filter_data, 60, 45) - # Remove mean value to match OM result - xrt.assert_allclose( - self.expected_grid.filter_rtp - self.expected_grid.filter_data.mean(), - rtp, - atol=1, + # Make a synthetic with only 2 known wavelengths + coordinates = vd.grid_coordinates((0, 10, 0, 10), spacing=0.05) + wavelength_high = 0.5 + wavelength_low = 10 + component_low = np.sin(2 * np.pi / wavelength_low * coordinates[0]) * np.cos( + 2 * np.pi / wavelength_low * coordinates[1] + ) + component_high = np.cos(2 * np.pi / wavelength_high * coordinates[0]) * np.sin( + 2 * np.pi / wavelength_high * coordinates[1] + ) + grid_low = xr.DataArray( + component_low, + coords={"x": coordinates[0][0, :], "y": coordinates[1][:, 0]}, + dims=("y", "x"), + ) + grid_high = xr.DataArray( + component_high, + coords={"x": coordinates[0][0, :], "y": coordinates[1][:, 0]}, + dims=("y", "x"), ) + grid = grid_low + grid_high + # Try to isolate the short wavelength component + high = gaussian_highpass(grid, wavelength=2) + # This is about a 10% error tolerance + npt.assert_allclose(high, grid_high, atol=0.14) diff --git a/pyproject.toml b/pyproject.toml index abc564de9..b85268a2d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -118,6 +118,7 @@ ignore = [ "RET504", # allow variable assignment only for return "PT001", # conventions for parenthesis on pytest.fixture "D200", # allow single line docstrings in their own line + "C420", # Replacing dict comprehension with `dict.fromkeys` is ugly ] [tool.ruff.lint.per-file-ignores]