Remove \partialz\partial\xi\_1,2 dss in hypsography / meshwarp,

Update dss_transform to support 3-component dss. modified: src/Topologies/dss.jl modified: ext/cuda/topologies_dss.jl
CliMA · May 2, 2024 · 8fd0a5d · 8fd0a5d
1 parent 9938bb6
commit 8fd0a5d
Show file tree

Hide file tree

Showing 4 changed files with 171 additions and 75 deletions.
diff --git a/ext/cuda/topologies_dss.jl b/ext/cuda/topologies_dss.jl
@@ -187,7 +187,8 @@ function Topologies.dss_transform!(
     perimeter::Topologies.Perimeter2D,
     scalarfidx::AbstractVector{Int},
     covariant12fidx::AbstractVector{Int},
-    contravariant12fidx::AbstractVector{Int},
+    covariant123fidx::AbstractVector{Int},
+    contravariant123fidx::AbstractVector{Int},
     localelems::AbstractVector{Int},
 )
     nlocalelems = length(localelems)
@@ -211,7 +212,8 @@ function Topologies.dss_transform!(
             perimeter,
             scalarfidx,
             covariant12fidx,
-            contravariant12fidx,
+            covariant123fidx,
+            contravariant123fidx,
             localelems,
         )
         auto_launch!(
@@ -235,7 +237,8 @@ function dss_transform_kernel!(
     perimeter::Topologies.Perimeter2D{Nq},
     scalarfidx::AbstractVector{Int},
     covariant12fidx::AbstractVector{Int},
-    contravariant12fidx::AbstractVector{Int},
+    covariant123fidx::AbstractVector{Int},
+    contravariant123fidx::AbstractVector{Int},
     localelems::AbstractVector{Int},
 ) where {FT <: AbstractFloat, Nq}
     gidx = threadIdx().x + (blockIdx().x - 1) * blockDim().x
@@ -257,41 +260,80 @@ function dss_transform_kernel!(
                 Topologies._get_idx(sizet_data, (level, ip, jp, fidx, elem))
             pperimeter_data[level, p, fidx, elem] = pdata[data_idx] * weight
         end
-        for fidx in covariant12fidx
-            data_idx1 =
-                Topologies._get_idx(sizet_data, (level, ip, jp, fidx, elem))
-            data_idx2 =
-                Topologies._get_idx(sizet_data, (level, ip, jp, fidx + 1, elem))
-            (idx11, idx12, idx21, idx22) =
-                Topologies._get_idx_metric(sizet_metric, (level, ip, jp, elem))
-            pperimeter_data[level, p, fidx, elem] =
+        for fidx in covariant123fidx
+                  data_idx1 = Topologies._get_idx(sizet_data, (level, ip, jp, fidx, elem))
+                data_idx2 =
+                    Topologies._get_idx(sizet_data, (level, ip, jp, fidx + 1, elem))
+                data_idx3 =
+                    Topologies._get_idx(sizet_data, (level, ip, jp, fidx + 2, elem))
                 (
-                    p∂ξ∂x[idx11] * pdata[data_idx1] +
-                    p∂ξ∂x[idx12] * pdata[data_idx2]
-                ) * weight
-            pperimeter_data[level, p, fidx + 1, elem] =
-                (
-                    p∂ξ∂x[idx21] * pdata[data_idx1] +
-                    p∂ξ∂x[idx22] * pdata[data_idx2]
-                ) * weight
+                    idx11,
+                    idx12,
+                    idx13,
+                    idx21,
+                    idx22,
+                    idx23,
+                    idx31,
+                    idx32,
+                    idx33,
+                ) = Topologies._get_idx_metric(sizet_metric, (level, ip, jp, elem))
+                # Covariant to physical transformation
+                pperimeter_data[level, p, fidx, elem] =
+                    (
+                        p∂ξ∂x[idx11] * pdata[data_idx1] +
+                        p∂ξ∂x[idx12] * pdata[data_idx2] +
+                        p∂ξ∂x[idx13] * pdata[data_idx3]
+                    ) * weight
+                pperimeter_data[level, p, fidx + 1, elem] =
+                    (
+                        p∂ξ∂x[idx21] * pdata[data_idx1] +
+                        p∂ξ∂x[idx22] * pdata[data_idx2] +
+                        p∂ξ∂x[idx23] * pdata[data_idx3]
+                    ) * weight
+                pperimeter_data[level, p, fidx + 2, elem] =
+                    (
+                        p∂ξ∂x[idx31] * pdata[data_idx1] +
+                        p∂ξ∂x[idx32] * pdata[data_idx2] +
+                        p∂ξ∂x[idx33] * pdata[data_idx3]
+                    ) * weight
         end
-        for fidx in contravariant12fidx
-            data_idx1 =
-                Topologies._get_idx(sizet_data, (level, ip, jp, fidx, elem))
-            data_idx2 =
-                Topologies._get_idx(sizet_data, (level, ip, jp, fidx + 1, elem))
-            (idx11, idx12, idx21, idx22) =
-                Topologies._get_idx_metric(sizet_metric, (level, ip, jp, elem))
-            pperimeter_data[level, p, fidx, elem] =
-                (
-                    p∂x∂ξ[idx11] * pdata[data_idx1] +
-                    p∂x∂ξ[idx21] * pdata[data_idx2]
-                ) * weight
-            pperimeter_data[level, p, fidx + 1, elem] =
+        for fidx in contravariant123fidx
+                data_idx1 = Topologies._get_idx(sizet_data, (level, ip, jp, fidx, elem))
+                data_idx2 =
+                    Topologies._get_idx(sizet_data, (level, ip, jp, fidx + 1, elem))
+                data_idx3 =
+                    Topologies._get_idx(sizet_data, (level, ip, jp, fidx + 2, elem))
                 (
-                    p∂x∂ξ[idx12] * pdata[data_idx1] +
-                    p∂x∂ξ[idx22] * pdata[data_idx2]
-                ) * weight
+                    idx11,
+                    idx12,
+                    idx13,
+                    idx21,
+                    idx22,
+                    idx23,
+                    idx31,
+                    idx32,
+                    idx33,
+                ) = Topologies._get_idx_metric(sizet_metric, (level, ip, jp, elem))
+                # Contravariant to physical transformation
+                pperimeter_data[level, p, fidx, elem] =
+                    (
+                        p∂x∂ξ[idx11] * pdata[data_idx1] +
+                        p∂x∂ξ[idx21] * pdata[data_idx2] +
+                        p∂x∂ξ[idx31] * pdata[data_idx3]
+                    ) * weight
+                pperimeter_data[level, p, fidx + 1, elem] =
+                    (
+                        p∂x∂ξ[idx12] * pdata[data_idx1] +
+                        p∂x∂ξ[idx22] * pdata[data_idx2] +
+                        p∂x∂ξ[idx32] * pdata[data_idx3]
+                    ) * weight
+                pperimeter_data[level, p, fidx + 2, elem] =
+                    (
+                        p∂x∂ξ[idx13] * pdata[data_idx1] +
+                        p∂x∂ξ[idx23] * pdata[data_idx2] +
+                        p∂x∂ξ[idx33] * pdata[data_idx3]
+                    ) * weight
+
         end
     end
     return nothing
@@ -306,7 +348,8 @@ function Topologies.dss_untransform!(
     perimeter::Topologies.Perimeter2D{Nq},
     scalarfidx::AbstractVector{Int},
     covariant12fidx::AbstractVector{Int},
-    contravariant12fidx::AbstractVector{Int},
+    covariant123fidx::AbstractVector{Int},
+    contravariant123fidx::AbstractVector{Int},
     localelems::AbstractVector{Int},
 ) where {Nq}
     nlocalelems = length(localelems)
@@ -328,7 +371,8 @@ function Topologies.dss_untransform!(
             perimeter,
             scalarfidx,
             covariant12fidx,
-            contravariant12fidx,
+            covariant123fidx,
+            contravariant123fidx,
             localelems,
         )
         auto_launch!(
@@ -351,7 +395,8 @@ function dss_untransform_kernel!(
     perimeter::Topologies.Perimeter2D{Nq},
     scalarfidx::AbstractVector{Int},
     covariant12fidx::AbstractVector{Int},
-    contravariant12fidx::AbstractVector{Int},
+    covariant123fidx::AbstractVector{Int},
+    contravariant123fidx::AbstractVector{Int},
     localelems::AbstractVector{Int},
 ) where {FT <: AbstractFloat, Nq}
     gidx = threadIdx().x + (blockIdx().x - 1) * blockDim().x
@@ -371,33 +416,73 @@ function dss_untransform_kernel!(
                 Topologies._get_idx(sizet_data, (level, ip, jp, fidx, elem))
             pdata[data_idx] = pperimeter_data[level, p, fidx, elem]
         end
-        for fidx in covariant12fidx
+        for fidx in covariant123fidx
             data_idx1 =
                 Topologies._get_idx(sizet_data, (level, ip, jp, fidx, elem))
             data_idx2 =
                 Topologies._get_idx(sizet_data, (level, ip, jp, fidx + 1, elem))
-            (idx11, idx12, idx21, idx22) =
-                Topologies._get_idx_metric(sizet_metric, (level, ip, jp, elem))
-            pdata[data_idx1] =
-                p∂x∂ξ[idx11] * pperimeter_data[level, p, fidx, elem] +
-                p∂x∂ξ[idx12] * pperimeter_data[level, p, fidx + 1, elem]
-            pdata[data_idx2] =
-                p∂x∂ξ[idx21] * pperimeter_data[level, p, fidx, elem] +
-                p∂x∂ξ[idx22] * pperimeter_data[level, p, fidx + 1, elem]
+            data_idx3 =
+                Topologies._get_idx(sizet_data, (level, ip, jp, fidx + 2, elem))
+                    (
+                        idx11,
+                        idx12,
+                        idx13,
+                        idx21,
+                        idx22,
+                        idx23,
+                        idx31,
+                        idx32,
+                        idx33,
+                ) = Topologies._get_idx_metric(sizet_metric, (level, ip, jp, elem))
+                pdata[data_idx1] =
+                    p∂x∂ξ[idx11] * pperimeter_data[level, p, fidx, elem] +
+                    p∂x∂ξ[idx12] *
+                    pperimeter_data[level, p, fidx + 1, elem] +
+                    p∂x∂ξ[idx13] * pperimeter_data[level, p, fidx + 2, elem]
+                pdata[data_idx2] =
+                    p∂x∂ξ[idx21] * pperimeter_data[level, p, fidx, elem] +
+                    p∂x∂ξ[idx22] *
+                    pperimeter_data[level, p, fidx + 1, elem] +
+                    p∂x∂ξ[idx23] * pperimeter_data[level, p, fidx + 2, elem]
+                pdata[data_idx3] =
+                    p∂x∂ξ[idx31] * pperimeter_data[level, p, fidx, elem] +
+                    p∂x∂ξ[idx32] *
+                    pperimeter_data[level, p, fidx + 1, elem] +
+                    p∂x∂ξ[idx33] * pperimeter_data[level, p, fidx + 2, elem]
         end
-        for fidx in contravariant12fidx
-            data_idx1 =
-                Topologies._get_idx(sizet_data, (level, ip, jp, fidx, elem))
-            data_idx2 =
-                Topologies._get_idx(sizet_data, (level, ip, jp, fidx + 1, elem))
-            (idx11, idx12, idx21, idx22) =
-                Topologies._get_idx_metric(sizet_metric, (level, ip, jp, elem))
-            pdata[data_idx1] =
-                p∂ξ∂x[idx11] * pperimeter_data[level, p, fidx, elem] +
-                p∂ξ∂x[idx21] * pperimeter_data[level, p, fidx + 1, elem]
-            pdata[data_idx2] =
-                p∂ξ∂x[idx12] * pperimeter_data[level, p, fidx, elem] +
-                p∂ξ∂x[idx22] * pperimeter_data[level, p, fidx + 1, elem]
+        for fidx in contravariant123fidx
+                    data_idx1 =
+                        Topologies._get_idx(sizet_data, (level, ip, jp, fidx, elem))
+                    data_idx2 =
+                        Topologies._get_idx(sizet_data, (level, ip, jp, fidx + 1, elem))
+                    data_idx3 =
+                        Topologies._get_idx(sizet_data, (level, ip, jp, fidx + 2, elem))
+                    (
+                        idx11,
+                        idx12,
+                        idx13,
+                        idx21,
+                        idx22,
+                        idx23,
+                        idx31,
+                        idx32,
+                        idx33,
+                    ) = Topologies._get_idx_metric(sizet_metric, (level, ip, jp, elem))
+                    pdata[data_idx1] =
+                        p∂ξ∂x[idx11] * pperimeter_data[level, p, fidx, elem] +
+                        p∂ξ∂x[idx21] *
+                        pperimeter_data[level, p, fidx + 1, elem] +
+                        p∂ξ∂x[idx31] * pperimeter_data[level, p, fidx + 2, elem]
+                    pdata[data_idx2] =
+                        p∂ξ∂x[idx12] * pperimeter_data[level, p, fidx, elem] +
+                        p∂ξ∂x[idx22] *
+                        pperimeter_data[level, p, fidx + 1, elem] +
+                        p∂ξ∂x[idx32] * pperimeter_data[level, p, fidx + 2, elem]
+                    pdata[data_idx3] =
+                        p∂ξ∂x[idx13] * pperimeter_data[level, p, fidx, elem] +
+                        p∂ξ∂x[idx23] *
+                        pperimeter_data[level, p, fidx + 1, elem] +
+                        p∂ξ∂x[idx33] * pperimeter_data[level, p, fidx + 2, elem]
         end
     end
     return nothing

diff --git a/src/Topologies/dss.jl b/src/Topologies/dss.jl
@@ -25,6 +25,8 @@ struct DSSBuffer{S, G, D, A, B, VI}
     "field id for all scalar fields stored in the `data` array"
     scalarfidx::VI
     "field id for all covariant123vector fields stored in the `data` array"
+    covariant12fidx::VI
+    "field id for all covariant12fidx fields stored in the `data` array"
     covariant123fidx::VI
     "field id for all contravariant123vector fields stored in the `data` array"
     contravariant123fidx::VI
@@ -108,7 +110,7 @@ function create_dss_buffer(
         internal_elems = DA(topology.internal_elems)
         perimeter_elems = DA(topology.perimeter_elems)
     end
-    scalarfidx, covariant123fidx, contravariant123fidx = Int[], Int[], Int[]
+    scalarfidx, covariant12fidx, covariant123fidx, contravariant123fidx = Int[], Int[], Int[], Int[]
     supportedvectortypes = Union{
         Geometry.UVector,
         Geometry.VVector,
@@ -191,6 +193,7 @@ function create_dss_buffer(
         send_buf_idx,
         recv_buf_idx,
         scalarfidx,
+        covariant12fidx,
         covariant123fidx,
         contravariant123fidx,
         internal_elems,
@@ -242,7 +245,7 @@ function dss_transform!(
     localelems::AbstractVector{Int},
 )
     if !isempty(localelems)
-        (; scalarfidx, covariant123fidx, contravariant123fidx, perimeter_data) =
+        (; scalarfidx, covariant12fidx, covariant123fidx, contravariant123fidx, perimeter_data) =
             dss_buffer
         (; ∂ξ∂x, ∂x∂ξ) = local_geometry
         dss_transform!(
@@ -254,6 +257,7 @@ function dss_transform!(
             weight,
             perimeter,
             scalarfidx,
+            covariant12fidx,
             covariant123fidx,
             contravariant123fidx,
             localelems,
@@ -292,7 +296,7 @@ function dss_untransform!(
     perimeter::Perimeter2D,
     localelems::AbstractVector{Int},
 )
-    (; scalarfidx, covariant123fidx, contravariant123fidx, perimeter_data) =
+    (; scalarfidx, covariant12fidx, covariant123fidx, contravariant123fidx, perimeter_data) =
         dss_buffer
     (; ∂ξ∂x, ∂x∂ξ) = local_geometry
     dss_untransform!(
@@ -303,6 +307,7 @@ function dss_untransform!(
         ∂x∂ξ,
         perimeter,
         scalarfidx,
+        covariant12fidx,
         covariant123fidx,
         contravariant123fidx,
         localelems,
@@ -320,6 +325,7 @@ end
         weight::DataLayouts.IJFH,
         perimeter::AbstractPerimeter,
         scalarfidx::Vector{Int},
+        covariant12fidx::Vector{Int},
         covariant123fidx::Vector{Int},
         contravariant123fidx::Vector{Int},
         localelems::Vector{Int},
@@ -336,6 +342,7 @@ Arguments:
 - `weight`: local dss weights for horizontal space
 - `perimeter`: perimeter iterator
 - `scalarfidx`: field index for scalar fields in the data layout
+- `covariant12fidx`: field index for Covariant12 vector fields in the data layout
 - `covariant123fidx`: field index for Covariant123 vector fields in the data layout
 - `localelems`: list of local elements to perform transformation operations on
 
@@ -350,6 +357,7 @@ function dss_transform!(
     weight::DataLayouts.IJFH,
     perimeter::Perimeter2D{Nq},
     scalarfidx::Vector{Int},
+    covariant12fidx::Vector{Int},
     covariant123fidx::Vector{Int},
     contravariant123fidx::Vector{Int},
     localelems::Vector{Int},
@@ -465,6 +473,7 @@ end
         ∂x∂ξ::Union{DataLayouts.IJFH, DataLayouts.VIJFH},
         perimeter::AbstractPerimeter,
         scalarfidx::Vector{Int},
+        covariant12fidx::Vector{Int},
         covariant123fidx::Vector{Int},
         contravariant123fidx::Vector{Int},
         localelems::Vector{Int},
@@ -493,6 +502,7 @@ function dss_untransform!(
     ∂x∂ξ::Union{DataLayouts.VIJFH, DataLayouts.IJFH},
     perimeter::Perimeter2D{Nq},
     scalarfidx::Vector{Int},
+    covariant12fidx::Vector{Int},
     covariant123fidx::Vector{Int},
     contravariant123fidx::Vector{Int},
     localelems::Vector{Int},

diff --git a/test/Spaces/ddss1_cs.jl b/test/Spaces/ddss1_cs.jl
@@ -12,7 +12,7 @@ import ClimaCore:
     DataLayouts
 
 @testset "DSS on Equiangular Cubed Sphere mesh (ne = 3, serial run)" begin
-    device = ClimaComms.device() #ClimaComms.CUDADevice()
+    device = ClimaComms.device() 
     context = ClimaComms.SingletonCommsContext(device)
     context_cpu =
         ClimaComms.SingletonCommsContext(ClimaComms.CPUSingleThreaded())