Sample dataarray (#53)

* add sample dataarray

* add changelog

Author: mathause

CHANGELOG.md

@@ -17,6 +17,7 @@ mplotutils now uses the MIT license instead of GPL-3.0 ([#51](https://github.com
  * Allow passing list of axes to ``set_map_layout``, renamed the files and extended
    the test coverage ([#42](https://github.com/mathause/mplotutils/pull/42)
    and [#43](https://github.com/mathause/mplotutils/pull/43)).
+ * Add function to create `xr.DataArray` sample data ([#53](https://github.com/mathause/mplotutils/pull/53)).
 
 ### Bug fixes
 

mplotutils/cartopy_utils.py

@@ -14,6 +14,24 @@
 def sample_data_map(nlons, nlats):
     """Returns `lons`, `lats`, and fake `data`
 
+    Parameters
+    ----------
+    nlons : int
+        Number of longitude grid cells.
+    nlats : int
+        Number of latitude grid cells.
+
+    Returns
+    -------
+    lon : ndarray
+        Array of the longitude coordiantes.
+    lat : ndarray
+        Array of the latitude coordiantes.
+    data : ndarray
+        Sample data.
+
+    Notes
+    -----
     adapted from:
     http://scitools.org.uk/cartopy/docs/v0.15/examples/axes_grid_basic.html
     """
@@ -32,6 +50,34 @@ def sample_data_map(nlons, nlats):
     return lon, lat, data
 
 
+def sample_dataarray(nlon, nlat):
+    """Returns `lons`, `lats`, and fake `data`
+
+    Parameters
+    ----------
+    nlons : int
+        Number of longitude grid cells.
+    nlats : int
+        Number of latitude grid cells.
+
+    Returns
+    -------
+    data : xr.DataArray
+        Sample data with coordiantes.
+
+    Notes
+    -----
+    adapted from:
+    http://scitools.org.uk/cartopy/docs/v0.15/examples/axes_grid_basic.html
+    """
+
+    import xarray as xr
+
+    lon, lat, data = sample_data_map(nlons=nlon, nlats=nlat)
+
+    return xr.DataArray(data, dims=("lat", "lon"), coords={"lon": lon, "lat": lat})
+
+
 def cyclic_dataarray(obj, coord="lon"):
     """Add a cyclic coordinate point to a DataArray or Dataset along a dimension.
 

mplotutils/tests/test_sample_data_map.py

@@ -1,22 +1,13 @@
 import numpy as np
+import pytest
+import xarray as xr
 
-from mplotutils import sample_data_map
+from mplotutils import sample_data_map, sample_dataarray
 
 
-def test_data_shape():
-
-    nlons = 10
-    nlats = 20
-
-    lon, lat, data = sample_data_map(nlons, nlats)
-
-    assert len(lon) == nlons
-    assert len(lat) == nlats
-
-    assert data.shape == (nlats, nlons)
-
-    nlons = 5
-    nlats = 10
+@pytest.mark.parametrize("nlons", [5, 10])
+@pytest.mark.parametrize("nlats", [10, 20])
+def test_data_shape(nlons, nlats):
 
     lon, lat, data = sample_data_map(nlons, nlats)
 
@@ -40,3 +31,13 @@ def test_lon():
 
     expected_lon = np.arange(0, 351, 10)
     assert np.allclose(lon, expected_lon)
+
+
+@pytest.mark.parametrize("nlon", [5, 10])
+@pytest.mark.parametrize("nlat", [10, 20])
+def test_sample_dataarray(nlon, nlat):
+
+    data = sample_dataarray(nlon, nlat)
+
+    assert isinstance(data, xr.DataArray)
+    assert data.shape == (nlat, nlon)