Mid-Level Enzyme + Oceananigans Integration Test

[wsmoses]

The simple test code from @christophernhill, now successfully differentiating!

Not sure what the correct answer should be, so someone should check this.

This probably makes sense to add a test somewhere at least, and also can be used as a starting point for future development.

using Enzyme
using Oceananigans
using KernelAbstractions:  @index, @kernel
using Oceananigans.Utils: launch!
using Oceananigans.Architectures: device

using EnzymeCore

EnzymeCore.EnzymeRules.inactive_type(::Type{<:Oceananigans.Grids.AbstractGrid}) = true

arch=CPU()
FT=Float64

N = 100
topo = (Periodic, Flat, Flat)
grid = RectilinearGrid(arch, FT, topology=topo, size=(N), halo=2, x=(-1, 1), y=(-1, 1), z=(-1, 1))

function del21d!(d2buf,fld::Field)
 g  = fld.grid
 Nx = g.Nx
 for i=1:Nx
  d2buf[i] = fld[i-1,1,1]+fld[i+1,1,1]-2fld[i,1,1]
 end
 return
end

@kernel function del21d_k!(d2buf_k,fld)
  i,j,k = @index(Global,NTuple)
  @inbounds d2buf_k[i,j,k] = fld[i-1,j,k]+fld[i+1,j,k]-2fld[i,j,k]
end

# 2. halo
function halo1d!(fld::Field)
        g  = fld.grid
        Hx = g.Hx
        Nx = g.Nx
        for i=1:Hx
         fld[i-Hx,1,1]=fld[Nx-Hx+i,1,1]
         fld[Nx+i,1,1]=fld[i      ,1,1]
        end
        return nothing
end

# 3. simple model
function diffuse1d_model!(jcost,fld)
  grid=fld.grid
  d2buf=ones(grid.Nx)
  arch=grid.architecture
  d2buf_k = CenterField(grid)
  k=1.0
  dt=0.1
  nsteps=50
  for i in 1:nsteps
    del21d!(d2buf,fld)
    ### kernel style
    event=launch!(arch,grid,:xyz,del21d_k!,d2buf_k,fld)
    for j in 1:fld.grid.Nx
     d2buf[i] = d2buf_k[i]
    end
    for j in 1:fld.grid.Nx
      fld[j,1,1] = fld[j,1,1] + k*d2buf[j]*dt
    end
    halo1d!(fld)
  end
  jcost[1]=fld[15,1,1].*fld[15,1,1]
  return nothing
end

c     = CenterField(grid)
c2    = CenterField(grid)
f(x, y, z) = exp( -50((x-grid.xᶠᵃᵃ[1])/grid.Lx-0.5)^2 )
set!(c,f)
halo1d!(c)
set!(c2,f)
halo1d!(c2)
j=[0.]
diffuse1d_model!(j,c)

bc=CenterField(grid)
set!(c,f)
set!(bc,0)
j=[0.]
bj=[1.]
autodiff(Reverse, diffuse1d_model!,Duplicated(j,bj), Duplicated(c,bc) )

@show bj
@show bc

[glwagner]

The next thing we're trying is

grid = RectilinearGrid(size=128, z=(-0.5, 0.5), topology=(Flat, Flat, Bounded))
diffusion = VerticalScalarDiffusivity(κ=1.0)

model = HydrostaticFreeSurfaceModel(; grid,
                                    tracers = :c,
                                    buoyancy = nothing,
                                    velocities = PrescribedVelocityFields(),
                                    closure = diffusion)

"""
    set_diffusivity!(model, diffusivity)

Change diffusivity of model to `diffusivity`.
"""
function set_diffusivity!(model, diffusivity)
    closure = VerticalScalarDiffusivity(; κ=diffusivity)
    names = tuple(:c) # tracernames(model.tracers)
    closure = with_tracers(names, closure)
    model.closure = closure
    return nothing
end

function set_initial_condition!(model, amplitude)
    # Set initial condition
    width = 0.1
    cᵢ(x, y, z) = amplitude * exp(-z^2 / (2width^2))
    set!(model, c=cᵢ)
    return nothing
end

function stable_diffusion!(model, amplitude, diffusivity)
    set_diffusivity!(model, diffusivity)
    set_initial_condition!(model, amplitude)

    # Do time-stepping
    Nx, Ny, Nz = size(model.grid)
    κ_max = maximum_diffusivity
    Δz = 2π / Nz
    Δt = 1e-1 * Δz^2 / κ_max

    model.clock.time = 0
    model.clock.iteration = 0

    for n = 1:10
        time_step!(model, Δt; euler=true)
    end

    # Compute scalar metric
    c = model.tracers.c

    # Hard way
    # c² = c^2
    # sum_c² = sum(c²)

    # Another way to compute it
    c² = Array(interior(c).^2)
    sum_c² = sum(c²)

    return sum_c²::Float64
end

We made a repo for all this stuff:

https://github.com/glwagner/Enzymanigans.jl/tree/main

check out stable_diffusion.jl.