Coverage for python/gsfit/database_writers/rtgsfit_mdsplus/poisson_matrix.py: 0%
48 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-14 15:25 +0000
1"""
2This module contains routines for discretizing the left-hand side of the Grad-Shafranov equation,
3which involves applying a modified Laplacian operator to the poloidal flux function.
5The resulting matrix is then factorized into a product of lower and upper triangular matrices (LU decomposition),
6which accelerates the Gaussian elimination process used in the RTGSFIT code.
7"""
9import numpy as np
10import numpy.typing as npt
11from scipy.linalg import lu
14def poisson_matrix(
15 r_vec: npt.NDArray[np.float64],
16 z_vec: npt.NDArray[np.float64],
17) -> npt.NDArray[np.float64]:
18 """
19 Construct the Poisson matrix used in the RTGSFIT code for discretizing the Grad-Shafranov equation.
21 This implementation follows the finite difference scheme described in Equation (45) of:
22 Moret, J.-M., et al. "Tokamak equilibrium reconstruction code LIUQE and its real time implementation."
23 Fusion Engineering and Design 91 (2015): 1-15.
25 The discretized equation has the form:
27 ψ_{i+1,j} + ψ_{i-1,j} + a_j * ψ_{i,j+1} + b_j * ψ_{i,j-1} - c_j * ψ_{i,j}
29 where:
30 - ψ is the poloidal flux function,
31 - a_j = (dz / dr)^2 * r_j / (r_j - dr / 2),
32 - b_j = (dz / dr)^2 * r_j / (r_j + dr / 2),
33 - c_j = 2 + a_j + b_j
35 Let:
36 - n_r = length of `r_vec` (number of radial grid points),
37 - n_z = length of `z_vec` (number of vertical grid points).
39 The resulting Poisson matrix has dimensions (n_r * n_z) x (n_r * n_z), corresponding to the total
40 number of grid points in the 2D domain.
42 Parameters:
43 r_vec (np.ndarray): 1D array of radial grid coordinates.
44 z_vec (np.ndarray): 1D array of vertical grid coordinates.
46 Returns:
47 np.ndarray: The assembled Poisson matrix representing the discretized operator.
48 """
50 n_r = len(r_vec)
51 n_z = len(z_vec)
52 n_grid = n_r * n_z
53 dr = r_vec[1] - r_vec[0]
54 dz = z_vec[1] - z_vec[0]
56 a = (dz / dr) ** 2 * r_vec / (r_vec + dr / 2)
57 b = (dz / dr) ** 2 * r_vec / (r_vec - dr / 2)
58 c = 2 + a + b
60 poiss_matrix = np.zeros((n_grid, n_grid), dtype=np.float64)
62 idx = -1
63 for i in range(n_z):
64 for j in range(n_r):
65 idx += 1
66 if i == 0 or i == n_z - 1 or j == 0 or j == n_r - 1:
67 poiss_matrix[idx, idx] = 1 # Dirichlet BC
68 else:
69 poiss_matrix[idx, idx] = -c[j] # coefficient of ψ_{i,j}
70 poiss_matrix[idx, idx + 1] = a[j] # coefficient of ψ_{i,j+1}
71 poiss_matrix[idx, idx - 1] = b[j] # coefficient of ψ_{i,j-1}
72 poiss_matrix[idx, idx + n_r] = 1 # coefficient of ψ_{i+1,j}
73 poiss_matrix[idx, idx - n_r] = 1 # coefficient of ψ_{i-1,j}
75 return poiss_matrix
78def poisson_matrix_lup_bands(
79 poiss_matrix: npt.NDArray[np.float64],
80 n_r: int,
81) -> tuple[
82 npt.NDArray[np.int64],
83 npt.NDArray[np.float64],
84 npt.NDArray[np.float64],
85]:
86 """
87 Perform LU decomposition of the Poisson matrix with banded storage.
89 This function uses the `scipy.linalg.lu` function to decompose the Poisson matrix
90 into a lower triangular matrix L and an upper triangular matrix U, along with a
91 permutation matrix P.
93 We then extract the non-zero bands of L and U, which are used in the RTGSFIT code
94 for efficient Gaussian elimination.
95 We don't store the main diagonal of L as it is always 1.
97 Note that it's not garunteed that L with have a bandwidth of n_r and U will have a bandwidth of n_r + 2,
98 so we should check this if we use a new grid.
100 Parameters:
101 poisson_matrix (np.ndarray): The Poisson matrix to be decomposed.
103 Returns:
104 tuple: A tuple containing the permutation matrix P, lower triangular matrix L,
105 and upper triangular matrix U.
106 """
107 perm, lower, upper = lu(poiss_matrix)
108 # Invert the permutation matrix P to get it in the same form
109 # as the Matlab lu function.
110 perm = perm.T
111 # Pad the matrices so we can extract the bands more easily.
112 lower_pad = np.pad(lower, ((0, 0), (n_r, 0)), mode="constant", constant_values=0)
113 upper_pad = np.pad(upper, ((0, 0), (0, n_r + 1)), mode="constant", constant_values=0)
114 n_grid = np.shape(poiss_matrix)[0]
115 lower_band = np.zeros((n_grid, n_r), dtype=np.float64)
116 upper_band = np.zeros((n_grid, n_r + 2), dtype=np.float64)
117 for ii in range(n_grid):
118 lower_band[ii, :] = lower_pad[ii, ii : ii + n_r]
119 upper_band[ii, :] = upper_pad[ii, ii : ii + n_r + 2]
120 # We precompute the inverse of the main diagonal of U to avoid repeated division later.
121 # In Gaussian elimination, the final step involves dividing by the diagonal elements of U.
122 # By inverting them here, we save computation time during back-substitution.
123 # For reference, see:
124 # Jardin, S. (2010). Computational Methods in Plasma Physics, CRC Press, p. 54.
125 upper_band[:, 0] = 1 / upper_band[:, 0]
126 _, perm_idx = np.nonzero(perm)
127 return perm_idx, lower_band.flatten(), upper_band.flatten()
130def compute_lup_bands(
131 r_vec: npt.NDArray[np.float64], z_vec: npt.NDArray[np.float64]
132) -> tuple[npt.NDArray[np.float64], npt.NDArray[np.float64], npt.NDArray[np.int64]]:
133 """
134 Compute the permutation indices and lower/upper bands of the Poisson matrix.
136 This function constructs the Poisson matrix using the provided radial and vertical grid vectors,
137 performs LU decomposition, and extracts the permutation indices, lower band, and upper band.
139 Parameters:
140 r_vec (np.ndarray): 1D array of radial grid coordinates.
141 z_vec (np.ndarray): 1D array of vertical grid coordinates.
143 Returns:
144 tuple: A tuple containing the permutation indices, lower band, and upper band.
145 """
146 poisson_mat = poisson_matrix(r_vec, z_vec)
147 perm_idx, lower_band, upper_band = poisson_matrix_lup_bands(poisson_mat, n_r=len(r_vec))
148 return lower_band, upper_band, perm_idx
151if __name__ == "__main__":
152 # Example usage
153 r_vec = np.linspace(110.000e-3, 1.070e0, 33)
154 z_vec = np.linspace(-960.000e-3, 960.000e-3, 65)
155 poisson_mat = poisson_matrix(r_vec, z_vec)
156 perm_idx, lower_band, upper_band = poisson_matrix_lup_bands(poisson_mat, n_r=len(r_vec))