[inch] A slightly better solution to the N_samples < N_ranks issue

krivenko · Jun 7, 2024 · fc67518 · fc67518
1 parent d7feede
commit fc67518
Showing 1 changed file with 12 additions and 16 deletions.
diff --git a/src/inchworm.jl b/src/inchworm.jl
@@ -175,17 +175,16 @@ function inchworm_step(expansion::Expansion,
                                         td.rand_params,
                                         seq_type=seq_type) do seq
                 skip!(seq, first(N_range) - 1, exact=true)
-                res::SectorBlockMatrix = rank_weight * contour_integral(
+                res::SectorBlockMatrix = (rank_weight == 0.0) ?
+                deepcopy(zero_sector_block_matrix) :
+                rank_weight * contour_integral(
                     t -> eval(td.topologies, t),
                     c,
                     trans,
                     init = deepcopy(zero_sector_block_matrix),
                     seq = seq,
                     N = length(N_range)
                 )
-                if rank_weight == 0.0
-                    res = deepcopy(zero_sector_block_matrix)
-                end
                 @timeit tmr "MPI all_reduce" begin
                 all_reduce!(res, +)
                 end # tmr
@@ -277,17 +276,16 @@ function inchworm_step_bare(expansion::Expansion,
             contrib_mean, contrib_std =
             mean_std_from_randomization(d, td.rand_params, seq_type=seq_type) do seq
                 skip!(seq, first(N_range) - 1, exact=true)
-                res::SectorBlockMatrix = rank_weight * contour_integral(
+                res::SectorBlockMatrix = (rank_weight == 0.0) ?
+                deepcopy(zero_sector_block_matrix) :
+                rank_weight * contour_integral(
                     t -> eval(td.topologies, t),
                     c,
                     trans,
                     init = deepcopy(zero_sector_block_matrix),
                     seq = seq,
                     N = length(N_range)
                 )
-                if rank_weight == 0.0
-                    res = deepcopy(zero_sector_block_matrix)
-                end
                 @timeit tmr "MPI all_reduce" begin
                 all_reduce!(res, +)
                 end # tmr
@@ -584,17 +582,16 @@ function diff_inchworm_step!(expansion::Expansion,
                                         td.rand_params,
                                         seq_type=seq_type) do seq
                 skip!(seq, first(N_range) - 1, exact=true)
-                res::SectorBlockMatrix = rank_weight * contour_integral(
+                res::SectorBlockMatrix = (rank_weight == 0.0) ?
+                deepcopy(zero_sector_block_matrix) :
+                rank_weight * contour_integral(
                     t -> eval(td.topologies, t),
                     c,
                     trans,
                     init = deepcopy(zero_sector_block_matrix),
                     seq = seq,
                     N = length(N_range)
                 )
-                if rank_weight == 0.0
-                    res = deepcopy(zero_sector_block_matrix)
-                end
                 @timeit tmr "MPI all_reduce" begin
                 all_reduce!(res, +)
                 end # tmr
@@ -865,17 +862,16 @@ function correlator_2p(expansion::Expansion,
                                         td.rand_params,
                                         seq_type=seq_type) do seq
                 skip!(seq, first(N_range) - 1, exact=true)
-                res::ComplexF64 = rank_weight * contour_integral(
+                res::ComplexF64 = (rank_weight == 0.0) ?
+                deepcopy(zero_sector_block_matrix) :
+                rank_weight * contour_integral(
                     t -> tr(eval(td.topologies, t)),
                     grid.contour,
                     trans,
                     init = ComplexF64(0),
                     seq = seq,
                     N = length(N_range)
                 )
-                if rank_weight == 0.0
-                    res = deepcopy(zero_sector_block_matrix)
-                end
                 @timeit tmr "MPI all_reduce" begin
                 MPI.Allreduce(res, +, MPI.COMM_WORLD)
                 end # tmr