Add regridded_snapshot!

And reduce allocations in regridding by preallocting the target field

NEWS.md

@@ -1,6 +1,12 @@
 ClimaUtilities.jl Release Notes
 ===============================
 
+v0.2.0
+------
+
+- Reduced alloctions in regridding. Now method `regridded_snapshot!`. PR
+  [#72](https://github.com/CliMA/ClimaUtilities.jl/pull/72)
+
 v0.1.9
 ------
 

docs/src/datahandling.md

@@ -102,4 +102,5 @@ ClimaUtilities.DataHandling.available_dates
 ClimaUtilities.DataHandling.previous_time
 ClimaUtilities.DataHandling.next_time
 ClimaUtilities.DataHandling.regridded_snapshot
+ClimaUtilities.DataHandling.regridded_snapshot!
 ```

ext/DataHandlingExt.jl

@@ -299,8 +299,6 @@ The `time` has to be available in the `data_handler`.
 
 `regridded_snapshot` potentially modifies the internal state of `data_handler` and it might be a very
 expensive operation.
-
-TODO: Add `regridded_snapshot!`
 """
 function DataHandling.regridded_snapshot(
     data_handler::DataHandler,
@@ -346,4 +344,9 @@ function DataHandling.regridded_snapshot(data_handler::DataHandler)
     return DataHandling.regridded_snapshot(data_handler, Dates.DateTime(0))
 end
 
+function DataHandling.regridded_snapshot!(out, data_handler, time)
+    out .= DataHandling.regridded_snapshot!(data_handler, time)
+    return nothing
+end
+
 end

ext/TimeVaryingInputsExt.jl

@@ -15,7 +15,7 @@ import ClimaUtilities.TimeVaryingInputs: NearestNeighbor, LinearInterpolation
 
 import ClimaUtilities.DataHandling
 import ClimaUtilities.DataHandling:
-    DataHandler, previous_time, next_time, regridded_snapshot, available_times
+    DataHandler, previous_time, next_time, regridded_snapshot!, available_times
 
 # Ideally, we should be able to split off the analytic part in a different
 # extension, but precompilation stops working when we do so
@@ -61,6 +61,7 @@ struct InterpolatingTimeVaryingInput23D{
     M <: AbstractInterpolationMethod,
     CC <: ClimaComms.AbstractCommsContext,
     R <: Tuple,
+    RR,
 } <: AbstractTimeVaryingInput
     """Object that has all the information on how to deal with files, data, and so on.
        Having to deal with files, it lives on the CPU."""
@@ -75,6 +76,9 @@ struct InterpolatingTimeVaryingInput23D{
     """Range of times over which the interpolator is defined. range is always defined on the
     CPU. Used by the in() function."""
     range::R
+
+    """Preallocated spaces where to do regridding"""
+    preallocated_regridded_fields::RR
 end
 
 """
@@ -96,11 +100,20 @@ function TimeVaryingInputs.TimeVaryingInput(
         error("DataHandler does not contain temporal data")
     issorted(available_times) || error("Can only interpolate with sorted times")
     range = (available_times[begin], available_times[end])
+
+    # TODO: Generalize the number of _regridded_fields depending on the interpolation
+    # stencil
+    preallocated_regridded_fields = (
+        zero(axes(data_handler.target_space)),
+        zero(axes(data_handler.target_space)),
+    )
+
     return InterpolatingTimeVaryingInput23D(
         data_handler,
         method,
         context,
         range,
+        preallocated_regridded_fields,
     )
 end
 
@@ -154,9 +167,9 @@ function TimeVaryingInputs.evaluate!(
 
     # The closest regridded_snapshot could be either the previous or the next one
     if (time - t0) <= (t1 - time)
-        dest .= regridded_snapshot(itp.data_handler, t0)
+        regridded_snapshot!(dest, itp.data_handler, t0)
     else
-        dest .= regridded_snapshot(itp.data_handler, t1)
+        regridded_snapshot!(dest, itp.data_handler, t1)
     end
     return nothing
 end
@@ -179,9 +192,13 @@ function TimeVaryingInputs.evaluate!(
     t0, t1 =
         previous_time(itp.data_handler, time), next_time(itp.data_handler, time)
     coeff = (time - t0) / (t1 - t0)
-    dest .=
-        (1 - coeff) .* regridded_snapshot(itp.data_handler, t0) .+
-        coeff .* regridded_snapshot(itp.data_handler, t1)
+
+    field_prev = itp.preallocated_regridded_fields[1]
+    field_next = itp.preallocated_regridded_fields[2]
+    regridded_snapshot!(field_prev, itp.data_handler, t0)
+    regridded_snapshot!(field_next, itp.data_handler, t1)
+
+    dest .= (1 - coeff) .* field_prev .+ coeff .* field_next
     return nothing
 end
 

src/DataHandling.jl

@@ -53,4 +53,6 @@ function next_time end
 
 function regridded_snapshot end
 
+function regridded_snapshot! end
+
 end