allocating curl, divergence functions

SpeedyWeather · Sep 11, 2023 · 883548e · 883548e
1 parent 1ab9bf4
commit 883548e
Show file tree

Hide file tree

Showing 2 changed files with 57 additions and 33 deletions.
diff --git a/src/SpeedyTransforms/SpeedyTransforms.jl b/src/SpeedyTransforms/SpeedyTransforms.jl
@@ -27,7 +27,9 @@ export  SpectralTransform,
 export  get_nlat_half
 
 # GRADIENTS
-export  curl!,
+export  curl,
+        divergence,
+        curl!,
         divergence!,
         UV_from_vor!,
         UV_from_vordiv!,

diff --git a/src/SpeedyTransforms/spectral_gradients.jl b/src/SpeedyTransforms/spectral_gradients.jl
@@ -1,12 +1,5 @@
 """
-    curl!(  curl::LowerTriangularMatrix,
-            u::LowerTriangularMatrix,
-            v::LowerTriangularMatrix,
-            S::SpectralTransform;
-            flipsign::Bool=false,
-            add::Bool=false,
-            )
-
+$(TYPEDSIGNATURES)
 Curl of a vector `u,v` written into `curl`, `curl = ∇×(u,v)`.
 `u,v` are expected to have a 1/coslat-scaling included, then `curl` is not scaled.
 `flipsign` option calculates -∇×(u,v) instead. `add` option calculates `curl += ∇×(u,v)` instead.
@@ -27,14 +20,7 @@ function curl!( curl::LowerTriangularMatrix,
 end
 
 """
-    divergence!(div::LowerTriangularMatrix,
-                u::LowerTriangularMatrix,
-                v::LowerTriangularMatrix,
-                S::SpectralTransform{NF};
-                flipsign::Bool=false,
-                add::Bool=false,
-                )
-
+$(TYPEDSIGNATURES)
 Divergence of a vector `u,v` written into `div`, `div = ∇⋅(u,v)`. 
 `u,v` are expected to have a 1/coslat-scaling included, then `div` is not scaled.
 `flipsign` option calculates -∇⋅(u,v) instead. `add` option calculates `div += ∇⋅(u,v)` instead.
@@ -55,12 +41,7 @@ function divergence!(   div::LowerTriangularMatrix,
 end
 
 """
-    _divergence!(   kernel,
-                    div::LowerTriangularMatrix,
-                    u::LowerTriangularMatrix,
-                    v::LowerTriangularMatrix,
-                    S::SpectralTransform)
-
+$(TYPEDSIGNATURES)
 Generic divergence function of vector `u`,`v` that writes into the output into `div`.
 Generic as it uses the kernel `kernel` such that curl, div, add or flipsign
 options are provided through `kernel`, but otherwise a single function is used."""
@@ -107,11 +88,57 @@ function _divergence!(  kernel,
 end
 
 """
-    UV_from_vor!(   U::LowerTriangularMatrix,
-                    V::LowerTriangularMatrix,
-                    vor::LowerTriangularMatrix,
-                    S::SpectralTransform)
+$(TYPEDSIGNATURES)
+Divergence (∇⋅) of two vector components `u,v` which need to have size (n+1)xn,
+the last row will be set to zero in the returned `LowerTriangularMatrix`.
+This function requires both `u,v` to be transforms of fields that are scaled with
+`1/cos(lat)`. An example usage is therefore
+
+    RingGrids.scale_coslat⁻¹!(u_grid)
+    RingGrids.scale_coslat⁻¹!(v_grid)
+    u = spectral(u_grid)
+    v = spectral(v_grid)
+    div = divergence(u,v)
+    div_grid = gridded(div)
+"""
+function divergence(u::LowerTriangularMatrix,
+                    v::LowerTriangularMatrix)
 
+    @assert size(u) == size(v) "Size $(size(u)) and $(size(v)) incompatible."
+
+    S = SpectralTransform(u)
+    div = similar(u)
+    divergence!(div,u,v,S,add=false,flipsign=false)
+    return div
+end
+
+"""
+$(TYPEDSIGNATURES)
+Curl (∇×) of two vector components `u,v` of size (n+1)xn, the last row
+will be set to zero in the returned `LowerTriangularMatrix`. This function
+requires both `u,v` to be transforms of fields that are scaled with
+`1/cos(lat)`. An example usage is therefore
+
+    RingGrids.scale_coslat⁻¹!(u_grid)
+    RingGrids.scale_coslat⁻¹!(v_grid)
+    u = spectral(u_grid)
+    v = spectral(v_grid)
+    vor = curl(u,v)
+    vor_grid = gridded(div)
+"""
+function curl(  u::LowerTriangularMatrix,
+                v::LowerTriangularMatrix)
+
+    @assert size(u) == size(v) "Size $(size(u)) and $(size(v)) incompatible."
+
+    S = SpectralTransform(u)
+    vor = similar(u)
+    curl!(vor,u,v,S,add=false,flipsign=false)
+    return vor
+end
+
+"""
+$(TYPEDSIGNATURES)
 Get U,V (=(u,v)*coslat) from vorticity ζ spectral space (divergence D=0)
 Two operations are combined into a single linear operation. First, invert the
 spherical Laplace ∇² operator to get stream function from vorticity. Then
@@ -178,12 +205,7 @@ function UV_from_vor!(  U::LowerTriangularMatrix{Complex{NF}},
 end
 
 """
-    UV_from_vordiv!(U::LowerTriangularMatrix,
-                    V::LowerTriangularMatrix,
-                    vor::LowerTriangularMatrix,
-                    div::LowerTriangularMatrix,
-                    S::SpectralTransform)
-
+$(TYPEDSIGNATURES)
 Get U,V (=(u,v)*coslat) from vorticity ζ and divergence D in spectral space.
 Two operations are combined into a single linear operation. First, invert the
 spherical Laplace ∇² operator to get stream function from vorticity and