Repository: fatiando/harmonica
Pull request #22 (2018-11-07 01:52:51)

Refactor normal gravity 
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.

**Reminders**

- [x] Run `make format` and `make check` to make sure the code follows the style guide.
- [x] Add tests for new features or tests that would have caught the bug that you're fixing.
- [x] Add new public functions/methods/classes to `doc/api/index.rst`.
- [x] Write detailed docstrings for all functions/methods.
- [x] If adding new functionality, add an example to the docstring, gallery, and/or tutorials.


leouieda: > I'm a little bit confused about priorities: code readability, easier testing, good styling.

Ideally they are all priorities. This is the real challenge of software development. The coding itself is not that hard. But coming up with a good design that can be well tested and understood is really hard.

> These refactors give me an idea how to improve future PRs.

That's why I've been asking for you reviews :) Actually, @aguspesce might want to keep an eye on these things as well.

Also, please feel free to disagree with my design choices. They're not always good choices and I've come to regret many made in the original fatiando repo (or even in the recent Verde).

Patch:
diff --git a/doc/api/index.rst b/doc/api/index.rst
index 714a726..7bd7a58 100644
--- a/doc/api/index.rst
+++ b/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
 ---------
diff --git a/examples/normal_gravity.py b/examples/normal_gravity.py
index e37aa89..43c063c 100644
--- a/examples/normal_gravity.py
+++ b/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()
diff --git a/harmonica/__init__.py b/harmonica/__init__.py
index 4854318..546a693 100644
--- a/harmonica/__init__.py
+++ b/harmonica/__init__.py
@@ -8,7 +8,7 @@ from .ellipsoid import (
     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
diff --git a/harmonica/gravity_corrections.py b/harmonica/gravity_corrections.py
new file mode 100644
index 0000000..8bfc5b7
--- /dev/null
+++ b/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
diff --git a/harmonica/normal_gravity.py b/harmonica/normal_gravity.py
deleted file mode 100644
index c4b7520..0000000
--- a/harmonica/normal_gravity.py
+++ /dev/null
@@ -1,75 +0,0 @@
-"""
-Gravity of ellipsoids models.
-"""
-import numpy as np
-from .ellipsoid import get_ellipsoid
-
-
-def normal_gravity(latitude, height):
-    """
-    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 height.
-    ``height``.
-    It makes use of the closed form expression of [LiGotze2001]_ and the
-    current reference ellipsoid defined by :func:`harmonica.set_ellipsoid`.
-
-    Parameters
-    ----------
-    latitude : float or numpy array
-        The latitude where the normal gravity will be computed (in degrees)
-    height : float or numpy array
-        The ellipsoidal (geometric) height of computation point (in meters).
-        (geometric) heights for geophysical purposes.
-
-    Returns
-    -------
-    gamma : float or numpy array
-        The computed normal gravity (in mGal).
-
-    """
-    ellipsoid = get_ellipsoid()
-    ecc = ellipsoid.linear_eccentricity
-    omega2 = ellipsoid.angular_velocity**2
-
-    big_w, q_l, q_0, b_l, cosbeta_l2, sinbeta_l2 = _normal_grav_factors(
-        ellipsoid, latitude, height)
-    term1 = ellipsoid.geocentric_grav_const / (b_l**2 + ecc**2)
-    term2 = 0.5 * sinbeta_l2 - 1 / 6
-    term2 *= ellipsoid.semimajor_axis**2 * ecc * q_l * omega2 / \
-        ((b_l**2 + ecc**2) * q_0)
-    term3 = -cosbeta_l2 * b_l * omega2
-    gamma = (term1 + term2 + term3) / big_w
-    # Convert gamma from SI to mGal
-    gamma *= 1e5
-    return gamma
-
-
-def _normal_grav_factors(ellipsoid, latitude, height):
-    ecc = ellipsoid.linear_eccentricity
-    semiminor = ellipsoid.semiminor_axis
-    rl2, zl2 = _normal_grav_r_z(ellipsoid, latitude, height)
-    big_d = (rl2 - zl2) / ecc**2
-    big_r = (rl2 + zl2) / ecc**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 - ecc**2 * cosbeta_l2)
-    q_0 = 0.5 * ((1 + 3 * (semiminor / ecc)**2) * np.arctan2(ecc, semiminor)
-                 - 3 * semiminor / ecc)
-    q_l = 3 * (1 + (b_l / ecc)**2) * (1 - b_l / ecc * np.arctan2(ecc, b_l)) - 1
-    big_w = np.sqrt((b_l**2 + ecc**2 * sinbeta_l2) / (b_l**2 + ecc**2))
-    return big_w, q_l, q_0, b_l, cosbeta_l2, sinbeta_l2
-
-
-def _normal_grav_r_z(ellipsoid, latitude, height):
-    semimajor = ellipsoid.semimajor_axis
-    semiminor = ellipsoid.semiminor_axis
-    lat = np.deg2rad(latitude)
-    coslat, sinlat = np.cos(lat), np.sin(lat)
-    beta = np.arctan2(semiminor * sinlat / coslat, semimajor)
-    cosbeta, sinbeta = np.cos(beta), np.sin(beta)
-    zl2 = (semiminor * sinbeta + height * sinlat)**2
-    rl2 = (semimajor * cosbeta + height * coslat)**2
-    return rl2, zl2


Test patch:
diff --git a/harmonica/tests/test_normal_gravity.py b/harmonica/tests/test_normal_gravity.py
index b23e81d..f9b739e 100644
--- a/harmonica/tests/test_normal_gravity.py
+++ b/harmonica/tests/test_normal_gravity.py
@@ -4,9 +4,8 @@ Testing normal gravity calculation.
 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():