Repository: jhardenberg/smmregrid (language: python)

Issue:
Dealing with levels with different masks
So far `smmregrid` cannot work with variables provided with different masks. This can be a problem when working with oceanic data, which have vertical levels with different bathymetry. 
We do not have a solution yet but this is to keep track of the issue. 


Discussion / hints:

Fix patch:
diff --git a/smmregrid/__init__.py b/smmregrid/__init__.py
index f8f5803..09075a6 100644
--- a/smmregrid/__init__.py
+++ b/smmregrid/__init__.py
@@ -1,1 +1,3 @@
-from .regrid import Regridder, cdo_generate_weights, regrid
+from .regrid import Regridder, regrid
+from .cdo_weights import cdo_generate_weights
+from .util import find_vert_coord
\ No newline at end of file
diff --git a/smmregrid/cdo_weights.py b/smmregrid/cdo_weights.py
new file mode 100644
index 0000000..e48b2be
--- /dev/null
+++ b/smmregrid/cdo_weights.py
@@ -0,0 +1,260 @@
+"""CDO-based generation of weights"""
+
+import os
+import sys
+import logging
+import tempfile
+import subprocess
+from multiprocessing import Process, Manager
+import numpy
+import xarray
+from .util import find_vert_coord
+from .weights import compute_weights_matrix3d, compute_weights_matrix, mask_weights, check_mask
+
+
+def worker(wlist, nnn, *args, **kwargs):
+    """Run a worker process"""
+    wlist[nnn] = cdo_generate_weights2d(*args, **kwargs).compute()
+
+def cdo_generate_weights(source_grid, target_grid, method="con", extrapolate=True,
+                         remap_norm="fracarea", remap_area_min=0.0, icongridpath=None,
+                         gridpath=None, extra=None, vert_coord=None, cdo="cdo", nproc=1):
+    """Generate the weights using CDO, handling both 2D and 3D cases"""
+
+    # Check if there is a vertical coordinate for 3d oceanic data
+    if not vert_coord:
+        vert_coord = find_vert_coord(source_grid)
+        logging.warning('vert_coord is %s',str(vert_coord))
+
+    if not vert_coord: # Are we 2D? Use default method
+        weights =cdo_generate_weights2d(
+            source_grid,
+            target_grid,
+            method=method,
+            extrapolate=extrapolate,
+            remap_norm=remap_norm,
+            remap_area_min=remap_area_min,
+            icongridpath=icongridpath,
+            gridpath=gridpath,
+            extra=extra,
+            cdo=cdo,
+            nproc=nproc)
+
+        # Precompute destination weights mask
+        weights_matrix = compute_weights_matrix(weights)
+        weights = mask_weights(weights, weights_matrix, vert_coord)
+        masked = int(check_mask(weights, vert_coord))
+        masked_xa = xarray.DataArray(masked, name="dst_grid_masked")
+
+        return xarray.merge([weights, masked_xa])
+
+    else:  # we are 3D
+        if extra:
+        # make sure extra is a flat list if it is not already
+            if not isinstance(extra, list):
+                extra = [extra]
+        else:
+            extra = []
+
+        if isinstance(source_grid, str):
+            sgrid = xarray.open_dataset(source_grid)
+        else:
+            sgrid = source_grid
+
+        nvert = sgrid[vert_coord].values.size
+        #print(nvert)
+
+        # for lev in range(0, nvert):
+        mgr = Manager()
+
+        # dictionaries are shared, so they have to be passed as functions
+        wlist= mgr.list(range(nvert))
+
+        num_blocks, remainder = divmod(nvert, nproc)
+        num_blocks = num_blocks + (0 if remainder == 0 else 1)
+
+        blocks = numpy.array_split(numpy.arange(nvert), num_blocks)
+        for block in blocks:
+            processes = []
+            for lev in block:
+                logging.info("Generating level: %s", str(lev))
+                extra2 = [f"-sellevidx,{lev+1}"]
+                ppp = Process(target=worker,
+                            args=(wlist, lev, source_grid, target_grid),
+                            kwargs=dict(method=method,
+                                        extrapolate=extrapolate,
+                                        remap_norm=remap_norm,
+                                        remap_area_min=remap_area_min,
+                                        icongridpath=icongridpath,
+                                        gridpath=gridpath,
+                                        extra=extra + extra2,
+                                        cdo=cdo,
+                                        nproc=nproc))
+                ppp.start()
+                processes.append(ppp)
+
+            for proc in processes:
+                proc.join()
+
+        weights = weightslist_to_3d(wlist, vert_coord)
+
+        # Precompute destination weights mask
+        weights_matrix = compute_weights_matrix3d(weights, vert_coord)
+        weights = mask_weights(weights, weights_matrix, vert_coord)
+        masked = check_mask(weights, vert_coord)
+        masked = [int(x) for x in masked]  # convert to list of int
+        masked_xa = xarray.DataArray(masked, coords={vert_coord: range(0, len(masked))}, name="dst_grid_masked")
+
+        return xarray.merge([weights, masked_xa])
+
+
+def cdo_generate_weights2d(source_grid, target_grid, method="con", extrapolate=True,
+                           remap_norm="fracarea", remap_area_min=0.0, icongridpath=None,
+                           gridpath=None, extra=None, cdo="cdo", nproc=1):
+    """
+    Generate weights for regridding using CDO
+
+    Available weight generation methods are:
+
+     * bic: SCRIP Bicubic
+     * bil: SCRIP Bilinear
+     * con: SCRIP First-order conservative
+     * con2: SCRIP Second-order conservative
+     * dis: SCRIP Distance-weighted average
+     * laf: YAC Largest area fraction
+     * ycon: YAC First-order conservative
+     * nn: Nearest neighbour
+
+    Run ``cdo gen${method} --help`` for details of each method
+
+    Args:
+        source_grid (xarray.DataArray): Source grid
+        target_grid (xarray.DataArray): Target grid
+            description
+        method (str): Regridding method - default conservative
+        extrapolate (bool): Extrapolate output field
+        remap_norm (str): Normalisation method for conservative methods
+        remap_area_min (float): Minimum destination area fraction
+        gridpath (str): where to store downloaded grids
+        icongridpath (str): location of ICON grids (e.g. /pool/data/ICON)
+        extra: command(s) to apply to source grid before weight generation (can be a list)
+        cdo: the command to launch cdo ["cdo"]
+        nproc: number of processes to use for weight generation (NOT USED!)
+
+        cdo: path to cdo binary
+    Returns:
+        :obj:`xarray.Dataset` with regridding weights
+    """
+
+    supported_methods = ["bic", "bil", "con", "con2", "dis", "laf", "nn", "ycon"]
+    if method not in supported_methods:
+        raise ValueError('The remap method provided is not supported!')
+    if remap_norm not in ["fracarea", "destarea"]:
+        raise ValueError('The remap normalization provided is not supported!')
+
+    # Make some temporary files that we'll feed to CDO
+    weight_file = tempfile.NamedTemporaryFile()
+
+    if isinstance(source_grid, str):
+        sgrid = source_grid
+    else:
+        source_grid_file = tempfile.NamedTemporaryFile()
+        source_grid.to_netcdf(source_grid_file.name)
+        sgrid = source_grid_file.name
+
+    if isinstance(target_grid, str):
+        tgrid = target_grid
+    else:
+        target_grid_file = tempfile.NamedTemporaryFile()
+        target_grid.to_netcdf(target_grid_file.name)
+        tgrid = target_grid_file.name
+
+    # Setup environment
+    env = os.environ
+    if extrapolate:
+        env["REMAP_EXTRAPOLATE"] = "on"
+    else:
+        env["REMAP_EXTRAPOLATE"] = "off"
+
+    env["CDO_REMAP_NORM"] = remap_norm
+    env["REMAP_AREA_MIN"] = "%f" % (remap_area_min)
+
+    if gridpath:
+        env["CDO_DOWNLOAD_PATH"] = gridpath
+    if icongridpath:
+        env["CDO_ICON_GRIDS"] = icongridpath
+
+    try:
+        # Run CDO
+        if extra:
+            # make sure extra is a flat list if it is not already
+            if not isinstance(extra, list):
+                extra = [extra]
+
+            subprocess.check_output(
+                [
+                    cdo,
+                    "gen%s,%s" % (method, tgrid)
+                ] + extra +
+                [
+                    sgrid,
+                    weight_file.name,
+                ],
+                stderr=subprocess.PIPE,
+                env=env,
+            )
+        else:
+            subprocess.check_output(
+                [
+                    cdo,
+                    "gen%s,%s" % (method, tgrid),
+                    sgrid,
+                    weight_file.name,
+                ],
+                stderr=subprocess.PIPE,
+                env=env,
+            )
+
+        # Grab the weights file it outputs as a xarray.Dataset
+        weights = xarray.open_dataset(weight_file.name, engine="netcdf4")
+        return weights
+
+    except subprocess.CalledProcessError as err:
+        # Print the CDO error message
+        logging.error(err.output.decode(), file=sys.stderr)
+        raise
+
+    finally:
+        # Clean up the temporary files
+        if not isinstance(source_grid, str):
+            source_grid_file.close()
+        if not isinstance(target_grid, str):
+            target_grid_file.close()
+        weight_file.close()
+
+
+def weightslist_to_3d(ds_list, vert_coord='lev'):
+    """
+    Function to combine a list of 2D cdo weights into a 3D one adding a vertical coordinate lev
+    """
+    # CDO 2.2.0 fix
+    if "numLinks" in ds_list[0].dims:
+        links_dim = "numLinks"
+    else:
+        links_dim = "num_links"
+    
+    dim_values = range(len(ds_list))
+    nl = [ds.src_address.size for ds in ds_list]
+    nl0 = max(nl)
+    nlda = xarray.DataArray(nl, coords={vert_coord: range(0, len(nl))}, name="link_length")
+    new_array = []
+    varlist = ["src_address", "dst_address", "remap_matrix", "src_grid_imask", "dst_grid_imask"]
+    ds0 = ds_list[0].drop_vars(varlist)
+    for x, d in zip(ds_list, dim_values):
+        nl1 = x.src_address.size
+#        xplist = [x[vname].pad(num_links=(0, nl0-nl1), mode='constant', constant_values=0)
+        xplist = [x[vname].pad(**{links_dim: (0, nl0-nl1), "mode": 'constant', "constant_values": 0})
+                  for vname in varlist ]
+        xp = xarray.merge(xplist)
+        new_array.append(xp.assign_coords({vert_coord: d}))
+    return xarray.merge([nlda, ds0, xarray.concat(new_array, vert_coord)])
diff --git a/smmregrid/checker.py b/smmregrid/checker.py
index 63437fe..1b88825 100644
--- a/smmregrid/checker.py
+++ b/smmregrid/checker.py
@@ -1,9 +1,12 @@
-import xarray as xr
+"""Module for cdo-smmregrid comparison"""
+
 import numpy as np
 import xarray as xr
+from cdo import Cdo
+import xarray as xr
 from smmregrid import cdo_generate_weights, Regridder
 
-from cdo import Cdo
+
 cdo = Cdo()
 
 
@@ -21,7 +24,7 @@ def find_var(xfield):
     return myvar
 
 
-def check_cdo_regrid(finput, ftarget, method='con', access='Dataset'):
+def check_cdo_regrid(finput, ftarget, method='con', access='Dataset', vert_coord=None):
     """Given a file to be interpolated finput over the ftarget grid,
     check if the output of the last variable is the same as produced
     by CDO remap command. This function is used for tests."""
@@ -52,7 +55,7 @@ def check_cdo_regrid(finput, ftarget, method='con', access='Dataset'):
     #interpolator = Regridder(weights=wfield)
 
     # method with automatic creation of weights
-    interpolator = Regridder(source_grid=finput, target_grid=ftarget, method=method)
+    interpolator = Regridder(source_grid=finput, target_grid=ftarget, method=method, vert_coord=vert_coord)
     rfield = interpolator.regrid(xfield)
 
     if access == 'Dataset':
diff --git a/smmregrid/esmf_weights.py b/smmregrid/esmf_weights.py
new file mode 100644
index 0000000..44c1d02
--- /dev/null
+++ b/smmregrid/esmf_weights.py
@@ -0,0 +1,96 @@
+"""ESMF-based generation of weights"""
+
+import tempfile
+import subprocess
+import xarray
+
+def esmf_generate_weights(
+    source_grid,
+    target_grid,
+    method="bilinear",
+    extrap_method="nearestidavg",
+    norm_type="dstarea",
+    line_type=None,
+    pole=None,
+    ignore_unmapped=False,
+):
+    """Generate regridding weights with ESMF
+
+    https://www.earthsystemcog.org/projects/esmf/regridding
+
+    Args:
+        source_grid (:obj:`xarray.Dataarray`): Source grid. If masked the mask
+            will be used in the regridding
+        target_grid (:obj:`xarray.Dataarray`): Target grid. If masked the mask
+            will be used in the regridding
+        method (str): ESMF Regridding method, see ``ESMF_RegridWeightGen --help``
+        extrap_method (str): ESMF Extrapolation method, see ``ESMF_RegridWeightGen --help``
+
+    Returns:
+        :obj:`xarray.Dataset` with regridding information from
+            ESMF_RegridWeightGen
+    """
+    # Make some temporary files that we'll feed to ESMF
+    source_file = tempfile.NamedTemporaryFile(suffix=".nc")
+    target_file = tempfile.NamedTemporaryFile(suffix=".nc")
+    weight_file = tempfile.NamedTemporaryFile(suffix=".nc")
+
+    rwg = "ESMF_RegridWeightGen"
+
+    if "_FillValue" not in source_grid.encoding:
+        source_grid.encoding["_FillValue"] = -1e20
+
+    if "_FillValue" not in target_grid.encoding:
+        target_grid.encoding["_FillValue"] = -1e20
+
+    try:
+        source_grid.to_netcdf(source_file.name)
+        target_grid.to_netcdf(target_file.name)
+
+        command = [
+            rwg,
+            "--source",
+            source_file.name,
+            "--destination",
+            target_file.name,
+            "--weight",
+            weight_file.name,
+            "--method",
+            method,
+            "--extrap_method",
+            extrap_method,
+            "--norm_type",
+            norm_type,
+            # '--user_areas',
+            "--no-log",
+            "--check",
+        ]
+
+        if isinstance(source_grid, xarray.DataArray):
+            command.extend(["--src_missingvalue", source_grid.name])
+        if isinstance(target_grid, xarray.DataArray):
+            command.extend(["--dst_missingvalue", target_grid.name])
+        if ignore_unmapped:
+            command.extend(["--ignore_unmapped"])
+        if line_type is not None:
+            command.extend(["--line_type", line_type])
+        if pole is not None:
+            command.extend(["--pole", pole])
+
+        out = subprocess.check_output(args=command, stderr=subprocess.PIPE)
+        print(out.decode("utf-8"))
+
+        weights = xarray.open_dataset(weight_file.name, engine="netcdf4")
+        # Load so we can delete the temp file
+        return weights.load()
+
+    except subprocess.CalledProcessError as err:
+        print(err)
+        print(err.output.decode("utf-8"))
+        raise
+
+    finally:
+        # Clean up the temporary files
+        source_file.close()
+        target_file.close()
+        weight_file.close()
diff --git a/smmregrid/regrid.py b/smmregrid/regrid.py
index bdf367f..8eac180 100644
--- a/smmregrid/regrid.py
+++ b/smmregrid/regrid.py
@@ -33,270 +33,23 @@
 multiple datasets.
 """
 
-from .dimension import remove_degenerate_axes
 
-import dask.array
 import math
-import os
-import sparse
-import subprocess
 import sys
-import tempfile
+import logging
 import xarray
 import numpy
+import sparse
+import dask.array
+from .dimension import remove_degenerate_axes
+from .cdo_weights import cdo_generate_weights
+from .util import find_vert_coord
+from .weights import compute_weights_matrix3d, compute_weights_matrix, mask_weights, check_mask
 
 
-def cdo_generate_weights(
-    source_grid,
-    target_grid,
-    method="con",
-    extrapolate=True,
-    remap_norm="fracarea",
-    remap_area_min=0.0,
-    icongridpath=None,
-    gridpath=None,
-    extra=None,
-    cdo="cdo"
-):
-    """
-    Generate weights for regridding using CDO
-
-    Available weight generation methods are:
-
-     * bic: SCRIP Bicubic
-     * bil: SCRIP Bilinear
-     * con: SCRIP First-order conservative
-     * con2: SCRIP Second-order conservative
-     * dis: SCRIP Distance-weighted average
-     * laf: YAC Largest area fraction
-     * ycon: YAC First-order conservative
-     * nn: Nearest neighbour
-
-    Run ``cdo gen${method} --help`` for details of each method
-
-    Args:
-        source_grid (xarray.DataArray): Source grid
-        target_grid (xarray.DataArray): Target grid
-            description
-        method (str): Regridding method - default conservative
-        extrapolate (bool): Extrapolate output field
-        remap_norm (str): Normalisation method for conservative methods
-        remap_area_min (float): Minimum destination area fraction
-        gridpath (str): where to store downloaded grids
-        icongridpath (str): location of ICON grids (e.g. /pool/data/ICON)
-        extra: command(s) to apply to source grid before weight generation (can be a list)
-        cdo: path to cdo binary
-    Returns:
-        :obj:`xarray.Dataset` with regridding weights
-    """
-
-    supported_methods = ["bic", "bil", "con", "con2", "dis", "laf", "nn", "ycon"]
-    if method not in supported_methods:
-        raise Exception
-    if remap_norm not in ["fracarea", "destarea"]:
-        raise Exception
-
-    # Make some temporary files that we'll feed to CDO
-    weight_file = tempfile.NamedTemporaryFile()
-
-    if type(source_grid) == str:
-        sgrid = source_grid
-    else:
-        source_grid_file = tempfile.NamedTemporaryFile()
-        source_grid.to_netcdf(source_grid_file.name)
-        sgrid = source_grid_file.name
-
-    if type(target_grid) == str:
-        tgrid = target_grid
-    else:
-        target_grid_file = tempfile.NamedTemporaryFile()
-        target_grid.to_netcdf(target_grid_file.name)
-        tgrid = target_grid_file.name
-
-    # Setup environment
-    env = os.environ
-    if extrapolate:
-        env["REMAP_EXTRAPOLATE"] = "on"
-    else:
-        env["REMAP_EXTRAPOLATE"] = "off"
-
-    env["CDO_REMAP_NORM"] = remap_norm
-    env["REMAP_AREA_MIN"] = "%f" % (remap_area_min)
-
-    if gridpath:
-        env["CDO_DOWNLOAD_PATH"] = gridpath
-    if icongridpath:
-        env["CDO_ICON_GRIDS"] = icongridpath
-
-    try:
-        # Run CDO
-        if extra:
-            # make sure extra is a flat list if it is not already
-            if not isinstance(extra, list):
-                extra = [extra]
-
-            subprocess.check_output(
-                [
-                    cdo,
-                    "gen%s,%s" % (method, tgrid)
-                ] + extra +
-                [
-                    sgrid,
-                    weight_file.name,
-                ],
-                stderr=subprocess.PIPE,
-                env=env,
-            )
-        else:
-            subprocess.check_output(
-                [
-                    cdo,
-                    "gen%s,%s" % (method, tgrid),
-                    sgrid,
-                    weight_file.name,
-                ],
-                stderr=subprocess.PIPE,
-                env=env,
-            )
-
-        # Grab the weights file it outputs as a xarray.Dataset
-        weights = xarray.open_dataset(weight_file.name, engine="netcdf4")
-        return weights
-
-    except subprocess.CalledProcessError as e:
-        # Print the CDO error message
-        print(e.output.decode(), file=sys.stderr)
-        raise
-
-    finally:
-        # Clean up the temporary files
-        if not type(source_grid) == str:
-            source_grid_file.close()
-        if not type(target_grid) == str:
-            target_grid_file.close()
-        weight_file.close()
-
-
-def esmf_generate_weights(
-    source_grid,
-    target_grid,
-    method="bilinear",
-    extrap_method="nearestidavg",
-    norm_type="dstarea",
-    line_type=None,
-    pole=None,
-    ignore_unmapped=False,
-):
-    """Generate regridding weights with ESMF
-
-    https://www.earthsystemcog.org/projects/esmf/regridding
-
-    Args:
-        source_grid (:obj:`xarray.Dataarray`): Source grid. If masked the mask
-            will be used in the regridding
-        target_grid (:obj:`xarray.Dataarray`): Target grid. If masked the mask
-            will be used in the regridding
-        method (str): ESMF Regridding method, see ``ESMF_RegridWeightGen --help``
-        extrap_method (str): ESMF Extrapolation method, see ``ESMF_RegridWeightGen --help``
-
-    Returns:
-        :obj:`xarray.Dataset` with regridding information from
-            ESMF_RegridWeightGen
-    """
-    # Make some temporary files that we'll feed to ESMF
-    source_file = tempfile.NamedTemporaryFile(suffix=".nc")
-    target_file = tempfile.NamedTemporaryFile(suffix=".nc")
-    weight_file = tempfile.NamedTemporaryFile(suffix=".nc")
-
-    rwg = "ESMF_RegridWeightGen"
-
-    if "_FillValue" not in source_grid.encoding:
-        source_grid.encoding["_FillValue"] = -1e20
-
-    if "_FillValue" not in target_grid.encoding:
-        target_grid.encoding["_FillValue"] = -1e20
-
-    try:
-        source_grid.to_netcdf(source_file.name)
-        target_grid.to_netcdf(target_file.name)
-
-        command = [
-            rwg,
-            "--source",
-            source_file.name,
-            "--destination",
-            target_file.name,
-            "--weight",
-            weight_file.name,
-            "--method",
-            method,
-            "--extrap_method",
-            extrap_method,
-            "--norm_type",
-            norm_type,
-            # '--user_areas',
-            "--no-log",
-            "--check",
-        ]
-
-        if isinstance(source_grid, xarray.DataArray):
-            command.extend(["--src_missingvalue", source_grid.name])
-        if isinstance(target_grid, xarray.DataArray):
-            command.extend(["--dst_missingvalue", target_grid.name])
-        if ignore_unmapped:
-            command.extend(["--ignore_unmapped"])
-        if line_type is not None:
-            command.extend(["--line_type", line_type])
-        if pole is not None:
-            command.extend(["--pole", pole])
-
-        out = subprocess.check_output(args=command, stderr=subprocess.PIPE)
-        print(out.decode("utf-8"))
-
-        weights = xarray.open_dataset(weight_file.name, engine="netcdf4")
-        # Load so we can delete the temp file
-        return weights.load()
-
-    except subprocess.CalledProcessError as e:
-        print(e)
-        print(e.output.decode("utf-8"))
-        raise
-
-    finally:
-        # Clean up the temporary files
-        source_file.close()
-        target_file.close()
-        weight_file.close()
-
-
-def compute_weights_matrix(weights):
-    """
-    Convert the weights from CDO/ESMF to a numpy array
-    """
-    w = weights
-    if w.title.startswith("ESMF"):
-        # ESMF style weights
-        src_address = w.col - 1
-        dst_address = w.row - 1
-        remap_matrix = w.S
-        w_shape = (w.sizes["n_a"], w.sizes["n_b"])
-
-    else:
-        # CDO style weights
-        src_address = w.src_address - 1
-        dst_address = w.dst_address - 1
-        remap_matrix = w.remap_matrix[:, 0]
-        w_shape = (w.sizes["src_grid_size"], w.sizes["dst_grid_size"])
-
-    # Create a sparse array from the weights
-    sparse_weights_delayed = dask.delayed(sparse.COO)(
-        [src_address.data, dst_address.data], remap_matrix.data, shape=w_shape
-    )
-    sparse_weights = dask.array.from_delayed(
-        sparse_weights_delayed, shape=w_shape, dtype=remap_matrix.dtype
-    )
-
-    return sparse_weights
+# default spatial dimensions and vertical coordinates
+default_space_dims = ['i', 'j', 'x', 'y', 'lon', 'lat', 'longitude', 'latitude',
+                     'cell', 'cells', 'ncells', 'values', 'value', 'nod2', 'pix']
 
 
 def apply_weights(source_data, weights, weights_matrix=None, masked=True, space_dims=None):
@@ -307,22 +60,22 @@ def apply_weights(source_data, weights, weights_matrix=None, masked=True, space_
         source_data (xarray.DataArray): Source dataset
         weights (xarray.DataArray): CDO weights information
         masked (bool): if the DataArray is masked
-        space_dims (list): dimensions on which the interpolation has to be done (e.g. ['lon', 'lat'])
+        space_dims (list): dimensions on which the interpolation has to be done (e.g.['lon', 'lat'])
 
     Returns:
         xarray.DataArray: Regridded version of the source dataset
     """
 
-    # Understand immediately if we need to retunr something or not
+    # Understand immediately if we need to return something or not
     # This is done if we have bounds variables
     if ("bnds" in source_data.name or "bounds" in source_data.name):
 
         # we keep time bounds, and we ignore all the rest
         if 'time' in source_data.name:
-            # print('original ' + source_data.name)
+            logging.info('original %s', source_data.name)
             return source_data
         else:
-            # print('empty ' + source_data.name)
+            logging.info('empty %s', source_data.name)
             return xarray.DataArray(data=None)
 
     # Alias the weights dataset from CDO
@@ -336,12 +89,13 @@ def apply_weights(source_data, weights, weights_matrix=None, masked=True, space_
     # array, multiplied by the weights matrix, then unstacked back into a 2d
     # array
 
-    if w.title.startswith("ESMF"):
+    #if w.title.startswith("ESMF"):
+    if "S" in w.variables:
         # ESMF style weights
-        src_address = w.col - 1
-        dst_address = w.row - 1
-        remap_matrix = w.S
-        w_shape = (w.sizes["n_a"], w.sizes["n_b"])
+        #src_address = w.col - 1
+        #dst_address = w.row - 1
+        #remap_matrix = w.S
+        #w_shape = (w.sizes["n_a"], w.sizes["n_b"])
 
         dst_grid_shape = w.dst_grid_dims.values
         dst_grid_center_lat = w.yc_b.data.reshape(dst_grid_shape[::-1])
@@ -353,10 +107,10 @@ def apply_weights(source_data, weights, weights_matrix=None, masked=True, space_
 
     else:
         # CDO style weights
-        src_address = w.src_address - 1
-        dst_address = w.dst_address - 1
-        remap_matrix = w.remap_matrix[:, 0]
-        w_shape = (w.sizes["src_grid_size"], w.sizes["dst_grid_size"])
+        #src_address = w.src_address - 1
+        #dst_address = w.dst_address - 1
+        #remap_matrix = w.remap_matrix[:, 0]
+        #w_shape = (w.sizes["src_grid_size"], w.sizes["dst_grid_size"])
 
         dst_grid_shape = w.dst_grid_dims.values
         dst_grid_center_lat = w.dst_grid_center_lat.data.reshape(
@@ -369,18 +123,17 @@ def apply_weights(source_data, weights, weights_matrix=None, masked=True, space_
         )
 
         dst_mask = w.dst_grid_imask
-        src_mask = w.src_grid_imask
+        #src_mask = w.src_grid_imask
 
         axis_scale = 180.0 / math.pi  # Weight lat/lon in radians
 
     # Dimension on which we can produce the interpolation
-    if space_dims is None: 
-        space_dims = ['i', 'j', 'x', 'y', 'lon', 'lat', 'longitude', 'latitude',
-                         'cell', 'cells', 'ncells', 'values', 'value', 'nod2', 'pix']
+    if space_dims is None:
+        space_dims = default_space_dims
     
     if not any(x in source_data.dims for x in space_dims):
-        print("None of dimensions on which we can interpolate is found in the DataArray. Does your DataArray include any of these?")
-        print(space_dims)
+        logging.error("None of dimensions on which we can interpolate is found in the DataArray. Does your DataArray include any of these?")
+        logging.error(space_dims)
         sys.exit('Dimensions mismatch')
 
     # Find dimensions to keep
@@ -388,8 +141,8 @@ def apply_weights(source_data, weights, weights_matrix=None, masked=True, space_
 
     kept_shape = list(source_data.shape[0:nd])
     kept_dims = list(source_data.dims[0:nd])
-    # print(kept_dims)
-    # print(kept_dims)
+    logging.info('Dimension kept:')
+    logging.info(kept_dims)
 
     if weights_matrix is None:
         weights_matrix = compute_weights_matrix(weights)
@@ -406,6 +159,7 @@ def apply_weights(source_data, weights, weights_matrix=None, masked=True, space_
     dask.array.ma.set_fill_value(source_array, 1e20)
     source_array = dask.array.ma.fix_invalid(source_array)
     source_array = dask.array.ma.filled(source_array)
+
     target_dask = dask.array.tensordot(source_array, weights_matrix, axes=1)
 
     # define and compute the new mask
@@ -438,6 +192,7 @@ def apply_weights(source_data, weights, weights_matrix=None, masked=True, space_
         },
         name=source_data.name,
     )
+
     target_da.coords["lat"] = xarray.DataArray(dst_grid_center_lat, dims=["i", "j"])
     target_da.coords["lon"] = xarray.DataArray(dst_grid_center_lon, dims=["i", "j"])
 
@@ -473,7 +228,6 @@ def apply_weights(source_data, weights, weights_matrix=None, masked=True, space_
 
     return target_da
 
-
 class Regridder(object):
     """Set up the regridding operation
 
@@ -490,36 +244,64 @@ class Regridder(object):
         source_grid (:class:`coecms.grid.Grid` or :class:`xarray.DataArray`): Source grid / sample dataset
         target_grid (:class:`coecms.grid.Grid` or :class:`xarray.DataArray`): Target grid / sample dataset
         weights (:class:`xarray.Dataset`): Pre-computed interpolation weights
+        vert_coord (str): Name of the vertical coordinate. 
+                          If provided, 3D weights are generated (default: None)
+        method (str): Method to use for interpolation (default: 'con')
+        space_dims (list): list of dimensions to interpolate (default: None)
+        transpose (bool): transpose the output so that the vertical coordinate is 
+                          just before other spatial coords (dafault: True)
     """
 
-    def __init__(self, source_grid=None, target_grid=None, weights=None, method='con', space_dims=None, cdo="cdo"):
-
-        if (source_grid is None or target_grid is None) and weights is None:
-            raise Exception(
+    def __init__(self, source_grid=None, target_grid=None, weights=None,
+                 method='con', space_dims=None, vert_coord=None, transpose=True,
+                 cdo='cdo'):
+        
+        if (source_grid is None or target_grid is None) and (weights is None):
+            raise ValueError(
                 "Either weights or source_grid/target_grid must be supplied"
             )
-
+        
+        self.transpose = transpose
+    
         # Is there already a weights file?
         if weights is not None:
-
             if not isinstance(weights, xarray.Dataset):
                 self.weights = xarray.open_mfdataset(weights)
             else:
                 self.weights = weights
+            
+            if not vert_coord:
+                self.vert_coord = find_vert_coord(self.weights)
+            else:
+                self.vert_coord = vert_coord
         else:
+
+            # Check if there is a vertical coordinate for 3d oceanic data
+            if not vert_coord:
+                self.vert_coord = find_vert_coord(source_grid)
+            else:
+                self.vert_coord = vert_coord
+
             # Generate the weights with CDO
-            # _source_grid = identify_grid(source_grid)
-            # _target_grid = identify_grid(target_grid)
-            self.weights = cdo_generate_weights(source_grid, target_grid, method=method, cdo=cdo)
-            #sys.exit('Missing capability of creating weights...')
+            self.weights = cdo_generate_weights(source_grid, target_grid, method=method,
+                                                vert_coord=self.vert_coord, cdo=cdo)
+
+        if self.vert_coord:
+            self.weights_matrix = compute_weights_matrix3d(self.weights, self.vert_coord)
+        else: 
+            self.weights_matrix = compute_weights_matrix(self.weights)
 
-        self.weights_matrix = compute_weights_matrix(self.weights)
+        # this section is used to create a target mask initializing the CDO weights (both 2d and 3d)
+        if "dst_grid_masked" in self.weights.variables:  # has a destination mask been precomputed?
+            self.masked = self.weights.dst_grid_masked.data  # ok, let's use it
+        else:
+            # compute the destination mask now
+            self.weights = mask_weights(self.weights, self.weights_matrix, self.vert_coord)
+            self.masked = check_mask(self.weights, self.vert_coord)
 
-        # this section is used to create a target mask initializing the CDO weights
-        self.weights = mask_weigths(self.weights, self.weights_matrix)
-        self.masked = check_mask(self.weights)
         self.space_dims = space_dims
 
+
     def regrid(self, source_data):
         """Regrid ``source_data`` to match the target grid
 
@@ -535,51 +317,102 @@ def regrid(self, source_data):
         if isinstance(source_data, xarray.Dataset):
 
             # apply the regridder on each DataArray
-            out = source_data.map(self.regrid, keep_attrs=True)
+            out = source_data.map(self.regrid_array, keep_attrs=True)
 
             # clean from degenerated variables
             degen_vars = [var for var in out.data_vars if out[var].dims == ()]
             return out.drop_vars(degen_vars)
-            # else:
-            #    sys.exit("source data has different mask, this can lead to unexpected" \
-            #        "results due to the format in which weights are generated. Aborting...")
-
+        
         elif isinstance(source_data, xarray.DataArray):
 
-            # print('DataArray access!')
-            return apply_weights(
-                source_data, self.weights, weights_matrix=self.weights_matrix, 
-                masked=self.masked, space_dims=self.space_dims
-            )
+            return self.regrid_array(source_data)
+
         else:
-            sys.exit('Cannot process this source_data, sure it is xarray?')
+            sys.exit('The object provided is not a Xarray object!')
 
+    def regrid_array(self, source_data):
 
-def mask_weigths(weights, weights_matrix):
-    """This functions precompute the mask for the target interpolation
-    Takes as input the weights from CDO and the precomputed weights matrix
-    Return the target mask"""
+        """Regridding selection through 2d and 3d arrays"""
+        
+        if (self.vert_coord and self.vert_coord in source_data.coords):
+            # if this is a 3D we specified the vertical coord and it has it
+            return self.regrid3d(source_data)
+        else:
+            return self.regrid2d(source_data)
 
-    src_mask = weights.src_grid_imask.data
-    target_mask = dask.array.tensordot(src_mask, weights_matrix, axes=1)
-    target_mask = dask.array.where(target_mask < 0.5, 0, 1)
-    weights.dst_grid_imask.data = target_mask
-    return weights
+    def regrid3d(self, source_data):
+        """Regrid ``source_data`` to match the target grid - 3D version
 
+        Args:
+            source_data (:class:`xarray.DataArray`): Source
+            variable
 
-def check_mask(weights):
-    """This check if the target mask is empty or full and
-    return a bool to be passed to the regridder"""
+        Returns:
+            :class:`xarray.DataArray` with a regridded
+            version of the source variable
+        """
 
-    w = weights.dst_grid_imask
-    v = w.sum()/len(w)
-    if v == 1:
-        return False
-    else:
-        return True
+        logging.info('3D DataArray access!')
 
+        # CDO 2.2.0 fix
+        if "numLinks" in self.weights.dims:
+            links_dim = "numLinks"
+        else:
+            links_dim = "num_links"
+        
+        # this is necessary to remove lev-bounds, temporary hack since they should
+        # be treated in a smarter way
+        if ("bnds" in source_data.name or "bounds" in source_data.name):
+            return source_data
 
-def regrid(source_data, target_grid=None, weights=None, cdo="cdo"):
+        data3d_list = []
+        for lev in range(0, source_data.coords[self.vert_coord].values.size):
+            xa = source_data.isel(**{self.vert_coord: lev})
+            wa = self.weights.isel(**{self.vert_coord: lev})
+            nl = wa.link_length.values
+            wa = wa.isel(**{links_dim: slice(0, nl)})
+            wm = self.weights_matrix[lev]
+            mm = self.masked[lev]
+            data3d_list.append(apply_weights(
+                xa, wa, weights_matrix=wm,
+                masked=mm, space_dims=self.space_dims)
+            )
+        data3d = xarray.concat(data3d_list, dim=self.vert_coord)
+
+        if self.transpose:
+            # Make sure that the vertical dimension is just before the spatial ones
+            if self.space_dims is None:
+                space_dims = default_space_dims
+            else:
+                space_dims = self.space_dims
+            dims = list(data3d.dims)
+            index = min([i for i, s in enumerate(dims) if s in space_dims])
+            dimst= dims[1:index] + [dims[0]] + dims[index:]
+            data3d = data3d.transpose(*dimst)
+
+            return data3d
+        else:
+            sys.exit('Cannot process this source data, are you sure it is an xarray?')
+
+    def regrid2d(self, source_data):
+        """Regrid ``source_data`` to match the target grid, 2D version
+
+        Args:
+            source_data (:class:`xarray.DataArray`): Source
+            variable
+
+        Returns:
+            :class:`xarray.DataArray` with a regridded
+            version of the source variable
+        """
+        logging.info('2D DataArray access!')
+        return apply_weights(
+            source_data, self.weights, weights_matrix=self.weights_matrix,
+            masked=self.masked, space_dims=self.space_dims
+        )
+
+
+def regrid(source_data, target_grid=None, weights=None, vert_coord=None, transpose=True, cdo='cdo'):
     """
     A simple regrid. Inefficient if you are regridding more than one dataset
     to the target grid because it re-generates the weights each time you call
@@ -590,11 +423,25 @@ def regrid(source_data, target_grid=None, weights=None, cdo="cdo"):
     Args:
         source_data (:class:`xarray.DataArray`): Source variable
         target_grid (:class:`coecms.grid.Grid` or :class:`xarray.DataArray`): Target grid / sample variable
+        vert_coord (str): Name of the vertical coordinate. 
+                          If provided, 3D weights are generated (default: None)
+        weights (:class:`xarray.Dataset`): Pre-computed interpolation weights
+        transpose (bool): If True, transpose the output so that the vertical
+                          coordinate is just before the other spatial coordinates (default: True)
+
         cdo (path): path of cdo executable ["cdo"]
     Returns:
         :class:`xarray.DataArray` with a regridded version of the source variable
     """
 
-    regridder = Regridder(source_data, target_grid=target_grid, weights=weights, cdo=cdo)
-
+    regridder = Regridder(source_data, target_grid=target_grid, weights=weights, 
+                          vert_coord=vert_coord, cdo=cdo, transpose=transpose)
     return regridder.regrid(source_data)
+
+
+# def combine_2d_to_3d(array_list, dim_name, dim_values):
+#     """
+#     Function to combine a list of 2D xarrays into a 3D one adding a vertical coordinate lev
+#     """
+#     new_array = [x.assign_coords({dim_name: d}) for x, d in zip(array_list, dim_values)]
+#     return xarray.concat(new_array, dim_name)
\ No newline at end of file
diff --git a/smmregrid/util.py b/smmregrid/util.py
new file mode 100644
index 0000000..184441d
--- /dev/null
+++ b/smmregrid/util.py
@@ -0,0 +1,17 @@
+"""Utils module"""
+
+import xarray
+
+default_vert_coords = ['lev', 'nz1', 'nz']
+
+def find_vert_coord(xfield):
+    """
+    Find a vertical coordinate among defaults
+    Used to define if we need the 3d interpolation with adaptive mask
+    """
+
+    if isinstance(xfield, str):
+        xfield = xarray.open_dataset(xfield)
+    for coord in default_vert_coords:
+        if coord in xfield.coords:
+            return coord
\ No newline at end of file
diff --git a/smmregrid/weights.py b/smmregrid/weights.py
new file mode 100644
index 0000000..98cc7bc
--- /dev/null
+++ b/smmregrid/weights.py
@@ -0,0 +1,123 @@
+"""Weight calculation utilities"""
+
+import dask.array
+import sparse
+
+
+def compute_weights_matrix3d(weights, vert_coord='lev'):
+    """
+    Convert the weights from CDO to a list of numpy arrays
+    """
+
+    # CDO 2.2.0 fix
+    if "numLinks" in weights.dims:
+        links_dim = "numLinks"
+    else:
+        links_dim = "num_links"
+
+    sparse_weights = []
+    nvert = weights[vert_coord].values.size
+    for i in range(0, nvert):
+        w = weights.loc[{vert_coord: i}]
+        nl = w.link_length.values
+        w = w.isel(**{links_dim: slice(0, nl)})
+        sparse_weights.append(compute_weights_matrix(w))
+    
+    return sparse_weights
+
+
+def compute_weights_matrix(weights):
+    """
+    Convert the weights from CDO/ESMF to a numpy array
+    """
+    w = weights
+    #if w.title.startswith("ESMF"):
+    if "S" in w.variables:
+        # ESMF style weights
+        src_address = w.col - 1
+        dst_address = w.row - 1
+        remap_matrix = w.S
+        w_shape = (w.sizes["n_a"], w.sizes["n_b"])
+
+    else:
+        # CDO style weights
+        src_address = w.src_address - 1
+        dst_address = w.dst_address - 1
+        remap_matrix = w.remap_matrix[:, 0]
+        w_shape = (w.sizes["src_grid_size"], w.sizes["dst_grid_size"])
+   
+    # Create a sparse array from the weights
+    sparse_weights_delayed = dask.delayed(sparse.COO)(
+        [src_address.data, dst_address.data], remap_matrix.data, shape=w_shape
+    )
+    sparse_weights = dask.array.from_delayed(
+        sparse_weights_delayed, shape=w_shape, dtype=remap_matrix.dtype
+    )
+
+    return sparse_weights
+
+
+def mask_tensordot(src_mask, weights_matrix):
+
+    """Apply tensor dot product to source mask to return destination mask"""
+
+    target_mask = dask.array.tensordot(src_mask, weights_matrix, axes=1)
+    target_mask = dask.array.where(target_mask < 0.5, 0, 1)
+    return target_mask
+
+
+def mask_weights(weights, weights_matrix, vert_coord=None):
+    """This functions precompute the mask for the target interpolation
+    Takes as input the weights from CDO and the precomputed weights matrix
+    Return the target mask: handle the 3d case"""
+
+    src_mask = weights.src_grid_imask
+    if vert_coord is not None:
+        for nlev in range(len(weights[vert_coord])):
+            mask = src_mask.loc[{vert_coord: nlev}].data
+            weights['dst_grid_imask'].loc[{vert_coord: nlev}] = mask_tensordot(mask, weights_matrix[nlev])
+    else:
+        mask = src_mask.data
+        weights['dst_grid_imask'].data = mask_tensordot(mask, weights_matrix)
+
+    return weights
+
+
+def check_mask(weights, vert_coord=None):
+    """Check if the target mask is empty or full and
+    return a bool to be passed to the regridder.
+    Handle the 3d case (5x time faster)"""
+
+    wdst = weights['dst_grid_imask']
+    if vert_coord is not None:
+        check = wdst.mean(dim=tuple(dim for dim in wdst.dims if dim != vert_coord))
+        out = (check != 1).data.tolist()
+    else:
+        check = wdst.mean()
+        out = (check != 1).data
+
+    return out
+
+
+def check_mask_old(weights, vert_coord=None):
+    """Old version: check if the target mask is empty or full and
+    return a bool to be passed to the regridder.
+    Handle the 3d case"""
+
+    if vert_coord is not None:
+        result = []
+        for nlev in range(len(weights[vert_coord])):
+            w = weights['dst_grid_imask'].loc[{vert_coord: nlev}]
+            v = w.sum()/len(w)
+            if v == 1:
+                result.append(False)
+            else:
+                result.append(True)
+        return result
+    else:
+        w = weights['dst_grid_imask']
+        v = w.sum()/len(w)
+        if v == 1:
+            return False
+        else:
+            return True
\ No newline at end of file


Test patch:
diff --git a/.github/workflows/mambatest.yml b/.github/workflows/mambatest.yml
index 8142f4b..d1e4557 100644
--- a/.github/workflows/mambatest.yml
+++ b/.github/workflows/mambatest.yml
@@ -40,14 +40,17 @@ jobs:
     - name: Lint with flake8
       run: |
         # install flake8
-        conda install flake8
+        python -m pip install flake8
         # stop the build if there are Python syntax errors or undefined names
         flake8 . --count --select=E9,F63,F7,F82 --show-source --statistics
         # exit-zero treats all errors as warnings. The GitHub editor is 127 chars wide
         flake8 . --count --exit-zero --max-complexity=10 --max-line-length=127 --statistics
-    - name: Test with pytest
+    - name: Install pytest
       run: |
-        conda install pytest
+         python -m pip install pytest
+    - name: Run tests
+      run: |
+        python -m pytest
         
   coverage:
     needs: build
diff --git a/tests/basic_test.py b/tests/basic_test.py
index b53970d..6708640 100644
--- a/tests/basic_test.py
+++ b/tests/basic_test.py
@@ -7,6 +7,18 @@
 indir = 'tests/data'
 tfile = os.path.join(indir, 'r360x180.nc')
 
+# test for pressure levels on gaussian grid (3D)
+@pytest.mark.parametrize("method", ['con', 'nn'])
+def test_nemo_3d(method):
+    fff = check_cdo_regrid(os.path.join(indir, 'so3d-nemo.nc'), tfile, method=method)
+    assert fff is True
+
+# test for pressure levels on gaussian grid (3D)
+@pytest.mark.parametrize("method", ['con', 'nn'])
+def test_fesom_3d(method):
+    fff = check_cdo_regrid(os.path.join(indir, 'temp3d-fesom.nc'), tfile, method=method)
+    assert fff is True
+
 # test for gaussian reduced grid (only nn)
 @pytest.mark.parametrize("method", ['nn', 'con'])
 def test_healpix(method):
@@ -43,8 +55,14 @@ def test_curivilinear(method):
     fff = check_cdo_regrid(os.path.join(indir, 'onlytos-ipsl.nc'), tfile, method=method)
     assert fff is True
 
-# test for pressure levels on gaussian grid
+# test for pressure levels on gaussian grid (2D, level-by-level)
 @pytest.mark.parametrize("method", ['con', 'nn', 'bil'])
-def test_plev_gaussian(method):
+def test_levbylev_plev_gaussian(method):
     fff = check_cdo_regrid(os.path.join(indir, 'ua-ecearth.nc'), tfile, method=method)
     assert fff is True
+
+# test for pressure levels on gaussian grid (3D)
+@pytest.mark.parametrize("method", ['con', 'nn', 'bil'])
+def test_full_plev_gaussian(method):
+    fff = check_cdo_regrid(os.path.join(indir, 'ua-ecearth.nc'), tfile, method=method, vert_coord="plev")
+    assert fff is True
\ No newline at end of file
diff --git a/tests/data/so3d-nemo.nc b/tests/data/so3d-nemo.nc
new file mode 100644
index 0000000..b2a63f9
Binary files /dev/null and b/tests/data/so3d-nemo.nc differ
diff --git a/tests/data/temp3d-fesom.nc b/tests/data/temp3d-fesom.nc
new file mode 100644
index 0000000..1090279
Binary files /dev/null and b/tests/data/temp3d-fesom.nc differ