Advanced (iterative) LB initialization for scenarios

lbstep.py

@@ -4,6 +4,7 @@ import numpy as np
 from lbmpy.boundaries.boundaryhandling import LatticeBoltzmannBoundaryHandling
 from lbmpy.creationfunctions import switch_to_symbolic_relaxation_rates_for_omega_adapting_methods, \
     create_lb_function, update_with_default_parameters
+from lbmpy.macroscopic_value_kernels import create_advanced_velocity_setter_collision_rule
 from lbmpy.simplificationfactory import create_simplification_strategy
 from lbmpy.stencils import get_stencil
 from pystencils.datahandling.serial_datahandling import SerialDataHandling
@@ -125,6 +126,9 @@ class LatticeBoltzmannStep:
         self._vtk_writer = None
         self.time_steps_run = 0
 
+        self._velocity_init_kernel = None
+        self._velocity_init_vel_backup = None
+
     @property
     def boundary_handling(self):
         """Boundary handling instance of the scenario. Use this to change the boundary setup"""
@@ -248,6 +252,85 @@ class LatticeBoltzmannStep:
     def write_vtk(self):
         self.vtk_writer(self.time_steps_run)
 
+    def run_iterative_initialization(self, velocity_relaxation_rate=1.0, convergence_threshold=1e-5, max_steps=5000,
+                                     check_residuum_after=100):
+        """Runs Advanced initialization of velocity field through iteration procedure.
+
+        Usually the pdfs are initialized in equilibrium with given density and velocity. Higher order moments are
+        set to their equilibrium values. This routine also initializes the higher order moments and the density field
+        using an iterative routine. For scenarios with high relaxation rates this might take long to converge.
+        For details, see Mei, Luo, Lallemand and Humieres: Consistent initial conditions for LBM simulations, 2005.
+
+        Args:
+            velocity_relaxation_rate: relaxation rate for the velocity moments - determines convergence behaviour
+                                      of the initialization scheme, should be in the range of the other relaxation
+                                      rate(s) otherwise the scheme could get unstable
+            convergence_threshold: The residuum is computed as average difference between prescribed and calculated
+                                   velocity field. If (residuum < convergence_threshold) the function returns
+                                   successfully.
+            max_steps: stop if not converged after this number of steps
+            check_residuum_after: the residuum criterion is tested after this number of steps
+
+        Returns:
+            tuple (residuum, steps_run) if successful or raises ValueError if not converged
+        """
+        dh = self.data_handling
+
+        def on_first_call():
+            self._velocity_init_vel_backup = 'velocity_init_vel_backup'
+            dh.add_array_like(self._velocity_init_vel_backup, self.velocity_data_name)
+            velocity_field = dh.fields[self.velocity_data_name]
+            collision_rule = create_advanced_velocity_setter_collision_rule(self.method, velocity_field,
+                                                                            velocity_relaxation_rate)
+
+            kernel = create_lb_function(collision_rule=collision_rule, field_name=self._pdf_arr_name,
+                                        temporary_field_name=self._tmp_arr_name,
+                                        optimization={'symbolic_field': dh.fields[self._pdf_arr_name]})
+            self._velocity_init_kernel = kernel
+
+        def make_velocity_backup():
+            for b in dh.iterate():
+                np.copyto(b[self._velocity_init_vel_backup], b[self.velocity_data_name])
+
+        def restore_velocity_backup():
+            for b in dh.iterate():
+                np.copyto(b[self.velocity_data_name], b[self._velocity_init_vel_backup])
+
+        def compute_residuum():
+            residuum = 0
+            for b in dh.iterate(ghost_layers=False, inner_ghost_layers=False):
+                residuum = np.average(np.abs(b[self._velocity_init_vel_backup] - b[self.velocity_data_name]))
+            reduce_result = dh.reduce_float_sequence([residuum, 1.0], 'sum', all_reduce=True)
+            return reduce_result[0] / reduce_result[1]
+
+        if self._velocity_init_kernel is None:
+            on_first_call()
+
+        make_velocity_backup()
+        outer_iterations = max_steps // check_residuum_after
+        global_residuum = None
+        steps_run = 0
+        for outer_iteration in range(outer_iterations):
+            for i in range(check_residuum_after):
+                steps_run += 1
+                self._sync()
+                self._boundary_handling(**self.kernel_params)
+                self._data_handling.run_kernel(self._velocity_init_kernel, **self.kernel_params)
+                self._data_handling.swap(self._pdf_arr_name, self._tmp_arr_name)
+            self._data_handling.run_kernel(self._getterKernel, **self.kernel_params)
+            global_residuum = compute_residuum()
+            if np.isnan(global_residuum) or global_residuum < convergence_threshold:
+                break
+
+        assert global_residuum is not None
+        converged = global_residuum < convergence_threshold
+        if not converged:
+            restore_velocity_backup()
+            raise ValueError("Iterative initialization did not converge after %d steps.\n"
+                             "Current residuum is %s" % (steps_run, global_residuum))
+
+        return global_residuum, steps_run
+
     def _compile_macroscopic_setter_and_getter(self):
         lb_method = self.method
         dim = lb_method.dim

macroscopic_value_kernels.py

@@ -105,7 +105,8 @@ def compile_macroscopic_values_setter(lb_method, quantities_to_set, pdf_arr=None
             fixed_kernel_parameters[quantity_name] = value
             at_least_one_field_input = True
             num_components = cqc.conserved_quantities[quantity_name]
-            field = Field.create_from_numpy_array(quantity_name, value, index_dimensions=0 if num_components <= 1 else 1)
+            field = Field.create_from_numpy_array(quantity_name, value,
+                                                  index_dimensions=0 if num_components <= 1 else 1)
             if num_components == 1:
                 value = field(0)
             else:
@@ -143,52 +144,39 @@ def compile_macroscopic_values_setter(lb_method, quantities_to_set, pdf_arr=None
     return setter
 
 
-def create_advanced_velocity_setter_collision_rule(lb_method, velocity_array, velocity_relaxation_rate=0.8):
+def create_advanced_velocity_setter_collision_rule(method, velocity_field: Field, velocity_relaxation_rate=0.8):
+    """Advanced initialization of velocity field through iteration procedure.
 
-    velocity_field = Field.create_from_numpy_array('vel_input', velocity_array, index_dimensions=1)
+    by Mei, Luo, Lallemand and Humieres: Consistent initial conditions for LBM simulations, 2005
 
-    cqc = lb_method.conserved_quantity_computation
+    Args:
+        method: lattice boltzmann method object
+        velocity_field: pystencils field
+        velocity_relaxation_rate: relaxation rate for the velocity moments - determines convergence behaviour
+                                of the initialization scheme
+
+    Returns:
+        LB collision rule
+    """
+    cqc = method.conserved_quantity_computation
     density_symbol = cqc.defined_symbols(order=0)[1]
     velocity_symbols = cqc.defined_symbols(order=1)[1]
 
     # density is computed from pdfs
-    eq_input_from_pdfs = cqc.equilibrium_input_equations_from_pdfs(lb_method.pre_collision_pdf_symbols)
+    eq_input_from_pdfs = cqc.equilibrium_input_equations_from_pdfs(method.pre_collision_pdf_symbols)
     eq_input_from_pdfs = eq_input_from_pdfs.new_filtered([density_symbol])
     # velocity is read from input field
-    vel_symbols = [velocity_field(i) for i in range(lb_method.dim)]
+    vel_symbols = [velocity_field(i) for i in range(method.dim)]
     eq_input_from_field = cqc.equilibrium_input_equations_from_init_values(velocity=vel_symbols)
     eq_input_from_field = eq_input_from_field.new_filtered(velocity_symbols)
     # then both are merged together
     eq_input = eq_input_from_pdfs.new_merged(eq_input_from_field)
 
     # set first order relaxation rate
-    lb_method = deepcopy(lb_method)
-    lb_method.set_first_moment_relaxation_rate(velocity_relaxation_rate)
+    method = deepcopy(method)
+    method.set_first_moment_relaxation_rate(velocity_relaxation_rate)
 
-    simplification_strategy = create_simplification_strategy(lb_method)
-    new_collision_rule = simplification_strategy(lb_method.get_collision_rule(eq_input))
+    simplification_strategy = create_simplification_strategy(method)
+    new_collision_rule = simplification_strategy(method.get_collision_rule(eq_input))
 
     return new_collision_rule
-
-
-def compile_advanced_velocity_setter(method, velocity_array, velocity_relaxation_rate=0.8, pdf_arr=None,
-                                     field_layout='numpy', optimization={}):
-    """
-    Advanced initialization of velocity field through iteration procedure according to
-    Mei, Luo, Lallemand and Humieres: Consistent initial conditions for LBM simulations, 2005
-
-    :param method:
-    :param velocity_array: array with velocity field
-    :param velocity_relaxation_rate: relaxation rate for the velocity moments - determines convergence behaviour
-                                   of the initialization scheme
-    :param pdf_arr: optional numpy array for pdf field - used to get optimal loop structure for kernel
-    :param field_layout: layout of the pdf field if pdf_arr was not given
-    :param optimization: dictionary with optimization hints
-    :return: stream-collide update function
-    """
-    from lbmpy.updatekernels import create_stream_pull_collide_kernel
-    from lbmpy.creationfunctions import create_lb_ast, create_lb_function
-    new_collision_rule = create_advanced_velocity_setter_collision_rule(method, velocity_array, velocity_relaxation_rate)
-    update_rule = create_stream_pull_collide_kernel(new_collision_rule, pdf_arr, generic_layout=field_layout)
-    ast = create_lb_ast(update_rule, optimization)
-    return create_lb_function(ast, optimization)

session.py

+from pystencils.sympy_gmpy_bug_workaround import *
 import sympy as sp
 import numpy as np
 from lbmpy.scenarios import *