Repository: fatiando/verde
Pull request #282 (2020-08-14 20:10:05)

Add function to convert grid to Dataset
Add a function that converts `coordinates` and `data` as `np.array`s to
`xarray.Dataset`. This would make handling grids easier, with the option of
efficiently saving them as netCDF files. Make `BaseGridder.grid` method to use
this function instead of building the `xarray.Dataset` by itself. Add new
`check_data_names` function that converts a `str` into a `tuple` with a single
`str`. Make `BaseGridder.grid`, `BaseGridder.scatter` and `BaseGridder.profile`
able to take a single `str` for the `data_names` argument. Add test function
for default `data_names`. Change `data_names` on all examples and tutorials to
a single `str`.


**Reminders**:

- [ ] Run `make format` and `make check` to make sure the code follows the style guide.
- [ ] Add tests for new features or tests that would have caught the bug that you're fixing.
- [ ] Add new public functions/methods/classes to `doc/api/index.rst` and the base `__init__.py` file for the package.
- [ ] Write detailed docstrings for all functions/classes/methods. It often helps to design better code if you write the docstrings first.
- [ ] If adding new functionality, add an example to the docstring, gallery, and/or tutorials.
- [ ] Add your full name, affiliation, and [ORCID](https://orcid.org) (optional) to the `AUTHORS.md` file (if you haven't already) in case you'd like to be listed as an author on the [Zenodo](https://zenodo.org/communities/fatiando) archive of the next release.


santisoler: Remove the `attrs` parameter. Users can add those later on the `xr.Dataset`. We might want to add metadata `attrs` that could be used by xarray and GMT when plotting, like units.
santisoler: @leouieda I tried to expalin a little bit better the difference between dims and extra_coords_name, although it might need more work. I added another example that uses extra_coords. Let me know what do you think about this.
santisoler: Thanks @leouieda for the review! I left a response to one of your questions. Let me know what do you think.

I still need to make `BaseGridder.grid` method to use this function. Should we do it here or on a separate PR?
santisoler: @leouieda If you have a spare time, please take a look at this and let me know if you would change something.
I hope I didn't make any backward breaking change. :crossed_fingers: 

Patch:
diff --git a/examples/project_grid.py b/examples/project_grid.py
index 5e80b49..2ac9aff 100644
--- a/examples/project_grid.py
+++ b/examples/project_grid.py
@@ -51,7 +51,7 @@ data = checkerboard.grid(
     region=region,
     spacing=spacing,
     projection=projection,
-    data_names=["checkerboard"],
+    data_names="checkerboard",
     dims=("latitude", "longitude"),
 )
 print("Geographic grid:")
diff --git a/examples/scipygridder.py b/examples/scipygridder.py
index 3bbd8e8..fbc57ce 100644
--- a/examples/scipygridder.py
+++ b/examples/scipygridder.py
@@ -58,7 +58,7 @@ grid = grd.grid(
     spacing=spacing,
     projection=projection,
     dims=["latitude", "longitude"],
-    data_names=["bathymetry_m"],
+    data_names="bathymetry_m",
 )
 print("Generated geographic grid:")
 print(grid)
diff --git a/examples/spline.py b/examples/spline.py
index bf5e272..19e42de 100644
--- a/examples/spline.py
+++ b/examples/spline.py
@@ -66,7 +66,7 @@ grid_full = chain.grid(
     spacing=spacing,
     projection=projection,
     dims=["latitude", "longitude"],
-    data_names=["temperature"],
+    data_names="temperature",
 )
 grid = vd.distance_mask(
     coordinates, maxdist=3 * spacing * 111e3, grid=grid_full, projection=projection
diff --git a/examples/spline_cv.py b/examples/spline_cv.py
index 9818d93..a880fdd 100644
--- a/examples/spline_cv.py
+++ b/examples/spline_cv.py
@@ -59,7 +59,7 @@ grid_full = spline.grid(
     spacing=spacing,
     projection=projection,
     dims=["latitude", "longitude"],
-    data_names=["temperature"],
+    data_names="temperature",
 )
 grid = vd.distance_mask(
     coordinates, maxdist=3 * spacing * 111e3, grid=grid_full, projection=projection
diff --git a/tutorials/chain.py b/tutorials/chain.py
index f4d1c46..f0ac27e 100644
--- a/tutorials/chain.py
+++ b/tutorials/chain.py
@@ -106,7 +106,7 @@ grid = chain.grid(
     spacing=spacing,
     projection=projection,
     dims=["latitude", "longitude"],
-    data_names=["bathymetry"],
+    data_names="bathymetry",
 )
 print(grid)
 
@@ -140,7 +140,7 @@ grid_trend = chain.named_steps["trend"].grid(
     spacing=spacing,
     projection=projection,
     dims=["latitude", "longitude"],
-    data_names=["bathymetry"],
+    data_names="bathymetry",
 )
 print(grid_trend)
 
diff --git a/tutorials/model_evaluation.py b/tutorials/model_evaluation.py
index 5e9674d..8a82526 100644
--- a/tutorials/model_evaluation.py
+++ b/tutorials/model_evaluation.py
@@ -91,7 +91,7 @@ grid = spline.grid(
     spacing=spacing,
     projection=projection,
     dims=["latitude", "longitude"],
-    data_names=["temperature"],
+    data_names="temperature",
 )
 print(grid)
 
diff --git a/tutorials/model_selection.py b/tutorials/model_selection.py
index ebcbe2c..cb0571a 100644
--- a/tutorials/model_selection.py
+++ b/tutorials/model_selection.py
@@ -128,7 +128,7 @@ grid = spline.grid(
     spacing=spacing,
     projection=projection,
     dims=["latitude", "longitude"],
-    data_names=["temperature"],
+    data_names="temperature",
 )
 print(grid)
 
@@ -174,7 +174,7 @@ grid_default = spline_default.grid(
     spacing=spacing,
     projection=projection,
     dims=["latitude", "longitude"],
-    data_names=["temperature"],
+    data_names="temperature",
 )
 
 ########################################################################################
diff --git a/tutorials/overview.py b/tutorials/overview.py
index 9ff5f69..ea252ba 100644
--- a/tutorials/overview.py
+++ b/tutorials/overview.py
@@ -121,7 +121,7 @@ print(grid)
 # and "northing") and data variables ("scalars"). You can overwrite these names by
 # setting the ``dims`` and ``data_names`` arguments.
 
-grid = spline.grid(spacing=30, dims=["latitude", "longitude"], data_names=["gravity"])
+grid = spline.grid(spacing=30, dims=["latitude", "longitude"], data_names="gravity")
 print(grid)
 
 plt.figure()
diff --git a/tutorials/projections.py b/tutorials/projections.py
index 0d6f34b..008f9f1 100644
--- a/tutorials/projections.py
+++ b/tutorials/projections.py
@@ -60,7 +60,7 @@ spline = vd.Spline().fit(filter_coords, filter_bathy)
 ########################################################################################
 # If we now call :meth:`verde.Spline.grid` we'll get back a grid evenly spaced in
 # projected Cartesian coordinates.
-grid = spline.grid(spacing=spacing * 111e3, data_names=["bathymetry"])
+grid = spline.grid(spacing=spacing * 111e3, data_names="bathymetry")
 print("Cartesian grid:")
 print(grid)
 
@@ -105,7 +105,7 @@ grid_geo = spline.grid(
     spacing=spacing,
     projection=projection,
     dims=["latitude", "longitude"],
-    data_names=["bathymetry"],
+    data_names="bathymetry",
 )
 print("Geographic grid:")
 print(grid_geo)
diff --git a/tutorials/weights.py b/tutorials/weights.py
index 46569e0..073f845 100644
--- a/tutorials/weights.py
+++ b/tutorials/weights.py
@@ -243,7 +243,7 @@ grid = spline.grid(
     spacing=spacing,
     projection=projection,
     dims=["latitude", "longitude"],
-    data_names=["velocity"],
+    data_names="velocity",
 )
 # Avoid showing interpolation outside of the convex hull of the data points.
 grid = vd.convexhull_mask(coordinates, grid=grid, projection=projection)
@@ -261,7 +261,7 @@ grid_unweighted = spline_unweighted.grid(
     spacing=spacing,
     projection=projection,
     dims=["latitude", "longitude"],
-    data_names=["velocity"],
+    data_names="velocity",
 )
 grid_unweighted = vd.convexhull_mask(
     coordinates, grid=grid_unweighted, projection=projection
diff --git a/verde/__init__.py b/verde/__init__.py
index f80b3e3..9f59575 100644
--- a/verde/__init__.py
+++ b/verde/__init__.py
@@ -15,7 +15,7 @@ from .coordinates import (
     longitude_continuity,
 )
 from .mask import distance_mask, convexhull_mask
-from .utils import variance_to_weights, maxabs, grid_to_table
+from .utils import variance_to_weights, maxabs, grid_to_table, make_xarray_grid
 from .io import load_surfer
 from .distances import median_distance
 from .blockreduce import BlockReduce, BlockMean
diff --git a/verde/base/base_classes.py b/verde/base/base_classes.py
index cc28f16..18f599b 100644
--- a/verde/base/base_classes.py
+++ b/verde/base/base_classes.py
@@ -3,13 +3,13 @@ Base classes for all gridders.
 """
 from abc import ABCMeta, abstractmethod
 
-import xarray as xr
 import pandas as pd
 from sklearn.base import BaseEstimator
 from sklearn.model_selection import BaseCrossValidator
 
 from ..coordinates import grid_coordinates, profile_coordinates, scatter_points
-from .utils import check_data, score_estimator
+from .utils import check_data, check_data_names, score_estimator
+from ..utils import make_xarray_grid
 
 
 # Pylint doesn't like X, y scikit-learn argument names.
@@ -380,9 +380,9 @@ class BaseGridder(BaseEstimator):
             order: northing dimension, easting dimension.
             **NOTE: This is an exception to the "easting" then
             "northing" pattern but is required for compatibility with xarray.**
-        data_names : list or None
+        data_names : str, list or None
             The name(s) of the data variables in the output grid. Defaults to
-            ``['scalars']`` for scalar data,
+            ``'scalars'`` for scalar data,
             ``['east_component', 'north_component']`` for 2D vector data, and
             ``['east_component', 'north_component', 'vertical_component']`` for
             3D vector data.
@@ -415,17 +415,24 @@ class BaseGridder(BaseEstimator):
             data = check_data(
                 self.predict(project_coordinates(coordinates, projection))
             )
+        # Get names for data and any extra coordinates
+        data_names = check_data_names(data_names)
         data_names = get_data_names(data, data_names)
-        coords = {dims[1]: coordinates[0][0, :], dims[0]: coordinates[1][:, 0]}
-        # Add extra coordinates to xr.Dataset as non-dimension coordinates
         extra_coords_names = self._get_extra_coords_names(coordinates)
-        for name, extra_coord in zip(extra_coords_names, coordinates[2:]):
-            coords[name] = (dims, extra_coord)
-        attrs = {"metadata": "Generated by {}".format(repr(self))}
-        data_vars = {
-            name: (dims, value, attrs) for name, value in zip(data_names, data)
-        }
-        return xr.Dataset(data_vars, coords=coords, attrs=attrs)
+        # Create xarray.Dataset
+        dataset = make_xarray_grid(
+            coordinates,
+            data,
+            data_names,
+            dims=dims,
+            extra_coords_names=extra_coords_names,
+        )
+        # Add metadata as attrs
+        metadata = "Generated by {}".format(repr(self))
+        dataset.attrs["metadata"] = metadata
+        for array in dataset:
+            dataset[array].attrs["metadata"] = metadata
+        return dataset
 
     def scatter(
         self,
@@ -469,9 +476,9 @@ class BaseGridder(BaseEstimator):
             following order: northing dimension, easting dimension.
             **NOTE: This is an exception to the "easting" then
             "northing" pattern but is required for compatibility with xarray.**
-        data_names : list of None
+        data_names : str, list or None
             The name(s) of the data variables in the output dataframe. Defaults
-            to ``['scalars']`` for scalar data,
+            to ``'scalars'`` for scalar data,
             ``['east_component', 'north_component']`` for 2D vector data, and
             ``['east_component', 'north_component', 'vertical_component']`` for
             3D vector data.
@@ -499,6 +506,7 @@ class BaseGridder(BaseEstimator):
             data = check_data(
                 self.predict(project_coordinates(coordinates, projection))
             )
+        data_names = check_data_names(data_names)
         data_names = get_data_names(data, data_names)
         columns = [(dims[0], coordinates[1]), (dims[1], coordinates[0])]
         extra_coords_names = self._get_extra_coords_names(coordinates)
@@ -572,9 +580,9 @@ class BaseGridder(BaseEstimator):
             following order: northing dimension, easting dimension.
             **NOTE: This is an exception to the "easting" then
             "northing" pattern but is required for compatibility with xarray.**
-        data_names : list of None
+        data_names : str, list or None
             The name(s) of the data variables in the output dataframe. Defaults
-            to ``['scalars']`` for scalar data,
+            to ``'scalars'`` for scalar data,
             ``['east_component', 'north_component']`` for 2D vector data, and
             ``['east_component', 'north_component', 'vertical_component']`` for
             3D vector data.
@@ -610,6 +618,7 @@ class BaseGridder(BaseEstimator):
         # profile but Cartesian distances.
         if projection is not None:
             coordinates = project_coordinates(coordinates, projection, inverse=True)
+        data_names = check_data_names(data_names)
         data_names = get_data_names(data, data_names)
         columns = [
             (dims[0], coordinates[1]),
diff --git a/verde/base/utils.py b/verde/base/utils.py
index 0adf31e..8cee240 100644
--- a/verde/base/utils.py
+++ b/verde/base/utils.py
@@ -104,6 +104,30 @@ def check_data(data):
     return data
 
 
+def check_data_names(data_names):
+    """
+    Check the *data_names* argument and make sure it's a tuple.
+    If ``data_names`` is a single string, return it as a tuple with a single
+    element.
+
+    This is the default form accepted by gridders and functions that require
+    the ``data_names`` argument.
+
+    Examples
+    --------
+
+    >>> check_data_names("dummy")
+    ('dummy',)
+    >>> check_data_names(("component_x", "component_y"))
+    ('component_x', 'component_y')
+    >>> check_data_names(["dummy"])
+    ['dummy']
+    """
+    if isinstance(data_names, str):
+        data_names = (data_names,)
+    return data_names
+
+
 def check_coordinates(coordinates):
     """
     Check that the given coordinate arrays are what we expect them to be.
diff --git a/verde/utils.py b/verde/utils.py
index ac6639c..3c40859 100644
--- a/verde/utils.py
+++ b/verde/utils.py
@@ -6,6 +6,7 @@ import functools
 import dask
 import numpy as np
 import pandas as pd
+import xarray as xr
 from scipy.spatial import cKDTree  # pylint: disable=no-name-in-module
 
 try:
@@ -18,7 +19,8 @@ try:
 except ImportError:
     numba = None
 
-from .base.utils import check_data, n_1d_arrays
+from .base.utils import check_data, check_data_names, n_1d_arrays
+from .coordinates import check_coordinates
 
 
 def dispatch(function, delayed=False, client=None):
@@ -207,6 +209,152 @@ def maxabs(*args, nan=True):
     return npmax(absolute)
 
 
+def make_xarray_grid(
+    coordinates,
+    data,
+    data_names,
+    dims=("northing", "easting"),
+    extra_coords_names=None,
+):
+    """
+    Create an :class:`xarray.Dataset` grid from numpy arrays
+
+    This functions creates an :class:`xarray.Dataset` out of 2d gridded data
+    including easting and northing coordinates, any extra coordinates (like
+    upward elevation, time, etc) and data arrays.
+
+    Use this to transform the outputs of :func:`verde.grid_coordinates` and
+    the ``predict`` method of a gridder into an :class:`xarray.Dataset`.
+
+    .. note::
+
+        This is a utility function to help create 2D grids (i.e., grids with
+        two ``dims`` coordinates). For arbitrary N-dimensional arrays, use
+        :mod:`xarray` directly.
+
+    Parameters
+    ----------
+    coordinates : tuple of arrays
+        Arrays with coordinates of each point in the grid. Each array must
+        contain values for a dimension in the order: easting, northing,
+        vertical, etc. All arrays must be 2d and need to have the same *shape*.
+        These coordinates can be generated through
+        :func:`verde.grid_coordinates`.
+    data : array or tuple of arrays
+        Array or tuple of arrays with data values on each point in the grid.
+        Each array must contain values for a dimension in the same order as
+        the coordinates. All arrays need to have the same *shape*.
+    data_names : str, list or None
+        The name(s) of the data variables in the output grid.
+    dims : list or None
+        The names of the northing and easting data dimensions, respectively,
+        in the output grid. Must be defined in the following order: northing
+        dimension, easting dimension.
+        **NOTE: This is an exception to the "easting" then
+        "northing" pattern but is required for compatibility with xarray.**
+        The easting and northing coordinates in the :class:`xarray.Dataset`
+        will have the same names as the passed dimensions.
+    extra_coords_names : str or list
+        Name or list of names for any additional coordinates besides the
+        easting and northing ones. Ignored if coordinates has
+        only two elements. The extra coordinates are non-index coordinates of
+        the grid array.
+
+    Returns
+    -------
+    grid : :class:`xarray.Dataset`
+        A 2D grid with one or more data variables.
+
+    Examples
+    --------
+
+    >>> import numpy as np
+    >>> import verde as vd
+    >>> # Create the coordinates of the regular grid
+    >>> coordinates = vd.grid_coordinates((-10, -6, 8, 10), spacing=2)
+    >>> # And some dummy data for each point of the grid
+    >>> data = np.ones_like(coordinates[0])
+    >>> # Create the grid
+    >>> grid = make_xarray_grid(coordinates, data, data_names="dummy")
+    >>> print(grid)
+    <xarray.Dataset>
+    Dimensions:   (easting: 3, northing: 2)
+    Coordinates:
+      * easting   (easting) float64 -10.0 -8.0 -6.0
+      * northing  (northing) float64 8.0 10.0
+    Data variables:
+        dummy     (northing, easting) float64 1.0 1.0 1.0 1.0 1.0 1.0
+
+    >>> # Create a grid with an extra coordinate
+    >>> coordinates = vd.grid_coordinates(
+    ...     (-10, -6, 8, 10), spacing=2, extra_coords=5
+    ... )
+    >>> # And some dummy data for each point of the grid
+    >>> data = np.ones_like(coordinates[0])
+    >>> # Create the grid
+    >>> grid = make_xarray_grid(
+    ...     coordinates, data, data_names="dummy", extra_coords_names="upward"
+    ... )
+    >>> print(grid)
+    <xarray.Dataset>
+    Dimensions:   (easting: 3, northing: 2)
+    Coordinates:
+      * easting   (easting) float64 -10.0 -8.0 -6.0
+      * northing  (northing) float64 8.0 10.0
+        upward    (northing, easting) float64 5.0 5.0 5.0 5.0 5.0 5.0
+    Data variables:
+        dummy     (northing, easting) float64 1.0 1.0 1.0 1.0 1.0 1.0
+
+    """
+    coordinates = check_coordinates(coordinates)
+    data = check_data(data)
+    data_names = check_data_names(data_names)
+    # dims is like shape with order (rows, cols) for the array
+    # so the first element is northing and second is easting
+    coords = {dims[1]: coordinates[0][0, :], dims[0]: coordinates[1][:, 0]}
+    # Extra coordinates are handled like 2D data arrays with
+    # the same dims and the data.
+    if coordinates[2:]:
+        extra_coords_names = _check_extra_coords_names(coordinates, extra_coords_names)
+        for name, extra_coord in zip(extra_coords_names, coordinates[2:]):
+            coords[name] = (dims, extra_coord)
+    # Generate object
+    if len(data) != len(data_names):
+        raise ValueError(
+            "Invalid data_names '{}'. ".format(data_names)
+            + "Number of data names must match the number of "
+            + "data arrays ('{}').".format(len(data))
+        )
+    data_vars = {name: (dims, value) for name, value in zip(data_names, data)}
+    return xr.Dataset(data_vars, coords)
+
+
+def _check_extra_coords_names(coordinates, extra_coords_names):
+    """
+    Sanity checks for extra_coords_names against coordinates
+
+    Assuming that there are extra coordinates on the ``coordinates`` tuple.
+    """
+    # Convert to tuple if it's a str
+    if isinstance(extra_coords_names, str):
+        extra_coords_names = (extra_coords_names,)
+    # Check if it's not None
+    if extra_coords_names is None:
+        raise ValueError(
+            "Invalid extra_coords_names equal to None. "
+            + "When passing one or more extra coordinate, "
+            + "extra_coords_names cannot be None."
+        )
+    # Check if there are the same number of extra_coords than extra_coords_name
+    if len(coordinates[2:]) != len(extra_coords_names):
+        raise ValueError(
+            "Invalid extra_coords_names '{}'. ".format(extra_coords_names)
+            + "Number of extra coordinates names must match the number of "
+            + "additional coordinates ('{}').".format(len(coordinates[2:]))
+        )
+    return extra_coords_names
+
+
 def grid_to_table(grid):
     """
     Convert a grid to a table with the values and coordinates of each point.


Test patch:
diff --git a/verde/tests/test_base.py b/verde/tests/test_base.py
index 9a148b6..7a0ee1e 100644
--- a/verde/tests/test_base.py
+++ b/verde/tests/test_base.py
@@ -146,6 +146,23 @@ def test_basegridder():
     npt.assert_allclose(prof.distance, coordinates_true[0][0, :])
 
 
+def test_basegridder_data_names():
+    """
+    Test default values for data_names
+    """
+    region = (0, 10, -10, -5)
+    shape = (50, 30)
+    angular, linear = 2, 100
+    coordinates = scatter_points(region, 1000, random_state=0)
+    data = angular * coordinates[0] + linear
+    grd = PolyGridder().fit(coordinates, data)
+    grid = grd.grid(region, shape)
+    prof = grd.profile((0, -10), (10, -10), 30)
+    # Check if default name for data_names was applied correctly
+    assert "scalars" in grid
+    assert "scalars" in prof
+
+
 def test_basegridder_projection():
     "Test basic functionality of BaseGridder when passing in a projection"
 
diff --git a/verde/tests/test_utils.py b/verde/tests/test_utils.py
index 5442606..3553c33 100644
--- a/verde/tests/test_utils.py
+++ b/verde/tests/test_utils.py
@@ -10,7 +10,14 @@ from scipy.spatial import cKDTree  # pylint: disable=no-name-in-module
 import pytest
 
 from ..coordinates import grid_coordinates
-from ..utils import parse_engine, dummy_jit, kdtree, grid_to_table, partition_by_sum
+from ..utils import (
+    parse_engine,
+    dummy_jit,
+    kdtree,
+    grid_to_table,
+    partition_by_sum,
+    make_xarray_grid,
+)
 from .. import utils
 
 
@@ -89,3 +96,116 @@ def test_partition_by_sum_fails_no_partitions():
     with pytest.raises(ValueError) as error:
         partition_by_sum(np.arange(10), 8)
     assert "Could not find partition points" in str(error)
+
+
+def test_build_grid():
+    """
+    Check if xarray.Dataset is correctly created
+    """
+    region = (-10, -5, 6, 10)
+    spacing = 1
+    coordinates = grid_coordinates(region, spacing=spacing)
+    data = np.ones_like(coordinates[0])
+    grid = make_xarray_grid(coordinates, data, data_names="dummy")
+    npt.assert_allclose(grid.easting, [-10, -9, -8, -7, -6, -5])
+    npt.assert_allclose(grid.northing, [6, 7, 8, 9, 10])
+    npt.assert_allclose(grid.dummy, 1)
+    assert grid.dummy.shape == (5, 6)
+    # Change dims
+    grid = make_xarray_grid(
+        coordinates, data, data_names="dummy", dims=("latitude", "longitude")
+    )
+    npt.assert_allclose(grid.longitude, [-10, -9, -8, -7, -6, -5])
+    npt.assert_allclose(grid.latitude, [6, 7, 8, 9, 10])
+    npt.assert_allclose(grid.dummy, 1)
+    assert grid.dummy.shape == (5, 6)
+
+
+def test_build_grid_multiple_data():
+    """
+    Check if xarray.Dataset with multiple data is correctly created
+    """
+    region = (-10, -5, 6, 10)
+    spacing = 1
+    coordinates = grid_coordinates(region, spacing=spacing)
+    data_arrays = tuple(i * np.ones_like(coordinates[0]) for i in range(1, 4))
+    data_names = list("data_{}".format(i) for i in range(1, 4))
+    dataset = make_xarray_grid(coordinates, data_arrays, data_names=data_names)
+    npt.assert_allclose(dataset.easting, [-10, -9, -8, -7, -6, -5])
+    npt.assert_allclose(dataset.northing, [6, 7, 8, 9, 10])
+    for i in range(1, 4):
+        npt.assert_allclose(dataset["data_{}".format(i)], i)
+        assert dataset["data_{}".format(i)].shape == (5, 6)
+
+
+def test_build_grid_extra_coords():
+    """
+    Check if xarray.Dataset with extra coords is correctly created
+    """
+    region = (-10, -5, 6, 10)
+    spacing = 1
+    extra_coords = [1, 2]
+    coordinates = grid_coordinates(region, spacing=spacing, extra_coords=extra_coords)
+    data = np.ones_like(coordinates[0])
+    dataset = make_xarray_grid(
+        coordinates,
+        data,
+        data_names="dummy",
+        extra_coords_names=["upward", "time"],
+    )
+    npt.assert_allclose(dataset.easting, [-10, -9, -8, -7, -6, -5])
+    npt.assert_allclose(dataset.northing, [6, 7, 8, 9, 10])
+    npt.assert_allclose(dataset.upward, 1)
+    npt.assert_allclose(dataset.time, 2)
+    npt.assert_allclose(dataset.dummy, 1)
+    assert dataset.dummy.shape == (5, 6)
+    assert dataset.upward.shape == (5, 6)
+    assert dataset.time.shape == (5, 6)
+
+
+def test_build_grid_invalid_names():
+    """
+    Check if errors are raise after invalid data names
+    """
+    region = (-10, -5, 6, 10)
+    spacing = 1
+    coordinates = grid_coordinates(region, spacing=spacing)
+    # Single data, multiple data_name
+    data = np.ones_like(coordinates[0])
+    with pytest.raises(ValueError):
+        make_xarray_grid(coordinates, data, data_names=["bla_1", "bla_2"])
+    # Multiple data, single data_name
+    data = tuple(i * np.ones_like(coordinates[0]) for i in (1, 2))
+    with pytest.raises(ValueError):
+        make_xarray_grid(coordinates, data, data_names="blabla")
+
+
+def test_build_grid_invalid_extra_coords():
+    """
+    Check if errors are raise after invalid extra coords
+    """
+    region = (-10, -5, 6, 10)
+    spacing = 1
+    # No extra coords, extra_coords_name should be ignored
+    coordinates = grid_coordinates(region, spacing=spacing)
+    data = np.ones_like(coordinates[0])
+    make_xarray_grid(coordinates, data, data_names="dummy", extra_coords_names="upward")
+    # Single extra coords, extra_coords_name equal to None
+    coordinates = grid_coordinates(region, spacing=spacing, extra_coords=1)
+    data = np.ones_like(coordinates[0])
+    with pytest.raises(ValueError):
+        make_xarray_grid(coordinates, data, data_names="dummy", extra_coords_names=None)
+    # Multiple extra coords, single extra_coords_name as a str
+    coordinates = grid_coordinates(region, spacing=spacing, extra_coords=[1, 2])
+    data = np.ones_like(coordinates[0])
+    with pytest.raises(ValueError):
+        make_xarray_grid(
+            coordinates, data, data_names="dummy", extra_coords_names="upward"
+        )
+    # Multiple extra coords, multiple extra_coords_name but not equal
+    coordinates = grid_coordinates(region, spacing=spacing, extra_coords=[1, 2, 3])
+    data = np.ones_like(coordinates[0])
+    with pytest.raises(ValueError):
+        make_xarray_grid(
+            coordinates, data, data_names="dummy", extra_coords_names=["upward", "time"]
+        )