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