Diffusion in a shrinking cylinder using FEM

cylinder_diffusion.py

import numpy as np
import scipy as sp
import matplotlib.pyplot as plt

from math import sqrt
from scipy.integrate import solve_ivp

def l1kw(a,b):
    """Lobatto quadrature using 2 knots (trapezoid-like)"""

    h = b - a

    w0 = (2*a + b) / 6
    w1 = (2*b + a) / 6

    return [a,b], [h*w0,h*w1]

def l2kw(a,b):
    """Lobatto quadrature using 3 knots (Simpson-like)"""

    h = b - a
    eta = a + h*(0.5 + h/(10*(a + b)))

    w0 = (3*a**2 + 6*a*b + b**2) / (12*(2*a + 3*b))
    w1 = 25*(a + b)**3/(12*(2*a + 3*b)*(3*a + 2*b))
    w2 = (a**2 + 6*a*b + 3*b**2)/(12*(3*a + 2*b))

    return [a,eta,b], [h*w0,h*w1,h*w2]

def g1kw(a,b):
    """Gauss quadrature rule with 1 knot"""
    eta = (2./3.)*(b**3 - a**3)/(b**2 - a**2)
    w = (b**2 - a**2)/2.0

    return [eta], [w]

def g2kw(a,b):
    """Gauss quadrature rule with 2 knots"""

    P2 = a**2 + 4*a*b + b**2
    R2 = a**2 + 7*a*b + b**2
    P4 = a**4 + 10*(a**3*b + a*b**3) + 28*a**2*b**2 + b**4
    P3 = a**3 + 4*(a**2*b + a*b**2) + b**3

    h = b - a

    eta1 = (6*P3 - h*sqrt(6*P4))/(10*P2)
    eta2 = (6*P3 + h*sqrt(6*P4))/(10*P2)

    w1 = 0.25*(a + b) - h*R2/(6*sqrt(6*P4))
    w2 = 0.25*(a + b) + h*R2/(6*sqrt(6*P4))

    return [eta1,eta2], [h*w1,h*w2]

def cumquad(x,y,qrule=g2kw):
    """Cumulative integral using cylindrical quadrature rule"""

    N = x.size

    s = np.empty_like(y)
    s[0] = 0.0

    for i in range(1,N):
        a, b = x[i-1:i+1]
        eta, w = qrule(a, b) # Quadrature rule

        lsum = 0
        for ei, wi in zip(eta,w):
            
            nb = (ei-a)/(b-a)
            na = 1 - nb

            yi = y[i-1]*na + y[i]*nb
            lsum += wi*yi

        s[i] = s[i-1] + lsum

    return s

def mass_matrix(x,qrule=None):
    """Mass matrix using either analytical or numerical quadrature"""

    N = x.size
    M = np.zeros((N,N))

    # Loop over elements
    for iel in range(N-1):

        a, b = x[iel:iel+2]
        h = b - a

        lm = np.zeros((2,2))

        if qrule is None:
            # Analytical integral
            lm[0,0] = (1./12.)*(-3*a**2 + 2*a*b + b**2)
            lm[0,1] = -(1./12.)*(a-b)*(a+b)
            lm[1,0] = lm[0,1]
            lm[1,1] = (1./12.)*(-a**2 - 2*a*b + 3*b**2)
        else:
            eta, w = qrule(a,b)

            for ei, wi in zip(eta,w):
                na, nb = 1.0 - (ei - a)/(b - a), (ei-a)/(b-a)
                lm[0,0] += wi*(na*na)
                lm[0,1] += wi*(na*nb)
                lm[1,1] += wi*(nb*nb)
        
            lm[1,0] = lm[0,1]

        # Insert the local element matrix into the global one
        M[iel:iel+2,iel:iel+2] += lm

    return M


def apply_stiffness_matrix(x,y,r,D,alpha,qrule=g2kw):
    """Applies the stiffness matrix to a vector of dofs

    Args:
      x (1d-array): The mesh (element edges)
      y (1d-array): Nodal values or DoFs
      r (1d-array): Nodal values of r in wet reference frame
      D (float): The diffusivity
      alpha (float): Shrinkage coefficient
      qrule, optional: A quadrature rule function

    Returns:
        The vector C.dot(y), where C is the stiffness matrix
    """

    def dzeta(u,r_over_xi):
        return D/(1 + alpha*u)**2 * (r_over_xi)**2

    N = x.size
    u = np.zeros_like(y)

    # Loop over elements
    for iel in range(N-1):

        a, b = x[iel:iel+2]
        h = b - a

        # Obtain quadrature rule fitted to element [a,b]
        eta, w = qrule(a,b)

        lb = np.zeros((2,2))

        # Loop over quadrature knots
        for ei, wi in zip(eta,w):

            # Basis functions
            na, nb = 1.0 - (ei-a)/(b-a), (ei-a)/(b-a)

            # Value at quadrature knots
            ye = na*y[iel] + nb*y[iel+1]
            re = na*r[iel] + nb*r[iel+1]

            # Diffusivity at the quadrature knots
            de = dzeta(ye, re/ei)

            # Basis functions derivative
            nax, nbx = -1.0/(b-a), 1.0/(b-a)

            # Quadrature contributions
            lb[0,0] += wi*(de*nax*nax)
            lb[0,1] += wi*(de*nax*nbx)
            lb[1,1] += wi*(de*nbx*nbx)

        # Symmetry
        lb[1,0] = lb[0,1]

        # Sum contribution of current element into global vector
        u[iel:iel+2] -= lb.dot(y[iel:iel+2])

    return u

def diffode(t,y,mlu_and_piv,D,alpha,x):
    """Diffusion PDE with full mass matrix

    The mass matrix M should be factorized beforehand
    """

    # Calculate wet radius
    r = np.sqrt(x**2 + 2*alpha*cumquad(x,y,qrule=g2kw))

    # Calculate right-hand side
    dydt = apply_stiffness_matrix(x,y,r,D,alpha,qrule=g2kw)
    dydt[-1] = 0

    # Solve linear system
    sp.linalg.lu_solve(mlu_and_piv,dydt,overwrite_b=True)
    
    return dydt

def run(N):

    U0 = 0.3
    Ustar = 0.1
    alpha = 1550./998.
    D = 1.e-10

    rwet = 0.001
    rdry = sqrt(rwet**2/(1+alpha*U0))

    tf = 6400
    teval = [0,30,60,120,240,480,960,1600,3200,6400]

    x = np.linspace(0,rdry,N)

    M = mass_matrix(x)
    
    # Dirichlet boundary at the surface
    M[-1,:] = 0
    M[-1,-1] = 1

    # Factorize mass matrix
    mlu_and_piv = sp.linalg.lu_factor(M)

    # Initial conditions
    y0 = U0*np.ones_like(x)
    y0[-1] = Ustar

    ode = lambda t, y: diffode(t,y,mlu_and_piv,D,alpha,x)

    sol = solve_ivp(ode,(0,tf),y0,
        method="BDF",
        t_eval=teval)

    # Calculate positions in post-processing step
    r = np.empty_like(sol.y)
    for i, ycol in enumerate(sol.y.T):
        r[:,i] = np.sqrt(x**2 + 2*alpha*cumquad(x,ycol))

    plt.plot(r,sol.y)

    labels = ["t = {}".format(te) for te in teval]
    plt.legend(labels)

    plt.xlabel(r"$r$")
    plt.ylabel(r"$u(r)$")
    plt.title("Diffusion in a shrinking cylinder")

    plt.show()

if __name__ == '__main__':

    run(N=41)