Move extrapolate kwarg to pressure intp constructor

Author: ph-kev

src/Remapping/interpolate_pressure.jl

@@ -30,6 +30,7 @@ struct PressureInterpolator{
     FACE <: Fields.Field,
     SPACE <: Spaces.AbstractSpace,
     LEVELS,
+    EXTRAPOLATE <: ClimaInterpolations.Interpolation1D.Extrapolate1D,
 }
     """A ClimaCore.Field representing pressure on center space. This field is
     defined on a space with height (z) as the vertical coordinate, not
@@ -51,6 +52,10 @@ struct PressureInterpolator{
     """A 1D vector of pressure coordinates to interpolate onto for every
     column"""
     pressure_levels::LEVELS
+
+    """Extrapolation condition when interpolating outside of the pressure
+    range"""
+    extrapolate::EXTRAPOLATE
 end
 
 """
@@ -130,15 +135,23 @@ function construct_pfull_grid(::Type{FT}, pressure_levels, device) where {FT}
 end
 
 """
-    PressureInterpolator(pfull_field::Fields.Field, pressure_levels)
+    PressureInterpolator(
+        pfull_field::Fields.Field,
+        pressure_levels;
+        extrapolate = ClimaInterpolations.Interpolation1D.Flat(),
+    )
 
 Construct a `PressureInterpolator` from `pfull_field`, a pressure field defined
 on a center space and `pressure_levels`, a vector of pressure levels to
 interpolate to.
 
 The pressure levels must be in ascending or descending order.
 """
-function PressureInterpolator(pfull_field::Fields.Field, pressure_levels)
+function PressureInterpolator(
+    pfull_field::Fields.Field,
+    pressure_levels;
+    extrapolate = ClimaInterpolations.Interpolation1D.Flat(),
+)
     if issorted(pressure_levels, rev = true)
         pressure_levels = sort(pressure_levels)
     end
@@ -150,7 +163,8 @@ function PressureInterpolator(pfull_field::Fields.Field, pressure_levels)
     pressure_space = construct_pressure_space(FT, space, pressure_levels)
     return PressureInterpolator(
         pfull_field,
-        pressure_space,
+        pressure_space;
+        extrapolate,
     )
 end
 
@@ -160,18 +174,22 @@ end
         pressure_space::Union{
             Spaces.AbstractFiniteDifferenceSpace,
             Spaces.ExtrudedFiniteDifferenceSpace,
-        },
+        };
+        extrapolate = ClimaInterpolations.Interpolation1D.Flat()
     )
 
 Construct a `PressureInterpolator` from `pfull_field`, a pressure field, and
 `pressure_space`, a space with pressure as the vertical coordinate.
+
+The default extrapolation behavior is flat extrapolation.
 """
 function PressureInterpolator(
     pfull_field::Fields.Field,
     pressure_space::Union{
         Spaces.AbstractFiniteDifferenceSpace,
         Spaces.ExtrudedFiniteDifferenceSpace,
-    },
+    };
+    extrapolate = ClimaInterpolations.Interpolation1D.Flat(),
 )
     axes(pfull_field).staggering isa Grids.CellCenter || error("The staggering of the
     pressure field must be cell center")
@@ -197,6 +215,7 @@ function PressureInterpolator(
         scratch_face_pressure_field,
         pressure_space,
         pressure_levels,
+        extrapolate,
     )
 end
 
@@ -258,21 +277,19 @@ end
 """
     interpolate_pressure(
         field::Fields.Field,
-        pfull_intp::PressureInterpolator;
-        extrapolate = ClimaInterpolations.Interpolation1D.Flat(),
+        pfull_intp::PressureInterpolator,
     )
 
 Vertically interpolate field onto a space identical to that of field, but with
 pressure as the vertical coordinate and return the interpolated field.
 """
 function interpolate_pressure(
     field::Fields.Field,
-    pfull_intp::PressureInterpolator;
-    extrapolate = ClimaInterpolations.Interpolation1D.Flat(),
+    pfull_intp::PressureInterpolator,
 )
     (; pfull_field) = pfull_intp
     dest = fill(one(eltype(pfull_field)), pfull_intp.pressure_space)
-    Remapping.interpolate_pressure!(dest, field, pfull_intp; extrapolate)
+    Remapping.interpolate_pressure!(dest, field, pfull_intp)
     return dest
 end
 
@@ -281,7 +298,6 @@ end
         dest::Fields.Field,
         field::Fields.Field,
         pfull_intp::PressureInterpolator;
-        extrapolate = ClimaInterpolations.Interpolation1D.Flat(),
     )
 
 Vertically interpolate `field` onto `dest` and return `nothing`.
@@ -291,10 +307,14 @@ The vertical coordinate of the space of `dest` must be in pressure.
 function interpolate_pressure!(
     dest::Fields.Field,
     field::Fields.Field,
-    pfull_intp::PressureInterpolator;
-    extrapolate = ClimaInterpolations.Interpolation1D.Flat(),
+    pfull_intp::PressureInterpolator,
 )
-    (; scratch_center_pressure_field, scratch_face_pressure_field, pressure_levels) =
+    (;
+        scratch_center_pressure_field,
+        scratch_face_pressure_field,
+        pressure_levels,
+        extrapolate,
+    ) =
         pfull_intp
     scratch_pfull_array = if axes(field).staggering isa Grids.CellCenter
         Fields.field2array(scratch_center_pressure_field)

test/Remapping/interpolate_pressure.jl

@@ -66,12 +66,12 @@ for FT in (Float32, Float64)
             pfull_field = fill(FT(1.0), space)
             pfull_field_array = Fields.field2array(pfull_field)
             pfull_field_array .= collect(10:-1:1)
-            pfull_levels = [0.0, 1.0, 5.5, 10.0, 12.0]
+            pressure_levels = [0.0, 1.0, 5.5, 10.0, 12.0]
 
             # Pressures should be in sorted order
             @test_throws ErrorException PressureInterpolator(pfull_field, [5.5, 10.0, 1.0])
 
-            pfull_intp = PressureInterpolator(pfull_field, pfull_levels)
+            pfull_intp = PressureInterpolator(pfull_field, pressure_levels)
             pfull_field_array = Fields.field2array(Remapping.pfull_field(pfull_intp))
 
             # Check the pressures along each column is sorted in descending
@@ -106,6 +106,22 @@ for FT in (Float32, Float64)
 
             interpolate_pressure!(dest2, dummy_field, pfull_intp)
             @test dest2 == dest
+
+            # Test linear extrapolation
+            pressure_levels = [-100.0, 100.0]
+            extrapolate = ClimaInterpolations.Interpolation1D.LinearExtrapolation()
+            linear_pfull_intp =
+                PressureInterpolator(pfull_field, pressure_levels; extrapolate)
+            dummy_field = fill(FT(1.0), space)
+            dummy_field_array = Fields.field2array(dummy_field)
+            dummy_field_array .= collect(10:-1:1)
+            dest = interpolate_pressure(dummy_field, linear_pfull_intp)
+            dest_array = Fields.field2array(dest)
+            @test dest_array == repeat(
+                [-100.0, 100.0],
+                1,
+                size(pfull_field_array, 2),
+            )
         end
     end