Clean up zeros

Project.toml

@@ -33,17 +33,18 @@ SeawaterPolynomials = "d496a93d-167e-4197-9f49-d3af4ff8fe40"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 StructArrays = "09ab397b-f2b6-538f-b94a-2f83cf4a842a"
+TimesDates = "bdfc003b-8df8-5c39-adcd-3a9087f5df4a"
 
 [weakdeps]
 Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
-Reactant = "3c362404-f566-11ee-1572-e11a4b42c853"
 Makie = "ee78f7c6-11fb-53f2-987a-cfe4a2b5a57a"
 MakieCore = "20f20a25-4f0e-4fdf-b5d1-57303727442b"
+Reactant = "3c362404-f566-11ee-1572-e11a4b42c853"
 
 [extensions]
 OceananigansEnzymeExt = "Enzyme"
-OceananigansReactantExt = "Reactant"
 OceananigansMakieExt = ["MakieCore", "Makie"]
+OceananigansReactantExt = "Reactant"
 
 [compat]
 Adapt = "4.1.1"
@@ -55,8 +56,8 @@ Distances = "0.10"
 DocStringExtensions = "0.8, 0.9"
 Enzyme = "0.13.14"
 FFTW = "1"
-Glob = "1.3"
 GPUArrays = "10.3"
+Glob = "1.3"
 IncompleteLU = "0.2"
 InteractiveUtils = "1.9"
 IterativeSolvers = "0.9"
@@ -83,8 +84,8 @@ julia = "1.9"
 CUDA_Runtime_jll = "76a88914-d11a-5bdc-97e0-2f5a05c973a2"
 DataDeps = "124859b0-ceae-595e-8997-d05f6a7a8dfe"
 Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
-Reactant = "3c362404-f566-11ee-1572-e11a4b42c853"
 MPIPreferences = "3da0fdf6-3ccc-4f1b-acd9-58baa6c99267"
+Reactant = "3c362404-f566-11ee-1572-e11a4b42c853"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 TimesDates = "bdfc003b-8df8-5c39-adcd-3a9087f5df4a"

src/DistributedComputations/distributed_architectures.jl

@@ -306,7 +306,7 @@ ranks(arch::Distributed) = ranks(arch.partition)
 child_architecture(arch::Distributed) = arch.child_architecture
 device(arch::Distributed)             = device(child_architecture(arch))
 
-zeros(FT, arch::Distributed, N...)         = zeros(FT, child_architecture(arch), N...)
+zeros(arch::Distributed, N...)             = zeros(child_architecture(arch), FT, N...)
 array_type(arch::Distributed)              = array_type(child_architecture(arch))
 sync_device!(arch::Distributed)            = sync_device!(arch.child_architecture)
 convert_to_device(arch::Distributed, arg)  = convert_to_device(child_architecture(arch), arg)

src/Grids/latitude_longitude_grid.jl

@@ -534,14 +534,14 @@ function allocate_metrics(grid::LatitudeLongitudeGrid)
         metric_size = (length(grid.Δλᶜᵃᵃ), length(grid.φᵃᶜᵃ))
     end
 
-    Δxᶠᶜ = OffsetArray(zeros(FT, arch, metric_size...), offsets...)
-    Δxᶜᶠ = OffsetArray(zeros(FT, arch, metric_size...), offsets...)
-    Δxᶠᶠ = OffsetArray(zeros(FT, arch, metric_size...), offsets...)
-    Δxᶜᶜ = OffsetArray(zeros(FT, arch, metric_size...), offsets...)
-    Azᶠᶜ = OffsetArray(zeros(FT, arch, metric_size...), offsets...)
-    Azᶜᶠ = OffsetArray(zeros(FT, arch, metric_size...), offsets...)
-    Azᶠᶠ = OffsetArray(zeros(FT, arch, metric_size...), offsets...)
-    Azᶜᶜ = OffsetArray(zeros(FT, arch, metric_size...), offsets...)
+    Δxᶠᶜ = OffsetArray(zeros(arch, FT, metric_size...), offsets...)
+    Δxᶜᶠ = OffsetArray(zeros(arch, FT, metric_size...), offsets...)
+    Δxᶠᶠ = OffsetArray(zeros(arch, FT, metric_size...), offsets...)
+    Δxᶜᶜ = OffsetArray(zeros(arch, FT, metric_size...), offsets...)
+    Azᶠᶜ = OffsetArray(zeros(arch, FT, metric_size...), offsets...)
+    Azᶜᶠ = OffsetArray(zeros(arch, FT, metric_size...), offsets...)
+    Azᶠᶠ = OffsetArray(zeros(arch, FT, metric_size...), offsets...)
+    Azᶜᶜ = OffsetArray(zeros(arch, FT, metric_size...), offsets...)
 
     if grid isa YRegularLLG
         Δyᶠᶜ = Δyᶠᶜᵃ(1, 1, 1, grid)

src/Grids/new_data.jl

@@ -64,7 +64,7 @@ with indices corresponding to a field on a `grid` of `size(grid)` and located at
 """
 function new_data(FT::DataType, arch, loc, topo, sz, halo_sz, indices=default_indices(length(loc)))
     Tx, Ty, Tz = total_size(loc, topo, sz, halo_sz, indices)
-    underlying_data = zeros(FT, arch, Tx, Ty, Tz)
+    underlying_data = zeros(arch, FT, Tx, Ty, Tz)
     indices = validate_indices(indices, loc, topo, sz, halo_sz)
     return offset_data(underlying_data, loc, topo, sz, halo_sz, indices)
 end

src/Grids/orthogonal_spherical_shell_grid.jl

@@ -860,10 +860,10 @@ function conformal_cubed_sphere_panel(filepath::AbstractString, architecture = C
     Tyᶠᶜ = total_length(loc_fc[2](), topology[2](), N[2], H[2])
     Tyᶜᶠ = total_length(loc_cf[2](), topology[2](), N[2], H[2])
 
-    λᶠᶜᵃ = offset_data(zeros(FT, architecture, Txᶠᶜ, Tyᶠᶜ), loc_fc, topology[1:2], N[1:2], H[1:2])
-    λᶜᶠᵃ = offset_data(zeros(FT, architecture, Txᶜᶠ, Tyᶜᶠ), loc_cf, topology[1:2], N[1:2], H[1:2])
-    φᶠᶜᵃ = offset_data(zeros(FT, architecture, Txᶠᶜ, Tyᶠᶜ), loc_fc, topology[1:2], N[1:2], H[1:2])
-    φᶜᶠᵃ = offset_data(zeros(FT, architecture, Txᶜᶠ, Tyᶜᶠ), loc_cf, topology[1:2], N[1:2], H[1:2])
+    λᶠᶜᵃ = offset_data(zeros(architecture, FT, Txᶠᶜ, Tyᶠᶜ), loc_fc, topology[1:2], N[1:2], H[1:2])
+    λᶜᶠᵃ = offset_data(zeros(architecture, FT, Txᶜᶠ, Tyᶜᶠ), loc_cf, topology[1:2], N[1:2], H[1:2])
+    φᶠᶜᵃ = offset_data(zeros(architecture, FT, Txᶠᶜ, Tyᶠᶜ), loc_fc, topology[1:2], N[1:2], H[1:2])
+    φᶜᶠᵃ = offset_data(zeros(architecture, FT, Txᶜᶠ, Tyᶜᶠ), loc_cf, topology[1:2], N[1:2], H[1:2])
 
     ## The vertical coordinates can come out of the regular rectilinear grid!
     Lz, z  = generate_coordinate(FT, topology, (Nξ, Nη, Nz), halo, z,  :z, 3, architecture)

src/Grids/zeros_and_ones.jl

@@ -1,10 +1,11 @@
 using Oceananigans.Architectures: device, AbstractArchitecture
 
-import KernelAbstractions: zeros
-
-zeros(FT, arch::AbstractArchitecture, N...) = zeros(device(arch), FT, N...)
+import KernelAbstractions
+import Base: zeros
 
+zeros(FT, arch::AbstractArchitecture, N...) = KernelAbstractions.zeros(device(arch), FT, N...)
 zeros(grid::AbstractGrid, N...) = zeros(eltype(grid), architecture(grid), N...)
 
 @inline Base.zero(grid::AbstractGrid) = zero(eltype(grid))
 @inline Base.one(grid::AbstractGrid)  = one(eltype(grid))
+

src/OutputReaders/field_time_series.jl

@@ -356,7 +356,7 @@ new_data(FT, grid, loc, indices, ::Nothing) = nothing
 # https://github.com/CliMA/ClimaOcean.jl/actions/runs/8804916198/job/24166354095)
 function new_data(FT, grid, loc, indices, Nt::Union{Int, Int64})
     space_size = total_size(grid, loc, indices)
-    underlying_data = zeros(FT, architecture(grid), space_size..., Nt)
+    underlying_data = zeros(architecture(grid), FT, space_size..., Nt)
     data = offset_data(underlying_data, grid, loc, indices)
     return data
 end