Fix some conversions

CliMA · Aug 29, 2024 · 515b92b · 515b92b
1 parent 78ac7bb
commit 515b92b
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 9 deletions.
diff --git a/ext/cuda/topologies_dss.jl b/ext/cuda/topologies_dss.jl
@@ -232,12 +232,13 @@ function dss_transform_kernel!(
         elem = localelems[localelemno]
         (ip, jp) = perimeter[p]
         loc = CI(ip, jp, 1, level, elem)
-        wt = weight[loc]
-        perimeter_data[CI(p, 1, 1, level, elem)] = Topologies.dss_transform(
+        src = Topologies.dss_transform(
             data[loc],
-            local_geometry[CI(ip, jp, 1, level, elem)],
-            wt,
+            local_geometry[loc],
+            weight[loc],
         )
+        perimeter_data[CI(p, 1, 1, level, elem)] =
+            Topologies.drop_vert_dim(eltype(perimeter_data), src)
     end
     return nothing
 end
@@ -291,10 +292,11 @@ function dss_untransform_kernel!(
         elem = localelems[localelemno]
         ip, jp = perimeter[p]
 
-        data[CI(ip, jp, 1, level, elem)] = Topologies.dss_untransform(
+        loc = CI(ip, jp, 1, level, elem)
+        data[loc] = Topologies.dss_untransform(
             eltype(data),
             perimeter_data[CI(p, 1, 1, level, elem)],
-            local_geometry[CI(ip, jp, 1, level, elem)],
+            local_geometry[loc],
         )
     end
     return nothing

diff --git a/src/Topologies/dss.jl b/src/Topologies/dss.jl
@@ -176,6 +176,15 @@ function dss_transform!(
     return nothing
 end
 
+# For DSS of Covariant123Vector, the third component is treated like a scalar
+# and is not transformed
+@inline drop_vert_dim(
+    ::Type{T},
+    x::Geometry.UVWVector,
+) where {T <: Geometry.UVVector} = Geometry.UVVector(x.u, x.v)
+@inline drop_vert_dim(::Type{T}, x::T) where {T} = x
+@inline drop_vert_dim(::Type{T}, x::Real) where {T <: Real} = x
+
 """
     function dss_transform!(
         ::ClimaComms.AbstractCPUDevice,
@@ -215,12 +224,14 @@ function dss_transform!(
         for (p, (ip, jp)) in enumerate(perimeter)
             for level in 1:nlevels
                 loc = CI(ip, jp, 1, level, elem)
-                wt = weight[loc]
-                perimeter_data[CI(p, 1, 1, level, elem)] = dss_transform(
+                src = dss_transform(
                     data[loc],
                     local_geometry[CI(ip, jp, 1, level, elem)],
-                    wt,
+                    weight[loc],
                 )
+                perimeter_data[CI(p, 1, 1, level, elem)] =
+                    drop_vert_dim(eltype(perimeter_data), src)
+
             end
         end
     end