Clean submodules for data handling and finite differences

chapman_enskog/chapman_enskog.py

@@ -2,10 +2,9 @@ import sympy as sp
 from collections import namedtuple
 from sympy.core.cache import cacheit
 from pystencils.cache import disk_cache
-from pystencils.derivative import full_diff_expand
+from pystencils.fd import expand_diff_full, Diff, DiffOperator, expand_diff_linear, \
+    expand_diff_products, normalize_diff_order
 from pystencils.sympyextensions import normalize_product, symmetric_product
-from pystencils.derivative import Diff, DiffOperator, expand_using_linearity, \
-    expand_using_product_rule, normalize_diff_order
 from lbmpy.moments import discrete_moment, moment_matrix, polynomial_to_exponent_representation, get_moment_indices, \
     non_aliased_moment
 from lbmpy.chapman_enskog.derivative import chapman_enskog_derivative_recombination, chapman_enskog_derivative_expansion
@@ -38,14 +37,14 @@ class ChapmanEnskogAnalysis:
 
         self.constants = constants
 
-        o_eps_moments1 = [expand_using_linearity(self._take_and_insert_moments(self.equations_by_order[1] * moment),
-                                                 constants=constants)
+        o_eps_moments1 = [expand_diff_linear(self._take_and_insert_moments(self.equations_by_order[1] * moment),
+                                             constants=constants)
                           for moment in moments_until_order1]
-        o_eps_moments2 = [expand_using_linearity(self._take_and_insert_moments(self.equations_by_order[1] * moment),
-                                                 constants=constants)
+        o_eps_moments2 = [expand_diff_linear(self._take_and_insert_moments(self.equations_by_order[1] * moment),
+                                             constants=constants)
                           for moment in moments_order2]
-        o_eps_sq_moments1 = [expand_using_linearity(self._take_and_insert_moments(self.equations_by_order[2] * moment),
-                                                    constants=constants)
+        o_eps_sq_moments1 = [expand_diff_linear(self._take_and_insert_moments(self.equations_by_order[2] * moment),
+                                                constants=constants)
                              for moment in moments_until_order1]
 
         self._equationsWithHigherOrderMoments = [self._ce_recombine(ord1 * self.epsilon + ord2 * self.epsilon ** 2)
@@ -63,7 +62,7 @@ class ChapmanEnskogAnalysis:
 
     def get_macroscopic_equations(self, substitute_higher_order_moments=False):
         if substitute_higher_order_moments:
-            return [full_diff_expand(e.subs(self.higher_order_moments), constants=self.constants)
+            return [expand_diff_full(e.subs(self.higher_order_moments), constants=self.constants)
                     for e in self._equationsWithHigherOrderMoments]
         else:
             return self._equationsWithHigherOrderMoments
@@ -383,7 +382,7 @@ def take_moments(eqn, pdf_to_moment_name=(('f', '\Pi'), ('\Omega f', '\\Upsilon'
         return result
 
     functions = sum(pdf_symbols, ())
-    eqn = expand_using_linearity(eqn, functions).expand()
+    eqn = expand_diff_linear(eqn, functions).expand()
 
     if eqn.func == sp.Mul:
         return handle_product(eqn)
@@ -401,7 +400,7 @@ def moment_selector(eq):
 
 
 def diff_expand_normalizer(eq):
-    return expand_using_product_rule(eq).expand()
+    return expand_diff_products(eq).expand()
 
 
 def chain_solve_and_substitute(assignments, unknown_selector, normalizing_func=diff_expand_normalizer):
@@ -487,12 +486,12 @@ def get_taylor_expanded_lb_equation(pdf_symbol_name="f", pdfs_after_collision_op
 
     functions = [pdf, collided_pdf]
     eq_4_5 = taylor_operator - dt * collided_pdf
-    applied_eq_4_5 = expand_using_linearity(DiffOperator.apply(eq_4_5, pdf, apply_to_constants=False), functions)
+    applied_eq_4_5 = expand_diff_linear(DiffOperator.apply(eq_4_5, pdf, apply_to_constants=False), functions)
 
     if shift:
         operator = ((dt / 2) * (dt_operator + c.dot(dx_operator))).expand()
-        op_times_eq_4_5 = expand_using_linearity(DiffOperator.apply(operator, applied_eq_4_5, apply_to_constants=False),
-                                                 functions).expand()
+        op_times_eq_4_5 = expand_diff_linear(DiffOperator.apply(operator, applied_eq_4_5, apply_to_constants=False),
+                                             functions).expand()
         op_times_eq_4_5 = normalize_diff_order(op_times_eq_4_5, functions)
         eq_4_7 = (applied_eq_4_5 - op_times_eq_4_5).subs(dt ** (taylor_order + 1), 0)
     else:
@@ -543,7 +542,7 @@ def chapman_enskog_ansatz(equation, time_derivative_orders=(1, 3), spatial_deriv
                                     for i in range(start_order, stop_order))
         max_expansion_order = max(max_expansion_order, stop_order)
     equation = equation.subs(subs_dict)
-    equation = expand_using_linearity(equation, functions=expanded_pdf_symbols).expand().collect(eps)
+    equation = expand_diff_linear(equation, functions=expanded_pdf_symbols).expand().collect(eps)
     result = {eps_order: equation.coeff(eps ** eps_order) for eps_order in range(1, 2 * max_expansion_order)}
     result[0] = equation.subs(eps, 0)
     return result
@@ -557,7 +556,7 @@ def match_to_navier_stokes(conservation_equations, rho=sp.Symbol("rho"), u=sp.sy
     def diff_simplify(eq):
         variables = eq.atoms(CeMoment)
         variables.update(funcs)
-        return expand_using_product_rule(expand_using_linearity(eq, variables)).expand()
+        return expand_diff_products(expand_diff_linear(eq, variables)).expand()
 
     def match_continuity_eq(continuity_eq):
         continuity_eq = diff_simplify(continuity_eq)

chapman_enskog/chapman_enskog_higher_order.py

 import sympy as sp
-from pystencils.derivative import full_diff_expand
+from pystencils.fd import expand_diff_full
 from lbmpy.chapman_enskog.chapman_enskog import CeMoment, Diff, take_moments
 from lbmpy.chapman_enskog.chapman_enskog import expanded_symbol
 from lbmpy.moments import moments_up_to_order, moments_of_order
@@ -61,7 +61,7 @@ def determine_higher_order_moments(epsilon_hierarchy, relaxation_rates, moment_c
     solvability_conditions = get_solvability_conditions(dim, order)
 
     def full_expand(expr):
-        return full_diff_expand(expr, constants=relaxation_rates)
+        return expand_diff_full(expr, constants=relaxation_rates)
 
     def tm(expr):
         expr = take_moments(expr, use_one_neighborhood_aliasing=True)

chapman_enskog/chapman_enskog_steady_state.py

 import sympy as sp
 import functools
-from pystencils.derivative import Diff, DiffOperator, expand_using_linearity, normalize_diff_order
-from pystencils.derivative import collect_derivatives, create_nested_diff
+from pystencils.fd import Diff, DiffOperator, expand_diff_linear, normalize_diff_order, \
+    collect_diffs, create_nested_diff
 from pystencils.sympyextensions import normalize_product, multidimensional_sum, kronecker_delta
 from lbmpy.chapman_enskog.chapman_enskog import LbMethodEqMoments, CeMoment, take_moments, insert_moments
 from lbmpy.chapman_enskog.chapman_enskog import expanded_symbol, chapman_enskog_ansatz, remove_higher_order_u
@@ -93,7 +93,7 @@ class SteadyStateChapmanEnskogAnalysis:
         for ce_eq, f_i in zip(chapman_enskog_hierarchy, self.f_syms):
             new_eq = -1 / self.collision_op_sym * (ce_eq - self.collision_op_sym * f_i)
             raw_hierarchy.append(new_eq)
-            new_eq = expand_using_linearity(new_eq.subs(substitution_dict), functions=self.f_syms + [self.force_sym])
+            new_eq = expand_diff_linear(new_eq.subs(substitution_dict), functions=self.f_syms + [self.force_sym])
             if new_eq:
                 substitution_dict[f_i] = new_eq
             inserted_hierarchy.append(new_eq)
@@ -163,7 +163,7 @@ class SteadyStateChapmanEnskogAnalysis:
 
             new_products.append(new_prod)
 
-        return normalize_diff_order(expand_using_linearity(sp.Add(*new_products), functions=self.physical_variables))
+        return normalize_diff_order(expand_diff_linear(sp.Add(*new_products), functions=self.physical_variables))
 
 
 # ----------------------------------------------------------------------------------------------------------------------
@@ -197,7 +197,7 @@ class SteadyStateChapmanEnskogAnalysisSRT:
 
         for i in range(2, len(self.scale_hierarchy)):
             eq = self.scale_hierarchy[i].subs(subs_dict)
-            eq = expand_using_linearity(eq, functions=expanded_pdf_symbols)
+            eq = expand_diff_linear(eq, functions=expanded_pdf_symbols)
             eq = normalize_diff_order(eq, functions=expanded_pdf_symbols)
             subs_dict[expanded_pdf_symbols[i]] = eq
             expanded_pdfs.append(eq)
@@ -230,10 +230,10 @@ class SteadyStateChapmanEnskogAnalysisSRT:
         # Continuity equation (mass transport)
         cont_eq = take_moments(recombined, max_expansion=(order + 1) * 2)
         cont_eq = handle_postcollision_values(cont_eq)
-        cont_eq = expand_using_linearity(cont_eq, constants=constants).expand().collect(dt)
+        cont_eq = expand_diff_linear(cont_eq, constants=constants).expand().collect(dt)
         self.continuity_equation_with_moments = cont_eq
         cont_eq = insert_moments(cont_eq, moment_computation, use_solvability_conditions=False)
-        cont_eq = expand_using_linearity(cont_eq, constants=constants).expand().collect(dt)
+        cont_eq = expand_diff_linear(cont_eq, constants=constants).expand().collect(dt)
         self.continuity_equation = cont_eq
 
         # Momentum equation (momentum transport)
@@ -242,10 +242,10 @@ class SteadyStateChapmanEnskogAnalysisSRT:
         for h in range(dim):
             mom_eq = take_moments(recombined * c[h], max_expansion=(order + 1) * 2)
             mom_eq = handle_postcollision_values(mom_eq)
-            mom_eq = expand_using_linearity(mom_eq, constants=constants).expand().collect(dt)
+            mom_eq = expand_diff_linear(mom_eq, constants=constants).expand().collect(dt)
             self.momentum_equations_with_moments.append(mom_eq)
             mom_eq = insert_moments(mom_eq, moment_computation, use_solvability_conditions=False)
-            mom_eq = expand_using_linearity(mom_eq, constants=constants).expand().collect(dt)
+            mom_eq = expand_diff_linear(mom_eq, constants=constants).expand().collect(dt)
             self.momentum_equations.append(mom_eq)
 
     def get_continuity_equation(self, order):
@@ -253,14 +253,14 @@ class SteadyStateChapmanEnskogAnalysisSRT:
             result = self.continuity_equation.subs(self.dt, 0)
         else:
             result = self.continuity_equation.coeff(self.dt ** order)
-        return collect_derivatives(result)
+        return collect_diffs(result)
 
     def get_momentum_equation(self, coordinate, order):
         if order == 0:
             result = self.momentum_equations[coordinate].subs(self.dt, 0)
         else:
             result = self.momentum_equations[coordinate].coeff(self.dt ** order)
-        return collect_derivatives(result)
+        return collect_diffs(result)
 
     def determine_viscosities(self, coordinate):
         """Matches the first order term of the momentum equation to Navier stokes.
@@ -288,7 +288,7 @@ class SteadyStateChapmanEnskogAnalysisSRT:
 
         first_order_terms = self.get_momentum_equation(coordinate, order=1)
         first_order_terms = remove_higher_order_u(first_order_terms)
-        first_order_terms = expand_using_linearity(first_order_terms, constants=[sp.Symbol("rho")])
+        first_order_terms = expand_diff_linear(first_order_terms, constants=[sp.Symbol("rho")])
 
         match_coeff_equations = []
         for diff in navier_stokes_ref.atoms(Diff):

chapman_enskog/derivative.py

 import sympy as sp
-from pystencils.derivative import Diff
+from pystencils.fd import Diff
 
 
 def chapman_enskog_derivative_expansion(expr, label, eps=sp.Symbol("epsilon"), start_order=1, stop_order=4):

lbstep.py

@@ -7,8 +7,7 @@ from lbmpy.creationfunctions import switch_to_symbolic_relaxation_rates_for_omeg
 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
-from pystencils import create_kernel, make_slice
+from pystencils import create_kernel, make_slice, create_data_handling
 from pystencils.slicing import SlicedGetter
 from pystencils.timeloop import TimeLoop
 
@@ -30,7 +29,8 @@ class LatticeBoltzmannStep:
         if data_handling is None:
             if domain_size is None:
                 raise ValueError("Specify either domain_size or data_handling")
-            data_handling = SerialDataHandling(domain_size, default_ghost_layers=1, periodicity=periodicity)
+            data_handling = create_data_handling(domain_size, default_ghost_layers=1,
+                                                 periodicity=periodicity, parallel=False)
 
         if 'stencil' not in method_parameters:
             method_parameters['stencil'] = 'D2Q9' if data_handling.dim == 2 else 'D3Q27'

phasefield/analytical.py

@@ -2,7 +2,7 @@ import sympy as sp
 from collections import defaultdict
 
 from pystencils.sympyextensions import multidimensional_sum as multi_sum, normalize_product, prod
-from pystencils.derivative import functional_derivative, expand_using_linearity, Diff, full_diff_expand
+from pystencils.fd import functional_derivative, expand_diff_linear, Diff, expand_diff_full
 
 order_parameter_symbol_name = "phi"
 surface_tension_symbol_name = "tau"
@@ -53,7 +53,7 @@ def free_energy_functional_3_phases(order_parameters=None, interface_width=inter
 
     f = f.subs(order_param_to_concentration_relation)
     if expand_derivatives:
-        f = expand_using_linearity(f, functions=order_parameters)
+        f = expand_diff_linear(f, functions=order_parameters)
 
     return f, transformation_matrix
 
@@ -198,7 +198,7 @@ def chemical_potentials_from_free_energy(free_energy, order_parameters=None):
         order_parameters.sort(key=lambda e: e.name)
         order_parameters = order_parameters[:-1]
     constants = [s for s in symbols if s not in order_parameters]
-    return sp.Matrix([expand_using_linearity(functional_derivative(free_energy, op), constants=constants)
+    return sp.Matrix([expand_diff_linear(functional_derivative(free_energy, op), constants=constants)
                       for op in order_parameters])
 
 
@@ -219,7 +219,7 @@ def substitute_laplacian_by_sum(eq, dim):
     functions = [d.args[0] for d in eq.atoms(Diff)]
     substitutions = {Diff(Diff(op)): sum(Diff(Diff(op, i), i) for i in range(dim))
                      for op in functions}
-    return full_diff_expand(eq.subs(substitutions))
+    return expand_diff_full(eq.subs(substitutions))
 
 
 def cosh_integral(f, var):
@@ -306,7 +306,7 @@ def force_from_pressure_tensor(pressure_tensor, functions=None):
 
     def force_component(b):
         r = -sum(Diff(pressure_tensor[a, b], a) for a in range(dim))
-        r = full_diff_expand(r, functions=functions)
+        r = expand_diff_full(r, functions=functions)
         return r
 
     return sp.Matrix([force_component(b) for b in range(dim)])

phasefield/experiments1D.py

@@ -2,7 +2,7 @@ import numpy as np
 import sympy as sp
 
 from pystencils import make_slice
-from pystencils.derivative import Diff
+from pystencils.fd import Diff
 
 
 def plot_status(phase_field_step, from_x=None, to_x=None):

phasefield/kerneleqs.py

 import sympy as sp
 from pystencils import Assignment
-from pystencils.finitedifferences import Discretization2ndOrder
+from pystencils.fd import Discretization2ndOrder
 from lbmpy.phasefield.analytical import chemical_potentials_from_free_energy, substitute_laplacian_by_sum, \
     force_from_phi_and_mu, symmetric_tensor_linearization, pressure_tensor_from_free_energy, force_from_pressure_tensor
 
@@ -58,7 +58,7 @@ def force_kernel_using_pressure_tensor(force_field, pressure_tensor_field, extra
 
 
 def cahn_hilliard_fd_eq(phase_idx, phi, mu, velocity, mobility, dx, dt):
-    from pystencils.finitedifferences import transient, advection, diffusion
+    from pystencils.fd.finitedifferences import transient, advection, diffusion
     cahn_hilliard = transient(phi, phase_idx) + advection(phi, velocity, phase_idx) - diffusion(mu, mobility, phase_idx)
     return Discretization2ndOrder(dx, dt)(cahn_hilliard)
 

phasefield/phasefieldstep.py

@@ -7,11 +7,10 @@ from lbmpy.lbstep import LatticeBoltzmannStep
 from lbmpy.phasefield.cahn_hilliard_lbm import cahn_hilliard_lb_method
 from lbmpy.phasefield.kerneleqs import mu_kernel, CahnHilliardFDStep, pressure_tensor_kernel, \
     force_kernel_using_pressure_tensor
-from pystencils import create_kernel
+from pystencils import create_kernel, create_data_handling
 from lbmpy.phasefield.analytical import chemical_potentials_from_free_energy, symmetric_tensor_linearization
 from pystencils.boundaries.boundaryhandling import FlagInterface
 from pystencils.boundaries.inkernel import add_neumann_boundary
-from pystencils.datahandling import SerialDataHandling
 from pystencils.assignment_collection.simplifications import sympy_cse_on_assignment_list
 from pystencils.slicing import make_slice, SlicedGetter
 
@@ -32,7 +31,7 @@ class PhaseFieldStep:
         target = optimization.get('target', 'cpu')
 
         if data_handling is None:
-            data_handling = SerialDataHandling(domain_size, periodicity=True)
+            data_handling = create_data_handling(domain_size, periodicity=True, parallel=False)
 
         self.free_energy = free_energy
         self.concentration_to_order_parameter = concentration_to_order_parameters

scenarios.py

@@ -59,8 +59,8 @@ def create_fully_periodic_flow(initial_velocity, periodicity_in_kernel=False, lb
         kwargs['optimization']['builtin_periodicity'] = (True, True, True)
 
     if data_handling is None:
-        data_handling = create_data_handling(parallel, domain_size, periodicity=not periodicity_in_kernel,
-                                             default_ghost_layers=1)
+        data_handling = create_data_handling(domain_size, periodicity=not periodicity_in_kernel,
+                                             default_ghost_layers=1, parallel=parallel)
     step = LatticeBoltzmannStep(data_handling=data_handling, name="periodic_scenario", lbm_kernel=lbm_kernel, **kwargs)
     for b in step.data_handling.iterate(ghost_layers=False):
         np.copyto(b[step.velocity_data_name], initial_velocity[b.global_slice])
@@ -84,7 +84,7 @@ def create_lid_driven_cavity(domain_size=None, lid_velocity=0.005, lbm_kernel=No
     """
     assert domain_size is not None or data_handling is not None
     if data_handling is None:
-        data_handling = create_data_handling(parallel, domain_size, periodicity=False, default_ghost_layers=1)
+        data_handling = create_data_handling(domain_size, periodicity=False, default_ghost_layers=1, parallel=parallel)
     step = LatticeBoltzmannStep(data_handling=data_handling, lbm_kernel=lbm_kernel, name="ldc", **kwargs)
 
     my_ubb = UBB(velocity=[lid_velocity, 0, 0][:step.method.dim])
@@ -99,22 +99,22 @@ def create_channel(domain_size=None, force=None, pressure_difference=None, u_max
                    duct=False, wall_boundary=NoSlip(), parallel=False, data_handling=None, **kwargs):
     """Create a channel scenario (2D or 3D).
     
-    :param domain_size: size of the simulation domain. First coordinate is the flow direction.
-    
-    The channel can be driven by one of the following methods. Please specify exactly one of the following parameters: 
-    :param force: Periodic channel, driven by a body force. Pass force in flow direction in lattice units here.
-    :param pressure_difference: Inflow and outflow are fixed pressure conditions, with the given pressure difference. 
-    :param u_max: Parabolic velocity profile prescribed at inflow, pressure boundary =1.0 at outflow.
-    
-    Geometry parameters:
-    :param diameter_callback: optional callback for channel with varying diameters. Only valid if duct=False.
-                             The callback receives x coordinate array and domain_size and returns a
-                             an array of diameters of the same shape
-    :param duct: if true the channel has rectangular instead of circular cross section
-    :param wall_boundary: instance of boundary class that should be set at the channel walls
-    :param parallel: True for distributed memory parallelization with walberla
-    :param data_handling: see documentation of :func:`create_fully_periodic_flow`
-    :param kwargs: all other keyword parameters are passed directly to scenario class.
+    The channel can be driven either by force, velocity inflow or pressure difference. Choose one and pass
+    exactly one of the parameters 'force', 'pressure_difference' or 'u_max'.
+
+    Args:
+        domain_size: size of the simulation domain. First coordinate is the flow direction.
+        force: Periodic channel, driven by a body force. Pass force in flow direction in lattice units here.
+        pressure_difference: Inflow and outflow are fixed pressure conditions, with the given pressure difference.
+        u_max: Parabolic velocity profile prescribed at inflow, pressure boundary =1.0 at outflow.
+        diameter_callback: optional callback for channel with varying diameters. Only valid if duct=False.
+                          The callback receives x coordinate array and domain_size and returns a
+                          an array of diameters of the same shape
+        duct: if true the channel has rectangular instead of circular cross section
+        wall_boundary: instance of boundary class that should be set at the channel walls
+        parallel: True for distributed memory parallelization with walberla
+        data_handling: see documentation of :func:`create_fully_periodic_flow`
+        kwargs: all other keyword parameters are passed directly to scenario class.
     """
     assert domain_size is not None or data_handling is not None
 
@@ -125,8 +125,8 @@ def create_channel(domain_size=None, force=None, pressure_difference=None, u_max
     if data_handling is None:
         dim = len(domain_size)
         assert dim in (2, 3)
-        data_handling = create_data_handling(parallel, domain_size, periodicity=periodicity[:dim],
-                                             default_ghost_layers=1)
+        data_handling = create_data_handling(domain_size, periodicity=periodicity[:dim],
+                                             default_ghost_layers=1, parallel=parallel)
 
     dim = data_handling.dim
     if force: