better particle bank checker in census loop

mcdc/kernel.py

@@ -1278,6 +1278,10 @@ def bank_rebalance(mcdc):
     idx_start, N_local, N = bank_scanning(mcdc["bank_source"], mcdc)
     idx_end = idx_start + N_local
 
+    # Abort if source bank is empty
+    if N == 0:
+        return
+
     distribute_work(N, mcdc)
 
     # Rebalance not needed if there is only one rank

mcdc/loop.py

@@ -90,7 +90,9 @@ def loop_fixed_source(data_arr, mcdc_arr):
             # Check and accordingly promote future particles to censused particle
             if kernel.get_bank_size(mcdc["bank_future"]) > 0:
                 kernel.check_future_bank(mcdc)
-
+            if idx_census > 0 and kernel.get_bank_size(mcdc["bank_source"]) == 0 and kernel.get_bank_size(mcdc['bank_census']) == 0 and kernel.get_bank_size(mcdc['bank_future']) == 0:
+                # No more particle to work on
+                 break
             # Loop over source particles
             seed_source = kernel.split_seed(seed_census, SEED_SPLIT_SOURCE)
             loop_source(seed_source, data, mcdc)
@@ -104,14 +106,11 @@ def loop_fixed_source(data_arr, mcdc_arr):
 
             # TODO: Output tally
 
-            # No particles censused?
-            if mcdc["bank_census"]["size"] == 0:
-                break
-
             # Manage particle banks: population control and work rebalance
             seed_bank = kernel.split_seed(seed_census, SEED_SPLIT_BANK)
             kernel.manage_particle_banks(seed_bank, mcdc)
 
+
         # Multi-batch closeout
         if mcdc["setting"]["N_batch"] > 1:
             # Reset banks