Refactor normal gravity function (#22)

Made some corrections to the docstring and tweaks to the example plot
(put the colorbar at the bottom and use pcolormesh instead of contourf
for faster plotting).

While I was at it, I moved the calculations back into a single function.
It was actually really hard to read the code with the two auxiliary
functions. Plus, there was no way of testing these functions
individually so it made no difference having them separated. In the end,
it's better to have more readable code than pleasing pylint. I added an
exception to pylint regarding this function.

I also changed the module name to `gravity_corrections.py` so we can
include Bouguer etc in there as well. These are really small functions
so there is no need to have them in separate modules.

doc/api/index.rst

@@ -7,13 +7,13 @@ API Reference
 
 .. currentmodule:: harmonica
 
-Input and Output
-----------------
+Gravity Corrections
+-------------------
 
 .. autosummary::
- :toctree: generated/
+  :toctree: generated/
 
- load_icgem_gdf
+ normal_gravity
 
 Reference Ellipsoids
 --------------------
@@ -26,13 +26,13 @@ Reference Ellipsoids
  get_ellipsoid
  print_ellipsoids
 
-Normal Gravity
---------------
+Input and Output
+----------------
 
 .. autosummary::
-  :toctree: generated/
+ :toctree: generated/
 
- normal_gravity
+ load_icgem_gdf
 
 Utilities
 ---------

examples/normal_gravity.py

@@ -1,38 +1,34 @@
 """
-Normal gravity calculation
-==========================
+Normal Gravity
+==============
 
-Gravitational disturbances allows geoscientists to infer the internal structure
-of the Earth. In order to compute a gravitational disturbance, we must define
-a normal gravity.
-Usually, the normal gravity is defined as the gravitational effect of the
-reference ellipsoid. [LiGotze2001]_ introduced a closed-form formula to compute
-the magntiude of the gradient of the gravitational potential generated by any
-reference ellipsoid.
-The :func:`harmonica.normal_gravity` function uses this closed-form formula to
-compute the gravitational effect of the reference ellipsoid given by
+Normal gravity is defined as the magnitude of the gradient of the gravity potential
+(gravitational + centrifugal) of a reference ellipsoid. Function
+:func:`harmonica.normal_gravity` implements a closed form solution [LiGotze2001]_ to
+calculate normal gravity at any latitude and height (it's symmetric in the longitudinal
+direction). The ellipsoid in the calculations used can be changed using
 :func:`harmonica.set_ellipsoid`.
 """
-import numpy as np
 import matplotlib.pyplot as plt
 import cartopy.crs as ccrs
 import harmonica as hm
 import verde as vd
 
-# Create a computation grid
+# Create a global computation grid with 1 degree spacing
 region = [0, 360, -90, 90]
-lon, lat = vd.grid_coordinates(region, spacing=1.0)
+longitude, latitude = vd.grid_coordinates(region, spacing=1)
 
-# Compute the gravitational effect of the WGS84 (default) ellipsoid
-# on its surface
-gamma = hm.normal_gravity(latitude=lat, height=0)
+# Compute normal gravity for the WGS84 ellipsoid (the default) on its surface
+gamma = hm.normal_gravity(latitude=latitude, height=0)
 
 # Make a plot of the normal gravity using Cartopy
-plt.figure(figsize=(6, 3))
+plt.figure(figsize=(9, 5))
 ax = plt.axes(projection=ccrs.Mollweide())
-ax.set_title("Normal gravity on the WGS84 ellipsoid")
+ax.set_title("Normal gravity of the WGS84 ellipsoid")
 ax.coastlines()
-plt.contourf(lon, lat, gamma, 150, transform=ccrs.PlateCarree())
-plt.colorbar(label="mGal")
+pc = ax.pcolormesh(longitude, latitude, gamma, transform=ccrs.PlateCarree())
+plt.colorbar(
+ pc, label="mGal", orientation="horizontal", aspect=50, pad=0.01, shrink=0.7
+)
 plt.tight_layout()
 plt.show()

harmonica/__init__.py

@@ -8,7 +8,7 @@
  ReferenceEllipsoid,
 )
 from .io import load_icgem_gdf
-from .normal_gravity import normal_gravity
+from .gravity_corrections import normal_gravity
 
 # Get the version number through versioneer
 __version__ = version.full_version

harmonica/gravity_corrections.py

@@ -0,0 +1,75 @@
+"""
+Gravity corrections like Normal Gravity and Bouguer corrections.
+"""
+import numpy as np
+
+from .ellipsoid import get_ellipsoid
+
+
+def normal_gravity(latitude, height): # pylint: disable=too-many-locals
+ """
+ Calculate normal gravity at any latitude and height.
+
+ Computes the magnitude of the gradient of the gravity potential (gravitational +
+ centrifugal) generated by a reference ellipsoid at the given latitude and
+ (geometric) height. Uses of a closed form expression [LiGotze2001]_ and the current
+ reference ellipsoid defined by :func:`harmonica.set_ellipsoid`.
+
+ Parameters
+ ----------
+ latitude : float or array
+ The (geodetic) latitude where the normal gravity will be computed (in degrees)
+ height : float or array
+ The ellipsoidal (geometric) height of computation point (in meters).
+
+ Returns
+ -------
+ gamma : float or array
+ The computed normal gravity (in mGal).
+
+ """
+ ell = get_ellipsoid()
+
+ latitude_radians = np.deg2rad(latitude)
+ sinlat = np.sin(latitude_radians)
+ coslat = np.cos(latitude_radians)
+
+ # The terms below follow the variable names in the Li and Goetze (2001) paper
+ beta = np.arctan2(ell.semiminor_axis * sinlat / coslat, ell.semimajor_axis)
+ zl2 = (ell.semiminor_axis * np.sin(beta) + height * sinlat) ** 2
+ rl2 = (ell.semimajor_axis * np.cos(beta) + height * coslat) ** 2
+ big_d = (rl2 - zl2) / ell.linear_eccentricity ** 2
+ big_r = (rl2 + zl2) / ell.linear_eccentricity ** 2
+ cosbeta_l2 = 0.5 * (1 + big_r) - np.sqrt(0.25 * (1 + big_r ** 2) - 0.5 * big_d)
+ sinbeta_l2 = 1 - cosbeta_l2
+ b_l = np.sqrt(rl2 + zl2 - ell.linear_eccentricity ** 2 * cosbeta_l2)
+ q_0 = 0.5 * (
+ (1 + 3 * (ell.semiminor_axis / ell.linear_eccentricity) ** 2)
+ * np.arctan2(ell.linear_eccentricity, ell.semiminor_axis)
+ - 3 * ell.semiminor_axis / ell.linear_eccentricity
+ )
+ q_l = (
+ 3
+ * (1 + (b_l / ell.linear_eccentricity) ** 2)
+ * (1 - b_l / ell.linear_eccentricity * np.arctan2(ell.linear_eccentricity, b_l))
+ - 1
+ )
+ big_w = np.sqrt(
+ (b_l ** 2 + ell.linear_eccentricity ** 2 * sinbeta_l2)
+ / (b_l ** 2 + ell.linear_eccentricity ** 2)
+ )
+
+ # Put together gamma using 3 terms
+ term1 = ell.geocentric_grav_const / (b_l ** 2 + ell.linear_eccentricity ** 2)
+ term2 = (0.5 * sinbeta_l2 - 1 / 6) * (
+ ell.semimajor_axis ** 2
+ * ell.linear_eccentricity
+ * q_l
+ * ell.angular_velocity ** 2
+ / ((b_l ** 2 + ell.linear_eccentricity ** 2) * q_0)
+ )
+ term3 = -cosbeta_l2 * b_l * ell.angular_velocity ** 2
+ gamma = (term1 + term2 + term3) / big_w
+
+ # Convert gamma from SI to mGal
+ return gamma * 1e5

harmonica/tests/test_normal_gravity.py

@@ -4,9 +4,8 @@
 import numpy as np
 import numpy.testing as npt
 
-from .. import normal_gravity
-from .. import set_ellipsoid, get_ellipsoid
-from ..ellipsoid import KNOWN_ELLIPSOIDS
+from ..gravity_corrections import normal_gravity
+from ..ellipsoid import KNOWN_ELLIPSOIDS, set_ellipsoid, get_ellipsoid
 
 
 def test_normal_gravity():
@@ -18,10 +17,8 @@ def test_normal_gravity():
  # Convert gamma to mGal
  gamma_pole = get_ellipsoid().gravity_pole * 1e5
  gamma_eq = get_ellipsoid().gravity_equator * 1e5
- npt.assert_allclose(gamma_pole, normal_gravity(-90, height),
- rtol=rtol)
- npt.assert_allclose(gamma_pole, normal_gravity(90, height),
- rtol=rtol)
+ npt.assert_allclose(gamma_pole, normal_gravity(-90, height), rtol=rtol)
+ npt.assert_allclose(gamma_pole, normal_gravity(90, height), rtol=rtol)
  npt.assert_allclose(gamma_eq, normal_gravity(0, height), rtol=rtol)
 
 
@@ -32,13 +29,16 @@ def test_normal_gravity_arrays():
  latitudes = np.array([-90, 90, 0])
  for ellipsoid_name in KNOWN_ELLIPSOIDS:
  with set_ellipsoid(ellipsoid_name):
- gammas = np.array([get_ellipsoid().gravity_pole,
- get_ellipsoid().gravity_pole,
- get_ellipsoid().gravity_equator])
+ gammas = np.array(
+ [
+ get_ellipsoid().gravity_pole,
+ get_ellipsoid().gravity_pole,
+ get_ellipsoid().gravity_equator,
+ ]
+ )
  # Convert gammas to mGal
  gammas *= 1e5
- npt.assert_allclose(gammas, normal_gravity(latitudes, heights),
- rtol=rtol)
+ npt.assert_allclose(gammas, normal_gravity(latitudes, heights), rtol=rtol)
 
 
 def test_no_zero_height():