Add WDMAM dataset to load_earth_magnetic_anomaly

pygmt/datasets/earth_magnetic_anomaly.py

@@ -5,33 +5,36 @@
 The grids are available in various resolutions.
 """
 from pygmt.datasets.load_remote_dataset import _load_remote_dataset
+from pygmt.exceptions import GMTInvalidInput
 from pygmt.helpers import kwargs_to_strings
 
 __doctest_skip__ = ["load_earth_magnetic_anomaly"]
 
 
 @kwargs_to_strings(region="sequence")
 def load_earth_magnetic_anomaly(
-    resolution="01d", region=None, registration=None, mag4km=False
+    resolution="01d", region=None, registration=None, data_source="emag2"
 ):
     r"""
     Load an Earth magnetic anomaly grid in various resolutions.
 
     The grids are downloaded to a user data directory
-    (usually ``~/.gmt/server/earth/earth_mag/`` or
-    ``~/.gmt/server/earth/earth_mag4km/``) the first time you invoke
+    (usually ``~/.gmt/server/earth/earth_mag/``,
+    ``~/.gmt/server/earth/earth_mag4km/``,
+    or ``~/.gmt/server/earth/earth_wdmam/``) the first time you invoke
     this function. Afterwards, it will load the grid from the data directory.
     So you'll need an internet connection the first time around.
 
     These grids can also be accessed by passing in the file name
     **@**\ *earth_mag_type*\_\ *res*\[_\ *reg*] to any grid plotting/processing
     function. *earth_mag_type* is the GMT name
-    for the dataset. The available options are **earth_mag** and
-    **earth_mag4km**. *res* is the grid resolution (see below), and *reg* is
-    grid registration type (**p** for pixel registration or **g** for gridline
-    registration).
+    for the dataset. The available options are **earth_mag**,
+    **earth_mag4km**, and **earth_wdmam**. *res* is the grid resolution
+    (see below), and *reg* is grid registration type (**p** for pixel
+    registration or **g** for gridline registration).
 
-    Refer to :gmt-datasets:`earth-mag.html` for more details.
+    Refer to :gmt-datasets:`earth-mag.html`
+    and :gmt-datasets:`earth-wdmam.html` for more details.
 
     Parameters
     ----------
@@ -52,12 +55,21 @@ def load_earth_magnetic_anomaly(
         ``"gridline"`` for gridline registration. Default is ``"gridline"``
         for all resolutions except ``"02m"`` which is ``"pixel"`` only.
 
-    mag4km : bool
-        Choose the data version to use. The default is ``False``, which is
-        observed at sea level over oceanic regions and has no data over land.
-        Set ``mag4km`` to ``True`` to use a version where all observations
-        are relative to an altitude of 4 km above the geoid and include data
-        over land.
+    data_source : str
+        Select the source of the magnetic anomaly data.
+
+        Available options:
+
+        - **emag2** : EMAG2 Global Earth Magnetic Anomaly Model [Default
+          option]. Only includes data is observed from sea level over
+          oceanic regions. See :gmt-datasets:`earth-mag.html`.
+
+        - **emag2_4km** : Use a version of EMAG2 where all observations
+          are relative to an altitude of 4 km above the geoid and include
+          data over land.
+
+        - **wdmam** : World Digital Magnetic Anomaly Map (WDMAM).
+          See :gmt-datasets:`earth-wdmam.html`
 
     Returns
     -------
@@ -87,14 +99,32 @@ def load_earth_magnetic_anomaly(
     ...     region=[120, 160, 30, 60],
     ...     registration="gridline",
     ... )
-    >>> # load the 20 arc-minutes grid of the mag4km dataset
+    >>> # load the 20 arc-minutes grid of the emag2_4km dataset
+    >>> grid = load_earth_magnetic_anomaly(
+    ...     resolution="20m", registration="gridline", data_source="emag2_4km"
+    ... )
+    >>> # load the 20 arc-minutes grid of the WDMAM dataset
     >>> grid = load_earth_magnetic_anomaly(
-    ...     resolution="20m", registration="gridline", mag4km=True
+    ...     resolution="20m", registration="gridline", data_source="wdmam"
     ... )
     """
-    dataset_prefix = "earth_mag4km_" if mag4km is True else "earth_mag_"
+    magnetic_anomaly_sources = {
+        "emag2": "earth_mag_",
+        "emag2_4km": "earth_mag4km_",
+        "wdmam": "earth_wdmam_",
+    }
+    if data_source not in magnetic_anomaly_sources:
+        raise GMTInvalidInput(
+            f"Invalid earth magnetic anomaly 'data_source' {data_source}, "
+            "valid values are 'emag2', 'emag2_4km', and 'wdmam'."
+        )
+    dataset_prefix = magnetic_anomaly_sources[data_source]
+    if data_source == "wdmam":
+        dataset_name = "earth_wdmam"
+    else:
+        dataset_name = "earth_magnetic_anomaly"
     grid = _load_remote_dataset(
-        dataset_name="earth_magnetic_anomaly",
+        dataset_name=dataset_name,
         dataset_prefix=dataset_prefix,
         resolution=resolution,
         region=region,

pygmt/datasets/load_remote_dataset.py

@@ -187,6 +187,24 @@ class GMTRemoteDataset(NamedTuple):
             "01m": Resolution(["pixel"], True),
         },
     ),
+    "earth_wdmam": GMTRemoteDataset(
+        title="WDMAM magnetic anomaly",
+        name="wdmam",
+        long_name="World Digital Magnetic Anomaly Map",
+        units="nT",
+        extra_attributes={"horizontal_datum": "WGS84"},
+        resolutions={
+            "01d": Resolution(["gridline", "pixel"], False),
+            "30m": Resolution(["gridline", "pixel"], False),
+            "20m": Resolution(["gridline", "pixel"], False),
+            "15m": Resolution(["gridline", "pixel"], False),
+            "10m": Resolution(["gridline", "pixel"], False),
+            "06m": Resolution(["gridline", "pixel"], False),
+            "05m": Resolution(["gridline", "pixel"], True),
+            "04m": Resolution(["gridline", "pixel"], True),
+            "03m": Resolution(["gridline"], True),
+        },
+    ),
 }
 
 

pygmt/helpers/testing.py

@@ -195,6 +195,9 @@ def download_test_data():
         # Earth vertical gravity gradient grids
         "@earth_vgg_01d_g",
         "@N00W030.earth_vgg_01m_p.nc",  # Specific grid for 01m test
+        # Earth WDMAM grids
+        "@earth_wdmam_01d_g",
+        "@S90E000.earth_wdmam_03m_g.nc",  # Specific grid for 03m test
         # Other cache files
         "@capitals.gmt",
         "@earth_relief_20m_holes.grd",

pygmt/tests/test_datasets_earth_magnetic_anomaly.py

@@ -68,20 +68,23 @@ def test_earth_mag_02m_without_region():
 
 def test_earth_mag_incorrect_resolution_registration():
     """
-    Test that an error is raised when trying to load a grid registration with
-    an unavailable resolution.
+    Test that an error is raised when trying to load a EMAG2 grid registration
+    with an unavailable resolution.
     """
     with pytest.raises(GMTInvalidInput):
         load_earth_magnetic_anomaly(
-            resolution="02m", region=[0, 1, 3, 5], registration="gridline", mag4km=False
+            resolution="02m",
+            region=[0, 1, 3, 5],
+            registration="gridline",
+            data_source="emag2_4km",
         )
 
 
 def test_earth_mag4km_01d():
     """
     Test some properties of the magnetic anomaly 4km 01d data.
     """
-    data = load_earth_magnetic_anomaly(resolution="01d", mag4km=True)
+    data = load_earth_magnetic_anomaly(resolution="01d", data_source="emag2_4km")
     assert data.name == "magnetic_anomaly"
     assert data.attrs["long_name"] == "Earth magnetic anomaly"
     assert data.attrs["units"] == "nT"
@@ -102,7 +105,7 @@ def test_earth_mag4km_01d_with_region():
         resolution="01d",
         region=[-10, 10, -5, 5],
         registration="gridline",
-        mag4km=True,
+        data_source="emag2_4km",
     )
     assert data.shape == (11, 21)
     npt.assert_allclose(data.lat, np.arange(-5, 6, 1))
@@ -125,3 +128,102 @@ def test_earth_mag_02m_default_registration():
     npt.assert_allclose(data.coords["lon"].data.max(), -9.01666667)
     npt.assert_allclose(data.min(), -231)
     npt.assert_allclose(data.max(), 131.79999)
+
+    data = load_earth_magnetic_anomaly(
+        resolution="05m",
+        region=[-115, -112, 4, 6],
+        registration="gridline",
+        data_source="emag2_4km",
+    )
+    assert data.shape == (25, 37)
+    assert data.lat.min() == 4
+    assert data.lat.max() == 6
+    assert data.lon.min() == -115
+    assert data.lon.max() == -112
+    npt.assert_allclose(data.min(), -128.40015)
+    npt.assert_allclose(data.max(), 76.80005)
+
+
+def test_earth_mag_01d_wdmam():
+    """
+    Test some properties of the WDMAM 01d data.
+    """
+    data = load_earth_magnetic_anomaly(
+        resolution="01d", registration="gridline", data_source="wdmam"
+    )
+    assert data.name == "wdmam"
+    assert data.attrs["long_name"] == "World Digital Magnetic Anomaly Map"
+    assert data.attrs["units"] == "nT"
+    assert data.attrs["horizontal_datum"] == "WGS84"
+    assert data.shape == (181, 361)
+    npt.assert_allclose(data.lat, np.arange(-90, 91, 1))
+    npt.assert_allclose(data.lon, np.arange(-180, 181, 1))
+    npt.assert_allclose(data.min(), -773.5)
+    npt.assert_allclose(data.max(), 1751.3)
+
+
+def test_earth_mag_01d_wdmam_with_region():
+    """
+    Test loading low-resolution WDMAM grid with 'region'.
+    """
+    data = load_earth_magnetic_anomaly(
+        resolution="01d",
+        region=[-10, 10, -5, 5],
+        registration="gridline",
+        data_source="wdmam",
+    )
+    assert data.shape == (11, 21)
+    npt.assert_allclose(data.lat, np.arange(-5, 6, 1))
+    npt.assert_allclose(data.lon, np.arange(-10, 11, 1))
+    npt.assert_allclose(data.min(), -103.900024)
+    npt.assert_allclose(data.max(), 102.19995)
+
+
+def test_earth_mag_03m_wdmam_with_region():
+    """
+    Test loading a subregion of high-resolution WDMAM data.
+    """
+    data = load_earth_magnetic_anomaly(
+        resolution="03m", region=[10, 13, -60, -58], data_source="wdmam"
+    )
+    assert data.gmt.registration == 0
+    assert data.shape == (41, 61)
+    assert data.lat.min() == -60
+    assert data.lat.max() == -58
+    assert data.lon.min() == 10
+    assert data.lon.max() == 13
+    npt.assert_allclose(data.min(), -639.7001)
+    npt.assert_allclose(data.max(), 629.6)
+
+
+def test_earth_mag_05m_wdmam_without_region():
+    """
+    Test loading a high-resolution WDMAM grid without passing 'region'.
+    """
+    with pytest.raises(GMTInvalidInput):
+        load_earth_magnetic_anomaly(
+            resolution="05m", registration="gridline", data_source="wdmam"
+        )
+
+
+def test_earth_mag_wdmam_incorrect_resolution_registration():
+    """
+    Test that an error is raised when trying to load a WDMAM grid registration
+    with an unavailable resolution.
+    """
+    with pytest.raises(GMTInvalidInput):
+        load_earth_magnetic_anomaly(
+            resolution="03m",
+            region=[0, 1, 3, 5],
+            registration="pixel",
+            data_source="wdmam",
+        )
+
+
+def test_earth_mag_data_source_error():
+    """
+    Test that an error is raised when an invalid argument is passed to
+    'data_source'.
+    """
+    with pytest.raises(GMTInvalidInput):
+        load_earth_magnetic_anomaly(resolution="01d", data_source="invalid")