Skip to content

Commit 78de681

Browse files
simone-silvestrisimone-silvestri
committed
start with adapting
1 parent abced34 commit 78de681

File tree

12 files changed

+389
-90
lines changed

12 files changed

+389
-90
lines changed

‎CondaPkg.toml‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11

22
[pip.deps]
3-
copernicusmarine = ">=2.0.0"
43
xarray = ">=2024.7.0"
5-
numpy = ">=2.0.0"
64
jax = ">=0.6"
5+
copernicusmarine = ">=2.0.0"
6+
numpy = "==2.3.1"
77
tensorflow = ">=2.17"

‎Manifest.toml‎

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1164,23 +1164,25 @@ version = "0.4.6"
11641164

11651165
[[deps.Roots]]
11661166
deps = ["Accessors", "CommonSolve", "Printf"]
1167-
git-tree-sha1 = "668e411c0616a70860249b4c96e5d35296631a1d"
1167+
git-tree-sha1 = "dc84d966d3f6240657501da6d43b8448da19f71e"
11681168
uuid = "f2b01f46-fcfa-551c-844a-d8ac1e96c665"
1169-
version = "2.2.8"
1169+
version = "2.2.9"
11701170

11711171
[deps.Roots.extensions]
11721172
RootsChainRulesCoreExt = "ChainRulesCore"
11731173
RootsForwardDiffExt = "ForwardDiff"
11741174
RootsIntervalRootFindingExt = "IntervalRootFinding"
11751175
RootsSymPyExt = "SymPy"
11761176
RootsSymPyPythonCallExt = "SymPyPythonCall"
1177+
RootsUnitfulExt = "Unitful"
11771178

11781179
[deps.Roots.weakdeps]
11791180
ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
11801181
ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
11811182
IntervalRootFinding = "d2bf35a9-74e0-55ec-b149-d360ff49b807"
11821183
SymPy = "24249f21-da20-56a4-8eb1-6a02cf4ae2e6"
11831184
SymPyPythonCall = "bc8888f7-b21e-4b7c-a06a-5d9c9496438c"
1185+
Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
11841186

11851187
[[deps.Rotations]]
11861188
deps = ["LinearAlgebra", "Quaternions", "Random", "StaticArrays"]
@@ -1524,9 +1526,9 @@ version = "1.3.0+0"
15241526

15251527
[[deps.libaec_jll]]
15261528
deps = ["Artifacts", "JLLWrappers", "Libdl"]
1527-
git-tree-sha1 = "f5733a5a9047722470b95a81e1b172383971105c"
1529+
git-tree-sha1 = "1aa23f01927b2dac46db77a56b31088feee0a491"
15281530
uuid = "477f73a3-ac25-53e9-8cc3-50b2fa2566f0"
1529-
version = "1.1.3+0"
1531+
version = "1.1.4+0"
15301532

15311533
[[deps.libblastrampoline_jll]]
15321534
deps = ["Artifacts", "Libdl"]

‎anderson_acceleration.jl‎

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,40 @@
1+
f(x) = sin(x) + atan(x)
2+
x0 = 1
3+
4+
k_max = 100 # Maximum number of iterations.
5+
tol_res = 1e-6 # Tolerance on the residual.
6+
m = 3 # Parameter m.
7+
8+
x = [x0, f(x0)] # Vector of iterates x.
9+
g = f(x) - x # Vector of residuals.
10+
11+
G_k = g[2] - g[1] # Matrix of increments in residuals.
12+
X_k = x[2] - x[1] # Matrix of increments in x.
13+
14+
k = 2
15+
while k < k_max && abs(g[k]) > tol_res
16+
m_k = min(k, m);
17+
18+
# Solve the optimization problem by QR decomposition.
19+
[Q, R] = qr(G_k)
20+
gamma_k = R \ (Q' * g(k))
21+
22+
# Compute new iterate and new residual.
23+
push!(x, x[k] + g[k] - (X_k + G_k) * gamma_k);
24+
push!(g, f(x[k + 1]) - x[k + 1]);
25+
26+
# Update increment matrices with new elements.
27+
X_k = [X_k..., x[k + 1] - x[k]]
28+
G_k = [G_k..., g[k + 1] - g[k]]
29+
30+
n = size(X_k, 2)
31+
if n > m_k
32+
X_k = X_k[:, n - m_k + 1:end];
33+
G_k = G_k[:, n - m_k + 1:end];
34+
end
35+
36+
k = k + 1;
37+
end
38+
39+
% Prints result: Computed fixed point 2.013444 after 9 iterations
40+
fprintf("Computed fixed point %f after %d iterations\n", x(end), k);

‎experiments/veros_forced_simulation.jl‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
1+
using ClimaOcean
2+
using PythonCall
3+
using Oceananigans
4+
5+
VerosModule = Base.get_extension(ClimaOcean, :ClimaOceanPythonCallExt)
6+
VerosModule.remove_outputs(:global_4deg)
7+
8+
ocean = VerosModule.veros_ocean_simulation("global_4deg", :GlobalFourDegreeSetup)
9+
VerosModule.veros_settings_set!(ocean, "dt_tracer", 1800.0)
10+
11+
atmos = JRA55PrescribedAtmosphere(; backend = JRA55NetCDFBackend(10))
12+
radiation = Radiation()
13+
coupled_model = OceanSeaIceModel(ocean, nothing; atmosphere=atmos, radiation)
14+

‎ext/ClimaOceanPythonCallExt/ClimaOceanPythonCallExt.jl‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
1+
module ClimaOceanPythonCallExt
2+
3+
using ClimaOcean
4+
using CondaPkg
5+
using PythonCall
6+
using Oceananigans
7+
using Oceananigans.DistributedComputations: @root
8+
9+
using Dates: DateTime
10+
11+
include("clima_ocean_copernicus.jl")
12+
include("clima_ocean_veros.jl")
13+
14+
end # module ClimaOceanPythonCallExt

‎ext/ClimaOceanPythonCallExt.jl‎ renamed to ‎ext/ClimaOceanPythonCallExt/clima_ocean_copernicus.jl‎

Lines changed: 1 addition & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,3 @@
1-
module ClimaOceanPythonCallExt
2-
3-
using ClimaOcean
4-
using CondaPkg
5-
using PythonCall
6-
using Oceananigans
7-
using Oceananigans.DistributedComputations: @root
8-
9-
using Dates: DateTime
101
using ClimaOcean.DataWrangling.Copernicus: CopernicusMetadata
112

123
import ClimaOcean.DataWrangling: download_dataset
@@ -95,6 +86,4 @@ function depth_bounds_kw(z)
9586
minimum_depth = - z[2]
9687
maximum_depth = - z[1]
9788
return (; minimum_depth, maximum_depth)
98-
end
99-
100-
end # module ClimaOceanPythonCallExt
89+
end

‎ext/ClimaOceanPythonCallExt/clima_ocean_veros.jl‎

Lines changed: 144 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,144 @@
1+
using CondaPkg
2+
3+
using ClimaOcean.OceanSeaIceModels: reference_density, heat_capacity, SeaIceSimulation
4+
5+
import Oceananigans.Fields: set!
6+
import Oceananigans.TimeSteppers: time_step!
7+
8+
import ClimaOcean.OceanSeaIceModels: OceanSeaIceModel
9+
10+
"""
11+
install_veros()
12+
13+
Install the Copernicus Marine CLI using CondaPkg.
14+
Returns a NamedTuple containing package information if successful.
15+
"""
16+
function install_veros()
17+
CondaPkg.add("veros"; channel = "conda-forge")
18+
cli = CondaPkg.which("veros")
19+
@info "... the veros CLI has been installed at $(cli)."
20+
return cli
21+
end
22+
23+
struct VerosOceanSimulation{S}
24+
setup :: S
25+
end
26+
27+
time_step!(sim::VerosOceanSimulation, Δt) = sim.setup.step()
28+
29+
function remove_outputs(setup::Symbol)
30+
rm("$(setup).averages.nc", force=true)
31+
rm("$(setup).energy.nc", force=true)
32+
rm("$(setup).overturning.nc", force=true)
33+
rm("$(setup).snapshot.nc", force=true)
34+
return nothing
35+
end
36+
37+
const Field2D = Field{<:Any, <:Any, <:Nothing}
38+
39+
function set!(field::Field2D, pyarray::Py, k=pyconvert(Int, pyarray.shape[2]))
40+
array = PyArray(pyarray)
41+
Nx, Ny, Nz = size(array)
42+
set!(field, view(array, 3:Nx-2, 3:Ny-2, k, 1))
43+
return field
44+
end
45+
46+
function veros_ocean_simulation(setup, setup_name)
47+
setups = pyimport("veros.setups." * setup)
48+
setup = @eval $setups.$setup_name()
49+
50+
# instantiate the setup
51+
setup.setup()
52+
53+
return VerosOceanSimulation(setup)
54+
end
55+
56+
function surface_grid(setup::VerosOceanSimulation)
57+
58+
xf = Array(PyArray(setup.state.variables.xu))
59+
yf = Array(PyArray(setup.state.variables.yu))
60+
61+
xc = Array(PyArray(setup.state.variables.xt))
62+
yc = Array(PyArray(setup.state.variables.yt))
63+
64+
xf = xf[2:end-2]
65+
yf = yf[2:end-2]
66+
67+
xc = xc[3:end-2]
68+
yc = yc[3:end-2]
69+
70+
xf[1] = xf[2] - 2xc[1]
71+
yf[1] = sign(yf[2]) * (yf[2] - 2yc[1])
72+
73+
TX = if xf[1] == 0 && xf[end] == 360
74+
Periodic
75+
else
76+
Bounded
77+
end
78+
79+
Nx = length(xc)
80+
Ny = length(yc)
81+
82+
return LatitudeLongitudeGrid(size=(Nx, Ny), longitude=xf, latitude=yf, topology=(TX, Bounded, Flat))
83+
end
84+
85+
function veros_set!(ocean::VerosOceanSimulation, v, x)
86+
s = ocean.setup
87+
pyexec("""
88+
with setup.state.variables.unlock():
89+
setup.state.variables.__setattr__(y, t)
90+
""", Main, (y=v, t=x, setup=s))
91+
end
92+
93+
function veros_settings_set!(ocean::VerosOceanSimulation, v, x)
94+
s = ocean.setup
95+
pyexec("""
96+
with setup.state.settings.unlock():
97+
setup.state.settings.__setattr__(y, t)
98+
""", Main, (y=v, t=x, setup=s))
99+
end
100+
101+
function OceanSeaIceModel(ocean::VerosOceanSimulation, sea_ice=nothing;
102+
atmosphere = nothing,
103+
radiation = Radiation(),
104+
clock = Clock(time=0),
105+
ocean_reference_density = 1020.0,
106+
ocean_heat_capacity = 3998.0,
107+
sea_ice_reference_density = reference_density(sea_ice),
108+
sea_ice_heat_capacity = heat_capacity(sea_ice),
109+
interfaces = nothing)
110+
111+
if sea_ice isa SeaIceSimulation
112+
if !isnothing(sea_ice.callbacks)
113+
pop!(sea_ice.callbacks, :stop_time_exceeded, nothing)
114+
pop!(sea_ice.callbacks, :stop_iteration_exceeded, nothing)
115+
pop!(sea_ice.callbacks, :wall_time_limit_exceeded, nothing)
116+
pop!(sea_ice.callbacks, :nan_checker, nothing)
117+
end
118+
end
119+
120+
# Contains information about flux contributions: bulk formula, prescribed fluxes, etc.
121+
if isnothing(interfaces) && !(isnothing(atmosphere) && isnothing(sea_ice))
122+
interfaces = ComponentInterfaces(atmosphere, ocean, sea_ice;
123+
ocean_reference_density,
124+
ocean_heat_capacity,
125+
sea_ice_reference_density,
126+
sea_ice_heat_capacity,
127+
radiation)
128+
end
129+
130+
arch = CPU()
131+
132+
ocean_sea_ice_model = OceanSeaIceModel(arch,
133+
clock,
134+
atmosphere,
135+
sea_ice,
136+
ocean,
137+
interfaces)
138+
139+
# Make sure the initial temperature of the ocean
140+
# is not below freezing and above melting near the surface
141+
initialization_update_state!(ocean_sea_ice_model)
142+
143+
return ocean_sea_ice_model
144+
end

‎ext/ClimaOceanPythonCallExt/veros_state_exchanger.jl‎

Lines changed: 84 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,84 @@
1+
import ClimaOcean.OceanSeaIceModels.InterfaceComputations:
2+
state_exchanger,
3+
sea_ice_ocean_interface,
4+
atmosphere_ocean_interface,
5+
initialize!,
6+
get_ocean_state,
7+
ocean_surface_fluxes,
8+
get_radiative_forcing,
9+
fill_up_net_fluxes!
10+
11+
12+
mutable struct VerosStateExchanger{G, OST, AST, AEX}
13+
exchange_grid :: G
14+
exchange_ocean_state :: OST
15+
exchange_atmosphere_state :: AST
16+
atmosphere_exchanger :: AEX
17+
end
18+
19+
mutable struct ExchangeOceanState{FC, CF, CC}
20+
u :: FC
21+
v :: CF
22+
T :: CC
23+
S :: CC
24+
end
25+
26+
ExchangeOceanState(grid) = ExchangeOceanState(Field{Face, Center, Nothing}(grid),
27+
Field{Center, Face, Nothing}(grid),
28+
Field{Center, Center, Nothing}(grid),
29+
Field{Center, Center, Nothing}(grid))
30+
31+
function state_exchanger(ocean::VerosOceanSimulation, atmosphere)
32+
exchange_grid = surface_grid(ocean)
33+
exchange_ocean_state = ExchangeOceanState(exchange_grid)
34+
exchange_atmosphere_state = ExchangeAtmosphereState(exchange_grid)
35+
36+
atmosphere_exchanger = AtmosphereExchanger(atmosphere, exchange_grid)
37+
38+
return VerosStateExchanger(exchange_grid,
39+
exchange_ocean_state,
40+
exchange_atmosphere_state,
41+
atmosphere_exchanger)
42+
end
43+
44+
atmosphere_ocean_interface(ocean::VerosOceanSimulation, args...) =
45+
atmosphere_ocean_interface(surface_grid(ocean), args...)
46+
47+
sea_ice_ocean_interface(ocean::VerosOceanSimulation, args...) =
48+
sea_ice_ocean_interface(surface_grid(ocean), args...)
49+
50+
initialize!(exchanger::VerosStateExchanger, atmosphere) =
51+
initialize!(exchanger.atmosphere_exchanger, exchanger.exchange_grid, atmosphere)
52+
53+
@inline function get_ocean_state(ocean::VerosOceanSimulation, exchanger::VerosStateExchanger)
54+
u = exchanger.exchange_ocean_state.u
55+
v = exchanger.exchange_ocean_state.v
56+
T = exchanger.exchange_ocean_state.T
57+
S = exchanger.exchange_ocean_state.S
58+
59+
set!(u, ocean.setup.state.variables.u)
60+
set!(v, ocean.setup.state.variables.v)
61+
set!(T, ocean.setup.state.variables.temp)
62+
set!(S, ocean.setup.state.variables.salt)
63+
64+
return (; u, v, T, S)
65+
end
66+
67+
@inline function ocean_surface_fluxes(ocean::VerosOceanSimulation)
68+
grid = surface_grid(ocean)
69+
u = Field{Face, Center, Nothing}(grid)
70+
v = Field{Center, Face, Nothing}(grid)
71+
T = Field{Center, Center, Nothing}(grid)
72+
S = Field{Center, Center, Nothing}(grid)
73+
74+
return (; u, v, T, S)
75+
end
76+
77+
@inline get_radiative_forcing(ocean::VerosOceanSimulation) = nothing
78+
79+
function fill_up_net_fluxes!(ocean::VerosOceanSimulation, net_ocean_fluxes)
80+
veros_set!(ocean, "surface_taux", net_ocean_fluxes.u)
81+
veros_set!(ocean, "surface_tauy", net_ocean_fluxes.v)
82+
83+
return nothing
84+
end

0 commit comments

Comments
 (0)