preprocess.py 6.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174
  1. import numpy as np
  2. from scipy.linalg import block_diag
  3. def shape(x, p):
  4. """generate p shape functions and its first order derivative
  5. (order p-1) at the given location x"""
  6. A = np.array([np.linspace(-1, 1, p)]).T**np.arange(p)
  7. C = np.linalg.inv(A).T
  8. x = np.array([x]).T
  9. shp = C.dot((x**np.arange(p)).T)
  10. shpx = C[:, 1::1].dot((x**np.arange(p - 1) * np.arange(1, p)).T)
  11. return shp, shpx
  12. def forcing(x):
  13. f = 1
  14. return f
  15. def bc(case, t=None):
  16. # boundary condition
  17. if case == 0:
  18. # advection-diffusion
  19. bc = [0, 0]
  20. if case == 1:
  21. # simple convection
  22. # bc = np.sin(2*np.pi*t)
  23. # adjoint boundary
  24. bc = [0, 1]
  25. return bc
  26. class discretization(object):
  27. """Given the problem statement, construct the discretization matrices"""
  28. def __init__(self, coeff, mesh, enrich=None):
  29. self.mesh = mesh
  30. self.coeff = coeff
  31. self.enrich = enrich
  32. # p is the number of basis functions
  33. self.numBasisFuncs = coeff.pOrder + 1
  34. self.tau_pos = coeff.tauPlus
  35. self.tau_neg = coeff.tauMinus
  36. self.conv = coeff.convection
  37. self.kappa = coeff.diffusion
  38. self.n_ele = len(mesh) - 1
  39. self.dist = self.distGen()
  40. # shape function and gauss quadrature
  41. self.gqOrder = 5 * self.numBasisFuncs
  42. self.xi, self.wi = np.polynomial.legendre.leggauss(self.gqOrder)
  43. self.shp, self.shpx = shape(self.xi, self.numBasisFuncs)
  44. # enrich the space if the enrich argument is given
  45. if enrich is not None:
  46. self.shpEnrich, self.shpxEnrich = shape(
  47. self.xi, self.numBasisFuncs + enrich)
  48. def distGen(self):
  49. dist = np.zeros(self.n_ele)
  50. for i in range(1, self.n_ele + 1):
  51. dist[i - 1] = self.mesh[i] - self.mesh[i - 1]
  52. return dist
  53. def matGen(self):
  54. self.lhsGen()
  55. self.eleGen()
  56. self.interfaceGen()
  57. def lhsGen(self):
  58. """Generate matrices associated with left hand side"""
  59. # elemental forcing vector F
  60. F = np.zeros(self.numBasisFuncs * self.n_ele)
  61. for i in range(1, self.n_ele + 1):
  62. f = self.dist[i - 1] / 2 * self.shp.dot(
  63. self.wi * forcing(self.mesh[i - 1] + 1 / 2 * (1 + self.xi) * self.dist[i - 1]))
  64. F[(i - 1) * self.numBasisFuncs:i * self.numBasisFuncs] = f
  65. F[0] += (self.conv + self.tau_pos) * bc(0)[0]
  66. F[-1] += (-self.conv + self.tau_neg) * bc(0)[1]
  67. # L, easy in 1d
  68. L = np.zeros(self.n_ele - 1)
  69. # R, easy in 1d
  70. R = np.zeros(self.numBasisFuncs * self.n_ele)
  71. R[0] = bc(0)[0]
  72. R[-1] = -bc(0)[1]
  73. self.F, self.L, self.R = F, L, R
  74. def eleGen(self):
  75. """Generate matrices associated with elements"""
  76. a = 1 / self.kappa * self.shp.dot(np.diag(self.wi).dot(self.shp.T))
  77. A = np.repeat(self.dist, self.numBasisFuncs) / \
  78. 2 * block_diag(*[a] * (self.n_ele))
  79. b = (self.shpx.T * np.ones((self.gqOrder, self.numBasisFuncs))
  80. ).T.dot(np.diag(self.wi).dot(self.shp.T))
  81. B = block_diag(*[b] * (self.n_ele))
  82. d = self.shp.dot(np.diag(self.wi).dot(self.shp.T))
  83. # assemble global D
  84. d_face = np.zeros((self.numBasisFuncs, self.numBasisFuncs))
  85. d_face[0, 0] = self.tau_pos
  86. d_face[-1, -1] = self.tau_neg
  87. D = np.repeat(self.dist, self.numBasisFuncs) / 2 * \
  88. block_diag(*[d] * (self.n_ele)) + \
  89. block_diag(*[d_face] * (self.n_ele))
  90. self.A, self.B, self.D = A, B, D
  91. def interfaceGen(self):
  92. """Generate matrices associated with interfaces"""
  93. tau_pos, tau_neg, conv = self.tau_pos, self.tau_neg, self.conv
  94. # elemental h
  95. h = np.zeros((2, 2))
  96. h[0, 0], h[-1, -1] = -conv - tau_pos, conv - tau_neg
  97. # mappinng matrix
  98. map_h = np.zeros((2, self.n_ele), dtype=int)
  99. map_h[:, 0] = np.arange(2)
  100. for i in np.arange(1, self.n_ele):
  101. map_h[:, i] = np.arange(
  102. map_h[2 - 1, i - 1], map_h[2 - 1, i - 1] + 2)
  103. # assemble H and eliminate boundaries
  104. H = np.zeros((self.n_ele + 1, self.n_ele + 1))
  105. for i in range(self.n_ele):
  106. for j in range(2):
  107. m = map_h[j, i]
  108. for k in range(2):
  109. n = map_h[k, i]
  110. H[m, n] += h[j, k]
  111. H = H[1:self.n_ele][:, 1:self.n_ele]
  112. # elemental e
  113. e = np.zeros((self.numBasisFuncs, 2))
  114. e[0, 0], e[-1, -1] = -conv - tau_pos, conv - tau_neg
  115. # mapping matrix
  116. map_e_x = np.arange(self.numBasisFuncs * self.n_ele,
  117. dtype=int).reshape(self.n_ele, self.numBasisFuncs).T
  118. map_e_y = map_h
  119. # assemble global E
  120. E = np.zeros((self.numBasisFuncs * self.n_ele, self.n_ele + 1))
  121. for i in range(self.n_ele):
  122. for j in range(self.numBasisFuncs):
  123. m = map_e_x[j, i]
  124. for k in range(2):
  125. n = map_e_y[k, i]
  126. E[m, n] += e[j, k]
  127. E = E[:, 1:self.n_ele]
  128. # elemental c
  129. c = np.zeros((self.numBasisFuncs, 2))
  130. c[0, 0], c[-1, -1] = -1, 1
  131. # assemble global C
  132. C = np.zeros((self.numBasisFuncs * self.n_ele, self.n_ele + 1))
  133. for i in range(self.n_ele):
  134. for j in range(self.numBasisFuncs):
  135. m = map_e_x[j, i]
  136. for k in range(2):
  137. n = map_e_y[k, i]
  138. C[m, n] += c[j, k]
  139. C = C[:, 1:self.n_ele]
  140. # elemental g
  141. g = np.zeros((2, self.numBasisFuncs))
  142. g[0, 0], g[-1, -1] = tau_pos, tau_neg
  143. # mapping matrix
  144. map_g_x = map_h
  145. map_g_y = np.arange(self.numBasisFuncs * self.n_ele,
  146. dtype=int).reshape(self.n_ele, self.numBasisFuncs).T
  147. # assemble global G
  148. G = np.zeros((self.n_ele + 1, self.numBasisFuncs * self.n_ele))
  149. for i in range(self.n_ele):
  150. for j in range(2):
  151. m = map_g_x[j, i]
  152. for k in range(self.numBasisFuncs):
  153. n = map_g_y[k, i]
  154. G[m, n] += g[j, k]
  155. G = G[1:self.n_ele, :]
  156. self.C, self.E, self.G, self.H = C, E, G, H