Merge pull request #810 from DHI/generic_concat_average

support average in generic concat

Author: ecomodeller

mikeio/generic.py

@@ -417,10 +417,7 @@ def concat(
         if keep == "last":
             if i < (len(infilenames) - 1):
                 dfs_n = DfsFileFactory.DfsGenericOpen(str(infilenames[i + 1]))
-                nf = dfs_n.FileInfo.TimeAxis.StartDateTime
-                next_start_time = datetime(
-                    nf.year, nf.month, nf.day, nf.hour, nf.minute, nf.second
-                )
+                next_start_time = dfs_n.FileInfo.TimeAxis.StartDateTime
                 dfs_n.Close()
 
             for timestep in range(n_time_steps):
@@ -478,6 +475,58 @@ def concat(
                 )  # get end time from current file
                 dfs_i.Close()
 
+        elif keep == "average":
+            ALPHA = 0.5  # averaging factor
+            last_file = i == (len(infilenames) - 1)
+            overlapping_with_next = False  # lets first asume no overlap
+
+            # Find the start time of next file
+            if not last_file:
+                dfs_n = DfsFileFactory.DfsGenericOpen(str(infilenames[i + 1]))
+                next_start_time = dfs_n.FileInfo.TimeAxis.StartDateTime
+            else:
+                next_start_time = datetime.max  # end of time ...
+
+            if i == 0:
+                timestep_n = 0  # have not read anything before
+
+            # lets start where we left off (if last file overlapped)
+            timestep = timestep_n
+            while timestep < n_time_steps:
+                current_time = start_time + timedelta(seconds=timestep * dt)
+                if current_time >= next_start_time:  # false if last file
+                    overlapping_with_next = True
+                    break
+                for item in range(n_items):
+                    itemdata = dfs_i.ReadItemTimeStep(item + 1, timestep)
+                    d = itemdata.Data
+                    darray = d.astype(np.float32)
+                    dfs_o.WriteItemTimeStepNext(0, darray)
+
+                timestep += 1
+
+            timestep_n = 0  # have not read anything from next file yet
+
+            if not overlapping_with_next:
+                dfs_n.Close()
+                continue  # next file
+
+            # Otherwhise read overlapping part
+            while timestep < n_time_steps:
+                for item in range(n_items):
+                    itemdata_i = dfs_i.ReadItemTimeStep(item + 1, timestep)
+                    itemdata_n = dfs_n.ReadItemTimeStep(item + 1, timestep_n)
+                    d_i = itemdata_i.Data
+                    d_n = itemdata_n.Data
+                    d_av = d_i * ALPHA + d_n * (1 - ALPHA)
+                    darray = d_av.astype(np.float32)
+                    dfs_o.WriteItemTimeStepNext(0, (darray))
+                timestep += 1
+                timestep_n += 1
+
+            # Close next file before opening it again
+            dfs_n.Close()
+
     dfs_o.Close()
 
 

tests/test_generic.py

@@ -268,7 +268,7 @@ def test_concat_keep(tmp_path: Path) -> None:
     test keep arguments of concatenation function
     """
     # added keep arguments to test
-    keep_args = ["first", "last"]
+    keep_args = ["first", "last", "average"]
 
     infiles = [
         "tests/testdata/tide1.dfs1",
@@ -317,11 +317,53 @@ def test_concat_keep(tmp_path: Path) -> None:
                     .all()
                     .all()
                 )
+                av_out = (
+                    (
+                        0.5 * (df_first.loc[overlap_idx] + df_last.loc[overlap_idx])
+                        == df_o.loc[overlap_idx]
+                    )
+                    .eq(True)
+                    .all()
+                    .all()
+                )
 
                 if keep_arg == "first":
                     assert first_out, "overlap should be with first dataset"
                 elif keep_arg == "last":
                     assert last_out, "overlap should be with last dataset"
+                elif keep_arg == "average":
+                    assert av_out, "overlap should be average of datasets"
+
+
+def test_concat_average(tmp_path: Path) -> None:
+    # Test for multiple items?
+    g = mikeio.Grid1D(x=range(5))
+    t = pd.date_range(start="2020-01-01", periods=5, freq="D")
+    d = np.zeros((5, 5))
+    #     x x x o o
+    #           o o x o o
+    #                 o o x x x
+    da_1 = mikeio.DataArray(data=d, time=t, geometry=g)
+    da_2 = mikeio.DataArray(data=d + 1, time=t + pd.DateOffset(days=3), geometry=g)
+    da_3 = mikeio.DataArray(data=d + 2, time=t + pd.DateOffset(days=6), geometry=g)
+
+    files = [tmp_path / "test1.dfs1", tmp_path / "test2.dfs1", tmp_path / "test3.dfs1"]
+    da_1.to_dfs(files[0])
+    da_2.to_dfs(files[1])
+    da_3.to_dfs(files[2])
+
+    # concat
+    fp = tmp_path / "concat.dfs1"
+
+    mikeio.generic.concat(files, fp, keep="average")
+    ds = mikeio.read(fp)
+    da_x0 = ds[0].isel(x=0)
+
+    assert np.allclose(
+        da_x0.values,
+        np.array([0.0, 0.0, 0.0, 0.5, 0.5, 1.0, 1.5, 1.5, 2.0, 2.0, 2.0]),
+        atol=1e-6,
+    )
 
 
 def test_concat_non_equidistant_dfs0(tmp_path: Path) -> None: