WIP - Forced collisions variance reduction technique

docs/source/pythonapi/index.rst

@@ -87,8 +87,10 @@ Defining techniques
 Techniques are enabled by calling methods on the ``mcdc.simulation`` singleton:
 
 - ``mcdc.simulation.implicit_capture(active=True)``
+- ``mcdc.simulation.forced_collisions(cells=[], weight_thresholds=[], weight_targets=[])``
 - ``mcdc.simulation.weighted_emission(active=True, weight_target=1.0)``
 - ``mcdc.simulation.weight_roulette(weight_threshold=0.0, weight_target=1.0)``
+- ``mcdc.simulation.weight_windows(weight_windows, mesh=None, energy=None)``
 - ``mcdc.simulation.population_control(active=True)``
 
 Running

mcdc/code_factory/numba_objects_generator.py

@@ -1013,6 +1013,14 @@ def generate_mcdc_access(targets):
                     object_name, attribute_name, shape[1], shape[2], True
                 )
 
+            elif len(shape) == 4:
+                text_getter += _accessor_4d_element(
+                    object_name, attribute_name, shape[1], shape[2], shape[3]
+                )
+
+                text_setter += _accessor_4d_element(
+                    object_name, attribute_name, shape[1], shape[2], shape[3], True
+                )
             text_getter += _accessor_chunk(object_name, attribute_name)
             text_setter += _accessor_chunk(object_name, attribute_name, True)
 
@@ -1157,6 +1165,25 @@ def _accessor_3d_element(object_name, attribute_name, stride_2, stride_3, setter
     return text
 
 
+def _accessor_4d_element(
+    object_name, attribute_name, stride_2, stride_3, stride_4, setter=False
+):
+    text = f"@njit\n"
+    if setter:
+        text += f"def {attribute_name}(index_1, index_2, index_3, index_4, {object_name}, data, value):\n"
+    else:
+        text += f"def {attribute_name}(index_1, index_2, index_3, index_4, {object_name}, data):\n"
+    text += f'    offset = {object_name}["{attribute_name}_offset"]\n'
+    text += f'    stride_2 = {object_name}["{stride_2}"]\n'
+    text += f'    stride_3 = {object_name}["{stride_3}"]\n'
+    text += f'    stride_4 = {object_name}["{stride_4}"]\n'
+    if setter:
+        text += f"    data[offset + index_1 * stride_2 * stride_3 * stride_4 + index_2 * stride_3 * stride_4 + index_3 * stride_4 + index_4] = value\n\n\n"
+    else:
+        text += f"    return data[offset + index_1 * stride_2 * stride_3 * stride_4 + index_2 * stride_3 * stride_4 + index_3 * stride_4 + index_4]\n\n\n"
+    return text
+
+
 # ======================================================================================
 # Misc.
 # ======================================================================================

mcdc/mcdc_get/__init__.py

@@ -86,12 +86,16 @@
 
 import mcdc.mcdc_get.settings as settings
 
+import mcdc.mcdc_get.forced_collisions as forced_collisions
+
 import mcdc.mcdc_get.implicit_capture as implicit_capture
 
 import mcdc.mcdc_get.population_control as population_control
 
 import mcdc.mcdc_get.weight_roulette as weight_roulette
 
+import mcdc.mcdc_get.weight_windows as weight_windows
+
 import mcdc.mcdc_get.weighted_emission as weighted_emission
 
 import mcdc.mcdc_get.source as source

mcdc/mcdc_get/forced_collisions.py

@@ -0,0 +1,90 @@
+# The following is automatically generated by code_factory.py
+
+from numba import njit
+
+
+@njit
+def cell_IDs(index, forced_collisions, data):
+    offset = forced_collisions["cell_IDs_offset"]
+    return data[offset + index]
+
+
+@njit
+def cell_IDs_all(forced_collisions, data):
+    start = forced_collisions["cell_IDs_offset"]
+    size = forced_collisions["cell_IDs_length"]
+    end = start + size
+    return data[start:end]
+
+
+@njit
+def cell_IDs_last(forced_collisions, data):
+    start = forced_collisions["cell_IDs_offset"]
+    size = forced_collisions["cell_IDs_length"]
+    end = start + size
+    return data[end - 1]
+
+
+@njit
+def cell_IDs_chunk(start, length, forced_collisions, data):
+    start += forced_collisions["cell_IDs_offset"]
+    end = start + length
+    return data[start:end]
+
+
+@njit
+def threshold_weights(index, forced_collisions, data):
+    offset = forced_collisions["threshold_weights_offset"]
+    return data[offset + index]
+
+
+@njit
+def threshold_weights_all(forced_collisions, data):
+    start = forced_collisions["threshold_weights_offset"]
+    size = forced_collisions["threshold_weights_length"]
+    end = start + size
+    return data[start:end]
+
+
+@njit
+def threshold_weights_last(forced_collisions, data):
+    start = forced_collisions["threshold_weights_offset"]
+    size = forced_collisions["threshold_weights_length"]
+    end = start + size
+    return data[end - 1]
+
+
+@njit
+def threshold_weights_chunk(start, length, forced_collisions, data):
+    start += forced_collisions["threshold_weights_offset"]
+    end = start + length
+    return data[start:end]
+
+
+@njit
+def target_weights(index, forced_collisions, data):
+    offset = forced_collisions["target_weights_offset"]
+    return data[offset + index]
+
+
+@njit
+def target_weights_all(forced_collisions, data):
+    start = forced_collisions["target_weights_offset"]
+    size = forced_collisions["target_weights_length"]
+    end = start + size
+    return data[start:end]
+
+
+@njit
+def target_weights_last(forced_collisions, data):
+    start = forced_collisions["target_weights_offset"]
+    size = forced_collisions["target_weights_length"]
+    end = start + size
+    return data[end - 1]
+
+
+@njit
+def target_weights_chunk(start, length, forced_collisions, data):
+    start += forced_collisions["target_weights_offset"]
+    end = start + length
+    return data[start:end]

mcdc/mcdc_get/weight_windows.py

@@ -0,0 +1,80 @@
+# The following is automatically generated by code_factory.py
+
+from numba import njit
+
+
+@njit
+def energy_bounds(index, weight_windows, data):
+    offset = weight_windows["energy_bounds_offset"]
+    return data[offset + index]
+
+
+@njit
+def energy_bounds_all(weight_windows, data):
+    start = weight_windows["energy_bounds_offset"]
+    size = weight_windows["energy_bounds_length"]
+    end = start + size
+    return data[start:end]
+
+
+@njit
+def energy_bounds_last(weight_windows, data):
+    start = weight_windows["energy_bounds_offset"]
+    size = weight_windows["energy_bounds_length"]
+    end = start + size
+    return data[end - 1]
+
+
+@njit
+def energy_bounds_chunk(start, length, weight_windows, data):
+    start += weight_windows["energy_bounds_offset"]
+    end = start + length
+    return data[start:end]
+
+
+@njit
+def lower_weights(index_1, index_2, index_3, index_4, weight_windows, data):
+    offset = weight_windows["lower_weights_offset"]
+    stride_2 = weight_windows["Nx"]
+    stride_3 = weight_windows["Ny"]
+    stride_4 = weight_windows["Nz"]
+    return data[offset + index_1 * stride_2 * stride_3 * stride_4 + index_2 * stride_3 * stride_4 + index_3 * stride_4 + index_4]
+
+
+@njit
+def lower_weights_chunk(start, length, weight_windows, data):
+    start += weight_windows["lower_weights_offset"]
+    end = start + length
+    return data[start:end]
+
+
+@njit
+def target_weights(index_1, index_2, index_3, index_4, weight_windows, data):
+    offset = weight_windows["target_weights_offset"]
+    stride_2 = weight_windows["Nx"]
+    stride_3 = weight_windows["Ny"]
+    stride_4 = weight_windows["Nz"]
+    return data[offset + index_1 * stride_2 * stride_3 * stride_4 + index_2 * stride_3 * stride_4 + index_3 * stride_4 + index_4]
+
+
+@njit
+def target_weights_chunk(start, length, weight_windows, data):
+    start += weight_windows["target_weights_offset"]
+    end = start + length
+    return data[start:end]
+
+
+@njit
+def upper_weights(index_1, index_2, index_3, index_4, weight_windows, data):
+    offset = weight_windows["upper_weights_offset"]
+    stride_2 = weight_windows["Nx"]
+    stride_3 = weight_windows["Ny"]
+    stride_4 = weight_windows["Nz"]
+    return data[offset + index_1 * stride_2 * stride_3 * stride_4 + index_2 * stride_3 * stride_4 + index_3 * stride_4 + index_4]
+
+
+@njit
+def upper_weights_chunk(start, length, weight_windows, data):
+    start += weight_windows["upper_weights_offset"]
+    end = start + length
+    return data[start:end]

mcdc/mcdc_set/__init__.py

@@ -86,12 +86,16 @@
 
 import mcdc.mcdc_set.settings as settings
 
+import mcdc.mcdc_set.forced_collisions as forced_collisions
+
 import mcdc.mcdc_set.implicit_capture as implicit_capture
 
 import mcdc.mcdc_set.population_control as population_control
 
 import mcdc.mcdc_set.weight_roulette as weight_roulette
 
+import mcdc.mcdc_set.weight_windows as weight_windows
+
 import mcdc.mcdc_set.weighted_emission as weighted_emission
 
 import mcdc.mcdc_set.source as source

mcdc/mcdc_set/forced_collisions.py

@@ -0,0 +1,90 @@
+# The following is automatically generated by code_factory.py
+
+from numba import njit
+
+
+@njit
+def cell_IDs(index, forced_collisions, data, value):
+    offset = forced_collisions["cell_IDs_offset"]
+    data[offset + index] = value
+
+
+@njit
+def cell_IDs_all(forced_collisions, data, value):
+    start = forced_collisions["cell_IDs_offset"]
+    size = forced_collisions["cell_IDs_length"]
+    end = start + size
+    data[start:end] = value
+
+
+@njit
+def cell_IDs_last(forced_collisions, data, value):
+    start = forced_collisions["cell_IDs_offset"]
+    size = forced_collisions["cell_IDs_length"]
+    end = start + size
+    data[end - 1] = value
+
+
+@njit
+def cell_IDs_chunk(start, length, forced_collisions, data, value):
+    start += forced_collisions["cell_IDs_offset"]
+    end = start + length
+    data[start:end] = value
+
+
+@njit
+def threshold_weights(index, forced_collisions, data, value):
+    offset = forced_collisions["threshold_weights_offset"]
+    data[offset + index] = value
+
+
+@njit
+def threshold_weights_all(forced_collisions, data, value):
+    start = forced_collisions["threshold_weights_offset"]
+    size = forced_collisions["threshold_weights_length"]
+    end = start + size
+    data[start:end] = value
+
+
+@njit
+def threshold_weights_last(forced_collisions, data, value):
+    start = forced_collisions["threshold_weights_offset"]
+    size = forced_collisions["threshold_weights_length"]
+    end = start + size
+    data[end - 1] = value
+
+
+@njit
+def threshold_weights_chunk(start, length, forced_collisions, data, value):
+    start += forced_collisions["threshold_weights_offset"]
+    end = start + length
+    data[start:end] = value
+
+
+@njit
+def target_weights(index, forced_collisions, data, value):
+    offset = forced_collisions["target_weights_offset"]
+    data[offset + index] = value
+
+
+@njit
+def target_weights_all(forced_collisions, data, value):
+    start = forced_collisions["target_weights_offset"]
+    size = forced_collisions["target_weights_length"]
+    end = start + size
+    data[start:end] = value
+
+
+@njit
+def target_weights_last(forced_collisions, data, value):
+    start = forced_collisions["target_weights_offset"]
+    size = forced_collisions["target_weights_length"]
+    end = start + size
+    data[end - 1] = value
+
+
+@njit
+def target_weights_chunk(start, length, forced_collisions, data, value):
+    start += forced_collisions["target_weights_offset"]
+    end = start + length
+    data[start:end] = value