How to adjoint the barotropic gyre model from Oceananigans

[swilliamson7]

As of right now, the steps to get to an adjoint-able barotropic gyre model are

1. Create a "forward step" function. Right now the barotropic gyre model in Oceananigans takes in some initial inputs, and then just runs the model forward in its entirety to return the fields $u,$ $v,$ and $\eta$ at the end, if understood correctly. Ideally, we would instead take just one forward step:
   $$f(\mathbf{u}(t), \mathbf{v}(t), \mathbf{\eta}(t)) = (\mathbf{u}(t + \Delta t), \mathbf{v}(t + \Delta t), \mathbf{\eta}(t + \Delta t) )$$
   (not trying to say $f$ returns a vector, just returns the fields at the next time step.)

2. Once we have that forward function, we first want to apply Enzyme to see if it works. This shouldn't be difficult at this point (hopefully). This step will also let us see if there are issues with Enzyme + Oceananigans that would need to be worked out.

3. Then we can run a trial sensitivity analysis. Namely, per Patrick's suggestion we can choose the cost function $$J(\mathbf{u}(t_f), \mathbf{v}(t_f)) = \mathbf{u}^2(t_f) + \mathbf{v}^2(t_f) = \sum_{j, k} u_{jk}^2 + v_{jk}^2$$ (i.e. we sum $u^2 + v^2$ at every point on the grid and add them all together.) Then we can ask what happens to this quantity if we change the initial $u$ field a little bit. This will just be running the "adjoint equations" step-by-step backwards.

[swilliamson7]

Adding to this, the way I learned how to run the barotropic gyre model in Oceananigans is through the steps

mkdir blah_20220707
cd blah_20220707
wget https://julialang-s3.julialang.org/bin/linux/x64/1.7/julia-1.7.3-linux-x86_64.tar.gz (or whatever version of Julia you need, Mac/PC)

tar -xzvf julia-1.7.3-linux-x86_64.tar.gz
git clone https://github.com/CliMA/Oceananigans.jl.git
cd Oceananigans.jl/
export JULIA_DEPOT_PATH=pwd/.julia

../julia-1.7.3/bin/julia --project=@. -e 'using Pkg;Pkg.instantiate()'
../julia-1.7.3/bin/julia --project=@.

include("validation/barotropic_gyre/barotropic_gyre.jl")

All of the setup for the model exists in that script "barotropic_gyre.jl"

[swilliamson7]

@vchuravy maybe you would know how we can modify the Oceananigans code to complete my steps? There might be a better way to use the code than to explicitly create some forward step like I described in (1), I mainly just want something I can actually give to Enzyme

[NoraLoose]

@glwagner may also be able to help us here!

[glwagner]

I'm happy to help! You may not need to clone Oceananigans (unless you have to modify the source code). I can submit a PR if you like! There is a function that takes one time-step:

using Oceananigans: time_step!

# create model

time_step!(model, dt)

which I can include in the PR.

[swilliamson7]

@glwagner that would be helpful, thank you! Do I need to add you to the repo to allow you to create a PR?

[glwagner]

You don't! I can submit one from a fork.

[vchuravy]

So the first step would be to just focus on the CPU version of Oceanigans :) @wsmoses is slowly working towards custom rules so we might be able to support FFTs soon enough.

[swilliamson7]

@vchuravy What's an FFT? Also, I (briefly) tried taking Greg's function that does one step (thanks a bunch to Greg!) and differentiating it with Enzyme yesterday. I ran into a question, the velocities are stored in an offset array, is Enzyme compatible with this? Will I need to provide an "offset array" placeholder for Enzyme to put the gradient in? I might open a new issue here to discuss this, but I think that the function I have should be all I need to test Enzyme.

[glwagner]

FFT: Fast Fourier Transform. I can update the model configuration so that it doesn't use an FFT (I think one of the changes I made reverted to an FFT-based solver for the free surface). We can also make the model fully explicit if that would help.

[vchuravy]

Yeah OffsetArray should be fine, but you need to provide the same structural type as the adjoint.

[heimbach]

It seems to me that a fully explicit version might be the best/easiest starting point in terms of attempting to adjoint as you don’t have to deal with AD-related solver issues. From: "Gregory L. Wagner" ***@***.***> Reply-To: DJ4Earth/DJ-barotropic-gyre ***@***.***> Date: Friday, August 5, 2022 at 9:16 AM To: DJ4Earth/DJ-barotropic-gyre ***@***.***> Cc: Subscribed ***@***.***> Subject: Re: [DJ4Earth/DJ-barotropic-gyre] How to adjoint the barotropic gyre model from Oceananigans (Issue #1) FFT: Fast Fourier Transform. I can update the model configuration so that it doesn't use an FFT (I think one of the changes I made reverted to an FFT-based solver for the free surface). We can also make the model fully explicit if that would help. — Reply to this email directly, view it on GitHub<https://nam12.safelinks.protection.outlook.com/?url=https%3A%2F%2Fgithub.com%2FDJ4Earth%2FDJ-barotropic-gyre%2Fissues%2F1%23issuecomment-1206507680&data=05%7C01%7C%7Caac381cda951456fbabc08da76ed2463%7C31d7e2a5bdd8414e9e97bea998ebdfe1%7C0%7C0%7C637953058172349326%7CUnknown%7CTWFpbGZsb3d8eyJWIjoiMC4wLjAwMDAiLCJQIjoiV2luMzIiLCJBTiI6Ik1haWwiLCJXVCI6Mn0%3D%7C3000%7C%7C%7C&sdata=6u7pp%2FSV7Z2HYmh08DEdG79jF705ACtpLJpSrHy%2FJ7o%3D&reserved=0>, or unsubscribe<https://nam12.safelinks.protection.outlook.com/?url=https%3A%2F%2Fgithub.com%2Fnotifications%2Funsubscribe-auth%2FABHJNDOTFFD3YPP6OB32O6LVXUO2NANCNFSM55POIGLA&data=05%7C01%7C%7Caac381cda951456fbabc08da76ed2463%7C31d7e2a5bdd8414e9e97bea998ebdfe1%7C0%7C0%7C637953058172349326%7CUnknown%7CTWFpbGZsb3d8eyJWIjoiMC4wLjAwMDAiLCJQIjoiV2luMzIiLCJBTiI6Ik1haWwiLCJXVCI6Mn0%3D%7C3000%7C%7C%7C&sdata=UC5w%2BhM6Os63ShOYpiYJsmRhTeRIiq1GeYIviqqfSsQ%3D&reserved=0>. You are receiving this because you are subscribed to this thread.Message ID: ***@***.***>