Update spacings calls to enable array operations. (#4691)

* [Reactant] Add recursive specialization for BinaryOperation

* fmt

* fix

* fix

* Removed erroneous end, updated compat to take Reactant 2.149

* Cosmetic upgrades

* Update Project.toml

Co-authored-by: Navid C. Constantinou <[email protected]>

* Stacking spacings to fix some dimension mismatch errors in Reactant

* Setting up array muls

* Not a problem

---------

Co-authored-by: William S. Moses <[email protected]>
Co-authored-by: Billy Moses <[email protected]>
Co-authored-by: Gregory Wagner <[email protected]>
Co-authored-by: Navid C. Constantinou <[email protected]>

Author: 4 people

src/Operators/spacings_and_areas_and_volumes.jl

@@ -109,6 +109,8 @@ end
 @inline Δrᵃᵃᶠ(i, j, k, grid) = getspacing(k, grid.z.Δᵃᵃᶠ)
 
 @inline Δzᵃᵃᶜ(i, j, k, grid) = getspacing(k, grid.z.Δᵃᵃᶜ)
+@inline Δzᵃᵃᶜ(i::AbstractArray, j::AbstractArray, k::AbstractArray, grid) = permutedims(Base.stack(collect(Base.stack(collect(getspacing(k, grid.z.Δᵃᵃᶜ) for _ in j)) for _ in i)), (3,2,1)) # one part
+
 @inline Δzᵃᵃᶠ(i, j, k, grid) = getspacing(k, grid.z.Δᵃᵃᶠ)
 
 #####
@@ -156,7 +158,9 @@ end
 @inline Δλᶠᵃᵃ(i, j, k, grid::LLGX) = @inbounds grid.Δλᶠᵃᵃ
 
 @inline Δφᵃᶜᵃ(i, j, k, grid::LLG)  = @inbounds grid.Δφᵃᶜᵃ[j]
+@inline Δφᵃᶜᵃ(i::AbstractArray, j::AbstractArray, k::AbstractArray, grid::LLG) = Base.stack(collect(collect(transpose(Base.stack(collect(Δφᵃᶜᵃ(1, j, 1, grid) for _ in i)))) for _ in k))
 @inline Δφᵃᶠᵃ(i, j, k, grid::LLG)  = @inbounds grid.Δφᵃᶠᵃ[j]
+@inline Δφᵃᶠᵃ(i::AbstractArray, j::AbstractArray, k::AbstractArray, grid::LLG) = Base.stack(collect(collect(transpose(Base.stack(collect(Δφᵃᶠᵃ(1, j, 1, grid) for _ in i)))) for _ in k))
 @inline Δφᵃᶜᵃ(i, j, k, grid::LLGY) = @inbounds grid.Δφᵃᶜᵃ
 @inline Δφᵃᶠᵃ(i, j, k, grid::LLGY) = @inbounds grid.Δφᵃᶠᵃ
 
@@ -165,9 +169,15 @@ end
 ### Precomputed metrics
 
 @inline Δyᵃᶜᵃ(i, j, k, grid::LLGY) = grid.Δyᶠᶜᵃ
+@inline Δyᵃᶜᵃ(i::AbstractArray, j::AbstractArray, k::AbstractArray, grid::LLGY) = Base.stack(collect(Base.stack(collect(Base.stack(collect(Δyᵃᶜᵃ(1, 1, 1, grid) for _ in i)) for _ in j)) for _ in k)) # other part
+
 @inline Δyᵃᶠᵃ(i, j, k, grid::LLGY) = grid.Δyᶜᶠᵃ
+@inline Δyᵃᶠᵃ(i::AbstractArray, j::AbstractArray, k::AbstractArray, grid::LLGY) = Base.stack(collect(Base.stack(collect(Base.stack(collect(Δyᵃᶠᵃ(1, 1, 1, grid) for _ in i)) for _ in j)) for _ in k))
+
 @inline Δyᵃᶜᵃ(i, j, k, grid::LLG)  = @inbounds grid.Δyᶠᶜᵃ[j]
+#@inline Δyᵃᶜᵃ(i::AbstractArray, j::AbstractArray, k::AbstractArray, grid::LLG) = Base.stack(collect(collect(transpose(Base.stack(collect(Δyᵃᶜᵃ(1, j, 1, grid) for _ in i)))) for _ in k))
 @inline Δyᵃᶠᵃ(i, j, k, grid::LLG)  = @inbounds grid.Δyᶜᶠᵃ[j]
+#@inline Δyᵃᶠᵃ(i::AbstractArray, j::AbstractArray, k::AbstractArray, grid::LLG) = Base.stack(collect(collect(transpose(Base.stack(collect(Δyᵃᶠᵃ(1, j, 1, grid) for _ in i)))) for _ in k))
 
 ### On-the-fly metrics
 
@@ -189,9 +199,17 @@ end
 ### Pre computed metrics
 
 @inline Δxᶜᶠᵃ(i, j, k, grid::LLG) = @inbounds grid.Δxᶜᶠᵃ[i, j]
+@inline Δxᶜᶠᵃ(i::AbstractArray, j::AbstractArray, k::AbstractArray, grid::LLG) = Base.stack(collect(Δxᶜᶠᵃ(i, j, 1, grid) for _ in k))
+
 @inline Δxᶠᶜᵃ(i, j, k, grid::LLG) = @inbounds grid.Δxᶠᶜᵃ[i, j]
+@inline Δxᶠᶜᵃ(i::AbstractArray, j::AbstractArray, k::AbstractArray, grid::LLG) = Base.stack(collect(Δxᶠᶜᵃ(i, j, 1, grid) for _ in k))
+
 @inline Δxᶠᶠᵃ(i, j, k, grid::LLG) = @inbounds grid.Δxᶠᶠᵃ[i, j]
+@inline Δxᶠᶠᵃ(i::AbstractArray, j::AbstractArray, k::AbstractArray, grid::LLG) = Base.stack(collect(Δxᶠᶠᵃ(i, j, 1, grid) for _ in k))
+
 @inline Δxᶜᶜᵃ(i, j, k, grid::LLG) = @inbounds grid.Δxᶜᶜᵃ[i, j]
+@inline Δxᶜᶜᵃ(i::AbstractArray, j::AbstractArray, k::AbstractArray, grid::LLG) = Base.stack(collect(Δxᶜᶜᵃ(i, j, 1, grid) for _ in k))
+
 
 ### XRegularLLG with pre computed metrics
 
@@ -306,7 +324,9 @@ for L1 in (:ᶜ, :ᶠ), L2 in (:ᶜ, :ᶠ)
 
         @eval begin
             @inline $Axˡˡˡ(i, j, k, grid) = $Δyˡˡˡ(i, j, k, grid) * $Δzˡˡˡ(i, j, k, grid)
+            @inline $Axˡˡˡ(i::AbstractArray, j::AbstractArray, k::AbstractArray, grid) = $Δyˡˡˡ(i, j, k, grid) .* $Δzˡˡˡ(i, j, k, grid)
             @inline $Ayˡˡˡ(i, j, k, grid) = $Δxˡˡˡ(i, j, k, grid) * $Δzˡˡˡ(i, j, k, grid)
+            @inline $Ayˡˡˡ(i::AbstractArray, j::AbstractArray, k::AbstractArray, grid) = $Δxˡˡˡ(i, j, k, grid) .* $Δzˡˡˡ(i, j, k, grid)
 
             # For the moment the horizontal area is independent of `z`. This might change if
             # we want to implement deep atmospheres where Az is a function of z