split function added test

CliMA · Apr 12, 2024 · 5d12533 · 5d12533
1 parent e0c12eb
commit 5d12533
Show file tree

Hide file tree

Showing 2 changed files with 53 additions and 35 deletions.
diff --git a/src/Regridder.jl b/src/Regridder.jl
@@ -657,41 +657,7 @@ function truncate_dataset(
         ds = NCDataset(datafile, "r")
         dates = ds["time"][:]
 
-        # if the simulation start date is before our first date in the dataset
-        # leave the beginning of the truncated dataset to be first date available
-        if date_start < dates[1]
-            start_id = 1
-            # if the simulation start date is after the last date in the dataset
-            # start the truncated dataset at its last possible date
-        elseif date_start > last(dates)
-            start_id = length(dates)
-            # if the simulation start date falls within the range of the dataset
-            # find the closest date to the start date and truncate there 
-        else
-            (~, start_id) = findmin(x -> abs(x - date_start), dates)
-            # if the closest date is after the start date, add one more date before 
-            if dates[start_id] > date_start
-                start_id = start_id - 1
-            end
-        end
-
-        # if the simulation end date is before our first date in the dataset
-        # truncate the end of the dataset to be the first date
-        if date_end < dates[1]
-            end_id = 1
-            # if the simulation end date is after the last date in the dataset
-            # leave the end of the dataset as is 
-        elseif date_end > last(dates)
-            end_id = length(dates)
-            # if the simulation end date falls within the range of the dataset
-            # find the closest date to the end date and truncate there 
-        else
-            (~, end_id) = findmin(x -> abs(x - date_end), dates)
-            # if the closest date is before the end date, add one more date after
-            if dates[end_id] < date_end
-                end_id = end_id + 1
-            end
-        end
+        (start_id, end_id) = find_idx_bounding_dates(dates, date_start, date_end)
 
         ds_truncated = NCDataset(datafile_truncated, "c")
         ds_truncated = write(ds_truncated, view(ds, time = start_id:end_id))
@@ -703,4 +669,49 @@ function truncate_dataset(
     end
 end
 
+"""
+    find_idx_bounding_dates(dates, date_start, date_end) 
+
+Returns the index range from dates that contains date_start to date_end
+"""
+function find_idx_bounding_dates(dates, date_start, date_end)
+    # if the simulation start date is before our first date in the dataset
+    # leave the beginning of the truncated dataset to be first date available
+    if date_start < dates[1]
+        start_id = 1
+        # if the simulation start date is after the last date in the dataset
+        # start the truncated dataset at its last possible date
+    elseif date_start > last(dates)
+        start_id = length(dates)
+        # if the simulation start date falls within the range of the dataset
+        # find the closest date to the start date and truncate there 
+    else
+        (~, start_id) = findmin(x -> abs(x - date_start), dates)
+        # if the closest date is after the start date, add one more date before 
+        if dates[start_id] > date_start
+            start_id = start_id - 1
+        end
+    end
+
+    # if the simulation end date is before our first date in the dataset
+    # truncate the end of the dataset to be the first date
+    if date_end < dates[1]
+        end_id = 1
+        # if the simulation end date is after the last date in the dataset
+        # leave the end of the dataset as is 
+    elseif date_end > last(dates)
+        end_id = length(dates)
+        # if the simulation end date falls within the range of the dataset
+        # find the closest date to the end date and truncate there 
+    else
+        (~, end_id) = findmin(x -> abs(x - date_end), dates)
+        # if the closest date is before the end date, add one more date after
+        if dates[end_id] < date_end
+            end_id = end_id + 1
+        end
+    end
+
+    return (; start_id, end_id)
+end
+
 end # Module
diff --git a/test/regridder_tests.jl b/test/regridder_tests.jl
@@ -374,6 +374,13 @@ end
             @test new_dates[length(new_dates) - 1] <= date_end
         end
 
+        # check that truncation is indexing correctly 
+        all_data = ds["SST"][:, :, :] 
+        new_data = ds_truncated["SST"][:, :, :]
+        (start_id, end_id) = Regridder.find_idx_bounding_dates(dates, date_start, date_end)
+        @test new_data[:, :, 1] ≈ all_data[:, :, start_id]
+        @test new_data[:, :, length(new_dates)] ≈ all_data[:, :, end_id]
+
         close(ds_truncated)
     end