Add structured_values_from_bitmask

src/Remapping/distributed_remapping.jl

@@ -59,10 +59,126 @@ containing_pid(target_point::Geometry.LatLongZPoint, topology) = containing_pid(
 Return a bitmask for points in `target_points` in the given `topology` that belong to the
 process with process number `pid`.
 
-This mask can be used to extract the `target_points` relevant to the given `pid` (and much
-more). The mask is the same shape and size as the input `target_points`.
+This mask can be used to extract the `target_points` relevant to the given `pid`. The mask
+is the same shape and size as the input `target_points`, which makes it particularly useful.
 """
 function target_points_pid_bitmask(target_points, topology, pid)
     pid_point = point -> containing_pid(point, topology)
     return pid_point.(target_points) .== pid
 end
+
+"""
+    structured_values_from_bitmask(values, bitmask)
+
+Return an object that has the same shape as `bitmask` but with values taken from `values`.
+("Structure" here is intended as having the same shape)
+
+This function is such that
+```
+values = initial_values[bitmask]
+# Values becomes an array here
+structured_values = structured_values_from_bitmask(values, bitmask)
+
+structured_values == initial_values # => true
+# structured_values is again a the same type as `initial_values`
+```
+"""
+function structured_values_from_bitmask(values, bitmask)
+    # We need to create a return type that look values and not bitmask (because it is a BitMatrix).
+    # So, we create a similar return value and fill it with 0s.
+    structured_values = zero(eltype(values)) .* similar(bitmask, eltype(values))
+    # structured_values[bitmask] is a view on the indices we want to overwrite, with .=, we
+    # change them in place
+    structured_values[bitmask] .= values
+    return structured_values
+end
+
+# struct RemapperDirector5{T, T1, T2}
+#     comms_ctx::T
+#     all_target_points::T1
+
+#     # A dictionary that maps each pid to the bitmask for the target points associated to that
+#     # process id
+#     target_points_pid_bitmasks::T2
+# end
+
+# RemapperDirector = RemapperDirector5
+
+# function RemapperDirector(target_points, topology, comms_ctx)
+#     target_points_bitmasks = Dict(
+#         pid => target_points_pid_bitmask(target_points, topology, pid) for
+#         pid in 1:ClimaComms.nprocs(comms_ctx)
+#     )
+#     return RemapperDirector(comms_ctx, target_points, target_points_bitmasks)
+# end
+
+# struct RemapperWorker1{T, T1, T2}
+#     comms_ctx::T
+
+#     # Target points that are on the process where this object is defined
+#     # local_target_points is stored as a 1D array
+#     local_target_points::T1
+
+#     # bitmask that identifies which target points are on process where the object is
+#     # defined. local_target_points_bitmask has to maintain the desired shape.
+#     local_target_points_bitmask::T2
+# end
+
+# function containing_pids(point, topology)
+#     containing_elem = Meshes.containing_element(topology.mesh, point)
+#     return topology.elempid[containing_elem]
+# end
+
+# pid, nprocs = ClimaComms.init(comms_ctx)
+
+# if ClimaComms.iamroot(comms_ctx)
+#     println("I am $pid and there are $nprocs processes")
+# end
+
+
+# function distributed_space(
+#     (n1, n2),
+#     (x1periodic, x2periodic),
+#     (Nq, Nv, Nf);
+#     x1min = -2π,
+#     x1max = 2π,
+#     x2min = -2π,
+#     x2max = 2π,
+# )
+#     domain = Domains.RectangleDomain(
+#         Domains.IntervalDomain(
+#             Geometry.XPoint(x1min),
+#             Geometry.XPoint(x1max),
+#             periodic = x1periodic,
+#             boundary_names = x1periodic ? nothing : (:west, :east),
+#         ),
+#         Domains.IntervalDomain(
+#             Geometry.YPoint(x2min),
+#             Geometry.YPoint(x2max),
+#             periodic = x2periodic,
+#             boundary_names = x2periodic ? nothing : (:north, :south),
+#         ),
+#     )
+#     mesh = Meshes.RectilinearMesh(domain, n1, n2)
+#     topology = Topologies.Topology2D(comms_ctx, mesh, Meshes.elements(mesh))
+#     quad = Spaces.Quadratures.GLL{Nq}()
+#     space = Spaces.SpectralElementSpace2D(topology, quad)
+
+#     return space
+# end
+
+# space = distributed_space((4, 1), (false, false), (3, 1, 1))
+
+# target_points = [Geometry.XYPoint(-5.0, 5.0),
+#                  Geometry.XYPoint(5.0, 5.0),
+#                  ]
+
+# compute_list_pids = [containing_pids(p, space.topology) for p in target_points]
+
+# remapper_director = RemapperDirector(target_points, compute_list_pids) ? ClimaComms.iamroot(comms_ctx) : nothing
+
+# mask = compute_list_pids .== pid
+
+# compute_list_this_process = target_points[mask]
+
+# println(pid, compute_list_this_process)

test/Remapping/distributed_remapping.jl

@@ -13,6 +13,22 @@ atexit() do
     global_logger(prev_logger)
 end
 
+@testset "structured_values_from_bitmask" begin
+
+    initial_values = reshape(1:27, (3, 3, 3))
+    bitmask = initial_values .% 5 .== 0
+    values = initial_values[bitmask]
+
+    expected_values = initial_values .* 0
+    expected_values[5] = 5
+    expected_values[10] = 10
+    expected_values[15] = 15
+    expected_values[20] = 20
+    expected_values[25] = 25
+    @test expected_values ==
+          Remapping.structured_values_from_bitmask(values, bitmask)
+end
+
 @testset "3D sphere" begin
     vertdomain = Domains.IntervalDomain(
         Geometry.ZPoint(0.0),