Coupling with JAX-GCM (to Oceananigans or veros ocean model)

[glwagner]

PR #602 adds capabilities for coupling veros with PrescribedAtmosphere or with a prognostic atmosphere like SpeedyWeather.

In the same vein, it would be nice to develop capability to build coupled models with JAX-GCM as the atmosphere component.

This would enable coupled simulations with JAX-GCM and Oceananigans or JAX-GCM and veros.

cc @simone-silvestri @jvmadan @duncanwp

[glwagner]

@duncanwp @jvmadan would it also be possible to couple NeuralGCM to veros and Oceananigans via this approach? Or will the coupling with JAX-GCM look different than coupling with NeuralGCM?

cc @milankl who may also be interested in this work

[glwagner]

Also relevant is the coupled Speedy-Oceananigans simulation that we have now:

https://clima.github.io/ClimaOceanDocumentation/dev/literated/global_climate_simulation/

This uses the helper atmosphere_simulation to build a SpeedyWeather simulation when given a SpectralGrid:

ClimaOcean.jl/examples/global_climate_simulation.jl

Lines 72 to 76 in 9ccb555

	nlayers = 4
	spectral_grid = SpeedyWeather.SpectralGrid(; trunc=63, nlayers, Grid=FullClenshawGrid)
	atmosphere = atmosphere_simulation(spectral_grid; output=true)
	atmosphere.model.output.output_dt = Hour(3)
	nothing # hide

which is implemented here:

ClimaOcean.jl/ext/ClimaOceanSpeedyWeatherExt/speedy_atmosphere_simulations.jl

Line 63 in 9ccb555

function atmosphere_simulation(spectral_grid::SpeedyWeather.SpectralGrid; output=false)

[duncanwp]

Thanks for the context and examples @glwagner.

We're working on a simple coupler internally and have a couple of APIs we're trying out so it's useful to compare. We should be able to open them up soon but in the mean time the basic pattern is to have a flux model calculate the updated boundary conditions for the atmosphere and just run it one day at a time with the ocean's SST. The main interface to JCM would be Model.run_from_state: https://github.com/climate-analytics-lab/jax-gcm/blob/aae23a0e0ccf689d2c19f4a2d3b206276fc57a0e/jcm/model.py#L424

This is currently distinct from the NeuralGCM interface, which itself is currently being updated, but the pattern would be very similar.

Looping in @meteorologytoday who has been working on this for us.

[glwagner]

Thanks for the context and examples @glwagner.

We're working on a simple coupler internally and have a couple of APIs we're trying out so it's useful to compare. We should be able to open them up soon but in the mean time the basic pattern is to have a flux model calculate the updated boundary conditions for the atmosphere and just run it one day at a time with the ocean's SST. The main interface to JCM would be Model.run_from_state: https://github.com/climate-analytics-lab/jax-gcm/blob/aae23a0e0ccf689d2c19f4a2d3b206276fc57a0e/jcm/model.py#L424

This is currently distinct from the NeuralGCM interface, which itself is currently being updated, but the pattern would be very similar.

Looping in @meteorologytoday who has been working on this for us.

Thanks @duncanwp ! Are the air-sea fluxes then contained in the forcing input to run_from_state? eg

https://github.com/climate-analytics-lab/jax-gcm/blob/aae23a0e0ccf689d2c19f4a2d3b206276fc57a0e/jcm/model.py#L426

For veros I think there is a step function, which is called by our time_step! function, eg

ClimaOcean.jl/ext/ClimaOceanPythonCallExt/veros_ocean_simulation.jl

Line 34 in 61ff994

time_step!(ocean::VerosOceanSimulation, Δt) = ocean.setup.step(ocean.setup.state)

The coupling is performed by filling internal veros arrays that store the surface fluxes here:

ClimaOcean.jl/ext/ClimaOceanPythonCallExt/veros_state_exchanger.jl

Line 86 in 61ff994

function fill_net_fluxes!(ocean::VerosOceanSimulation, net_ocean_fluxes)

I guess veros does not support tracer fluxes yet, so we can only do momentum fluxes (am I reading that right @simone-silvestri ?)

[meteorologytoday]

Hi, thanks for including me. I think the wrapping of Jax-gcm and NeuralGCM into a component that can be coupled with ClimaOcean.jl should be fairly similar.

Any fully-jax style model will provide some sort of "step_function" that advances the input model state with a given forcing and returns the output and history, i.e., final_state, history = step_function(state, forcing, time). This is "stateless" design can be easily integrated into state-aware framework, such as CliMa, by using a wrapper to cache the current model state.

Currently, jax-gcm receives sea-surface temperature as forcing, and computes the fluxes. In the future, I believe they want to allow jax-gcm to be forced by heat/momentum/freshwater fluxes, such that the air-sea flux algorithm can be isolated, ideal for air-sea interaction research.

Hopefully, this makes sense.

Just some additional input: This is jax-gcm's run_from_state signature
https://github.com/climate-analytics-lab/jax-gcm/blob/aae23a0e0ccf689d2c19f4a2d3b206276fc57a0e/jcm/model.py#L424-L430

and the structure of ForcingData is here:
https://github.com/climate-analytics-lab/jax-gcm/blob/aae23a0e0ccf689d2c19f4a2d3b206276fc57a0e/jcm/forcing.py#L9-L17

which contains ice, snow, sea-surface temperature and such.

[glwagner]

Currently, jax-gcm receives sea-surface temperature as forcing, and computes the fluxes. In the future, I believe they want to allow jax-gcm to be forced by heat/momentum/freshwater fluxes, such that the air-sea flux algorithm can be isolated, ideal for air-sea interaction research.

Hopefully, this makes sense.

It makes a lot of sense! Thank you.

I think this exercise is most interesting if ClimaOcean.OceanSeaIceModel computes the interface state (fluxes, and optionally skin temperature, skin velocity, etc). This allows coupled configurations to be compared apples to apples because they use exactly the same flux/interface state formulation.

It could also be interesting to run JAX-GCM with run_from_state wherein the SST is the top layer of an Oceananigans ocean model. Presumably that would be pretty easy since you simply need to communicate the fluxes computed by JAX-GCM back to Oceananigans. But I don't think there's a point to using OceanSeaIceModel. One could use ocean = ClimaOcean.ocean_simulation(grid) and work from there.

[meteorologytoday]

@glwagner

1. I totally agree that there is a lot of value in separating the air-sea flux algorithm, and this is what we aim for 👏.
2. The internal coupler under development might put this into a component called FluxModel, because workflow-wise computing air-sea flux behaves pretty much like another coupled component.
3. And yes, it is definitely nice to use the existing framework to create a beta version ClimaOcean-JaxGcm coupled model, and go from there. I would be happy to provide help.

I also personally think it would be a fantastic example of a cross-language coupled model!

[glwagner]

Any fully-jax style model will provide some sort of "step_function" that advances the input model state with a given forcing and returns the output and history, i.e., final_state, history = step_function(state, forcing, time).

It looks like veros does something like this:

ClimaOcean.jl/ext/ClimaOceanPythonCallExt/veros_ocean_simulation.jl

Line 34 in 61ff994

time_step!(ocean::VerosOceanSimulation, Δt) = ocean.setup.step(ocean.setup.state)

Veros has objects called "setups" that store their state:

ClimaOcean.jl/ext/ClimaOceanPythonCallExt/veros_ocean_simulation.jl

Lines 84 to 89 in 61ff994

	function VerosOceanSimulation(setup, setup_name::Symbol)
	setups = pyimport("veros.setups." * setup)
	setup = @eval $setups.$setup_name()
	# instantiate the setup
	setup.setup()

I also personally think it would be a fantastic example of a cross-language coupled model!

Yes for sure. I think that python and julia communicate seamlessly (so it seems), so there is no real barrier to this. The GCM design is by far the more challenging problem.

Help me understand if I have this right: to start from the example in the README, we would have

from jcm.model import Model
from juliacall import Main as jl

# Create a model with default configuration
time_step = 30.0

model = Model(
    time_step=time_step,          # minutes
    layers=8,                 # vertical levels
    horizontal_resolution=31  # T31 spectral grid
)

jl.seval("using Oceananigans")
jl.seval("using ClimaOcean")
# etc
ocean = jl.ocean_simulation(grid) # build Oceananigans simulation

Next we run one time step with run_from_state:

# Because I don't know how else to generate an initial condition:
state = model.run(
    save_interval=10.0,  # save every 10 days
    total_time=1.0     # total simulation time in days
)

# Build `ForcingData`, including Oceananigans SST?
# may need to use the interpolation mechanism we've build for SpeedyWeather

# then
model.run_from_state(state, forcing, total_time=time_step / 86400)
jl.time_step!(ocean, time_step) # if using juliacall

the workflow can be managed by either julia or python via PythonCall: https://github.com/JuliaPy/PythonCall.jl

[meteorologytoday]

Hi @glwagner
I modify your code according to the current jax-gcm design:

import xarray as xr
from juliacall import Main as jl
from jcm.model import Model
from jcm.forcing import ForcingData 

# Simulation setup
total_steps = 360
coupling_time_step = 86400.0 # seconds

# Ocean initialization
jl.seval("using Oceananigans")
jl.seval("using ClimaOcean")
ocean_model = jl.ocean_simulation(grid)

# Atmosphere initialization
atmosphere_model = Model(
    time_step = 30.0, # minutes
    layers=8,                 # vertical levels
    horizontal_resolution=31  # T31 spectral grid
)
atmosphere_state = atmosphere_model._prepare_initial_modal_state()
atmosphere_forcing = ForcingData.zeros().copy(
    sea_surface_temperature = ocean.SST,
    # and other variables
)
ocean_forcing = dict()

atmosphere_history = []
# Running Model
for step in range(total_steps):
    print(f"Doing step {step:d} ... ")
    ocean_model.time_step!(ocean_model, ocean_forcing, coupling_time_step)
    atmosphere_state, _history = atmosphere_model.run_from_state(
        initial_state = atmosphere_state,
        forcing = atmosphere_forcing,
        total_time = coupling_time_step / 86400.0,
    )
    atmosphere_history.append(_history.to_xarray())
    
    # update forcing

    atmosphere_forcing = atmosphere_forcing.copy(
        sea_surface_temperature = ocean.SST,
        # and other variables. Please see https://github.com/climate-analytics-lab/jax-gcm/blob/main/jcm/forcing.py
    )

    # Upward = positive, downward = negative
    ocean_forcing = dict(
        net_surface_shortwave_radiation_flux = - _history.shortwave_rad.rsns,
        net_surface_longwave_radiation_flux = _history.surface_flux.rlns,
        freshwater_flux_precipitation = - (_history.condensation.precls + _history.convection.precnv) / 1e3,
        freshwater_flux_evaporation = _history.surface_flux.evap / 1e3,
        sensible_heat_flux = _history.physics.surface_flux.shf,
    )

atmosphere_history = xr.merge(atmosphere_history)

Of course there are some details about grid-remapping that we need to think about. But this should be enough to illustrate necessary blocks.

Add: Please be aware that jax-gcm dimension is (lon, lat) rather than (lat, lon). With T31 grid you can quickly check variables shape should be (nlon, nlat)=(96, 48)

[simone-silvestri]

This looks amazing! We would also need to grab the fluxes computed by jcm so that we can input them into the ocean_simulation to force the ocean.

[meteorologytoday]

@simone-silvestri Hi I just updated the code demo to include the extraction of surface fluxes from jax-gcm

[duncanwp]

Thanks @meteorologytoday !

For reference we're pretty settled on most of this API for v1, but also open to feedback + suggestions going forward. We would like a neater (public) way to generate an initial state, and this might happen in the constructor in future.

[glwagner]

Here's a demo of using juliacall to run an Oceananigans model from python:

https://github.com/NumericalEarth/pynanigans/blob/main/single_column_model.py

I suspect the reverse is similarly seamless. @duncanwp if you want (and depending on grid compatibility) you may be able to use ClimaOcean Metadata utilities for downloading data + Oceananigans for interpolating for state initialization. SpeedyWeather might have something too. We don't yet have the ability to build speedy grids natively in Oceananigans, but I think some of them could be supported.