diff --git a/examples/scripts/Project.toml b/examples/scripts/Project.toml index 03b35b927..9dfaf41d8 100644 --- a/examples/scripts/Project.toml +++ b/examples/scripts/Project.toml @@ -2,6 +2,7 @@ CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0" Checkpointing = "eb46d486-4f9c-4c3d-b445-a617f2a2f1ca" +Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" GeoMakie = "db073c08-6b98-4ee5-b6a4-5efafb3259c6" InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240" Interpolations = "a98d9a8b-a2ab-59e6-89dd-64a1c18fca59" diff --git a/examples/scripts/soil_heat_global_sensitivity.jl b/examples/scripts/soil_heat_global_sensitivity.jl new file mode 100644 index 000000000..c16eb8767 --- /dev/null +++ b/examples/scripts/soil_heat_global_sensitivity.jl @@ -0,0 +1,60 @@ +# currently only works in Julia 1.10 with Enzyme +using Terrarium + +using CUDA +using Dates +using Rasters, NCDatasets +using Statistics + +using CairoMakie, GeoMakie +using Enzyme, Checkpointing + +import RingGrids +import SpeedyWeather + +# run on GPU if available +arch = CUDA.functional() ? Terrarium.GPU() : Terrarium.CPU() + +# Load land-sea mask at ~1° resolution +land_sea_frac = convert.(Float32, dropdims(Raster("inputs/era5-land_land_sea_mask_N72.nc"), dims=Ti)) +land_sea_frac_field = RingGrids.FullGaussianField(Matrix(land_sea_frac), input_as=Matrix) +heatmap(land_sea_frac_field) + +# Load ERA-5 2 meter air temperature at ~1° resolution +Tair_raster = Raster("inputs/external/era5-land/2m_temperature/era5_land_2m_temperature_2023_N72.nc") +Tsurf_0 = convert.(Float32, replace_missing(Tair_raster, NaN)) .- 273.15f0 +# heatmap(Tair_raster[:,:,1]) + +# Set up grids +land_mask = land_sea_frac_field .> 0.5 # select only grid points with > 50% land + +# for now let's do it actually without the mask to keep the example simple +grid = ColumnRingGrid(arch, Float64, ExponentialSpacing(N=30), land_mask.grid) #, land_mask.grid, land_mask) +lon, lat = RingGrids.get_londlatds(grid.rings) + +# Construct input sources +Tair_forcing = InputSource(grid, rebuild(Tair_raster, name=:Tair)) +Tsurf_0 = Tair_raster[Ti(1)][findall(land_mask)] + +# Initial conditions +initializer = FieldInitializers( + # steady-ish state initial condition for temperature + temperature = (x,z) -> -1 - 0.01*z, + # fully saturated soil pores + saturation_water_ice = 1.0, +) +model = SoilModel(grid; initializer) +# constant surface temperature of 1°C +boundary_conditions = PrescribedSurfaceTemperature(:Tair) +integrator = initialize(model, ForwardEuler(), Tair_forcing; boundary_conditions) + +# spin up a little +@time run!(integrator, period=Day(5), Δt=900.0) + +# Enzyme prep +scheme = Revolve(1) +dintegrator = make_zero(integrator) +N_t = 200 + +# this uses checkpointing +autodiff(set_runtime_activity(Enzyme.Reverse), run!, Const, Duplicated(integrator, dintegrator), Const(scheme), Const(N_t))