Ocean in the Cloud

]up; add Oceananigans CuArrays AbstractPlotting#master https://github.com/JuliaPlots/WGLMakie.jl.git#sd-jsserve https://github.com/SimonDanisch/JSServe.jl Observables Distributed

using Distributed
procs = addprocs(1)
@everywhere using Statistics, Oceananigans, CuArrays
const global_simulation_array = RemoteChannel()

@everywhere function start_simulation(simulation_result)
    FT = Float32
    Nx, Ny, Nz = 128, 128, 64     # Number of grid points.
    Lx, Ly, Lz = 2000, 2000, 1000  # Domain size [meters].
    end_time = 4day
    Δt = 4  # Time step in seconds.

    grid = RegularCartesianGrid(FT, (Nx, Ny, Nz), (Lx, Ly, Lz))

    # Physical constants
    φ  = 58       # Latitude to run simulation at. Corresponds to Labrador Sea.
    ρ₀ = 1027    # Density of seawater [kg/m³]
    αᵥ = 2.07e-4  # Volumetric coefficient of thermal expansion for water [K⁻¹].
    cₚ = 4181.3   # Isobaric mass heat capacity [J / kg·K].

    planetary_constants = Earth(lat=φ)
    f, g = planetary_constants.f, planetary_constants.g  # Coriolis parameter and standard gravity.

    # Parameters for generating surface cooling disk.
    Rc = 600   # Radius of cooling disk [m].
    Ts = 20    # Surface temperature [°C].
    Q₀ = -800  # Cooling disk heat flux [W/m²].
    Q₁ = 1     # Noise added to cooling disk heat flux [W/m²].
    Ns = 5 * (f * Rc/Lz)  # Stratification or Brunt–Väisälä frequency [s⁻¹].
    ∂T∂z = Ns^2 / (g * αᵥ)  # Vertical temperature gradient [K/m].

    # Center horizontal coordinates so that (x₀,y₀) = (0,0) corresponds to the center
    # of the domain (and the cooling disk).
    x₀ = grid.xC .- mean(grid.xC)
    y₀ = grid.yC .- mean(grid.yC)

    # Generate surface heat flux field.
    Q = Float32.((Q₀ .+ Q₁ .* (0.5f0 .+ randn(Nx, Ny))) / (ρ₀ * cₚ))
    Q = CuArray(Q)

    # Set surface heat flux to zero outside of cooling disk of radius Rc.
    r₀² = Float32.(@. x₀*x₀ + y₀'*y₀')
    Q[findall(r₀² .> Rc^2)] .= 0f0;

    Tbcs = HorizontallyPeriodicBCs(
        top = BoundaryCondition(Flux, Q),
        bottom = BoundaryCondition(Gradient, ∂T∂z)
    )

    model = Model(
        float_type = FT,
        arch = GPU(),
        N = (Nx, Ny, Nz),
        L = (Lx, Ly, Lz),
        ν = 4e-2, κ = 4e-2,
        constants = planetary_constants,
        bcs = BoundaryConditions(T=Tbcs)
    )
    ε(μ) = μ * randn()
    T₀(x, y, z) = Ts + ∂T∂z * z + ε(1e-4)
    set_ic!(model; T=T₀)
    # model.tracers.T.data.parent .= CuArray(Float32.(T₀));
    Hx, Hy, Hz = model.grid.Hx, model.grid.Hy, model.grid.Hz
    temp(model) = Array(model.tracers.T.data.parent[1+Hx:Nx+Hx, 1+Hy:Ny+Hy, 1+Hz:Nz+Hz])
    while true
        time_step!(model; Nt = 2, Δt = Δt)
        put!(simulation_result, temp(model))
    end
end
f = @spawnat 2 start_simulation(global_simulation_array)

using WGLMakie, AbstractPlotting, JSServe, Observables
using JSServe.DOM
# create a signal that always contains our newest simulation state
const ocean_signal = Observable(take!(global_simulation_array))
taking_task = @async while true
    ocean_signal[] = take!(global_simulation_array)
    yield()
end

function dom_handler(session, request)
  scene3d = Scene(show_axis = false)
  linesegments!(scene3d, FRect3D(Vec3f0(0), Vec3f0(1)))
  brain_data = ocean_signal[]
  volume = Observables.async_latest(ocean_signal)
  planes = (:yz, :xz, :xy)
  three = nothing
  rs = LinRange.(0, size(brain_data) ./ 128, size(brain_data))
  sliders = ntuple(3) do i
    slid = JSServe.Slider(1:2:size(volume[], i), value = size(volume[], i) ÷ 2)
    idx = Observables.async_latest(Observables.observe(slid))
    plane = planes[i]; r = rs[i]
    indices = ntuple(3) do j
      planes[j] == plane ? 1 : (:)
    end
    ridx = Iterators.filter((1,2,3)) do j
      planes[j] != plane
    end
    heatm = heatmap!(
      scene3d, getindex.((rs,), ridx)..., volume[][indices...],
      colorrange = (19.90f0, 20.1f0),
      interpolate = true
      )[end]

    function transform_planes(idx, vol)
      transform!(heatm, (plane, r[idx]))
      indices = ntuple(3) do j
        planes[j] == plane ? idx : (:)
      end
      if checkbounds(Bool, vol, indices...)
        heatm[3][] = view(vol, indices...)
        three !== nothing && WGLMakie.redraw!(three)
      end
    end
    onany(transform_planes, idx, volume)
    transform_planes(idx[], volume[])
    slid
  end
  three, canvas = WGLMakie.three_display(session, scene3d)
  DOM.um(
    "h1", "Ocean Simulation:",
    DOM.div(DOM.div(sliders...), canvas)
  )
end

app = JSServe.Application(
    dom_handler,
    get(ENV, "WEBIO_SERVER_HOST_URL", "127.0.0.1"),
    parse(Int, get(ENV, "WEBIO_HTTP_PORT", "8081")),
    verbose = false
);

HTML("<a href=$(repr(JSServe.server_proxy_url[]))>Visit Simulation</a>")

sleep(100)