preprocess.py 6.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179
  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. # the following init are for the sake of simplicity
  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 lhsGen(self):
  54. """Generate matrices associated with left hand side"""
  55. if self.enrich is not None:
  56. numBasisFuncs = self.numBasisFuncs + self.enrich
  57. shp = self.shpEnrich
  58. else:
  59. numBasisFuncs = self.numBasisFuncs
  60. shp = self.shp
  61. # forcing vector F
  62. F = np.zeros(numBasisFuncs * self.n_ele)
  63. for i in range(1, self.n_ele + 1):
  64. f = self.dist[i - 1] / 2 \
  65. * shp.dot(self.wi * forcing(self.mesh[i - 1] +
  66. 1 / 2 * (1 + self.xi) *
  67. self.dist[i - 1]))
  68. F[(i - 1) * numBasisFuncs:i * numBasisFuncs] = f
  69. F[0] += (self.conv + self.tau_pos) * bc(0)[0]
  70. F[-1] += (-self.conv + self.tau_neg) * bc(0)[1]
  71. # L, easy in 1d
  72. L = np.zeros(self.n_ele - 1)
  73. # R, easy in 1d
  74. R = np.zeros(numBasisFuncs * self.n_ele)
  75. R[0] = bc(0)[0]
  76. R[-1] = -bc(0)[1]
  77. return F, L, R
  78. def eleGen(self):
  79. """Generate matrices associated with elements"""
  80. a = 1 / self.kappa * self.shp.dot(np.diag(self.wi).dot(self.shp.T))
  81. A = np.repeat(self.dist, self.numBasisFuncs) / \
  82. 2 * block_diag(*[a] * (self.n_ele))
  83. b = (self.shpx.T * np.ones((self.gqOrder, self.numBasisFuncs))
  84. ).T.dot(np.diag(self.wi).dot(self.shp.T))
  85. B = block_diag(*[b] * (self.n_ele))
  86. d = self.shp.dot(np.diag(self.wi).dot(self.shp.T))
  87. # assemble D
  88. d_face = np.zeros((self.numBasisFuncs, self.numBasisFuncs))
  89. d_face[0, 0] = self.tau_pos
  90. d_face[-1, -1] = self.tau_neg
  91. D = np.repeat(self.dist, self.numBasisFuncs) / 2 * \
  92. block_diag(*[d] * (self.n_ele)) + \
  93. block_diag(*[d_face] * (self.n_ele))
  94. return A, B, D
  95. def interfaceGen(self):
  96. """Generate matrices associated with interfaces"""
  97. tau_pos, tau_neg, conv = self.tau_pos, self.tau_neg, self.conv
  98. # elemental h
  99. h = np.zeros((2, 2))
  100. h[0, 0], h[-1, -1] = -conv - tau_pos, conv - tau_neg
  101. # mappinng matrix
  102. map_h = np.zeros((2, self.n_ele), dtype=int)
  103. map_h[:, 0] = np.arange(2)
  104. for i in np.arange(1, self.n_ele):
  105. map_h[:, i] = np.arange(
  106. map_h[2 - 1, i - 1], map_h[2 - 1, i - 1] + 2)
  107. # assemble H and eliminate boundaries
  108. H = np.zeros((self.n_ele + 1, self.n_ele + 1))
  109. for i in range(self.n_ele):
  110. for j in range(2):
  111. m = map_h[j, i]
  112. for k in range(2):
  113. n = map_h[k, i]
  114. H[m, n] += h[j, k]
  115. H = H[1:self.n_ele][:, 1:self.n_ele]
  116. # elemental e
  117. e = np.zeros((self.numBasisFuncs, 2))
  118. e[0, 0], e[-1, -1] = -conv - tau_pos, conv - tau_neg
  119. # mapping matrix
  120. map_e_x = np.arange(self.numBasisFuncs * self.n_ele,
  121. dtype=int).reshape(self.n_ele,
  122. self.numBasisFuncs).T
  123. map_e_y = map_h
  124. # assemble global E
  125. E = np.zeros((self.numBasisFuncs * self.n_ele, self.n_ele + 1))
  126. for i in range(self.n_ele):
  127. for j in range(self.numBasisFuncs):
  128. m = map_e_x[j, i]
  129. for k in range(2):
  130. n = map_e_y[k, i]
  131. E[m, n] += e[j, k]
  132. E = E[:, 1:self.n_ele]
  133. # elemental c
  134. c = np.zeros((self.numBasisFuncs, 2))
  135. c[0, 0], c[-1, -1] = -1, 1
  136. # assemble global C
  137. C = np.zeros((self.numBasisFuncs * self.n_ele, self.n_ele + 1))
  138. for i in range(self.n_ele):
  139. for j in range(self.numBasisFuncs):
  140. m = map_e_x[j, i]
  141. for k in range(2):
  142. n = map_e_y[k, i]
  143. C[m, n] += c[j, k]
  144. C = C[:, 1:self.n_ele]
  145. # elemental g
  146. g = np.zeros((2, self.numBasisFuncs))
  147. g[0, 0], g[-1, -1] = tau_pos, tau_neg
  148. # mapping matrix
  149. map_g_x = map_h
  150. map_g_y = np.arange(self.numBasisFuncs * self.n_ele,
  151. dtype=int).reshape(self.n_ele,
  152. self.numBasisFuncs).T
  153. # assemble global G
  154. G = np.zeros((self.n_ele + 1, self.numBasisFuncs * self.n_ele))
  155. for i in range(self.n_ele):
  156. for j in range(2):
  157. m = map_g_x[j, i]
  158. for k in range(self.numBasisFuncs):
  159. n = map_g_y[k, i]
  160. G[m, n] += g[j, k]
  161. G = G[1:self.n_ele, :]
  162. return C, E, G, H