tonik 0.0.10__tar.gz → 0.0.12__tar.gz

{tonik-0.0.10 → tonik-0.0.12}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: tonik
-Version: 0.0.10
+Version: 0.0.12
 Summary: Store time series data as HDF5 files and access them through an API.
 Project-URL: Homepage, https://tsc-tools.github.io/tonik
 Project-URL: Issues, https://github.com/tsc-tools/tonik/issues

{tonik-0.0.10 → tonik-0.0.12}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "tonik"
-version = "0.0.10"
+version = "0.0.12"
 authors = [
   { name="Yannik Behr", email="y.behr@gns.cri.nz" },
   { name="Christof Mueller", email="c.mueller@gns.cri.nz" }

{tonik-0.0.10 → tonik-0.0.12}/src/tonik/storage.py

@@ -155,11 +155,7 @@ class Path(object):
 
         xd_index = dict(datetime=slice(self.starttime, self.endtime))
         with xr.open_dataset(filename, group='original', engine=self.backend) as ds:
-            try:
-                rq = ds.loc[xd_index].load()
-            except KeyError:
-                ds = ds.sortby("datetime")
-                rq = ds.loc[xd_index].load()
+            rq = ds.loc[xd_index].load()
 
         # Stack features
         if stack_length is not None:

tonik-0.0.12/src/tonik/xarray2zarr.py

@@ -0,0 +1,78 @@
+import logging
+import os
+
+import xarray as xr
+
+logger = logging.getLogger(__name__)
+
+
+def merge_arrays(xds_old: xr.DataArray, xds_new: xr.DataArray) -> xr.DataArray:
+    """
+    Merge two xarray datasets with the same datetime index.
+
+    Parameters
+    ----------
+    xds_old : xr.DataArray
+        Old array.
+    xds_new : xr.DataArray
+        New array.
+
+    Returns
+    -------
+    xr.DataArray
+        Merged array.
+    """
+    xda_old = xds_old.drop_duplicates(
+        'datetime', keep='last')
+    xda_new = xds_new.drop_duplicates(
+        'datetime', keep='last')
+    xda_new = xda_new.combine_first(xda_old)
+    return xda_new
+
+
+def xarray2zarr(xds: xr.Dataset, path: str, mode: str = 'a'):
+    """
+    Write xarray dataset to zarr files.
+
+    Parameters
+    ----------
+    xds : xr.Dataset
+        Dataset to write.
+    path : str
+        Path to write the dataset.
+    mode : str, optional
+        Write mode, by default 'a'.
+
+    Returns
+    -------
+    None
+    """
+    for feature in xds.data_vars.keys():
+        fout = os.path.join(path, feature + '.zarr')
+        if not os.path.exists(fout) or mode == 'w':
+            xds[feature].to_zarr(
+                fout, group='original', mode='w')
+        else:
+            xds_existing = xr.open_zarr(fout, group='original')
+            if xds_existing.datetime[0] > xds.datetime[0] or xds_existing.datetime[-1] > xds.datetime[-1]:
+                xda_new = merge_arrays(xds_existing[feature], xds[feature])
+                xda_new.to_zarr(fout, group='original', mode='w')
+            else:
+                try:
+                    overlap = xds_existing.datetime.where(
+                        xds_existing.datetime == xds.datetime)
+                    if overlap.size > 0:
+                        xds[feature].loc[dict(datetime=overlap)].to_zarr(
+                            fout, group='original', mode='r+', region='auto')
+                        xds[feature].drop_sel(datetime=overlap).to_zarr(
+                            fout, group='original', mode='a', append_dim="datetime")
+                    else:
+                        xds[feature].to_zarr(
+                            fout, group='original', append_dim='datetime')
+                except Exception as e:
+                    msg = f"Appending {feature} to {fout} failed: {e}\n"
+                    msg += "Attempting to merge the two datasets."
+                    logger.error(msg)
+                    # remove duplicate datetime entries
+                    xda_new = merge_arrays(xds_existing[feature], xds[feature])
+                    xda_new.to_zarr(fout, group='original', mode='w')

{tonik-0.0.10 → tonik-0.0.12}/tests/test_save.py

@@ -161,18 +161,21 @@ def test_xarray2zarr_outofsequence(tmp_path_factory):
     """
     temp_dir = tmp_path_factory.mktemp('test_xarray2zarr')
     start = datetime(2022, 7, 18, 8, 0, 0)
+    middle = datetime(2022, 7, 18, 12, 0, 0)
     end = datetime(2022, 7, 19, 12, 0, 0)
-    xdf1 = generate_test_data(dim=1, ndays=1, tstart=start)
-    xdf2 = generate_test_data(dim=1, ndays=1, tstart=end)
+    xdf1 = generate_test_data(dim=1, intervals=3, tstart=start)
+    xdf2 = generate_test_data(dim=1, intervals=3, tstart=middle)
+    xdf3 = generate_test_data(dim=1, intervals=3, tstart=end)
     g = Storage('test_experiment', rootdir=temp_dir,
                 starttime=start, endtime=end + timedelta(days=1),
                 backend='zarr')
     c = g.get_substore('MDR', '00', 'HHZ')
-    c.save(xdf2)
+    c.save(xdf3)
     c.save(xdf1)
+    c.save(xdf2)
     xdf_test = c('rsam')
     np.testing.assert_array_equal(
-        xdf_test.datetime.values, xdf1.merge(xdf2).datetime.values)
+        xdf_test.datetime.values, xr.merge([xdf1, xdf2, xdf3]).datetime.values)
 
 
 def test_xarray2zarr_duplicates(tmp_path_factory):

tonik-0.0.10/src/tonik/xarray2zarr.py

@@ -1,37 +0,0 @@
-import logging
-import os
-
-import xarray as xr
-
-logger = logging.getLogger(__name__)
-
-
-def xarray2zarr(xds, path, mode='a'):
-    for feature in xds.data_vars.keys():
-        fout = os.path.join(path, feature + '.zarr')
-        if not os.path.exists(fout) or mode == 'w':
-            xds[feature].to_zarr(
-                fout, group='original', mode='w')
-        else:
-            xds_existing = xr.open_zarr(fout, group='original')
-            try:
-                overlap = xds_existing.datetime.where(
-                    xds_existing.datetime == xds.datetime)
-                if overlap.size > 0:
-                    xds[feature].loc[dict(datetime=overlap)].to_zarr(
-                        fout, group='original', mode='r+', region='auto')
-                    xds[feature].drop_sel(datetime=overlap).to_zarr(
-                        fout, group='original', mode='a', append_dim="datetime")
-                else:
-                    xds[feature].to_zarr(
-                        fout, group='original', append_dim='datetime')
-            except Exception as e:
-                msg = f"Appending {feature} to {fout} failed: {e}\n"
-                msg += "Attempting to merge the two datasets."
-                logger.error(msg)
-                # remove duplicate datetime entries
-                xda_existing = xds_existing[feature].drop_duplicates(
-                    'datetime', keep='last')
-                xda_new = xds[feature].drop_duplicates('datetime', keep='last')
-                xda_new = xda_new.combine_first(xda_existing)
-                xda_new.to_zarr(fout, group='original', mode='w')