Source code for pyamg.gallery.diffusion

"""Generate a diffusion stencil.

Supports isotropic diffusion (FE,FD), anisotropic diffusion (FE, FD), and
rotated anisotropic diffusion (FD).

The stencils include redundancy to maintain readability for simple cases (e.g.
isotropic diffusion).

"""
# pylint: disable=redefined-builtin

import numpy as np


[docs] def diffusion_stencil_2d(epsilon=1.0, theta=0.0, type='FE'): """Rotated Anisotropic diffusion in 2d of the form. -div Q A Q^T grad u = - (C^2 + eps S^2) u_xx - 2(1 - eps) C S u_xy - (eps C^2 + S^2) u_yy where C=cos(theta), S=sin(theta), and Q = [cos(theta) -sin(theta)] [sin(theta) cos(theta)] A = [1 0 ] [0 eps ] Parameters ---------- epsilon : float Anisotropic diffusion coefficient: -div A grad u, where A = [1 0; 0 epsilon]. The default is isotropic, epsilon=1.0. theta : float Rotation angle `theta` from the positive x-axis in radians. Defines -div Q A Q^T grad, where Q = [cos(`theta`) -sin(`theta`); sin(`theta`) cos(`theta`)]. The default is `theta` = 0.0. type : {'FE','FD'} Specifies the discretization as Q1 finite element (FE) or 2nd order finite difference (FD). Returns ------- array Diffusion stencil of size (3,3). See Also -------- stencil_grid, poisson Notes ----- Not all combinations are supported. The stencil is ordered with y varying first; see `stencil_grid` for more details. Examples -------- >>> import numpy as np >>> import scipy as sp >>> from pyamg.gallery.diffusion import diffusion_stencil_2d >>> from pyamg.gallery import stencil_grid >>> sten = diffusion_stencil_2d(epsilon=0.0001,theta=np.pi/6,type='FD') >>> print(sten) [[-0.2164847 -0.750025 0.2164847] [-0.250075 2.0002 -0.250075 ] [ 0.2164847 -0.750025 -0.2164847]] Consider a 2 x 4 grid ([x0, x1] x [y0, y1, y2, y3]). The first dimension of the stencil defines x. >>> nx, ny = (2, 4) >>> sten = diffusion_stencil_2d(epsilon=0.1, type='FD') >>> A = stencil_grid(sten, (nx, ny)).toarray() >>> print(sten) [[-0. -1. 0. ] [-0.1 2.2 -0.1] [ 0. -1. -0. ]] >>> print(A) [[ 2.2 -0.1 0. 0. -1. 0. 0. 0. ] [-0.1 2.2 -0.1 0. 0. -1. 0. 0. ] [ 0. -0.1 2.2 -0.1 0. 0. -1. 0. ] [ 0. 0. -0.1 2.2 0. 0. 0. -1. ] [-1. 0. 0. 0. 2.2 -0.1 0. 0. ] [ 0. -1. 0. 0. -0.1 2.2 -0.1 0. ] [ 0. 0. -1. 0. 0. -0.1 2.2 -0.1] [ 0. 0. 0. -1. 0. 0. -0.1 2.2]] """ eps = float(epsilon) # for brevity theta = float(theta) C = np.cos(theta) S = np.sin(theta) CS = C*S CC = C**2 SS = S**2 if type == 'FE': # FE approximation to # -div K grad u # using the weak form (K grad u, grad v) # see _symbolic_fe_helper() for more details a = (-1*eps - 1)*CC + (-1*eps - 1)*SS + (3*eps - 3)*CS b = (2*eps - 4)*CC + (-4*eps + 2)*SS c = (-1*eps - 1)*CC + (-1*eps - 1)*SS + (-3*eps + 3)*CS d = (-4*eps + 2)*CC + (2*eps - 4)*SS e = (8*eps + 8)*CC + (8*eps + 8)*SS stencil = np.array([[a, b, c], [d, e, d], [c, b, a]]) / 6.0 elif type == 'FD': # discretizing # -div Q A Q^T grad u # = - (C^2 + eps S^2) u_xx - 2(1 - eps) C S u_xy - (eps C^2 + S^2) u_yy # = - E u_xx - F u_xy - G u_yy # with hx = hy = h = 1 gives # = E (- u_{i-1,j} + 2 u_{i,j} - u_{i+1, j}) # + F/4 (- u_{i+1,j+1} + u_{i+1,j-1} + u_{i-1,j+1} - u_{i-1,j-1}) # + G (- u_{i,j-1} + 2 u_{i,j} - u_{i,j+1}) # For a stencil centered at [i,j] this leads to # [i-1,j+1] ---- [i,j+1] ---- [i+1,j+1] [ F/4] ---- [ -G ] ---- [-F/4] # | | | | | | # | | | | | | # | | | | | | # [i-1,j ] ---- [i,j ] ---- [i+1,j+1] = [ -E ] ---- [2 E + 2G] ---- [-E ] # | | | | | | # | | | | | | # | | | | | | # [i-1,j-1] ---- [i,j-1] ---- [i+1,j+1] [-F/4] ---- [ -G ] ---- [ F/4] # And the stencil, with y varying first: # stencil = [-F/4 -E F/4] # column 0: [i-1,j-1] [i-1,j] [i-1,j+1] # [-G 2E+2G -G] # column 1: [i,j-1] [i,j] [i,j+1] # [ F/4 -E -F/4] # column 2: [i+1,j-1] [i+1,j] [i+1,j+1] # [ | | ] # [0.5(eps-1) C S | -(C^2 + eps S^2) | 0.5(1-eps) C S ] # [_________________________________________________________] # [ | | ] # = [-(eps C^2 + S^2) | 2 (eps + 1) | -(eps C^2 + S^2)] # [_________________________________________________________] # [ | | ] # [0.5(1-eps) C S | -(C^2 + eps S^2) | 0.5(eps-1) C S ] # [ | | ] # = [a, b, c] # [d, e, d] # [c, b, a] a = 0.5*(eps - 1)*CS b = -(eps*SS + CC) c = -a d = -(eps*CC + SS) e = 2.0*(eps + 1) stencil = np.array([[a, b, c], [d, e, d], [c, b, a]]) else: raise ValueError('only stencil types "FE" and "FD" are supported') return stencil
def _symbolic_fe_helper(): """Generate the stencil for 2D FE using SymPy.""" from sympy import symbols, integrate, Matrix # noqa: PLC0415 from sympy.vector import CoordSys3D, gradient # noqa: PLC0415 C, S, eps = symbols('C S eps') N = CoordSys3D('N') x, y = N.x, N.y # Define the rotation and anisotropy Q = Matrix([[C, -S], [S, C]]) A = Matrix([[1, 0], [0, eps]]) K = Q @ A @ Q.T # Start with a reference element ordering: # [2 3] # [0 1] # And defeine four basis functions phi0 = (1-x)*(1-y) phi1 = x*(1-y) phi2 = (1-x)*y phi3 = x*y # Make space for a 3x3 stencil sten = np.empty((3, 3), dtype=object) # Define a weak form def a(phi_l, phi_r): """Define the weak form. weak form form -div K grad u = (K grad u, grad v) """ gradu = gradient(phi_l) gradv = gradient(phi_r) Kgradu = K @ Matrix([gradu.coeff(N.i), gradu.coeff(N.j)]) Kgradu = Kgradu[0]*N.i + Kgradu[1]*N.j I = integrate(Kgradu.dot(gradv), (N.x, 0, 1), (N.y, 0, 1)) return I # Consider a four element mesh to create the stencil at [4] # 2--5--8 # | | | # 1--4--7 # | | | # 0--3--6 sten[0, 0] = a(phi3, phi0) # 4-0 sten[0, 1] = a(phi3, phi2) + a(phi1, phi0) # 4-1 sten[0, 2] = a(phi1, phi2) # 4-2 sten[1, 0] = a(phi3, phi1) + a(phi2, phi0) # 4-3 sten[1, 1] = a(phi3, phi3) + a(phi1, phi1) \ + a(phi2, phi2) + a(phi0, phi0) # 4-4 sten[1, 2] = a(phi1, phi3) + a(phi0, phi2) # 4-5 sten[2, 0] = a(phi2, phi1) # 4-6 sten[2, 1] = a(phi2, phi3) + a(phi0, phi1) # 4-7 sten[2, 2] = a(phi0, phi3) # 4-8 # now set a, b, c, d, and e a = 6 * sten[0, 0] b = 6 * sten[0, 1] c = 6 * sten[0, 2] d = 6 * sten[1, 0] e = 6 * sten[1, 1] print(f'{a=}') print(f'{b=}') print(f'{c=}') print(f'{d=}') print(f'{e=}') def _symbolic_rotation_helper(): """Use SymPy to generate the 3D matrices for diffusion_stencil_3d.""" # pylint: disable=import-error,import-outside-toplevel from sympy import symbols, Matrix # noqa: PLC0415 cpsi, spsi = symbols('cpsi, spsi') cth, sth = symbols('cth, sth') cphi, sphi = symbols('cphi, sphi') Rpsi = Matrix([[cpsi, spsi, 0], [-spsi, cpsi, 0], [0, 0, 1]]) Rth = Matrix([[1, 0, 0], [0, cth, sth], [0, -sth, cth]]) Rphi = Matrix([[cphi, sphi, 0], [-sphi, cphi, 0], [0, 0, 1]]) Q = Rpsi * Rth * Rphi epsy, epsz = symbols('epsy, epsz') A = Matrix([[1, 0, 0], [0, epsy, 0], [0, 0, epsz]]) D = Q * A * Q.T for i in range(3): for j in range(3): print(f'D[{i}, {j}] = {D[i, j]}') def _symbolic_product_helper(): """Use SymPy to generate the 3D products for diffusion_stencil_3d.""" from sympy import symbols, Matrix # noqa: PLC0415 D11, D12, D13, D21, D22, D23, D31, D32, D33 =\ symbols('D11, D12, D13, D21, D22, D23, D31, D32, D33') D = Matrix([[D11, D12, D13], [D21, D22, D23], [D31, D32, D33]]) grad = Matrix([['dx', 'dy', 'dz']]).T div = grad.T a = div * D * grad print(a[0]) def diffusion_stencil_3d(epsilony=1.0, epsilonz=1.0, theta=0.0, phi=0.0, psi=0.0, type='FD'): """Rotated Anisotropic diffusion in 3d of the form. -div Q A Q^T grad u Q = Rpsi Rtheta Rphi Rpsi = [ c s 0 ] [-s c 0 ] [ 0 0 1 ] c = cos(psi) s = sin(psi) Rtheta = [ 1 0 0 ] [ 0 c s ] [ 0 -s c ] c = cos(theta) s = sin(theta) Rphi = [ c s 0 ] [-s c 0 ] [ 0 0 1 ] c = cos(phi) s = sin(phi) Here Euler Angles are used: http://en.wikipedia.org/wiki/Euler_angles This results in Q = [ cphi*cpsi - cth*sphi*spsi, cpsi*sphi + cphi*cth*spsi, spsi*sth] [ - cphi*spsi - cpsi*cth*sphi, cphi*cpsi*cth - sphi*spsi, cpsi*sth] [ sphi*sth, -cphi*sth, cth] A = [1 0 ] [0 epsy ] [0 0 epsz] D = Q A Q^T Parameters ---------- epsilony : float, optional Anisotropic diffusion coefficient in the y-direction where A = [1 0 0; 0 epsilon_y 0; 0 0 epsilon_z]. The default is isotropic, epsilon=1.0 epsilonz : float, optional Anisotropic diffusion coefficient in the z-direction where A = [1 0 0; 0 epsilon_y 0; 0 0 epsilon_z]. The default is isotropic, epsilon=1.0 theta : float, optional Euler rotation angle `theta` in radians. The default is 0.0. phi : float, optional Euler rotation angle `phi` in radians. The default is 0.0. psi : float, optional Euler rotation angle `psi` in radians. The default is 0.0. type : {'FE','FD'} Specifies the discretization as Q1 finite element (FE) or 2nd order finite difference (FD) Returns ------- stencil : numpy array A 3x3 diffusion stencil See Also -------- stencil_grid, poisson, _symbolic_rotation_helper, _symbolic_product_helper Notes ----- Not all combinations are supported. Examples -------- >>> import numpy as np >>> import scipy as sp >>> from pyamg.gallery.diffusion import diffusion_stencil_2d >>> sten = diffusion_stencil_2d(epsilon=0.0001,theta=np.pi/6,type='FD') >>> print(sten) [[-0.2164847 -0.750025 0.2164847] [-0.250075 2.0002 -0.250075 ] [ 0.2164847 -0.750025 -0.2164847]] """ epsy = float(epsilony) # for brevity epsz = float(epsilonz) # for brevity theta = float(theta) phi = float(phi) psi = float(psi) D = np.zeros((3, 3)) cphi = np.cos(phi) sphi = np.sin(phi) cth = np.cos(theta) sth = np.sin(theta) cpsi = np.cos(psi) spsi = np.sin(psi) # from _symbolic_rotation_helper D[0, 0] = epsy*(cphi*cth*spsi + cpsi*sphi)**2 + epsz*spsi**2*sth**2 +\ (cphi*cpsi - cth*sphi*spsi)**2 D[0, 1] = cpsi*epsz*spsi*sth**2 +\ epsy*(cphi*cpsi*cth - sphi*spsi)*(cphi*cth*spsi + cpsi*sphi) +\ (cphi*cpsi - cth*sphi*spsi)*(-cphi*spsi - cpsi*cth*sphi) D[0, 2] = -cphi*epsy*sth*(cphi*cth*spsi + cpsi*sphi) +\ cth*epsz*spsi*sth + sphi*sth*(cphi*cpsi - cth*sphi*spsi) D[1, 0] = cpsi*epsz*spsi*sth**2 +\ epsy*(cphi*cpsi*cth - sphi*spsi)*(cphi*cth*spsi + cpsi*sphi) +\ (cphi*cpsi - cth*sphi*spsi)*(-cphi*spsi - cpsi*cth*sphi) D[1, 1] = cpsi**2*epsz*sth**2 + epsy*(cphi*cpsi*cth - sphi*spsi)**2 +\ (-cphi*spsi - cpsi*cth*sphi)**2 D[1, 2] = -cphi*epsy*sth*(cphi*cpsi*cth - sphi*spsi) +\ cpsi*cth*epsz*sth + sphi*sth*(-cphi*spsi - cpsi*cth*sphi) D[2, 0] = -cphi*epsy*sth*(cphi*cth*spsi + cpsi*sphi) + cth*epsz*spsi*sth +\ sphi*sth*(cphi*cpsi - cth*sphi*spsi) D[2, 1] = -cphi*epsy*sth*(cphi*cpsi*cth - sphi*spsi) + cpsi*cth*epsz*sth +\ sphi*sth*(-cphi*spsi - cpsi*cth*sphi) D[2, 2] = cphi**2*epsy*sth**2 + cth**2*epsz + sphi**2*sth**2 stencil = np.zeros((3, 3, 3)) if type == 'FE': raise NotImplementedError('FE not implemented yet') if type == 'FD': # from _symbolic_product_helper # dx*(D11*dx + D21*dy + D31*dz) + # dy*(D12*dx + D22*dy + D32*dz) + # dz*(D13*dx + D23*dy + D33*dz) # # D00*dxx + # (D10+D01)*dxy + # (D20+D02)*dxz + # D11*dyy + # (D21+D12)*dyz + # D22*dzz i, j, k = (1, 1, 1) # dxx stencil[[i-1, i, i+1], j, k] += np.array([-1, 2, -1]) * D[0, 0] # dyy stencil[i, [j-1, j, j+1], k] += np.array([-1, 2, -1]) * D[1, 1] # dzz stencil[i, j, [k-1, k, k+1]] += np.array([-1, 2, -1]) * D[2, 2] L = np.array([-1, -1, 1, 1]) M = np.array([-1, 1, -1, 1]) # dxy stencil[i + L, j + M, k] \ += 0.25 * np.array([1, -1, -1, 1]) * (D[1, 0] + D[0, 1]) # dxz stencil[i + L, j, k + M] \ += 0.25 * np.array([1, -1, -1, 1]) * (D[2, 0] + D[0, 2]) # dyz stencil[i, j + L, k + M] \ += 0.25 * np.array([1, -1, -1, 1]) * (D[2, 1] + D[1, 2]) return stencil