Optimizing a hacky coupler between Oceananigans and WRF

[jakob-braga]

Hello I've recently wrote a bit of hacky code that couples Oceananigans to WRF. In essence this is what the code looks like within a single time step of Oceananignas:

• write the current ocean SST to the WRF "TSK" (Basically SST) variable in a restart file
• restart WRF from the restart file that I've edited, run it for the same time as the ocean models time step
• read in the last WRF output frame and read in the wind stress, update the wind stress boundary condition
• continue the ocean model

It seems to work pretty well, however clearly the writing and reading to files on every time step of the ocean model it a huge bottleneck and is going to seemingly cause problems running on GPU.

Has anyone done something like this before? If so were you able to overcome this bottleneck somehow?

[simone-silvestri]

Hi @jakob-braga. Interesting work, do you have a link to the code you have written to share?

As for coupling with an atmosphere, we have integrated Oceanigans with a prescribed atmosphere (not an atmospheric model) in ClimaOcean. You can check out how this coupling is implemented here. Essentially, we load the prescribed atmosphere using the FieldTimeSeries utility into a PrescribedAtmosphere type, which includes atmospheric properties such as temperature, winds, and humidity. This is then incorporated into an OceanSeaIceModel. During execution, the model computes the ocean time step and the atmosphere-ocean fluxes, which are used as top boundary conditions for ocean velocities and tracers.

We’re currently working on developing documentation and including examples to help users understand how to use ClimaOcean effectively.

That said, coupling a GPU model with a CPU model can be inefficient due to the overhead of memory transfers between the two architectures. It might be more effective to run everything on the CPU if memory transfer is the bottleneck. Additionally, it is possible to call Fortran code from Julia (see here), which could allow coupling without the need for disk read/write operations.

[glwagner]

@jakob-braga may be running an Oceananigans LES though, and wants to use WRF to compute the fluxes. So I'm not sure that ClimaOcean is the best way forward for that. @jakob-braga can you clarify the kind of application you are working on?

[jakob-braga]

@simone-silvestri Thanks for the insight, I will look into triggering the WRF Fortran code from within Julia, this could be helpful. I'm not sure the prescribed atmosphere is exactly what I am after in ClimaOcean (see below).

(@glwagner) I'm currently looking at two way coupling between the atmosphere and ocean. The code I have now runs an idealized Oceananigans LES alongside an idealized WRF LES in a way that lets me share some of the fluxes between them and allows them to evolve dynamically together, this gives me both atmospheric and ocean outputs that lets me understand the coupling better, especially in areas with a heterogeneous SST.

The project seems to work well right now (despite the IO bottleneck), but is currently a little hacky and rough, I will spend today cleaning it up a bit and make a repository to share later today or this weekend.

Thanks so much for the quick replies!

[jakob-braga]

@simone-silvestri @glwagner here is a link to the project as is for reference Oceananigans-WRF-Coupler.

[glwagner]

Is your WRF implementation on the GPU? If not then you will always be limited by CPU-GPU data transfer.

If you can run WRF on the GPU then getting this to work probably requires figuring out how to control your WRF simulation from julia (actually more specifically I think you need to be able to take a single WRF time-step from julia). If you are able to do that, then you might be able to setup a time-stepping algorithm that copies the Oceananigans-derived SST into the appropriate WRF variable in between time-steps.

[jakob-braga]

Thanks for the reply and insight in the issues on GPU-CPU data transfer, I will stick to using CPUs for the Oceananignans sim.

I will look into taking a single time step in WRF although this may be a bit daunting of a task.

[glwagner]

How do you run WRF right now?

The tricky part is namespaces. To make this work without writing data to disk, you need to have a common namespace for Oceananigans and WRF variables, right?

Note, a similar effort was undertaken for FV3 here: https://gmd.copernicus.org/articles/14/4401/2021/

I see there is something called wrf-python but this looks like its only for analyzing output, not running WRF.

[glwagner]

Are you tied to WRF? I think this project is feasible and indeed, with an atmospheric model capable of computing fluxes that requires only SST, maybe not even too difficult. But I think two critical requirements are 1) the atmospheric LES runs on GPUs and 2) the atmospheric model has a high-level interface so that performant coupling can be achieved without a heroic multi-year effort.

[jakob-braga]

Sorry, I'm not sure how but I just missed these replies. As of right now I am running WRF as an executable that gets built from compiling, which makes this not possible, I spent the last few days looking at trying to run Fortran code in Julia, but WRF is such a massive project that its not clear to me how to simply make a single time step. I may continue looking into this however when I have time. I have used wrf-python, form my understanding it is mostly just for post processing.

As of right now I would say I am unfortunately tied to WRF, I am a masters student half way through my studies and the primary part of my research is not building a coupler but looking at SST influence on the atmosphere and wind stress influence on the ocean. It took me a little bit to learn WRF and I am not sure I can commit to learning another model. What I have built allows me to get better results and allows me to talk more about the significance of coupling (even though it is quite slow).

After looking around I think the next easier step to speed this up is to see if I can save the intermediate files virtually to memory instead of actually writing to disk, which should give a decent speed up.

Thank you for the replies! If I have time to come back to this later, I do think using a different atmospheric model would make this a lot easier and will probably be worth it. If I have some free time before my masters is over I may also be able to give this a try.

[glwagner]

No worries at all! Thank you for a very clear description of the project --- it's challenging to say the least! I wish it wasn't so challenging, but this is software in 2024.

Curious to hear updates on what you figure out and the kind of performance you're able to achieve!

I'd also be curious to know what WRF developers might say if you posed the same question to them. But I think it is perfectly reasonable for you to stay the road you're currently on.