Merge pull request #4 from jagoosw/jsw/adapts

Add adapt methods for wind stress and heat exchange

Project.toml

@@ -1,7 +1,7 @@
 name = "Walrus"
 uuid = "bec5164f-1c7a-4c32-8768-be9354254d6a"
 authors = ["Jago Stong-Wright <[email protected]>"]
-version = "0.4.1"
+version = "0.4.2"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

src/Walrus.jl

@@ -10,11 +10,17 @@ export WindStress, WindStressBoundaryConditions, LogarithmicNeutralWind
 
 export SurfaceHeatExchange, SurfaceHeatExchangeBoundaryCondition
 
+using Adapt: adapt
+
+import Adapt: adapt_structure
+
 # This is how we will standarise function vs values, I will have to think of a way to deal with discrete vs continuous at some point
 struct ReturnValue{FT}
     value :: FT
 end
 
+adapt_structure(to, rv::ReturnValue) = ReturnValue(adapt(to, rv.value))
+
 @inline (return_value::ReturnValue)(args...) = return_value.value
 
 display_input(return_value::ReturnValue) = return_value.value

src/surface_heating.jl

@@ -4,13 +4,17 @@ export SurfaceHeatExchange, SurfaceHeatExchangeBoundaryCondition
 
 using Roots
 
+using Adapt: adapt
+
 using Oceananigans.BoundaryConditions: FluxBoundaryCondition
 
 using Walrus: ReturnValue, display_input
 using Walrus.WindStressModel: WindStress, 
                               LogarithmicNeutralWind, 
                               find_velocity_roughness_length
 
+import Adapt: adapt_structure
+
 struct SurfaceHeatExchange{WS, AT, LH, VP, FT} <: Function
                    wind_stress :: WS
                air_temperature :: AT
@@ -24,6 +28,18 @@ struct SurfaceHeatExchange{WS, AT, LH, VP, FT} <: Function
      stephan_boltzman_constant :: FT
 end
 
+adapt_structure(to, sh::SurfaceHeatExchange) = 
+    SurfaceHeatExchange(adapt(to, sh.wind_stress),
+                        adapt(to, sh.air_temperature),
+                        adapt(to, sh.latent_heat_vaporisation),
+                        adapt(to, sh.vapour_pressure),
+                        sh.water_specific_heat_capacity,
+                        sh.water_density,
+                        sh.air_specific_heat_capacity,
+                        sh.air_density,
+                        sh.air_water_mixing_ratio,
+                        sh.stephan_boltzman_constant)
+
 """
     SurfaceHeatExchange(; wind_stress,
                           air_temperature = 18, # °C

src/wind_stress.jl

@@ -3,10 +3,13 @@ module WindStressModel
 export WindStress, WindStressBoundaryConditions, LogarithmicNeutralWind
 
 using Roots
-using Walrus: ReturnValue, display_input
+
+using Adapt: adapt
 using Oceananigans.BoundaryConditions: FluxBoundaryCondition
 using Oceananigans.BuoyancyModels: g_Earth
+using Walrus: ReturnValue, display_input
 
+import Adapt: adapt_structure
 import Base: summary, show
 
 struct WindStress{WS, WD, DC, FT} <: Function
@@ -17,6 +20,12 @@ struct WindStress{WS, WD, DC, FT} <: Function
              water_density :: FT
 end
 
+adapt_structure(to, ws::WindStress) = WindStress(adapt(to, ws.reference_wind_speed),
+                                                 adapt(to, ws.reference_wind_direction),
+                                                 adapt(to, ws.drag_coefficient),
+                                                 ws.air_density,
+                                                 ws.water_density)
+
 """
     WindStress(; reference_wind_speed, 
                  reference_wind_direction,
@@ -242,6 +251,8 @@ This parameterisaion is described in [smith1988](@citet)
             gravity_acceleration :: FT = g_Earth
 end
 
+adapt_structure(to, dc::LogarithmicNeutralWind) = dc
+
 @inline velocity_roughness_length_roots(z₀, params) = 
     log(params.reference_height/z₀)/(params.κ * params.wind_speed) * 
         (params.ν * params.b + params.aᶜ * (params.κ * params.wind_speed / log(params.reference_height/z₀))^3) - z₀