opt: Parallelise calculation of coordinate arrays

There are functions which calculate radial and angular coordinates from Cartesian. Use Strided.jl to parallelise them.
musoke · Oct 20, 2023 · d557bdd · d557bdd
1 parent 3c1bb59
commit d557bdd
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 9 deletions.
diff --git a/Project.toml b/Project.toml
@@ -15,6 +15,7 @@ NPZ = "15e1cf62-19b3-5cfa-8e77-841668bca605"
 PencilArrays = "0e08944d-e94e-41b1-9406-dcf66b6a9d2e"
 PencilFFTs = "4a48f351-57a6-4416-9ec4-c37015456aae"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
+Strided = "5e0ebb24-38b0-5f93-81fe-25c709ecae67"
 
 [compat]
 AbstractFFTs = "0.5, 1.0"
@@ -28,4 +29,5 @@ NPZ = "0.4"
 PencilArrays = "0.11, 0.12, 0.13, 0.14, 0.15, 0.16, 0.17, 0.18"
 PencilFFTs = "0.12, 0.13, 0.14, 0.15"
 Statistics = "1.8, 1.9"
+Strided = "2.0.4"
 julia = "1.8, 1.9"
diff --git a/src/grids.jl b/src/grids.jl
@@ -1,3 +1,5 @@
+using Strided
+
 """
     AbstractGrids
 
@@ -342,15 +344,15 @@ true
 ```
 """
 function radius_spherical(x, y, z)
-    rs = (x^2 + y^2 + z^2)^(1 // 2)
+    sqrt(x^2 + y^2 + z^2)
 end
 
 function radius_spherical(grids)
-    radius_spherical.(grids.x, grids.y, grids.z)
+    @strided radius_spherical.(grids.x, grids.y, grids.z)
 end
 
 function radius_spherical(grids, r0)
-    radius_spherical.(grids.x .- r0[1], grids.y .- r0[2], grids.z .- r0[3])
+    @strided radius_spherical.(grids.x .- r0[1], grids.y .- r0[2], grids.z .- r0[3])
 end
 
 """
@@ -366,11 +368,11 @@ function polar_angle(x, y, z)
 end
 
 function polar_angle(grids)
-    polar_angle.(grids.x, grids.y, grids.z)
+    @strided polar_angle.(grids.x, grids.y, grids.z)
 end
 
 function polar_angle(grids, r0)
-    polar_angle.(grids.x .- r0[1], grids.y .- r0[2], grids.z .- r0[3])
+    @strided polar_angle.(grids.x .- r0[1], grids.y .- r0[2], grids.z .- r0[3])
 end
 
 
@@ -386,11 +388,11 @@ function azimuthal_angle(x, y, z)
 end
 
 function azimuthal_angle(grids)
-    azimuthal_angle.(grids.x, grids.y, grids.z)
+    @strided azimuthal_angle.(grids.x, grids.y, grids.z)
 end
 
 function azimuthal_angle(grids, r0)
-    azimuthal_angle.(grids.x .- r0[1], grids.y .- r0[2], grids.z .- r0[3])
+    @strided azimuthal_angle.(grids.x .- r0[1], grids.y .- r0[2], grids.z .- r0[3])
 end
 
 """
@@ -426,11 +428,11 @@ function radius_cylindrical(x, y, z)
 end
 
 function radius_cylindrical(grids)
-    radius_cylindrical.(grids.x, grids.y, grids.z)
+    @strided radius_cylindrical.(grids.x, grids.y, grids.z)
 end
 
 function radius_cylindrical(grids, r0)
-    radius_cylindrical.(grids.x .- r0[1], grids.y .- r0[2], grids.z .- r0[3])
+    @strided radius_cylindrical.(grids.x .- r0[1], grids.y .- r0[2], grids.z .- r0[3])
 end
 
 """