Remove Tullio dependency; take #2 (#3037)

* remove tullio dependency

* bugfix

* add abs to abstract operations

* solved the issue

* probably not super efficient

* indices for gridmetricoperations

* just remove tullio

* change timescale definition

* do not introduce indices now

* Update src/Advection/cell_advection_timescale.jl

Co-authored-by: Gregory L. Wagner <wagner.greg@gmail.com>

* Update src/Diagnostics/cfl.jl

* accurate_cell_advection_timescale -> cell_advection_timescale

---------

Co-authored-by: Gregory L. Wagner <wagner.greg@gmail.com>
Co-authored-by: Navid C. Constantinou <navidcy@users.noreply.github.com>

Manifest.toml

@@ -2,7 +2,7 @@
 
 julia_version = "1.8.5"
 manifest_format = "2.0"
-project_hash = "96eebad427fe7f986c00777f7738aa636a9e621d"
+project_hash = "31dd08f0370dcfe9232b9df8c2131e894bcff3ca"
 
 [[deps.AMGX]]
 deps = ["AMGX_jll", "CEnum", "CUDA", "JSON", "Libdl", "SparseArrays"]
@@ -156,12 +156,6 @@ version = "1.0.0"
 deps = ["Printf"]
 uuid = "ade2ca70-3891-5945-98fb-dc099432e06a"
 
-[[deps.DiffRules]]
-deps = ["IrrationalConstants", "LogExpFunctions", "NaNMath", "Random", "SpecialFunctions"]
-git-tree-sha1 = "c5b6685d53f933c11404a3ae9822afe30d522494"
-uuid = "b552c78f-8df3-52c6-915a-8e097449b14b"
-version = "1.12.2"
-
 [[deps.Distributed]]
 deps = ["Random", "Serialization", "Sockets"]
 uuid = "8ba89e20-285c-5b6f-9357-94700520ee1b"
@@ -431,12 +425,6 @@ git-tree-sha1 = "d3e32d2b0530d929a047ceab8799ae4204de6c88"
 uuid = "85f8d34a-cbdd-5861-8df4-14fed0d494ab"
 version = "0.12.12"
 
-[[deps.NaNMath]]
-deps = ["OpenLibm_jll"]
-git-tree-sha1 = "a7c3d1da1189a1c2fe843a3bfa04d18d20eb3211"
-uuid = "77ba4419-2d1f-58cd-9bb1-8ffee604a2e3"
-version = "1.0.1"
-
 [[deps.NetCDF_jll]]
 deps = ["Artifacts", "HDF5_jll", "JLLWrappers", "LibCURL_jll", "Libdl", "Pkg", "XML2_jll", "Zlib_jll"]
 git-tree-sha1 = "072f8371f74c3b9e1b26679de7fbf059d45ea221"
@@ -703,12 +691,6 @@ git-tree-sha1 = "94f38103c984f89cf77c402f2a68dbd870f8165f"
 uuid = "3bb67fe8-82b1-5028-8e26-92a6c54297fa"
 version = "0.9.11"
 
-[[deps.Tullio]]
-deps = ["ChainRulesCore", "DiffRules", "LinearAlgebra", "Requires"]
-git-tree-sha1 = "7871a39eac745697ee512a87eeff06a048a7905b"
-uuid = "bc48ee85-29a4-5162-ae0b-a64e1601d4bc"
-version = "0.3.5"
-
 [[deps.TupleTools]]
 git-tree-sha1 = "3c712976c47707ff893cf6ba4354aa14db1d8938"
 uuid = "9d95972d-f1c8-5527-a6e0-b4b365fa01f6"

Project.toml

@@ -36,7 +36,6 @@ SeawaterPolynomials = "d496a93d-167e-4197-9f49-d3af4ff8fe40"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 StructArrays = "09ab397b-f2b6-538f-b94a-2f83cf4a842a"
-Tullio = "bc48ee85-29a4-5162-ae0b-a64e1601d4bc"
 
 [compat]
 AMGX = "0.1.3"
@@ -62,7 +61,6 @@ PencilFFTs = "0.13.5, 0.14"
 Rotations = "1.0"
 SeawaterPolynomials = "0.3"
 StructArrays = "0.4, 0.5, 0.6"
-Tullio = "0.2, 0.3"
 julia = "1.6"
 
 [extras]

src/AbstractOperations/AbstractOperations.jl

@@ -76,7 +76,8 @@ include("show_abstract_operations.jl")
 # Make some operators!
 
 # Some operators:
-import Base: sqrt, sin, cos, exp, tanh, -, +, /, ^, *
+import Base: sqrt, sin, cos, exp, tanh, abs, -, +, /, ^, *
+import Base: abs
 
 @unary sqrt sin cos exp tanh abs
 @unary -

src/Advection/cell_advection_timescale.jl

@@ -1,5 +1,15 @@
 using Oceananigans.AbstractOperations: KernelFunctionOperation
 
+"""
+    cell_advection_timescale(grid, velocities)
+
+Return the advection timescale for `grid` with `velocities`. The advection timescale
+is the minimum over all `i, j, k` in the `grid` of
+
+```
+  1 / (|u(i, j, k)| / Δxᶠᶜᶜ(i, j, k) + |v(i, j, k)| / Δyᶜᶠᶜ(i, j, k) + |w(i, j, k)| / Δzᶜᶜᶠ(i, j, k))
+```
+"""
 function cell_advection_timescale(grid, velocities)
     u, v, w = velocities
     τ = KernelFunctionOperation{Center, Center, Center}(cell_advection_timescaleᶜᶜᶜ, grid, u, v, w)
@@ -11,7 +21,6 @@ end
     Δy = Δyᶜᶠᶜ(i, j, k, grid)
     Δz = Δzᶜᶜᶠ(i, j, k, grid)
 
-    return @inbounds min(Δx / abs(u[i, j, k]),
-                         Δy / abs(v[i, j, k]),
-                         Δz / abs(w[i, j, k])) 
+    inverse_timescale = @inbounds abs(u[i, j, k]) / Δx + abs(v[i, j, k]) / Δy + abs(w[i, j, k]) / Δz
+    return 1 / inverse_timescale
 end

src/CubedSpheres/CubedSpheres.jl

@@ -156,15 +156,15 @@ end
 ##### CFL for cubed sphere fields
 #####
 
-import Oceananigans.Diagnostics: accurate_cell_advection_timescale
+import Oceananigans.Advection: cell_advection_timescale
 
-function accurate_cell_advection_timescale(grid::ConformalCubedSphereGrid, velocities)
+function cell_advection_timescale(grid::ConformalCubedSphereGrid, velocities)
 
     min_timescale_on_faces = []
 
     for (face_index, grid_face) in enumerate(grid.faces)
         velocities_face = get_face(velocities, face_index)
-        min_timescale_on_face = accurate_cell_advection_timescale(grid_face, velocities_face)
+        min_timescale_on_face = cell_advection_timescale(grid_face, velocities_face)
         push!(min_timescale_on_faces, min_timescale_on_face)
     end
 
@@ -175,10 +175,11 @@ end
 ##### compute...
 #####
 
-import Oceananigans.Fields: compute!
 using Oceananigans.AbstractOperations: _compute!
 using Oceananigans.Fields: compute_at!
 
+import Oceananigans.Fields: compute!
+
 const CubedSphereComputedField{LX, LY, LZ} = Field{LX, LY, LZ,
                                                    <:AbstractOperation,
                                                    <:ConformalCubedSphereGrid} where {LX, LY, LZ}

src/Diagnostics/cfl.jl

@@ -12,7 +12,7 @@ struct CFL{D, S}
 end
 
 """
-    CFL(Δt [, timescale = Oceananigans.cell_advection_timescale])
+    CFL(Δt [, timescale = Oceananigans.Advection.cell_advection_timescale])
 
 Return an object for computing the Courant-Freidrichs-Lewy (CFL) number
 associated with time step `Δt` or `TimeStepWizard` and `timescale`.
@@ -78,36 +78,3 @@ julia> dcfl(model)
 ```
 """
 DiffusiveCFL(Δt) = CFL(Δt, cell_diffusion_timescale)
-
-#####
-##### Accurate CFL via reduction
-#####
-
-using CUDA, CUDAKernels, KernelAbstractions, Tullio
-
-using Oceananigans.Models
-using Oceananigans.Grids: halo_size
-using Oceananigans.Operators: Δxᶠᶜᶜ, Δyᶜᶠᶜ, Δzᶜᶜᶠ
-
-accurate_cell_advection_timescale(model) = accurate_cell_advection_timescale(model.grid, model.velocities)
-
-function accurate_cell_advection_timescale(grid, velocities)
-    Nx, Ny, Nz = size(grid)
-    Hx, Hy, Hz = halo_size(grid)
-
-    is = 1+Hx:Nx+Hx
-    js = 1+Hy:Ny+Hy
-    ks = 1+Hz:Nz+Hz
-
-    u = view(velocities.u.data.parent, is, js, ks)
-    v = view(velocities.v.data.parent, is, js, ks)
-    w = view(velocities.w.data.parent, is, js, ks)
-
-    min_timescale = minimum(
-        @tullio (min) timescale[k] := 1 / (  abs(u[i, j, k]) / Δxᶠᶜᶜ(i, j, k, grid)
-                                           + abs(v[i, j, k]) / Δyᶜᶠᶜ(i, j, k, grid)
-                                           + abs(w[i, j, k]) / Δzᶜᶜᶠ(i, j, k, grid))
-    )
-
-    return min_timescale
-end

test/test_diagnostics.jl

@@ -1,3 +1,5 @@
+include("dependencies_for_runtests.jl")
+
 using Oceananigans.Diagnostics
 using Oceananigans.Diagnostics: AbstractDiagnostic
 
@@ -34,7 +36,7 @@ function advective_cfl_diagnostic_is_correct_on_regular_grid(arch, FT)
     model = TestModel_RegularRectGrid(arch, FT)
 
     Δt = FT(1.3e-6)
-    Δx = FT(model.grid.Δxᶜᵃᵃ)
+    Δx = minimum_xspacing(model.grid)
     u₀ = FT(1.2)
     CFL_by_hand = Δt * u₀ / Δx
 
@@ -58,7 +60,7 @@ function advective_cfl_diagnostic_is_correct_on_vertically_stretched_grid(arch,
     return cfl(model) ≈ CFL_by_hand
 end
 
-function accurate_advective_cfl_on_regular_grid(arch, FT)
+function advective_timescale_cfl_on_regular_grid(arch, FT)
     model = TestModel_RegularRectGrid(arch, FT)
 
     Δt = FT(1.7)
@@ -75,12 +77,12 @@ function accurate_advective_cfl_on_regular_grid(arch, FT)
 
     CFL_by_hand = Δt * (abs(u₀) / Δx + abs(v₀) / Δy + abs(w₀) / Δz)
 
-    cfl = CFL(FT(Δt), Oceananigans.Diagnostics.accurate_cell_advection_timescale)
+    cfl = CFL(FT(Δt), Oceananigans.Advection.cell_advection_timescale)
 
     return cfl(model) ≈ CFL_by_hand
 end
 
-function accurate_advective_cfl_on_stretched_grid(arch, FT)
+function advective_timescale_cfl_on_stretched_grid(arch, FT)
     grid = RectilinearGrid(arch, size=(4, 4, 8), x=(0, 100), y=(0, 100), z=[k^2 for k in 0:8])
     model = NonhydrostaticModel(grid=grid)
 
@@ -100,12 +102,12 @@ function accurate_advective_cfl_on_stretched_grid(arch, FT)
 
     CFL_by_hand = Δt * (abs(u₀) / Δx + abs(v₀) / Δy + abs(w₀) / Δz_min)
 
-    cfl = CFL(FT(Δt), Oceananigans.Diagnostics.accurate_cell_advection_timescale)
+    cfl = CFL(FT(Δt), Oceananigans.Advection.cell_advection_timescale)
 
     return cfl(model) ≈ CFL_by_hand
 end
 
-function accurate_advective_cfl_on_lat_lon_grid(arch, FT)
+function advective_timescale_cfl_on_lat_lon_grid(arch, FT)
     grid = LatitudeLongitudeGrid(arch, size=(8, 8, 8), longitude=(-10, 10), latitude=(0, 45), z=(-1000, 0))
     model = HydrostaticFreeSurfaceModel(grid=grid, momentum_advection=VectorInvariant())
 
@@ -132,7 +134,7 @@ function accurate_advective_cfl_on_lat_lon_grid(arch, FT)
 
     CFL_by_hand = Δt * (abs(u₀) / Δx_min + abs(v₀) / Δy_min + abs(w₀) / Δz)
 
-    cfl = CFL(FT(Δt), Oceananigans.Diagnostics.accurate_cell_advection_timescale)
+    cfl = CFL(FT(Δt), Oceananigans.Advection.cell_advection_timescale)
 
     return cfl(model) ≈ CFL_by_hand
 end
@@ -172,9 +174,9 @@ end
                 @test diffusive_cfl_diagnostic_is_correct(arch, FT)
                 @test advective_cfl_diagnostic_is_correct_on_regular_grid(arch, FT)
                 @test advective_cfl_diagnostic_is_correct_on_vertically_stretched_grid(arch, FT)
-                @test accurate_advective_cfl_on_regular_grid(arch, FT)
-                @test accurate_advective_cfl_on_stretched_grid(arch, FT)
-                @test accurate_advective_cfl_on_lat_lon_grid(arch, FT)
+                @test advective_timescale_cfl_on_regular_grid(arch, FT)
+                @test advective_timescale_cfl_on_stretched_grid(arch, FT)
+                @test advective_timescale_cfl_on_lat_lon_grid(arch, FT)
             end
         end
     end