Repository: FESOM/pyfesom2 (language: python)

Issue:
fix case for different mesh diag versions



Discussion / hints:

Fix patch:
diff --git a/.github/workflows/pyfesom2_onPR.yml b/.github/workflows/pyfesom2_onPR.yml
index 2535736..f5fbad7 100644
--- a/.github/workflows/pyfesom2_onPR.yml
+++ b/.github/workflows/pyfesom2_onPR.yml
@@ -20,7 +20,7 @@ jobs:
       fail-fast: false
       matrix:
         os: ["ubuntu-latest", "macos-latest", "windows-latest"]
-        python-version: ["3.6", "3.7", "3.8"]
+        python-version: ["3.7", "3.8", "3.9"]
     steps:
       - uses: actions/checkout@v2
       - uses: conda-incubator/setup-miniconda@v2
diff --git a/ci/requirements-py37.yml b/ci/requirements-py37.yml
index 9ac7292..c990a13 100644
--- a/ci/requirements-py37.yml
+++ b/ci/requirements-py37.yml
@@ -17,9 +17,9 @@ dependencies:
   - ipython
   - pytest
   - cmocean
-  - zarr
-  - fsspec
+  - zarr>=2.8.3
+  - fsspec==2021.06.1
   - requests
-  - aiohttp 
+  - aiohttp
   - dask
   - xarray
diff --git a/pyfesom2/accessor.py b/pyfesom2/accessor.py
index c48c359..fdc9de6 100644
--- a/pyfesom2/accessor.py
+++ b/pyfesom2/accessor.py
@@ -159,7 +159,7 @@ def __init__(self, lon: ArrayLike, lat: ArrayLike, faces: ArrayLike):
 
 def select_bbox(xr_obj: Union[xr.DataArray, xr.Dataset],
                 bbox: BoundingBox,
-                faces: Optional[ArrayLike] = None) -> xr.Dataset:
+                coords_dataset: Optional[xr.Dataset] = None) -> xr.Dataset:
     """Returns subset Dataset or DataArray for bounding box.
 
     This method uses triangulation indices in faces (as argument or as coordinate in a dataset) to select nodes
@@ -175,38 +175,44 @@ def select_bbox(xr_obj: Union[xr.DataArray, xr.Dataset],
     bbox
         Bounding box can be specified as as sequence of size 4 (lists or tuple or array) containing bounds
         from lower-left to upper-right of longitudes and latitudes. For instance: (xmin, ymin, xmax, ymax).
-    faces
-        For Datasets containing faces as coordinate information the argument is not necessary.
-        For DataArrays faces argument, defining indices of faces defining triangles, is necessary.
-
+    coords_dataset
+        if xr_obj is a fesom2 Dataset, this argument is not necessary.
+        if xr_obj is a xr.DataArray: all metadata needed for selection cannot be often contained in a dataarray,
+        this can be provided as a (coordinate) dataset containing necessary additional coordinate information
+        using this argument.
     Returns
     -------
 
     """
     from .ut import cut_region
-    faces = getattr(xr_obj, "faces", faces)
-    if faces is None:
-        raise ValueError(f"When passing a dataset it needs have faces in coords, or "
-                         f"faces need to be passed explicitly.\n"
-                         f"When passing a data array, argument faces can't be None,"
-                         f"faces must be indices[nelem,3] that define triangles.")
-
-    mesh = SimpleMesh(xr_obj.lon, xr_obj.lat, faces)
+
+    if isinstance(xr_obj, xr.Dataset):
+        lats, lons = xr_obj.lat, xr_obj.lon
+        faces = xr_obj.faces
+        ret = xr_obj
+    elif isinstance(xr_obj, xr.DataArray):
+        if coords_dataset is None:
+            raise ValueError(f"Selection on a dataarray needs coords_dataset argument containing lon,lat, faces as"
+                             f" coordinates.")
+        coords_dataset = coords_dataset[['lon', 'lat', 'faces']] # in case it is not just get needed coords and dims
+        lats, lons = coords_dataset.lat, coords_dataset.lon
+        faces = coords_dataset.faces
+        ret = xr.merge([xr_obj, coords_dataset])
+
+    mesh = SimpleMesh(lons, lats, faces)
     bbox = np.asarray(bbox)
     # cut region takes xmin, xmax, ymin, ymax
-    cut_faces, _ = cut_region(mesh, [bbox[0], bbox[2], bbox[1], bbox[3]])
+    cut_faces, cut_indices = cut_region(mesh, [bbox[0], bbox[2], bbox[1], bbox[3]])
     cut_faces = np.asarray(cut_faces)
     uniq, inv_index = np.unique(cut_faces.ravel(), return_inverse=True)
     new_faces = inv_index.reshape(cut_faces.shape)
-    ret = xr_obj.isel(nod2=uniq)
-    if isinstance(xr_obj, xr.DataArray):
-        ret = ret.to_dataset()
-    ret = ret.assign_coords({'faces': (('nelem', 'three'), new_faces)})
+    ret = ret.isel(nod2=uniq, nelem=cut_indices)
+    ret['faces'] = (('nelem', 'three'), new_faces)
     return ret
 
 
 def select_region(xr_obj: Union[xr.DataArray, xr.Dataset], region: Region,
-                  faces: Optional[ArrayLike] = None) -> xr.Dataset:
+                  coords_dataset: Optional[xr.DataArray] = None) -> xr.Dataset:
     """Returns a FESOM data subset for specified arbitrary region.
 
     This method uses vectorized Shapely's vectorized routines to find nodes contained by specified polygon. Faces
@@ -239,20 +245,27 @@ def select_region(xr_obj: Union[xr.DataArray, xr.Dataset], region: Region,
         raise ValueError(f"Supplied region data can be a sequence of (minlon, minlat, maxlon, maxlat) or "
                          f"a Shapely's Polygon. This {region} is not supported.")
 
-    faces = getattr(xr_obj, "faces", faces)
-    if faces is None:
-        raise ValueError(f"When passing a dataset it needs have faces in coords, or"
-                         f"faces need to be passed explicitly.\n"
-                         f"When passing a data array, argument faces can't be None,"
-                         f"faces must be indices[nelem,3] that define triangles.")
-    faces = np.asarray(faces)
+    if isinstance(xr_obj, xr.Dataset):
+        lats, lons = xr_obj.lat, xr_obj.lon
+        faces = np.asarray(xr_obj.faces)
+        nelem = np.asarray(xr_obj.nelem)
+        ret = xr_obj
+    elif isinstance(xr_obj, xr.DataArray):
+        if coords_dataset is None:
+            raise ValueError(f"Selection on a dataarray needs coords_dataset argument containing lon,lat, faces as"
+                             f" coordinates.")
+        coords_dataset = coords_dataset[['lon', 'lat', 'faces']] # in case it is not just get needed coords and dims
+        lats, lons = coords_dataset.lat, coords_dataset.lon
+        faces = np.asarray(coords_dataset.faces)
+        nelem = np.asarray(coords_dataset.nelem)
+        ret = xr.merge([xr_obj, coords_dataset])
 
     # buffer is necessry to facilitte floating point comparisions
     # buffer can be thought as tolerance around region in degrees
     # its value should be at least precision of data type of lats, lons (np.finfo)
     region = region.buffer(1e-6)
     prep_region = prep(region)
-    selection = vectorized_contains(prep_region, np.asarray(xr_obj.lon), np.asarray(xr_obj.lat))
+    selection = vectorized_contains(prep_region, np.asarray(lons), np.asarray(lats))
     if np.count_nonzero(selection) == 0:
         warnings.warn('No points in domain are within region, returning original data.')
         return xr_obj
@@ -260,10 +273,11 @@ def select_region(xr_obj: Union[xr.DataArray, xr.Dataset], region: Region,
     selection = selection[faces]
     face_mask = np.all(selection, axis=1)
     cut_faces = faces[face_mask]
-    cut_faces = np.array(cut_faces, ndmin=1)
+    cut_faces = np.asarray(cut_faces)
+    cut_indices = nelem[face_mask]
     uniq, inv_index = np.unique(cut_faces.ravel(), return_inverse=True)
     new_faces = inv_index.reshape(cut_faces.shape)
-    ret = xr_obj.isel(nod2=uniq)
+    ret = ret.isel(nod2=uniq, nelem=cut_indices)
 
     if 'faces' in ret.coords:
         ret = ret.drop_vars('faces')
@@ -272,11 +286,7 @@ def select_region(xr_obj: Union[xr.DataArray, xr.Dataset], region: Region,
         warnings.warn("No found points for the region are contained in dataset's triangulation (faces), "
                       "returning object without faces.")
         return ret  # no faces in coords
-
-    if isinstance(xr_obj, xr.DataArray):
-        ret = ret.to_dataset()
-
-    ret = ret.assign_coords({'faces': (('nelem', 'three'), new_faces)})
+    ret['faces'] = (('nelem', 'three'), new_faces)
     return ret
 
 
@@ -355,7 +365,7 @@ def select_points(xr_obj: Union[xr.Dataset, xr.DataArray],
     else:
         raise NotImplementedError('tolerance is currently not supported.')
 
-    other_dims = {k: xr.DataArray(np.array(v, ndmin=1), dims=sel_dim) for k, v in other_dims.items()}
+    other_dims = {k: xr.DataArray(np.asarray(v), dims=sel_dim) for k, v in other_dims.items()}
     ret_obj = xr_obj.isel(nod2=xr.DataArray(ind, dims=sel_dim)).sel(**other_dims, method=method)
 
     # from faces, which will not be useful in returned dataset
@@ -603,10 +613,12 @@ def select(self, method: str = 'nearest', tolerance: float = None, region: Optio
         xr.Dataset
         """
 
-        sel_obj = self._xrobj.to_dataset()
-        sel_obj = sel_obj.assign_coords({'faces': (self._context_dataset.faces.dims,
-                                                   self._context_dataset.faces.values)})
-        tree = self._context_dataset.pyfesom2._tree
+        coords_dataset = self._context_dataset[['lon', 'lat', 'faces']]
+        sel_obj = xr.merge([self._xrobj, coords_dataset])
+        if "lat" in indexers or "lon" in indexers:
+            tree = self._context_dataset.pyfesom2._tree
+        else:
+            tree = None
         sel_obj = select(sel_obj, method=method, tolerance=tolerance, region=region, path=path, tree=tree,
                          **indexers)
         return sel_obj
diff --git a/pyfesom2/config.py b/pyfesom2/config.py
new file mode 100644
index 0000000..165a57d
--- /dev/null
+++ b/pyfesom2/config.py
@@ -0,0 +1,4 @@
+import os.path
+USER_CONFIG_DIR = os.path.expanduser(os.path.join("~", ".pyfesom2"))
+USER_CACHE_DIR = os.path.join(USER_CONFIG_DIR, "cached_data")
+
diff --git a/pyfesom2/datasets.py b/pyfesom2/datasets.py
index 3836180..d658964 100644
--- a/pyfesom2/datasets.py
+++ b/pyfesom2/datasets.py
@@ -1,12 +1,14 @@
+import os
 import warnings
 from typing import Sequence
 from typing import Tuple, Optional
 
+import fsspec
 import numpy as np
 import xarray as xr
 
-from . import load_mesh
-from .ut import get_no_cyclic
+from .config import USER_CACHE_DIR
+from .load_mesh_data import load_mesh
 
 cmip6_grids = {
     'AWI-CM-LR': {
@@ -18,9 +20,6 @@
     'AWI-CM-HR': {
         "path_url"   : "https://swift.dkrz.de/v1/dkrz_035d8f6ff058403bb42f8302e6badfbc/pyfesom2/cmip6-grids/zarr/awicm-hr",
         "Dataset URL": "https://swiftbrowser.dkrz.de/public/dkrz_035d8f6ff058403bb42f8302e6badfbc/pyfesom2/cmip6-grids/zarr/awicm-hr"},
-    "A01"      : {
-        "path_url"   : "https://swift.dkrz.de/v1/dkrz_02942825-0cab-44f3-ad37-80fd5d2e37e3/FESOM2_data/A01",
-        "Dataset URL": "https://swiftbrowser.dkrz.de/public/dkrz_02942825-0cab-44f3-ad37-80fd5d2e37e3/FESOM2_data/A01"}
 }
 
 frontier_datasets = {
@@ -32,17 +31,21 @@
         "Dataset URL": "https://swiftbrowser.dkrz.de/public/dkrz_035d8f6ff058403bb42f8302e6badfbc/pyfesom2/frontier/rossby42_level_aceess"},
     "ROSSBY4.2_spatial": {
         "path_url"   : "https://swift.dkrz.de/v1/dkrz_035d8f6ff058403bb42f8302e6badfbc/pyfesom2/frontier/rossby42_spatial_aceess",
-        "Dataset URL": "https://swiftbrowser.dkrz.de/public/dkrz_035d8f6ff058403bb42f8302e6badfbc/pyfesom2/frontier/rossby42_spatial_aceess"}}
+        "Dataset URL": "https://swiftbrowser.dkrz.de/public/dkrz_035d8f6ff058403bb42f8302e6badfbc/pyfesom2/frontier/rossby42_spatial_aceess"},
+    "A01"              : {
+        "path_url"   : "https://swift.dkrz.de/v1/dkrz_02942825-0cab-44f3-ad37-80fd5d2e37e3/FESOM2_data/A01",
+        "Dataset URL": "https://swiftbrowser.dkrz.de/public/dkrz_02942825-0cab-44f3-ad37-80fd5d2e37e3/FESOM2_data/A01"}
+}
 
 all_datasets = {
     "CORE"   : {
-        "path_url"   : "https://swift.dkrz.de/v1/dkrz_02942825-0cab-44f3-ad37-80fd5d2e37e3/FESOM2_data/LCORE",
-        "Dataset URL": "https://swiftbrowser.dkrz.de/public/dkrz_02942825-0cab-44f3-ad37-80fd5d2e37e3/FESOM2_data/LCORE",
-        "var_list"   : ["temp", "salt", "a_ice", "m_ice", "ssh", "sst", "mesh"]},
+        "path_url"   : "https://swift.dkrz.de/v1/dkrz_035d8f6ff058403bb42f8302e6badfbc/pyfesom2/tutorial/core2",
+        "Dataset URL": "https://swiftbrowser.dkrz.de/public/dkrz_035d8f6ff058403bb42f8302e6badfbc/pyfesom2/tutorial/core2",
+        "group"      : "variables"},
     "pi-grid": {
         'path_url'   : "https://swift.dkrz.de/v1/dkrz_035d8f6ff058403bb42f8302e6badfbc/pyfesom2/tutorial/pi-grid",
-        "var_list"   : ['a_ice', 'm_ice', 'temp', 'u', 'v', 'w', 'mesh'],
-        "Dataset URL": "https://swiftbrowser.dkrz.de/public/dkrz_035d8f6ff058403bb42f8302e6badfbc/pyfesom2/tutorial/pi-grid"},
+        "Dataset URL": "https://swiftbrowser.dkrz.de/public/dkrz_035d8f6ff058403bb42f8302e6badfbc/pyfesom2/tutorial/pi-grid",
+        "group"      : "variables"},
     **cmip6_grids,
     **frontier_datasets
 }
@@ -54,7 +57,8 @@ class RemoteZarrDataset:
     Dataset is only loaded on .load() for dataset contaning more variables
     """
 
-    def __init__(self, path_url: str, var_list: Optional[Sequence] = None, consolidated: bool = True, **kwargs):
+    def __init__(self, path_url: str, var_list: Optional[Sequence] = None, group: Optional[str] = None,
+                 consolidated: bool = True, **kwargs):
         """Initializes a remote zarr dataset.
 
         Dataset is loaded only on .load() method.
@@ -72,33 +76,81 @@ def __init__(self, path_url: str, var_list: Optional[Sequence] = None, consolida
         """
         self.path_url = path_url  # can also be local path remove fsspec in that case
         self.var_list = var_list
+        self.group = group
         self.is_consolidated = consolidated
         self.dset_attrs = kwargs
         self._ds = None
+        self._fs = None
+        self.cache_path = os.path.join(USER_CACHE_DIR, self.path_url.split("/")[-1])
+        self._is_zarr_cache = False
 
-    @property
-    def merged_dataset(self):
+    def _merged_dataset(self):
         """Merges data variables from remote url
 
         Returns
         -------
             xr.Dataset
         """
-        import fsspec
         if self._ds is None:
             if self.var_list is not None:
                 urls = [self.path_url + "/" + var for var in self.var_list]
-                dataset_list = [xr.open_zarr(fsspec.get_mapper(url), consolidated=self.is_consolidated) for url in urls]
+                dataset_list = [
+                    # xr.open_zarr(zarr.LRUStoreCache(self._fs.get_mapper(self.path_url), max_size=None),
+                    #             group=self.group, consolidated=self.is_consolidated) for url in
+                    # urls]
+                    xr.open_zarr(self._fs.get_mapper(self.path_url),
+                                 group=self.group, consolidated=self.is_consolidated) for url in
+                    urls]
                 self._ds = xr.merge(dataset_list)
             else:
-                self._ds = xr.open_zarr(self.path_url, consolidated=True)
+                # store = zarr.LRUStoreCache(self._fs.get_mapper(self.path_url), max_size=None)  # unlimited size
+                data_url = self.cache_path if self._is_zarr_cache else self.path_url
+                store = self._fs.get_mapper(data_url)
+                self._ds = xr.open_zarr(store, group=self.group, consolidated=self.is_consolidated)
 
             self._ds.attrs.update(self.dset_attrs)
         return self._ds
 
-    def load(self):
-        return self.merged_dataset
-
+    def load(self, cached=True):
+        if not cached:
+            self._fs = fsspec.filesystem("http")
+        else:
+            if ".zattrs" in fsspec.get_mapper(self.cache_path):
+                self._is_zarr_cache = True
+                self._fs = fsspec.get_filesystem_class('file')()
+            else:
+                self._fs = fsspec.filesystem("filecache", target_protocol='http',
+                                             cache_storage=self.cache_path, expiry_time=False, check_files=False)
+        return self._merged_dataset()
+
+    def download(self, root_path=None, **zarr_copy_kwargs):
+        import zarr
+        if not os.path.exists(os.path.join(self.cache_path, ".zattrs")):
+            self.clear_cache()  # if we download entire dataset we dont need other caches
+        mapper = fsspec.get_mapper(self.path_url)
+        output_path = root_path if root_path is not None else self.cache_path
+        try:
+            zarr.copy_all(zarr.open_consolidated(mapper), zarr.open_group(output_path), **zarr_copy_kwargs)
+            if self.is_consolidated:
+                import requests
+                metadata = requests.get(self.path_url + "/.zmetadata")
+                with open(output_path + "/.zmetadata", 'wt') as f:
+                    f.write(metadata.text)
+        #return f"Dataset downloaded at: {output_path}"
+        except zarr.CopyError:
+            warnings.warn(f'Cache directory {self.cache_path} for this dataset already contains entire data, if any errors persist use .clear_cache() method to clear the cache and re-download using .download() method.')
+        return self.load()
+
+    def clear_cache(self):
+        import shutil
+        try:
+            shutil.rmtree(self.cache_path)
+        except OSError:
+            warnings.warn(f'No cache dir {self.cache_path} found, skipping deletion.')
+
+
+# TODO; generate the datasets automatically from dictionary keys
+# one option is to use user friendly keys in datasets and serattr on sys.modules[__name__] using a function
 
 core = RemoteZarrDataset(**all_datasets['CORE'])
 tutorial_dataset = RemoteZarrDataset(**all_datasets['pi-grid'])
@@ -140,8 +192,7 @@ def fesom_mesh_to_xr(path: str, alpha: int = 0, beta: int = 0, gamma: int = 0) -
 
     """
     mesh = load_mesh(path, abg=[alpha, beta, gamma])
-    ncyclic_inds = get_no_cyclic(mesh, mesh.elem)
-    triangles = mesh.elem[ncyclic_inds]
+    triangles = mesh.elem
     nz_values = np.absolute(mesh.zlev)
     coords_dataset = xr.Dataset(coords={'lon'  : ('nod2', mesh.x2),
                                         'lat'  : ('nod2', mesh.y2),
diff --git a/pyfesom2/diagnostics.py b/pyfesom2/diagnostics.py
index f64269a..aff681b 100644
--- a/pyfesom2/diagnostics.py
+++ b/pyfesom2/diagnostics.py
@@ -236,7 +236,14 @@ def hovm_data(data, mesh, meshdiag=None, runid="fesom", mask=None):
         data = add_timedim(data)
 
     diag = get_meshdiag(mesh, meshdiag, runid)
-    nod_area = diag.rename_dims({"nl": "nz1", "nod_n": "nod2"}).nod_area
+    if 'nod_n' in diag.dims:
+        nod_area = diag.rename_dims({"nl": "nz1", "nod_n": "nod2"}).nod_area[:-1,:]
+    else:
+        nod_area = diag.nod_area[:-1,:]
+        nod_area = nod_area.rename({"nz": "nz1"})
+        nod_area = nod_area.assign_coords({'nz1':data.nz1.values})
+        
+    print(nod_area)
     nod_area.load()
     if mask is not None:
         nod_area = nod_area[:, mask]
@@ -244,12 +251,12 @@ def hovm_data(data, mesh, meshdiag=None, runid="fesom", mask=None):
 
     if isinstance(data, xr.DataArray):
         nod_area = nod_area.where(nod_area != 0)
-        hdg_total = (data * nod_area[:-1, :].T).sum(dim="nod2")
-        hdg_variable = hdg_total / (nod_area[:-1, :].T).sum(axis=0)
+        hdg_total = (data * nod_area.T).sum(dim="nod2")
+        hdg_variable = hdg_total / (nod_area.T).sum(axis=0)
         hdg_variable = hdg_variable.compute()
     else:
-        hdg_total = (data * nod_area[:-1, :].T.data).sum(axis=1)
-        hdg_variable = hdg_total / (nod_area[:-1, :].T).sum(axis=0).data
+        hdg_total = (data * nod_area.T.data).sum(axis=1)
+        hdg_variable = hdg_total / (nod_area.T).sum(axis=0).data
 
     return hdg_variable
 


Test patch:
diff --git a/tests/test_accessor.py b/tests/test_accessor.py
index cf286a8..c3b15b6 100644
--- a/tests/test_accessor.py
+++ b/tests/test_accessor.py
@@ -137,7 +137,7 @@ def test_select_bbox(dataset, bbox):
     with pytest.raises(ValueError):
         select_bbox(dataset[data_var], bbox)
 
-    sda = select_bbox(dataset[data_var], bbox, faces=dataset.faces)
+    sda = select_bbox(dataset[data_var], bbox, coords_dataset=dataset)
     slon_min, slat_min, slon_max, slat_max = (sda.lon.min(), sda.lat.min(),
                                               sda.lon.max(), sda.lon.max())
     assert slon_min >= bbox[0] and slon_max <= bbox[2]
@@ -183,7 +183,7 @@ def test_select_region(dataset, region):
         # dataarrays cannot contain faces so ValueError
         sda = select_region(test_data_array, region)
 
-    sda = select_region(test_data_array, region, faces=dataset.faces)
+    sda = select_region(test_data_array, region, coords_dataset=dataset)
     # returned region selection on data array is a dataset
     assert isinstance(sda, xr.Dataset)