`ECCOPrescribedAtmosphere` causes error when time indices are interpolated in `OceanSeaIceModel`

[xkykai]

Here I'll outline a (perhaps not so) MWE which demonstrates typical use of OceanSeaIceModel using ECCOPrescribedAtmosphere

using ClimaOcean
using ClimaOcean.DataWrangling.ECCO
using Oceananigans
using Oceananigans.Units
using Oceananigans.Operators: Δx, Δy
using Dates
using CUDA

catke_closure = ClimaOcean.Oceans.default_ocean_closure()

arch = GPU()

Nx = 720 # longitudinal direction 
Ny = 360 # meridional direction 
Nz = 100

z_faces = ExponentialDiscretization(Nz, -6000, 0; scale=1800)
z_surf = z_faces(Nz)

grid = TripolarGrid(arch;
                    size = (Nx, Ny, Nz),
                    z = z_faces,
                    halo = (7, 7, 7))

bottom_height = regrid_bathymetry(grid; minimum_depth=15, major_basins=1, interpolation_passes=55)
grid = ImmersedBoundaryGrid(grid, GridFittedBottom(bottom_height); active_cells_map=true)

tracer_advection   = WENO(order=7)
momentum_advection = WENOVectorInvariant(order=5)
free_surface       = SplitExplicitFreeSurface(grid; cfl=0.8, fixed_Δt=40minutes)

@inline Δ²ᵃᵃᵃ(i, j, k, grid, lx, ly, lz) =  2 * (1 / (1 / Δx(i, j, k, grid, lx, ly, lz)^2 + 1 / Δy(i, j, k, grid, lx, ly, lz)^2))
@inline geometric_νhb(i, j, k, grid, lx, ly, lz, clock, fields, λ) = Δ²ᵃᵃᵃ(i, j, k, grid, lx, ly, lz)^2 / λ

horizontal_viscosity = HorizontalScalarBiharmonicDiffusivity(ν=geometric_νhb, discrete_form=true, parameters=25days)

closure = (catke_closure, horizontal_viscosity)

start_date = DateTime(1992, 1, 1)
end_date = start_date + Year(2)
simulation_period = Dates.value(Second(end_date - start_date))

ECCO_dir = joinpath(homedir(), "ECCO_data")
mkpath(ECCO_dir)

@inline mask(x, y, z, t) = z >= z_surf - 1
Smetadata = Metadata(:salinity; dataset=ECCO4Monthly(), dir=ECCO_dir, start_date, end_date)
FS = DatasetRestoring(Smetadata, grid; rate = 1/30days, mask, time_indices_in_memory = 10)

ocean = ocean_simulation(grid; Δt=1minutes,
                        momentum_advection,
                        tracer_advection,
                        timestepper = :SplitRungeKutta3,
                        free_surface,
                        forcing = (; S = FS),
                        closure)

set!(ocean.model, T=Metadatum(:temperature; dataset=ECCO4Monthly(), date=start_date, dir=ECCO_dir),
                    S=Metadatum(:salinity;    dataset=ECCO4Monthly(), date=start_date, dir=ECCO_dir))

sea_ice = sea_ice_simulation(grid, ocean; dynamics=nothing)

set!(sea_ice.model, h=Metadatum(:sea_ice_thickness;     dataset=ECCO4Monthly(), date=start_date, dir=ECCO_dir),
                    ℵ=Metadatum(:sea_ice_concentration; dataset=ECCO4Monthly(), date=start_date, dir=ECCO_dir))

@info "Setting up prescribed atmosphere (ECCO4 forcing)"
atmosphere = ECCOPrescribedAtmosphere(arch; dir=ECCO_dir, start_date, end_date, time_indices_in_memory=24)

radiation  = Radiation()

@info "Built atmosphere model $(atmosphere)"

omip = OceanSeaIceModel(ocean, sea_ice; atmosphere, radiation)

At the last stage of building the OceanSeaIceModel with atmosphere=ECCOPrescribedAtmosphere, we hit a scalar indexing disallowed error:

julia> omip = OceanSeaIceModel(ocean, sea_ice; atmosphere, radiation)
ERROR: Scalar indexing is disallowed.
Invocation of getindex resulted in scalar indexing of a GPU array.
This is typically caused by calling an iterating implementation of a method.
Such implementations *do not* execute on the GPU, but very slowly on the CPU,
and therefore should be avoided.

If you want to allow scalar iteration, use `allowscalar` or `@allowscalar`
to enable scalar iteration globally or for the operations in question.
Stacktrace:
  [1] error(s::String)
    @ Base ./error.jl:44
  [2] errorscalar(op::String)
    @ GPUArraysCore ~/.julia/packages/GPUArraysCore/aNaXo/src/GPUArraysCore.jl:151
  [3] _assertscalar(op::String, behavior::GPUArraysCore.ScalarIndexing)
    @ GPUArraysCore ~/.julia/packages/GPUArraysCore/aNaXo/src/GPUArraysCore.jl:124
  [4] assertscalar(op::String)
    @ GPUArraysCore ~/.julia/packages/GPUArraysCore/aNaXo/src/GPUArraysCore.jl:112
  [5] getindex
    @ ~/.julia/packages/GPUArrays/nc02F/src/host/indexing.jl:50 [inlined]
  [6] first
    @ ./abstractarray.jl:452 [inlined]
  [7] interpolating_time_indices
    @ ~/.julia/packages/Oceananigans/t5Rdl/src/OutputReaders/field_time_series_indexing.jl:55 [inlined]
  [8] TimeInterpolator
    @ ~/.julia/packages/Oceananigans/t5Rdl/src/OutputReaders/field_time_series_indexing.jl:35 [inlined]
  [9] interpolate_state!(exchanger::ClimaOcean.OceanSeaIceModels.InterfaceComputations.ComponentExchanger{…}, grid::ImmersedBoundaryGrid{…}, atmosphere::PrescribedAtmosphere{…}, coupled_model::OceanSeaIceModel{…})
    @ ClimaOcean.Atmospheres ~/.julia/packages/ClimaOcean/tL4Wo/src/Atmospheres/interpolate_atmospheric_state.jl:63
 [10] update_state!(coupled_model::OceanSeaIceModel{Simulation{…}, PrescribedAtmosphere{…}, Simulation{…}, ComponentInterfaces{…}, Clock{…}, GPU{…}})
    @ ClimaOcean.OceanSeaIceModels ~/.julia/packages/ClimaOcean/tL4Wo/src/OceanSeaIceModels/time_step_ocean_sea_ice_model.jl:44
 [11] initialization_update_state!(model::OceanSeaIceModel{Simulation{…}, PrescribedAtmosphere{…}, Simulation{…}, ComponentInterfaces{…}, Clock{…}, GPU{…}})
    @ ClimaOcean.OceanSeaIceModels ~/.julia/packages/ClimaOcean/tL4Wo/src/OceanSeaIceModels/ocean_sea_ice_model.jl:66
 [12] OceanSeaIceModel(ocean::Simulation{…}, sea_ice::Simulation{…}; atmosphere::PrescribedAtmosphere{…}, radiation::Radiation{…}, clock::Clock{…}, ocean_reference_density::Float64, ocean_heat_capacity::Float64, sea_ice_reference_density::Float64, sea_ice_heat_capacity::Float64, interfaces::Nothing)
    @ ClimaOcean.OceanSeaIceModels ~/.julia/packages/ClimaOcean/tL4Wo/src/OceanSeaIceModels/ocean_sea_ice_model.jl:211
 [13] top-level scope
    @ REPL[69]:1
Some type information was truncated. Use `show(err)` to see complete types

This seems to be coming from time interpolation