Various interpolation changes

Manifest.toml

@@ -2,7 +2,7 @@
 
 julia_version = "1.10.2"
 manifest_format = "2.0"
-project_hash = "acc6dd5077921f9cacb46eb572b93736282df7cc"
+project_hash = "5864f7280ef526f4ebcf0f0ff9628f004cf1c488"
 
 [[deps.AbstractFFTs]]
 deps = ["LinearAlgebra"]
@@ -38,9 +38,9 @@ version = "0.1.37"
 
 [[deps.Adapt]]
 deps = ["LinearAlgebra", "Requires"]
-git-tree-sha1 = "cde29ddf7e5726c9fb511f340244ea3481267608"
+git-tree-sha1 = "6a55b747d1812e699320963ffde36f1ebdda4099"
 uuid = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
-version = "3.7.2"
+version = "4.0.4"
 weakdeps = ["StaticArrays"]
 
     [deps.Adapt.extensions]
@@ -61,9 +61,9 @@ version = "0.1.0"
 
 [[deps.BFloat16s]]
 deps = ["LinearAlgebra", "Printf", "Random", "Test"]
-git-tree-sha1 = "dbf84058d0a8cbbadee18d25cf606934b22d7c66"
+git-tree-sha1 = "2c7cc21e8678eff479978a0a2ef5ce2f51b63dff"
 uuid = "ab4f0b2a-ad5b-11e8-123f-65d77653426b"
-version = "0.4.2"
+version = "0.5.0"
 
 [[deps.Base64]]
 uuid = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f"
@@ -92,33 +92,38 @@ uuid = "179af706-886a-5703-950a-314cd64e0468"
 version = "0.1.3"
 
 [[deps.CUDA]]
-deps = ["AbstractFFTs", "Adapt", "BFloat16s", "CEnum", "CUDA_Driver_jll", "CUDA_Runtime_Discovery", "CUDA_Runtime_jll", "Crayons", "DataFrames", "ExprTools", "GPUArrays", "GPUCompiler", "KernelAbstractions", "LLVM", "LLVMLoopInfo", "LazyArtifacts", "Libdl", "LinearAlgebra", "Logging", "NVTX", "Preferences", "PrettyTables", "Printf", "Random", "Random123", "RandomNumbers", "Reexport", "Requires", "SparseArrays", "Statistics", "UnsafeAtomicsLLVM"]
-git-tree-sha1 = "95ac638373ac40e29c1b6d086a3698f5026ff6a6"
+deps = ["AbstractFFTs", "Adapt", "BFloat16s", "CEnum", "CUDA_Driver_jll", "CUDA_Runtime_Discovery", "CUDA_Runtime_jll", "Crayons", "DataFrames", "ExprTools", "GPUArrays", "GPUCompiler", "KernelAbstractions", "LLVM", "LLVMLoopInfo", "LazyArtifacts", "Libdl", "LinearAlgebra", "Logging", "NVTX", "Preferences", "PrettyTables", "Printf", "Random", "Random123", "RandomNumbers", "Reexport", "Requires", "SparseArrays", "StaticArrays", "Statistics"]
+git-tree-sha1 = "fdd9dfb67dfefd548f51000cc400bb51003de247"
 uuid = "052768ef-5323-5732-b1bb-66c8b64840ba"
-version = "5.1.2"
-weakdeps = ["ChainRulesCore", "SpecialFunctions"]
+version = "5.4.3"
 
     [deps.CUDA.extensions]
     ChainRulesCoreExt = "ChainRulesCore"
+    EnzymeCoreExt = "EnzymeCore"
     SpecialFunctionsExt = "SpecialFunctions"
 
+    [deps.CUDA.weakdeps]
+    ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+    EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869"
+    SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
+
 [[deps.CUDA_Driver_jll]]
-deps = ["Artifacts", "JLLWrappers", "LazyArtifacts", "Libdl", "Pkg"]
-git-tree-sha1 = "d01bfc999768f0a31ed36f5d22a76161fc63079c"
+deps = ["Artifacts", "JLLWrappers", "Libdl", "Pkg"]
+git-tree-sha1 = "325058b426c2b421e3d2df3d5fa646d72d2e3e7e"
 uuid = "4ee394cb-3365-5eb0-8335-949819d2adfc"
-version = "0.7.0+1"
+version = "0.9.2+0"
 
 [[deps.CUDA_Runtime_Discovery]]
 deps = ["Libdl"]
-git-tree-sha1 = "38f830504358e9972d2a0c3e5d51cb865e0733df"
+git-tree-sha1 = "f3b237289a5a77c759b2dd5d4c2ff641d67c4030"
 uuid = "1af6417a-86b4-443c-805f-a4643ffb695f"
-version = "0.2.4"
+version = "0.3.4"
 
 [[deps.CUDA_Runtime_jll]]
 deps = ["Artifacts", "CUDA_Driver_jll", "JLLWrappers", "LazyArtifacts", "Libdl", "TOML"]
-git-tree-sha1 = "9704e50c9158cf8896c2776b8dbc5edd136caf80"
+git-tree-sha1 = "afea94249b821dc754a8ca6695d3daed851e1f5a"
 uuid = "76a88914-d11a-5bdc-97e0-2f5a05c973a2"
-version = "0.10.1+0"
+version = "0.14.1+0"
 
 [[deps.ChainRulesCore]]
 deps = ["Compat", "LinearAlgebra"]
@@ -308,21 +313,21 @@ version = "6.2.1+6"
 
 [[deps.GPUArrays]]
 deps = ["Adapt", "GPUArraysCore", "LLVM", "LinearAlgebra", "Printf", "Random", "Reexport", "Serialization", "Statistics"]
-git-tree-sha1 = "85d7fb51afb3def5dcb85ad31c3707795c8bccc1"
+git-tree-sha1 = "a74c3f1cf56a3dfcdef0605f8cdb7015926aae30"
 uuid = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7"
-version = "9.1.0"
+version = "10.3.0"
 
 [[deps.GPUArraysCore]]
 deps = ["Adapt"]
-git-tree-sha1 = "2d6ca471a6c7b536127afccfa7564b5b39227fe0"
+git-tree-sha1 = "ec632f177c0d990e64d955ccc1b8c04c485a0950"
 uuid = "46192b85-c4d5-4398-a991-12ede77f4527"
-version = "0.1.5"
+version = "0.1.6"
 
 [[deps.GPUCompiler]]
-deps = ["ExprTools", "InteractiveUtils", "LLVM", "Libdl", "Logging", "Scratch", "TimerOutputs", "UUIDs"]
-git-tree-sha1 = "a846f297ce9d09ccba02ead0cae70690e072a119"
+deps = ["ExprTools", "InteractiveUtils", "LLVM", "Libdl", "Logging", "Preferences", "Scratch", "Serialization", "TOML", "TimerOutputs", "UUIDs"]
+git-tree-sha1 = "ab29216184312f99ff957b32cd63c2fe9c928b91"
 uuid = "61eb1bfa-7361-4325-ad38-22787b887f55"
-version = "0.25.0"
+version = "0.26.7"
 
 [[deps.Glob]]
 git-tree-sha1 = "97285bbd5230dd766e9ef6749b80fc617126d496"
@@ -439,19 +444,19 @@ version = "0.9.23"
 
 [[deps.LLVM]]
 deps = ["CEnum", "LLVMExtra_jll", "Libdl", "Preferences", "Printf", "Requires", "Unicode"]
-git-tree-sha1 = "839c82932db86740ae729779e610f07a1640be9a"
+git-tree-sha1 = "2470e69781ddd70b8878491233cd09bc1bd7fc96"
 uuid = "929cbde3-209d-540e-8aea-75f648917ca0"
-version = "6.6.3"
+version = "8.1.0"
 weakdeps = ["BFloat16s"]
 
     [deps.LLVM.extensions]
     BFloat16sExt = "BFloat16s"
 
 [[deps.LLVMExtra_jll]]
 deps = ["Artifacts", "JLLWrappers", "LazyArtifacts", "Libdl", "TOML"]
-git-tree-sha1 = "88b916503aac4fb7f701bb625cd84ca5dd1677bc"
+git-tree-sha1 = "597d1c758c9ae5d985ba4202386a607c675ee700"
 uuid = "dad2f222-ce93-54a1-a47d-0025e8a3acab"
-version = "0.0.29+0"
+version = "0.0.31+0"
 
 [[deps.LLVMLoopInfo]]
 git-tree-sha1 = "2e5c102cfc41f48ae4740c7eca7743cc7e7b75ea"
@@ -1006,9 +1011,9 @@ version = "0.2.1"
 
 [[deps.UnsafeAtomicsLLVM]]
 deps = ["LLVM", "UnsafeAtomics"]
-git-tree-sha1 = "bf2c553f25e954a9b38c9c0593a59bb13113f9e5"
+git-tree-sha1 = "4073c836c2befcb041e5fe306cb6abf621eb3140"
 uuid = "d80eeb9a-aca5-4d75-85e5-170c8b632249"
-version = "0.1.5"
+version = "0.2.0"
 
 [[deps.XML2_jll]]
 deps = ["Artifacts", "JLLWrappers", "Libdl", "Libiconv_jll", "Zlib_jll"]

Project.toml

@@ -1,7 +1,7 @@
 name = "Walrus"
 uuid = "bec5164f-1c7a-4c32-8768-be9354254d6a"
 authors = ["Jago Stong-Wright <[email protected]>"]
-version = "0.5.3"
+version = "0.5.4"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
@@ -11,7 +11,7 @@ Oceananigans = "9e8cae18-63c1-5223-a75c-80ca9d6e9a09"
 Roots = "f2b01f46-fcfa-551c-844a-d8ac1e96c665"
 
 [compat]
-Adapt = "3.7"
+Adapt = "3, 4"
 CUDA = "5"
 KernelAbstractions = "0.9.22"
 OceanBioME = "0.10.5"

src/interpolations.jl

@@ -8,35 +8,86 @@ using Oceananigans.Architectures: on_architecture, CPU
 
 import Adapt: adapt_structure
 
-struct SimpleInterpolation{R, V}
-     range :: R
-    values :: V
+struct SimpleInterpolation{R, V, M} <: Function
+               range :: R
+              values :: V
+  boundary_condition :: M
 end
 
 adapt_structure(to, itp::SimpleInterpolation) = SimpleInterpolation(adapt(to, itp.range),
                                                                     adapt(to, itp.values))
 
-function SimpleInterpolation(range::Array, values; arch = CPU())
-    x₀ = minimum(range)
+function coordinate_to_range(coordinate)
+    x₀ = minimum(coordinate)
+    x₁ = maximum(coordinate)
 
-    (range[2] - range[1] ≈ range[end] - range[end - 1]) || throw(ArgumentError("Interpolation range must be regularly spaced"))
+    (coordinate[2] - coordinate[1] ≈ coordinate[end] - coordinate[end - 1]) || throw(ArgumentError("Interpolation coordinate must be regularly spaced"))
 
-    Δx = range[2] - range[1]
+    Δx = coordinate[2] - coordinate[1]
+
+    return (; x₀, x₁, Δx)
+end
+
+function SimpleInterpolation(coordinate::Array, values; arch = CPU(), boundary_condition = Cyclic())
+    range = coordinate_to_range(coordinate)
+
+    return SimpleInterpolation(range, on_architecture(arch, values), boundary_condition)
+end
+
+function SimpleInterpolation(coordinates::Tuple, values; arch = CPU(), boundary_conditions = Cyclic())
+    if !(boundary_conditions isa Tuple)
+        boundary_conditions = tuple([boundary_conditions for n in 1:length(coordinates)]...)
+    end
+
+    length(coordinates) == length(boundary_conditions) || 
+        throw(ArgumentError("You must either provide one boundary condition to be used in all directions, or one for each direction"))
+
+    ranges = tuple([coordinate_to_range(coord) for coord in coordinates])
+
+    return SimpleInterpolation(ranges, on_architecture(arch, values), boundary_condition)
+end
+
+(itp::SimpleInterpolation)(X...) = itp(X..., itp.values, itp.range, itp.boundary_condition)
+
+struct Limited end
+struct Cyclic end
+
+@inline function (::Limited)(x, x₀, Δx, N)
+    # this will (hopeflly) result in an out of bounds error 
+    # but the whole point of this excersise was to make it kernel safe
+    # so I'm not erroring here
+    n₀ = floor(Int, (x - x₀) / Δx)
+
+    return x, n₀ + 1, n₀ + 2
+end
+
+@inline function (::Cyclic)(x, x₀, Δx, N)
+    n₀ = floor(Int, (x - x₀) / Δx)
+
+    n₁ = mod(n₀, N) + 1
+
+    n₂ = ifelse(n₁ == N, 1, n₁ + 1)
+
+    x = mod(x - x₀, N * Δx) + x₀
 
-    return SimpleInterpolation((; x₀, Δx), on_architecture(arch, values))
+    return x, n₁, n₂
 end
 
+# the terminal case
+function (itp::SimpleInterpolation)(x::Number, values, range, boundary_condition)
+    x₀ = range.x₀
+    x₁ = range.x₁
+    Δx = range.Δx
+    N  = length(values)
 
-function (itp::SimpleInterpolation)(x)
-    n₁ = floor(Int, (x - itp.range.x₀) / itp.range.Δx)
+    x, n₁, n₂ = boundary_condition(x, x₀, Δx, N)
 
-    x₁ = itp.range.x₀ + itp.range.Δx * n₁
-    x₂ = x₁ + itp.range.Δx
+    x₁ = x₀ + Δx * (n₁ - 1)
 
-    y₁ = @inbounds itp.values[n₁ + 1]
-    y₂ = @inbounds itp.values[n₁ + 2]
+    y₁ = @inbounds values[n₁]
+    y₂ = @inbounds values[n₂]
 
-    return y₁ + (x - x₁) * (y₂ - y₁) / (x₂ - x₁)
+    return y₁ + (x - x₁) * (y₂ - y₁) / Δx
 end
 
 end # module

src/radiative_transfer/homogeneous_body_heating.jl

@@ -73,7 +73,7 @@ function HomogeneousBodyHeating(; surface_flux,
                                   water_heat_capacity = 3991.0, # J K⁻¹ kg⁻¹
                                   water_density = 1026.0) # kg m⁻³
 
-    isa(surface_flux, Function) || (surface_flux = ReturnValue(surface_flux))
+    surface_flux = normalise_surface_function(surface_flux)
 
     return HomogeneousBodyHeating(water_attenuation_coefficient,
                                   water_heat_capacity,
@@ -86,15 +86,13 @@ end
     cₚ = heating.water_heat_capacity
     α = heating.water_attenuation_coefficient
 
-    x, y, _ = node(i, j, k, grid, Center(), Center(), Center())
+    surface_flux = get_value(heating.surface_flux, i, j, grid, clock)
 
     zᶠ = znode(i, j, k, grid, Center(), Center(), Face())
 
     zᶠ⁺ = znode(i, j, k + 1, grid, Center(), Center(), Face())
 
-    t = clock.time
-
-    return α * heating.surface_flux(x, y, t) * (exp(- α * abs(zᶠ⁺)) - exp(- α * abs(zᶠ))) / (ρ * cₚ)
+    return α * surface_flux * (exp(- α * abs(zᶠ⁺)) - exp(- α * abs(zᶠ))) / (ρ * cₚ)
 end
 
 

src/radiative_transfer/radiative_transfer.jl

@@ -19,7 +19,7 @@ using Oceananigans.Utils: launch!
 
 using KernelAbstractions.Extras: @unroll
 
-using Walrus: normalise_surface_function
+using Walrus: get_value, normalise_surface_function
 
 import Oceananigans.Biogeochemistry: update_biogeochemical_state!, update_tendencies!, AbstractBiogeochemistry
 

src/surface_heating.jl

@@ -157,16 +157,16 @@ function SurfaceHeatExchange(; wind_stress,
 end
 
 """
-    SurfaceHeatExchange(; wind_stress,
-                          air_temperature = 18, # °C
-                          latent_heat_vaporisation = EmpiricalLatentHeatVaporisation(),
-                          vapour_pressure = AugustRocheMagnusVapourPressure(),
-                          water_specific_heat_capacity = 3991., # J / K / kg
-                          water_density = 1026., # kg / m³
-                          air_specific_heat_capacity = 1003.5, # J / K / kg
-                          air_density = 1.204, # kg
-                          air_water_mixing_ratio = 0.001, # kg / kg
-                          stephan_boltzman_constant = 5.670374419e-8) # W / K⁴
+    SurfaceHeatExchangeBoundaryCondition(; wind_stress,
+                                           air_temperature = 18, # °C
+                                           latent_heat_vaporisation = EmpiricalLatentHeatVaporisation(),
+                                           vapour_pressure = AugustRocheMagnusVapourPressure(),
+                                           water_specific_heat_capacity = 3991., # J / K / kg
+                                           water_density = 1026., #kg / m³
+                                           air_specific_heat_capacity = 1003.5, # J / K / kg
+                                           air_density = 1.204, # kg
+                                           air_water_mixing_ratio = 0.001, # kg / kg
+                                           stephan_boltzman_constant = 5.670374419e-8) # W / K⁴
 
 A convenience constructor returning `SurfaceHeatExchange` as a boundary condition
 
@@ -181,7 +181,7 @@ Keyword Arguments
 - `air_specific_heat_capacit`: the specific heat capacity of air in J / K / kg
 - `air_density`: air density in kg / m³
 - `air_water_mixing_ratio`: water content of air in kg / kg
-- `stephan_boltzman_constant`: the Stephan-Boltzman constant in W / K⁴
+- `stephan_boltzman_constant`: the Stephan-Boltzman constant in W / m² / K⁴
 
 Example
 =======
@@ -232,16 +232,16 @@ end
 
     z₀ = find_velocity_roughness_length(drag_coefficient, wind_speed, 10, params)
 
-    ū = κ * wind_speed / log(10 / z₀)
+    ū = κ * wind_speed / log(2/z₀)
 
-    isfinite(ū) || (ū = 0)
+    ū = ifelse(isfinite(ū), ū, 0)
 
     Rᵣ = ū * z₀ / params.ν
     zₒₜ = min(1.15e-4, 5.5e-5 * Rᵣ ^ -0.6)
 
-    result = params.κ ^ 2 / (log(10/z₀) * log(10/zₒₜ)) # hmm this might be meant to be 2
+    result = params.κ ^ 2 / (log(2/z₀) * log(2/zₒₜ))
 
-    isfinite(result) || (result = 0) # this should only occur if wind speed is zero in which case stress is zero anyway
+    result = ifelse(isfinite(result), result, 0)# this should only occur if wind speed is zero in which case stress is zero anyway
 
     return result
 end
@@ -277,8 +277,6 @@ end
     ρʷ  = interface.water_density
     cₚʷ = interface.water_specific_heat_capacity
 
-    t = clock.time
-
     u = @inbounds model_fields.u[i, j, grid.Nz]
     v = @inbounds model_fields.v[i, j, grid.Nz]
     T = @inbounds model_fields.T[i, j, grid.Nz]
@@ -291,7 +289,7 @@ end
     uʷ = - wind_speed * sind(wind_direction)
     vʷ = - wind_speed * cosd(wind_direction)
 
-    relative_speed = √((uʷ - u)^2 + (vʷ - v)^2)
+    relative_speed = wind_speed#√((uʷ - u)^2 + (vʷ - v)^2)
 
     cʰ = Cʰ(interface.wind_stress.drag_coefficient, relative_speed)