From 515b92b458e1252366199d5a1893cd7b01fd8194 Mon Sep 17 00:00:00 2001
From: Charles Kawczynski <kawczynski.charles@gmail.com>
Date: Thu, 29 Aug 2024 10:27:55 -0400
Subject: [PATCH] Fix some conversions

---
 ext/cuda/topologies_dss.jl | 14 ++++++++------
 src/Topologies/dss.jl      | 17 ++++++++++++++---
 2 files changed, 22 insertions(+), 9 deletions(-)

diff --git a/ext/cuda/topologies_dss.jl b/ext/cuda/topologies_dss.jl
index c297ba97e9..6a5d22a9e3 100644
--- 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
index d8ddacbb70..06316dd162 100644
--- 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