Merge branch 'oldroydb' into 'master'

Oldroyd-B

Closes walberla/walberla#158

See merge request pycodegen/lbmpy!103

lbmpy/oldroydb.py

+import pystencils as ps
+import sympy as sp
+import numpy as np
+
+from pystencils.boundaries.boundaryconditions import Boundary
+from pystencils.stencil import inverse_direction_string, direction_string_to_offset
+
+
+class OldroydB:
+    def __init__(self, dim, u, tau, F, tauflux, tauface, lambda_p, eta_p, vof=True):
+        assert not ps.FieldType.is_staggered(u)
+        assert not ps.FieldType.is_staggered(tau)
+        assert not ps.FieldType.is_staggered(F)
+        assert ps.FieldType.is_staggered(tauflux)
+        assert ps.FieldType.is_staggered(tauface)
+        
+        self.dim = dim
+        self.u = u
+        self.tau = tau
+        self.F = F
+        self.tauflux = tauflux
+        self.tauface_field = tauface
+        self.lambda_p = lambda_p
+        self.eta_p = eta_p
+        
+        full_stencil = ["C"] + self.tauflux.staggered_stencil + \
+            list(map(inverse_direction_string, self.tauflux.staggered_stencil))
+        self.stencil = tuple(map(lambda d: tuple(ps.stencil.direction_string_to_offset(d, self.dim)), full_stencil))
+        full_stencil = ["C"] + self.tauface_field.staggered_stencil + \
+            list(map(inverse_direction_string, self.tauface_field.staggered_stencil))
+        self.force_stencil = tuple(map(lambda d: tuple(ps.stencil.direction_string_to_offset(d, self.dim)),
+                                   full_stencil))
+        
+        self.disc = ps.fd.FVM1stOrder(self.tau, self._flux(), self._source())
+        if vof:
+            self.vof = ps.fd.VOF(self.tauflux, self.u, self.tau)
+        else:
+            self.vof = None
+    
+    def _flux(self):
+        return [self.tau.center_vector.applyfunc(lambda t: t * self.u.center_vector[i]) for i in range(self.dim)]
+    
+    def _source(self):
+        gradu = sp.Matrix([[ps.fd.diff(self.u.center_vector[j], i) for j in range(self.dim)] for i in range(self.dim)])
+        gamma = gradu + gradu.transpose()
+        return self.tau.center_vector * gradu + gradu.transpose() * self.tau.center_vector + \
+            (self.eta_p * gamma - self.tau.center_vector) / self.lambda_p
+    
+    def tauface(self):
+        return ps.AssignmentCollection([ps.Assignment(self.tauface_field.staggered_vector_access(d),
+                                        (self.tau.center_vector + self.tau.neighbor_vector(d)) / 2)
+                                        for d in self.tauface_field.staggered_stencil])
+    
+    def force(self):
+        full_stencil = self.tauface_field.staggered_stencil + \
+            list(map(inverse_direction_string, self.tauface_field.staggered_stencil))
+        dtau = sp.Matrix([sum([sum([
+                          self.tauface_field.staggered_access(d, (i, j)) * direction_string_to_offset(d)[i]
+                          for i in range(self.dim)]) / sp.Matrix(direction_string_to_offset(d)).norm()
+                         for d in full_stencil]) for j in range(self.dim)])
+        A0 = sum([sp.Matrix(direction_string_to_offset(d)).norm() for d in full_stencil])
+        return ps.AssignmentCollection(ps.Assignment(self.F.center_vector, dtau / A0 * 2 * self.dim))
+    
+    def flux(self):
+        if self.vof is not None:
+            return self.vof
+        else:
+            return self.disc.discrete_flux(self.tauflux)
+    
+    def continuity(self):
+        cont = self.disc.discrete_continuity(self.tauflux)
+        tau_copy = sp.Matrix(self.dim, self.dim, lambda i, j: sp.Symbol("tau_old_%d_%d" % (i, j)))
+        tau_subs = {self.tau.center_vector[i, j]: tau_copy[i, j] for i in range(self.dim) for j in range(self.dim)}
+        return [ps.Assignment(tau_copy[i, j], self.tau.center_vector[i, j])
+                for i in range(self.dim) for j in range(self.dim)] + \
+               [ps.Assignment(a.lhs, a.rhs.subs(tau_subs)) for a in cont]
+
+
+class Flux(Boundary):
+    inner_or_boundary = True  # call the boundary condition with the fluid cell
+    single_link = False  # needs to be called for all directional fluxes
+    
+    def __init__(self, stencil, value=None):
+        self.stencil = stencil
+        self.value = value
+
+    def __call__(self, field, direction_symbol, **kwargs):
+        assert ps.FieldType.is_staggered(field)
+        
+        assert all([s == 0 for s in self.stencil[0]])
+        accesses = [field.staggered_vector_access(ps.stencil.offset_to_direction_string(d))
+                    for d in self.stencil[1:]]
+        conds = [sp.Equality(direction_symbol, d + 1) for d in range(len(accesses))]
+        
+        if self.value is None:
+            val = sp.Matrix(np.zeros(accesses[0].shape, dtype=int))
+        else:
+            val = self.value
+        
+        # use conditional
+        conditional = None
+        for a, c, d in zip(accesses, conds, self.stencil[1:]):
+            d = ps.stencil.offset_to_direction_string(d)
+            assignments = []
+            for i in range(len(a)):
+                fac = 1
+                if ps.FieldType.is_staggered_flux(field) and type(a[i]) is sp.Mul and a[i].args[0] == -1:
+                    fac = a[i].args[0]
+                    a[i] *= a[i].args[0]
+                assignments.append(ps.Assignment(a[i], fac * val[i]))
+            if len(assignments) > 0:
+                conditional = ps.astnodes.Conditional(ps.data_types.type_all_numbers(c, "int"),
+                                                      ps.astnodes.Block(assignments),
+                                                      conditional)
+        return [conditional]
+
+    def __hash__(self):
+        return hash((Flux, self.stencil, self.value))
+
+    def __eq__(self, other):
+        return type(other) == Flux and other.stencil == self.stencil and self.value == other.value
+
+
+class Extrapolation(Boundary):
+    inner_or_boundary = True  # call the boundary condition with the fluid cell
+    single_link = False  # needs to be called for all directional fluxes
+
+    def __init__(self, stencil, src_field, order):
+        self.stencil = stencil
+        self.src = src_field
+        if order == 0:
+            self.weights = (1,)
+        elif order == 1:
+            self.weights = (sp.Rational(3, 2), - sp.Rational(1, 2))
+        elif order == 2:
+            self.weights = (sp.Rational(15, 8), - sp.Rational(10, 8), sp.Rational(3, 8))
+        else:
+            raise NotImplementedError("weights are not known for extrapolation orders > 2")
+
+    def __call__(self, field, direction_symbol, **kwargs):
+        assert ps.FieldType.is_staggered(field)
+        
+        assert all([s == 0 for s in self.stencil[0]])
+        accesses = [field.staggered_vector_access(ps.stencil.offset_to_direction_string(d))
+                    for d in self.stencil[1:]]
+        conds = [sp.Equality(direction_symbol, d + 1) for d in range(len(accesses))]
+        
+        # use conditional
+        conditional = None
+        for a, c, o in zip(accesses, conds, self.stencil[1:]):
+            assignments = []
+            src = [self.src.neighbor_vector(tuple([-1 * n * i for i in o])) for n in range(len(self.weights))]
+            interp = self.weights[0] * src[0]
+            for i in range(1, len(self.weights)):
+                interp += self.weights[i] * src[i]
+            for i in range(len(a)):
+                fac = 1
+                if ps.FieldType.is_staggered_flux(field) and type(a[i]) is sp.Mul and a[i].args[0] == -1:
+                    fac = a[i].args[0]
+                    a[i] *= a[i].args[0]
+                assignments.append(ps.Assignment(a[i], fac * interp[i]))
+            if len(assignments) > 0:
+                conditional = ps.astnodes.Conditional(ps.data_types.type_all_numbers(c, "int"),
+                                                      ps.astnodes.Block(assignments),
+                                                      conditional)
+        return [conditional]
+
+    def __hash__(self):
+        return hash((Extrapolation, self.stencil, self.src, self.weights))
+
+    def __eq__(self, other):
+        return type(other) == Extrapolation and other.stencil == self.stencil and \
+            other.src == self.src and other.weights == self.weights
+
+
+class ForceOnBoundary(Boundary):
+    inner_or_boundary = False  # call the boundary condition with the boundary cell
+    single_link = False  # needs to be called for all directional fluxes
+    
+    def __init__(self, stencil, force_field):
+        self.stencil = stencil
+        self.force_field = force_field
+        
+        assert not ps.FieldType.is_staggered(force_field)
+
+    def __call__(self, face_stress_field, direction_symbol, **kwargs):
+        assert ps.FieldType.is_staggered(face_stress_field)
+        
+        assert all([s == 0 for s in self.stencil[0]])
+        accesses = [face_stress_field.staggered_vector_access(ps.stencil.offset_to_direction_string(d))
+                    for d in self.stencil[1:]]
+        conds = [sp.Equality(direction_symbol, d + 1) for d in range(len(accesses))]
+        
+        # use conditional
+        conditional = None
+        for a, c, o in zip(accesses, conds, self.stencil[1:]):
+            assignments = ps.Assignment(self.force_field.center_vector,
+                                        self.force_field.center_vector + 1 * a.transpose() * sp.Matrix(o))
+            conditional = ps.astnodes.Conditional(ps.data_types.type_all_numbers(c, "int"),
+                                                  ps.astnodes.Block(assignments),
+                                                  conditional)
+        return [conditional]
+
+    def __hash__(self):
+        return hash((ForceOnBoundary, self.stencil, self.force_field))
+
+    def __eq__(self, other):
+        return type(other) == ForceOnBoundary and other.stencil == self.stencil and \
+            other.force_field == self.force_field

lbmpy_tests/test_oldroydb.py

+import pystencils as ps
+from lbmpy.stencils import get_stencil
+from lbmpy.updatekernels import create_stream_pull_with_output_kernel
+from lbmpy import create_lb_update_rule, relaxation_rate_from_lattice_viscosity, ForceModel, Method, LBStencil
+from lbmpy.macroscopic_value_kernels import macroscopic_values_setter
+from pystencils.boundaries.boundaryhandling import BoundaryHandling
+from pystencils.boundaries.boundaryconditions import Boundary, Neumann, Dirichlet
+from lbmpy.boundaries.boundaryhandling import LatticeBoltzmannBoundaryHandling
+from lbmpy.boundaries import NoSlip
+
+from lbmpy.oldroydb import *
+import pytest
+
+
+# # Lattice Boltzmann with Finite-Volume Oldroyd-B
+# # taken from the electronic supplement of https://doi.org/10.1140/epje/s10189-020-00005-6,
+# # available at https://doi.org/10.24416/UU01-2AFZSW
+
+pytest.importorskip('scipy.optimize')
+def test_oldroydb():
+    import scipy.optimize
+
+    # ## Definitions
+
+    # In[2]:
+
+
+    L = (34, 34)
+    lambda_p = sp.Symbol("lambda_p")
+    eta_p = sp.Symbol("eta_p")
+
+    lb_stencil = LBStencil("D2Q9")
+    fv_stencil = LBStencil("D2Q9")
+    eta = 1-eta_p
+    omega = relaxation_rate_from_lattice_viscosity(eta)
+
+    f_pre = 0.00001
+
+
+    # ## Data structures
+
+    # In[3]:
+
+
+    dh = ps.create_data_handling(L, periodicity=(True, False), default_target=ps.Target.CPU)
+
+    opts = {'cpu_openmp': True, 
+            'cpu_vectorize_info': None,
+            'target': dh.default_target}
+
+    src = dh.add_array('src', values_per_cell=len(lb_stencil), layout='c')
+    dst = dh.add_array_like('dst', 'src')
+    ρ = dh.add_array('rho', layout='c', latex_name='\\rho')
+    u = dh.add_array('u', values_per_cell=dh.dim, layout='c')
+    tauface = dh.add_array('tau_face', values_per_cell=(len(fv_stencil)//2, dh.dim, dh.dim), latex_name='\\tau_f',
+                           field_type=ps.FieldType.STAGGERED, layout='c')
+
+    tau = dh.add_array('tau', values_per_cell=(dh.dim, dh.dim), layout='c', latex_name='\\tau')
+    tauflux = dh.add_array('j_tau', values_per_cell=(len(fv_stencil)//2, dh.dim, dh.dim), latex_name='j_\\tau',
+                           field_type=ps.FieldType.STAGGERED_FLUX, layout='c')
+    F = dh.add_array('F', values_per_cell=dh.dim, layout='c')
+
+    fluxbh = BoundaryHandling(dh, tauflux.name, fv_stencil, name="flux_boundary_handling",
+                              openmp=opts['cpu_openmp'], target=dh.default_target)
+    ubh = BoundaryHandling(dh, u.name, lb_stencil, name="velocity_boundary_handling",
+                           openmp=opts['cpu_openmp'], target=dh.default_target)
+    taufacebh = BoundaryHandling(dh, tauface.name, fv_stencil, name="tauface_boundary_handling",
+                                 openmp=opts['cpu_openmp'], target=dh.default_target)
+
+
+    # ## Solver
+
+    # In[4]:
+
+
+    collision = create_lb_update_rule(stencil=lb_stencil,
+                                      method=Method.TRT,
+                                      relaxation_rate=omega, 
+                                      compressible=True,
+                                      force_model=ForceModel.GUO, 
+                                      force=F.center_vector+sp.Matrix([f_pre,0]),
+                                      kernel_type='collide_only',
+                                      optimization={'symbolic_field': src})
+
+    stream = create_stream_pull_with_output_kernel(collision.method, src, dst, {'density': ρ, 'velocity': u})
+
+    lbbh = LatticeBoltzmannBoundaryHandling(collision.method, dh, src.name,
+                                            openmp=opts['cpu_openmp'], target=dh.default_target)
+
+    stream_kernel = ps.create_kernel(stream, **opts).compile()
+    collision_kernel = ps.create_kernel(collision, **opts).compile()
+
+
+    # In[5]:
+
+
+    ob = OldroydB(dh.dim, u, tau, F, tauflux, tauface, lambda_p, eta_p)
+    flux_kernel = ps.create_staggered_kernel(ob.flux(), **opts).compile()
+    tauface_kernel = ps.create_staggered_kernel(ob.tauface(), **opts).compile()
+    continuity_kernel = ps.create_kernel(ob.continuity(), **opts).compile()
+    force_kernel = ps.create_kernel(ob.force(), **opts).compile()
+
+
+    # ## Set up the simulation
+
+    # In[6]:
+
+
+    init = macroscopic_values_setter(collision.method, velocity=(0,)*dh.dim, 
+                                     pdfs=src.center_vector, density=ρ.center)
+    init_kernel = ps.create_kernel(init, ghost_layers=0).compile()
+
+    # no-slip for the fluid, no-flux for the stress
+    noslip = NoSlip()
+    noflux = Flux(fv_stencil)
+    nostressdiff = Flux(fv_stencil, tau.center_vector)
+
+    # put some good values into the boundaries so we can take derivatives
+    noforce = Neumann() # put the same stress into the boundary cells that is in the nearest fluid cell
+    noflow = Dirichlet((0,)*dh.dim) # put zero velocity into the boundary cells
+
+    lbbh.set_boundary(noslip, ps.make_slice[:, :4])
+    lbbh.set_boundary(noslip, ps.make_slice[:, -4:])
+    fluxbh.set_boundary(noflux, ps.make_slice[:, :4])
+    fluxbh.set_boundary(noflux, ps.make_slice[:, -4:])
+    ubh.set_boundary(noflow, ps.make_slice[:, :4])
+    ubh.set_boundary(noflow, ps.make_slice[:, -4:])
+    taufacebh.set_boundary(nostressdiff, ps.make_slice[:, :4])
+    taufacebh.set_boundary(nostressdiff, ps.make_slice[:, -4:])
+
+    for bh in lbbh, fluxbh, ubh, taufacebh:
+        assert len(bh._boundary_object_to_boundary_info) == 1, "Restart kernel to clear boundaries"
+
+    def init():
+        dh.fill(ρ.name, np.nan, ghost_layers=True, inner_ghost_layers=True)
+        dh.fill(ρ.name, 1)
+        dh.fill(u.name, np.nan, ghost_layers=True, inner_ghost_layers=True)
+        dh.fill(u.name, 0)
+        dh.fill(tau.name, np.nan, ghost_layers=True, inner_ghost_layers=True)
+        dh.fill(tau.name, 0)
+        dh.fill(tauflux.name, np.nan, ghost_layers=True, inner_ghost_layers=True)
+        dh.fill(tauface.name, np.nan, ghost_layers=True, inner_ghost_layers=True)
+        dh.fill(F.name, np.nan, ghost_layers=True, inner_ghost_layers=True)
+        dh.fill(F.name, 0) # needed for LB initialization
+    
+        sync_tau() # force calculation inside the initialization needs neighbor taus
+        dh.run_kernel(init_kernel)
+        dh.fill(F.name, np.nan)
+
+
+    # In[7]:
+
+
+    sync_pdfs = dh.synchronization_function([src.name]) # needed before stream, but after collision
+    sync_u = dh.synchronization_function([u.name]) # needed before continuity, but after stream
+    sync_tau = dh.synchronization_function([tau.name]) # needed before flux and tauface, but after continuity
+
+    def time_loop(steps, lambda_p_val, eta_p_val):
+        dh.all_to_gpu()
+        vmid = np.empty((2,steps//10+1))
+        sync_tau()
+        sync_u()
+        ubh()
+        i = -1
+        for i in range(steps):
+            dh.run_kernel(flux_kernel)
+            fluxbh() # zero the fluxes into/out of boundaries
+            dh.run_kernel(continuity_kernel, **{lambda_p.name: lambda_p_val, eta_p.name: eta_p_val})
+        
+            sync_tau()
+            dh.run_kernel(tauface_kernel) # needed for force
+            taufacebh()
+            dh.run_kernel(force_kernel)
+        
+            dh.run_kernel(collision_kernel, **{eta_p.name: eta_p_val})
+            sync_pdfs()
+            lbbh() # bounce-back populations into boundaries
+            dh.run_kernel(stream_kernel)
+            sync_u()
+            ubh() # need neighboring us for flux and continuity
+        
+            dh.swap(src.name, dst.name)
+        
+            if i % 10 == 0:
+                if u.name in dh.gpu_arrays:
+                    dh.to_cpu(u.name)
+                uu = dh.gather_array(u.name)
+                uu = uu[L[0]//2-1:L[0]//2+1, L[1]//2-1:L[1]//2+1, 0].mean()
+                if np.isnan(uu):
+                    raise Exception(f"NaN encountered after {i} steps")
+                vmid[:, i//10] = [i, uu]
+        sync_u()
+        dh.all_to_cpu()
+    
+        return vmid[:,:i//10+1]
+
+
+    # ## Analytical solution
+    # 
+    # comes from Waters and King, Unsteady flow of an elastico-viscous liquid, Rheologica Acta 9, 345–355 (1970).
+
+    # In[9]:
+
+
+    def N(n):
+        return (2*n-1)*np.pi
+
+    def Alpha_n(N, El, eta_p):
+        return 1+(1-eta_p)*El*N*N/4
+
+    def Beta_n(alpha_n,N, El):
+        return np.sqrt(np.abs(alpha_n*alpha_n - El*N*N))
+
+    def Gamma_n(N, El, eta_p):
+        return 1-(1+eta_p)*El*N*N/4
+
+    def G(alpha_n,beta_n,gamma_n,flag,T):
+        if(flag):
+            return ((1.0 - gamma_n/beta_n)*np.exp(-(alpha_n+beta_n)*T/2) +
+                    (1.0 + gamma_n/beta_n)*np.exp((beta_n-alpha_n)*T/2))
+        else:
+            return 2*np.exp(-alpha_n*T/2)*(np.cos(beta_n*T/2) + (gamma_n/beta_n)*np.sin(beta_n*T/2))
+
+    def W(T, El, eta_p):
+        W_ = 1.5
+        for n in range(1,1000):
+            N_ = N(n)
+            alpha_n = Alpha_n(N_, El, eta_p)
+        
+            if(alpha_n*alpha_n - El*N_*N_ < 0):
+                flag_ = False
+            else:
+                flag_ = True
+
+            beta_n = Beta_n(alpha_n,N_, El)
+            gamma_n = Gamma_n(N_, El, eta_p)
+            G_=G(alpha_n,beta_n,gamma_n,flag_,T)
+
+            W_ -= 24*(np.sin(N_/2)/(N_*N_*N_))*G_
+
+        return W_
+
+
+    # ## Run the simulation
+
+    # In[11]:
+
+
+    lambda_p_val = 3000
+    eta_p_val = 0.9
+
+    init()
+    vmid = time_loop(lambda_p_val*4, lambda_p_val, eta_p_val)
+
+    actual_width = sum(dh.gather_array(lbbh.flag_array_name)[L[0]//2,:] == 1)
+    uref = float(f_pre*actual_width**2/(8*(eta+eta_p)))
+
+    Wi = lambda_p_val*uref/(actual_width/2)
+    Re = uref*(actual_width/2)/(eta+eta_p)
+    El = float(Wi/Re)
+
+    pref = 1/W(vmid[0,-1]/lambda_p_val, El, eta_p_val)
+
+    El_measured, pref_measured = scipy.optimize.curve_fit(lambda a, b, c: W(a, b, eta_p_val)*c,
+                                                          vmid[0,:]/lambda_p_val, vmid[1,:]/vmid[1,-1],
+                                                          p0=(El, pref))[0]
+    measured_width = np.sqrt(4*lambda_p_val*float(eta+eta_p)/El_measured)
+    
+    print(f"El={El}, El_measured={El_measured}")
+    print(f"L={actual_width}, L_measured={measured_width}")
+
+    assert abs(measured_width - actual_width) < 1, "effective channel width differs significantly from defined width"
+
+    an = W(vmid[0,:]/lambda_p_val, El, eta_p_val)*pref
+    an_measured = W(vmid[0,:]/lambda_p_val, El_measured, eta_p_val)*pref_measured
+
+    diff = abs(vmid[1,:]/vmid[1,-1]-an_measured)/an_measured
+    assert diff[lambda_p_val//5:].max() < 0.03, "maximum velocity deviation is too large"
+
+#    from pystencils import plot as plt
+#
+#     plt.xlabel("$t$")
+#     plt.ylabel(r"$u_{max}/u_{max}^{Newtonian}$")
+#     plt.plot(vmid[0,:], vmid[1,:]/vmid[1,-1] if vmid[1,-1] != 0 else 0, label='FVM')
+#     plt.plot(vmid[0,:], np.ones_like(vmid[0,:]), 'k--', label='Newtonian')
+#
+#     plt.plot(vmid[0,:], an, label="analytic")
+#     plt.plot(vmid[0,:], an_measured, label="analytic, fit width")
+#     plt.legend()
+#
+#     if eta_p_val == 0.1:
+#         plt.ylim(0.9, 1.15)
+#     elif lambda_p_val == 9000:
+#         plt.ylim(0.8, 1.5)
+#     elif eta_p_val == 0.3:
+#         plt.ylim(0.8, 1.4)
+#     plt.show()