preprocess.py 4.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152
  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 = np.cos(2*np.pi*x)
  14. # f = 4*pi**2*sin(2*pi*x)
  15. f = 1
  16. return f
  17. def bc(case, t=None):
  18. # boundary condition
  19. if case == 0:
  20. # advection-diffusion
  21. bc = [0, 0]
  22. if case == 1:
  23. # simple convection
  24. # bc = np.sin(2*np.pi*t)
  25. # adjoint boundary
  26. bc = [0, 1]
  27. return bc
  28. def discretization(coeff, mesh):
  29. """Given the problem statement, construct the discretization matrice"""
  30. # p is the number of basis functions
  31. p = coeff.pOrder + 1
  32. tau_pos = coeff.tauPlus
  33. tau_neg = coeff.tauMinus
  34. # order of gauss quadrature
  35. gorder = 2 * p
  36. # shape function and gauss quadrature
  37. xi, wi = np.polynomial.legendre.leggauss(gorder)
  38. shp, shpx = shape(xi, p)
  39. con = coeff.convection
  40. kappa = coeff.diffusion
  41. n_ele = len(mesh) - 1
  42. dist = np.zeros(n_ele)
  43. # elemental forcing vector
  44. F = np.zeros(p * n_ele)
  45. for i in range(1, n_ele + 1):
  46. dist[i - 1] = mesh[i] - mesh[i - 1]
  47. f = dist[i - 1] / 2 * shp.dot(
  48. wi * forcing(mesh[i - 1] + 1 / 2 * (1 + xi) * dist[i - 1]))
  49. F[(i - 1) * p:(i - 1) * p + p] = f
  50. F[0] += (con + tau_pos) * bc(0)[0]
  51. F[-1] += (-con + tau_neg) * bc(0)[1]
  52. d = shp.dot(np.diag(wi).dot(shp.T))
  53. # assemble global D
  54. d_face = np.zeros((p, p))
  55. d_face[0, 0] = tau_pos
  56. d_face[-1, -1] = tau_neg
  57. D = np.repeat(dist, p) / 2 * block_diag(*
  58. [d] * (n_ele)) + block_diag(*[d_face] * (n_ele))
  59. a = 1 / kappa * shp.dot(np.diag(wi).dot(shp.T))
  60. A = np.repeat(dist, p) / 2 * block_diag(*[a] * (n_ele))
  61. b = (shpx.T * np.ones((gorder, p))).T.dot(np.diag(wi).dot(shp.T))
  62. B = block_diag(*[b] * (n_ele))
  63. # elemental h
  64. h = np.zeros((2, 2))
  65. h[0, 0], h[-1, -1] = -con - tau_pos, con - tau_neg
  66. # mappinng matrix
  67. map_h = np.zeros((2, n_ele), dtype=int)
  68. map_h[:, 0] = np.arange(2)
  69. for i in np.arange(1, n_ele):
  70. map_h[:, i] = np.arange(
  71. map_h[2 - 1, i - 1], map_h[2 - 1, i - 1] + 2)
  72. # assemble H and eliminate boundaries
  73. H = np.zeros((n_ele + 1, n_ele + 1))
  74. for i in range(n_ele):
  75. for j in range(2):
  76. m = map_h[j, i]
  77. for k in range(2):
  78. n = map_h[k, i]
  79. H[m, n] += h[j, k]
  80. H = H[1:n_ele][:, 1:n_ele]
  81. # elemental g
  82. g = np.zeros((2, p))
  83. g[0, 0], g[-1, -1] = tau_pos, tau_neg
  84. # mapping matrix
  85. map_g_x = map_h
  86. map_g_y = np.arange(p * n_ele, dtype=int).reshape(n_ele, p).T
  87. # assemble global G
  88. G = np.zeros((n_ele + 1, p * n_ele))
  89. for i in range(n_ele):
  90. for j in range(2):
  91. m = map_g_x[j, i]
  92. for k in range(p):
  93. n = map_g_y[k, i]
  94. G[m, n] += g[j, k]
  95. G = G[1:n_ele, :]
  96. # elemental e
  97. e = np.zeros((p, 2))
  98. e[0, 0], e[-1, -1] = -con - tau_pos, con - tau_neg
  99. # mapping matrix
  100. map_e_x = np.arange(p * n_ele, dtype=int).reshape(n_ele, p).T
  101. map_e_y = map_h
  102. # assemble global E
  103. E = np.zeros((p * n_ele, n_ele + 1))
  104. for i in range(n_ele):
  105. for j in range(p):
  106. m = map_e_x[j, i]
  107. for k in range(2):
  108. n = map_e_y[k, i]
  109. E[m, n] += e[j, k]
  110. E = E[:, 1:n_ele]
  111. # L, easy in 1d
  112. L = np.zeros(n_ele - 1)
  113. # elemental c
  114. c = np.zeros((p, 2))
  115. c[0, 0], c[-1, -1] = -1, 1
  116. # assemble global C
  117. C = np.zeros((p * n_ele, n_ele + 1))
  118. for i in range(n_ele):
  119. for j in range(p):
  120. m = map_e_x[j, i]
  121. for k in range(2):
  122. n = map_e_y[k, i]
  123. C[m, n] += c[j, k]
  124. C = C[:, 1:n_ele]
  125. # L, easy in 1d
  126. L = np.zeros(n_ele - 1)
  127. # R, easy in 1d
  128. R = np.zeros(p * n_ele)
  129. R[0] = bc(0)[0]
  130. R[-1] = -bc(0)[1]
  131. return A, B, C, D, E, F, G, H, L, R, m