diff --git a/examples/example02 - static_forward_solve_MASTU.ipynb b/examples/example02 - static_forward_solve_MASTU.ipynb index 9c02a5a..c333172 100644 --- a/examples/example02 - static_forward_solve_MASTU.ipynb +++ b/examples/example02 - static_forward_solve_MASTU.ipynb @@ -143,10 +143,28 @@ "outputs": [], "source": [ "from freegsnke import GSstaticsolver\n", + "GSStaticSolver = GSstaticsolver.NKGSsolver(eq)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The solver can alternatively be initialized while specifying the solver type. A DST solver (below) would be significantly faster than the default (but more accurate) LU sparse solver.\n", "\n", + "The solver can also be initialized with the argument `cache_greens` set to `False`, which will allow for larger grids to be run (due to lower memory requirements) at the expense of longer execution time when solving." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ "GSStaticSolver = GSstaticsolver.NKGSsolver(\n", " eq,\n", - " gs_operator_order=4, # use 2 for lower direct-solver setup cost\n", + " solver_type='DST',\n", + " #cache_greens=False,\n", ")" ] }, diff --git a/examples/example05a - nonlinear_and_linear_evolution_with_GS.ipynb b/examples/example05a - nonlinear_and_linear_evolution_with_GS.ipynb index c7f51d1..98a7a25 100644 --- a/examples/example05a - nonlinear_and_linear_evolution_with_GS.ipynb +++ b/examples/example05a - nonlinear_and_linear_evolution_with_GS.ipynb @@ -217,6 +217,7 @@ "- `min_dIy_dI`: threshold value below which passive structure normal modes are dropped. Modes with norm(d(Iy)/dI)<`min_dIy_dI` are dropped, which filters out modes that do not actually couple with the plasma.\n", "- `max_mode_frequency`: threshold value for characteristic frequencies above which passive structure normal modes are dropped (i.e. the fast modes).\n", "- **`n_linearization_workers`**: number of independent worker processes used to build the current-response (`dIydI`) and profile-parameter-response (`dIydtheta`) Jacobian columns. The default, `1`, retains serial construction. Values larger than one parallelise both the initial linearisation and later relinearisations without limiting the inner threading of ordinary time-step GS solves.\n", + "- `cache_myy`: determines whether the Myy matrix should be calculated once and saved, or re-calculated as needed. The default is `True`, setting to `False` reduces memory usage significantly, allowing for more resolved meshes to be used. Note that not caching Myy can have serious performance impact when evolving the plasma non-linearly (i.e. when passing `linear_only=False` to the `nlstepper`); for linear evolution no noticeable impact should be expected.\n", "\n", "Other customisable inputs are available, do see the documentation or the later notebook on \"Growth Rates\" for more details. For example, one may explicitly set your own resistance and inductance matrices for the tokamak, rather than the geometrical value calculated internally in FreeGSNKE.\n", "\n", @@ -241,6 +242,7 @@ " plasma_resistivity=1e-6,\n", " max_mode_frequency=10**2.5,\n", " n_linearization_workers=1, # set >1 to parallelise Jacobian construction\n", + " cache_myy=True,\n", ")" ] }, diff --git a/freegsnke/GSstaticsolver.py b/freegsnke/GSstaticsolver.py index 4b4e602..8e816d3 100644 --- a/freegsnke/GSstaticsolver.py +++ b/freegsnke/GSstaticsolver.py @@ -24,6 +24,8 @@ import freegs4e import numpy as np from freegs4e.gradshafranov import Greens +from freegs4e.gs_solver import GSDSTSolver, GSLUSolver +from freegs4e.multigrid import createMultigridSolver from . import nk_solver_H as nk_solver @@ -73,7 +75,11 @@ def __init__( l2_reg=1e-6, collinearity_reg=1e-6, seed=42, - gs_operator_order=4, + gs_operator_order=None, + cache_greens=True, + solver_type="LUsparse", + mg_kwargs=None, + max_batch_size=2_000_000_000, ): """ Initialise the Grad–Shafranov nonlinear solver. @@ -111,11 +117,33 @@ def __init__( • Krylov perturbation generation • Directional exploration in nonlinear solve - gs_operator_order : {2, 4}, optional (default=4) + gs_operator_order : {2, 4}, optional (default=None) Finite-difference order of the linear Grad-Shafranov operator. - Fourth order is more accurate; second order reduces sparse matrix - construction and factorisation costs when that accuracy trade-off - is acceptable. + Fourth order is more accurate; second order may reduce execution + time. Default is the highest value supported by `solver_type`. It + is STRONGLY recommended to leave the default value. + + cache_greens : bool, optional (default=True) + Determines whether the Greens function should be pre-calculated and + stored, or re-calculated every time the boundary flux is computed + (reducing memory usage by up to an order of magnitude). + + solver_type: str (default='LUsparse') + The type of linear solver to use for the GS equation. Supported + options are 'LUsparse', 'DST', 'multigrid'. 'LUsparse' supports + order 2,4. The 'DST' solver is faster, but only supports order 2. + Use of 'multigrid' is discouraged and is planned to be deprecated. + + mg_kwargs + dict with optional keyword arguments to pass to + `freegs4e.multigrid.createMultigridSolver` during multigrid solver + initialization. Ignored whenever `solver_type` != 'multigrid'. + + max_batch_size : int, float (default=2_000_000_000) + Sets a maximum size (number of entries) for the batches of the Greens function when calculating the + boundary flux, allowing for further reduction in memory footprint. Only applicable when + `cache_greens=False`, ignored otherwise. Float type values supported for convenience (to allow for + use of scientific notation, e.g. `max_batch_size=2e9`). Attributes ---------- @@ -164,14 +192,6 @@ def __init__( dZ = Z[0, 1] - Z[0, 0] self.dRdZ = dR * dZ - if gs_operator_order == 2: - gs_operator = freegs4e.gradshafranov.GSsparse - elif gs_operator_order == 4: - gs_operator = freegs4e.gradshafranov.GSsparse4thOrder - else: - raise ValueError("gs_operator_order must be either 2 or 4") - self.gs_operator_order = gs_operator_order - # nonlinear solver backend self.nksolver = nk_solver.nksolver( problem_dimension=self.nx * self.ny, @@ -179,16 +199,8 @@ def __init__( collinearity_reg=collinearity_reg, ) - # linear GS solver used inside nonlinear iteration - self.linear_GS_solver = freegs4e.multigrid.createVcycle( - nx, - ny, - gs_operator(eq.R[0, 0], eq.R[-1, 0], eq.Z[0, 0], eq.Z[0, -1]), - nlevels=1, - ncycle=1, - niter=2, - direct=True, - ) + # define the GS linear solver (del*Psi=RHS with fixed RHS) + self.configureLinearSolver(solver_type, gs_operator_order, mg_kwargs) # collect boundary grid indices for Dirichlet conditions bndry_indices = np.concatenate( @@ -201,12 +213,13 @@ def __init__( ) self.bndry_indices = bndry_indices - # Plasma current is confined inside the limiter, so only those Green - # columns contribute to the free-boundary condition. - self.plasma_source_mask = np.asarray( - eq.limiter_handler.mask_inside_limiter, dtype=bool - ) - self.greenfunc = self._build_boundary_green(self.plasma_source_mask) + # Cache Greens function if necessary + if cache_greens: + self.greenfunc = self._build_full_boundary_greens() + self._max_batch_size = None + else: + self.greenfunc = None + self._max_batch_size = int(max_batch_size) # Precompute geometric RHS coefficient # Comes from GS equation: @@ -216,7 +229,134 @@ def __init__( # random generator used for NK search direction exploration self.rng = np.random.default_rng(seed=seed) - def _build_boundary_green(self, source_mask): + def _build_full_boundary_greens(self): + """ + Calculates the Greens function giving the responses of boundary nodes to internal nodes: + Jtor(R',Z') → ψ_boundary(R,Z) + + Fills the array sequentially to optimize memory usage. + """ + + bndry_indices = self.bndry_indices + n_bndry_nodes = bndry_indices.shape[0] + + R_1D = self.R[:, 0] + Z_1D = self.Z[0, :] + + # Pre-allocate full array + greenfunc = np.empty( + (n_bndry_nodes, self.R.shape[0], self.R.shape[1]), + ) + + # Computation is performed in batches, to limit memory consumption + # Batches correspond to slices over ROWS of the Green's function + + num_batches = 16 # fine-tuned to balance memory vs. compute needs + batch_len = n_bndry_nodes // num_batches # number of rows in the batch + + for i in range(num_batches): + + start = i * batch_len + end = start + batch_len + end = ( + end if i != num_batches - 1 else n_bndry_nodes + ) # last batch gets the remainder + + # Fill up slice of greenfunc in-place. Applies dRdZ factor automatically. + Greens( + self.R[np.newaxis, :, :], + self.Z[np.newaxis, :, :], + R_1D[bndry_indices[:, 0]][start:end, np.newaxis, np.newaxis], + Z_1D[bndry_indices[:, 1]][start:end, np.newaxis, np.newaxis], + scale_factor=self.dRdZ, + out=greenfunc[start:end, :, :], + ) + + # filter out Greens(x,y;x,y), to prevent infinity/NaNs + greenfunc[start:end, bndry_indices[:, 0], bndry_indices[:, 1]] = 0 + + return greenfunc + + def _calculate_boundary_flux(self): + """ + Compute boundary flux via Green's function convolution: psi_bnd = ∫ G(R,Z; R',Z') Jtor(R',Z') dR'dZ' + + Implemented using tensor contraction. + + If Green's function was precomputed, uses the cached version. Otherwise, computes it on the fly. + + Returns + ------- + Boundary flux vector as a flattened array + """ + + # Implemented using tensor contraction: + # Contract: + # greenfunc axis (1,2) with jtor axis (0,1) + + if self.greenfunc is not None: + psi_bnd = np.tensordot(self.greenfunc, self.jtor, axes=([1, 2], [0, 1])) + + else: + + # Computation is performed in batches, to limit memory consumption + # Batches correspond to contiguous groups of FULL ROWS of the Green's function + # (only last dimension is sliced) + + bndry_indices = self.bndry_indices + n_bndry_nodes = bndry_indices.shape[0] + + R_1D = self.R[:, 0] + Z_1D = self.Z[0, :] + + psi_bnd = np.empty(n_bndry_nodes) + + # calculating psi_bnd in 16 batches reduces total RSS contribution to 0.25*greenfunc.nbytes + default_num_batches = 16 + + # determine maximum number of rows allowed for a batch + max_batch_len = self._max_batch_size // (self.nx * self.ny) + + # determine number of rows in a batch + batch_len = n_bndry_nodes // default_num_batches + batch_len = min(batch_len, max_batch_len) + batch_len = max(batch_len, 2) # ensure at least two rows are in a batch + + # calculate number of batches (adds one more if there is a remainder) + num_batches = (n_bndry_nodes - 1) // batch_len + 1 + + # pre-allocate re-usable buffer + greenfunc_buff = np.empty((batch_len, self.nx, self.ny)) + + for i in range(num_batches): + + start = i * batch_len + end = start + batch_len + end = min(end, n_bndry_nodes) # last batch can be smaller + + greenfunc = greenfunc_buff[: end - start] + + # Calculate greenfunc for these boundary nodes. Applies dRdZ factor automatically. + Greens( + self.R[np.newaxis, :, :], + self.Z[np.newaxis, :, :], + R_1D[bndry_indices[:, 0]][start:end, np.newaxis, np.newaxis], + Z_1D[bndry_indices[:, 1]][start:end, np.newaxis, np.newaxis], + scale_factor=self.dRdZ, + out=greenfunc, + ) + + # filter out values at boundary Greens(x,y;x,y), to prevent infinity/NaNs + greenfunc[:, bndry_indices[:, 0], bndry_indices[:, 1]] = 0 + + # weighted sum over the last two axes + psi_bnd[start:end] = np.tensordot( + greenfunc, self.jtor, axes=([1, 2], [0, 1]) + ) + + return psi_bnd + + def _build_masked_boundary_greens(self, source_mask): """Build the boundary Green matrix for a selected set of source points.""" source_indices = np.flatnonzero(source_mask) boundary_indices = np.ravel_multi_index( @@ -241,10 +381,62 @@ def _build_boundary_green(self, source_mask): greenfunc[np.flatnonzero(matches), positions[matches]] = 0.0 return np.ascontiguousarray(greenfunc * self.dRdZ) - def _boundary_flux_from_jtor(self, jtor): + def _calculate_masked_boundary_flux(self, jtor): """Return boundary flux from plasma current inside the limiter.""" return self.greenfunc @ jtor[self.plasma_source_mask] + def configureLinearSolver(self, solver_type, order, mg_kwargs): + """ + Creates and assigns the linear solver `self.linear_GS_solver` using the arguments provided. + + Also sets the attribute `self.gs_operator_order`. + + Parameters + ---------- + solver_type: str + The type of linear solver to use for the GS equation. Supported options are 'LUsparse', + 'DST', 'multigrid'. + order : int + Order of differential operators used in calculations. + Must be either 2 or 4. + mg_kwargs + dict with kwargs to pass to `freegs4e.multigrid.createMultigridSolver` during multigrid solver + initialization. Ignored whenever `solver_type` != 'multigrid'. + """ + + if solver_type == "LUsparse": + if order is None: + order = 4 + self.linear_GS_solver = GSLUSolver(self.R, self.Z, order=order) + + elif solver_type == "DST": + if order is None: + order = 2 + self.linear_GS_solver = GSDSTSolver(self.R, self.Z, order=order) + + elif solver_type == "multigrid": + + if order is None: + order = 4 + if mg_kwargs is None: + mg_kwargs = {} + elif not isinstance(mg_kwargs, dict): + raise TypeError("mg_kwargs needs to be of type dict") + + self.linear_GS_solver = createMultigridSolver( + self.R, + self.Z, + order, + **mg_kwargs, + ) + + else: + raise ValueError(f"Solver type {solver_type} not recognized") + + self.gs_operator_order = order + + return self.linear_GS_solver + def freeboundary(self, plasma_psi, tokamak_psi, profiles): """ Apply free-boundary Grad–Shafranov boundary conditions and compute @@ -319,22 +511,19 @@ def freeboundary(self, plasma_psi, tokamak_psi, profiles): self.rhs = self.rhs_before_jtor * self.jtor # ------------------------------------------------------------ - # Compute boundary flux via Green's function convolution - # - # psi_boundary = ∫ G(R,Z; R',Z') Jtor(R',Z') dR'dZ' - # - # Implemented as a matrix-vector product over source points inside the - # limiter, outside which the plasma current is identically zero. + # Calculate boundary flux (flat array) # ------------------------------------------------------------ - self.psi_boundary = np.zeros_like(self.R) - psi_bnd = self._boundary_flux_from_jtor(self.jtor) + psi_bnd = self._calculate_boundary_flux() # ------------------------------------------------------------ # Map flattened Green's solution back to boundary grid # ------------------------------------------------------------ + self.psi_boundary = np.zeros_like(self.R) + # Vertical boundaries self.psi_boundary[:, 0] = psi_bnd[: self.nx] self.psi_boundary[:, -1] = psi_bnd[self.nx : 2 * self.nx] + # Horizontal boundaries self.psi_boundary[0, 1 : self.ny - 1] = psi_bnd[ 2 * self.nx : 2 * self.nx + self.ny - 2 diff --git a/freegsnke/Myy_builder.py b/freegsnke/Myy_builder.py index 2892e9e..ffd10e9 100644 --- a/freegsnke/Myy_builder.py +++ b/freegsnke/Myy_builder.py @@ -1,5 +1,5 @@ """ -Defines the plasma_current Object, which handles the lumped parameter model +Defines the plasma_current Object, which handles the lumped parameter models used as an effective circuit equation for the plasma. Copyright 2025 UKAEA, UKRI-STFC, and The Authors, as per the COPYRIGHT and README files. @@ -15,96 +15,132 @@ it under the terms of the GNU Lesser General Public License as published by the Free Software Foundation, either version 3 of the License, or (at your option) any later version. - + You should have received a copy of the GNU Lesser General Public License -along with FreeGSNKE. If not, see . +along with FreeGSNKE. If not, see . """ +import abc + +import numexpr as ne import numpy as np from freegs4e.gradshafranov import Greens +from freegs4e.parallel import threaded_take + + +def make_Myy_handler( + domain_type, limiter_handler, layer_size=None, tolerance=None, cache_myy=None +): + """Returns an Myy_handler object of appropriate type given the domain_type specified. + + Parameters + ---------- + domain_type: str + Type of domain to define Myy over. Supported values: "reduced", "fft". + limiter_handler : FreeGSNKE limiter object, i.e. eq.limiter_handler + Sets the properties of the domain grid and those of the limiter + layer_size : int, optional + (for reduced domain only) + Used when recalculating myy. + A layer of layer_size pixels is added to envelop the mask defined by the + plasma. This 'broadened' mask defines the pixels included in the myy matrix + By default 5 + tolerance : int, optional + (for reduced domain only) + Used to check if myy needs recalculating. Myy is not recalculated if + the mask defined by the plasma region, broadened by tolerance pixels, + is fully contained in the domain of the current myy matrix, + By default 3 + + """ + + # TODO: perhaps replace `domain_type` with something more generic that directly distinguishes + # between cached and otf reduced Myy. This may be a good idea depending on how the GPU version + # is coupled. (options: "real_cached", "real_otf", "fft", "gpu") + + if domain_type == "reduced": + + # Apply defaults + layer_size = 5 if layer_size is None else layer_size + tolerance = 3 if tolerance is None else tolerance + cache_myy = True if cache_myy is None else cache_myy + + return Reduced_Myy_handler(limiter_handler, layer_size, tolerance, cache_myy) + + elif domain_type == "fft": + return FFT_Myy_handler(limiter_handler) + + else: + raise ValueError(f"Domain type {domain_type} not recognized.") + + +class Myy_handler(abc.ABC): + """Abstract base class for objects handling all operations which involve the Myy matrix, + i.e. the mututal inductance matrix of all domain grid points. + """ + + def make_1D(self, R, Z): + + dz = Z[0, 1] - Z[0, 0] + nZ = Z.shape[1] + Z_1D = np.arange(0, dz * nZ, dz) + R_1D = R[:, 0] + + return R_1D, Z_1D + + @abc.abstractmethod + def grid_greens(self, R_1D, Z_1D): + """Calculates and stores the Green's function values in the corresponding Myy domain. + + Parameters + ---------- + R_1D : np.ndarray + Array of dim=1 representing the R values of the grid in the appropriate Myy domain. + Z_1D : np.ndarray + Array of dim=1 representing the Z values of the grid in the appropriate Myy domain. + """ + pass + + @abc.abstractmethod + def check_Myy(self, hatIy): + """As defined in Reduced_Myy_handler. Definition enforced for compatibility with + non_linear_solver""" + pass + + @abc.abstractmethod + def force_build_Myy(self, hatIy): + """As defined in Reduced_Myy_handler. Definition enforced for compatibility with + non_linear_solver""" + pass + + @abc.abstractmethod + def dot(self, hatIy): + """Performs the product with a vector defined on the reduced plasma domain, i.e. inside the + limiter, independently of the domain of Myy. Returns a vector on the input domain. + + Parameters + ---------- + hatIy : np.ndarray + 1d vector on reduced plasma domain, e.g. inside the limiter + """ + pass + + @abc.abstractmethod + def _myy_dot(self, vector): + """Performs the dot product Myy@vector for a vector already projected to the domain + in which Myy is defined. Returns a vector on the input domain.""" + pass + -# class plasma_current: -# """Implements the plasma circuit equation in projection on $I_{y}^T$: - -# $$I_{y}^T/I_p (M_{yy} \dot{I_y} + M_{ye} \dot{I_e} + R_p I_y) = 0$$ -# """ - -# def __init__(self, plasma_pts, Rm1, P, plasma_resistance_1d, Mye): -# """Implements the object dealing with the plasma circuit equation in projection on $I_y$, -# I_y being the plasma toroidal current density distribution: - -# $$I_{y}^T/I_p (M_{yy} \dot{I_y} + M_{ye} \dot{I_e} + R_p I_y) = 0$$ - -# Parameters -# ---------- -# plasma_pts : freegsnke.limiter_handler.plasma_pts -# Domain points in the domain that are included in the evolutive calculations. -# A typical choice would be all domain points inside the limiter. Defaults to None. -# Rm1 : np.ndarray -# The diagonal matrix of all metal vessel resistances to the power of -1 ($R^{-1}$). -# P : np.ndarray -# Matrix used to change basis from normal mode currents to vessel metal currents. -# plasma_resistance_1d : np.ndarray -# Vector of plasma resistance values for all grid points in the reduced plasma domain. -# plasma_resistance_1d = 2pi resistivity R/dA for all plasma_pts -# Mye : np.ndarray -# Matrix of mutual inductances between plasma grid points and all vessel coils. - -# """ - -# self.plasma_pts = plasma_pts -# self.Rm1 = Rm1 -# self.P = P -# self.Mye = Mye -# self.Ryy = plasma_resistance_1d -# self.Myy_matrix = self.Myy() - -# def reset_modes(self, P): -# """Allows a reset of the attributes set up at initialization time following a change -# in the properties of the selected normal modes for the passive structures. - -# Parameters -# ---------- -# P : np.ndarray -# New change of basis matrix. -# """ -# self.P = P - - -# def Myy( -# plasma_pts, -# ): -# """Calculates the matrix of mutual inductances between all plasma grid points - -# Parameters -# ---------- -# plasma_pts : np.ndarray -# Array with R and Z coordinates of all the points inside the limiter - -# Returns -# ------- -# Myy : np.ndarray -# Array of mutual inductances between plasma grid points -# """ -# greenm = Greens( -# plasma_pts[:, np.newaxis, 0], -# plasma_pts[:, np.newaxis, 1], -# plasma_pts[np.newaxis, :, 0], -# plasma_pts[np.newaxis, :, 1], -# ) -# return 2 * np.pi * greenm - - -class Myy_handler: +class Reduced_Myy_handler(Myy_handler): """Object handling all operations which involve the Myy matrix, i.e. the mututal inductance matrix of all domain grid points. To reduce memory usage, the domain on which myy is built and stored is set adaptively, so to cover the plasma. This object handles this adaptive aspect. - """ - def __init__(self, limiter_handler, layer_size=5, tolerance=3): + def __init__(self, limiter_handler, layer_size=5, tolerance=3, cache_myy=True): """Instantiates the object Parameters @@ -135,39 +171,46 @@ def __init__(self, limiter_handler, layer_size=5, tolerance=3): self.idxs_mask_red = self.extract_index_mask(self.mask_inside_limiter_red) - self.gg = self.grid_greens( - self.reduce_rect_domain(limiter_handler.eqR), - self.reduce_rect_domain(limiter_handler.eqZ), - ) + R_red = self.reduce_rect_domain(limiter_handler.eqR) + Z_red = self.reduce_rect_domain(limiter_handler.eqZ) + R_1D, Z_1D = self.make_1D(R_red, Z_red) + + self.gg = self.grid_greens(R_1D, Z_1D) self.layer_size = layer_size self.tolerance = tolerance - def grid_greens(self, R, Z): - """Calculates and stores the green function values on the minimal rectangular - region that fully encompasses the limiter. Uses that the green functions are invariant - for vertical translations. + self.cache_myy = cache_myy + + def grid_greens(self, R_1D, Z_1D): + """Calculates the Green's function values on the minimal rectangular region that fully + encompasses the limiter. Uses that the green functions are invariant for vertical + translations. Parameters ---------- - R : np.ndarray - Like eq.R, but on the rectangular reduced domain, - i.e. self.reduce_rect_domain(limiter_handler.eqR) + R_1D : np.ndarray + Like eq.R, but on the rectangular reduced domain and only in 1D + i.e. self.make_1D(self.reduce_rect_domain(limiter_handler.eqR),...) Z : np.ndarray - Like eq.Z, but on the rectangular reduced domain - """ + Like eq.Z, but on the rectangular reduced domain and only in 1D - dz = Z[0, 1] - Z[0, 0] - nZ = np.shape(Z)[1] + Returns + ------- + ggreens: np.ndarray + Array of shape (nR,nR,nZ) containing the values of the Green's function in the reduced + rectangular domain. + """ ggreens = Greens( - R[:, 0][:, np.newaxis, np.newaxis], - dz * np.arange(nZ)[np.newaxis, np.newaxis, :], - R[:, 0][np.newaxis, :, np.newaxis], + R_1D[:, np.newaxis, np.newaxis], + Z_1D[np.newaxis, np.newaxis, :], + R_1D[np.newaxis, :, np.newaxis], 0, + scale_factor=2.0 * np.pi, ) - return 2 * np.pi * ggreens + return ggreens def build_mask_from_hatIy(self, hatIy, layer_size): """Builds the mask that will be used by build_myy_from_mask @@ -200,6 +243,7 @@ def build_Myy_from_mask(self, mask): i.e. the smallest rectangular domain around limiter mask (same size as self.mask_inside_limiter_red) """ + self.myy_mask_red = mask self.outside_myy_mask = np.logical_not(mask) @@ -207,16 +251,26 @@ def build_Myy_from_mask(self, mask): self.idxs_myy_mask_red = self.extract_index_mask(mask) - r_idxs = np.tile( - self.idxs_myy_mask_red[0][:, np.newaxis], - (1, nmask), - ) - dz_idxs = np.abs( - self.idxs_myy_mask_red[1][np.newaxis, :] - - self.idxs_myy_mask_red[1][:, np.newaxis] - ) + if self.cache_myy: + dz_idxs = self.idxs_myy_mask_red[1] + r_idxs = self.idxs_myy_mask_red[0] + + self.myy = np.empty((nmask, nmask)) + + d1, d2, d3 = self.gg.shape + d23 = d2 * d3 + + # important to keep this as a python loop, do not try to vectorize + for i in range(nmask): - self.myy = self.gg[r_idxs, r_idxs.T, dz_idxs] + idxs1 = r_idxs + idxs2 = r_idxs[i] + + idxs3 = np.abs(dz_idxs[i] - dz_idxs) + idcs = idxs1 * d23 + idxs2 * d3 + idxs3 + + # same as self.myy[i] = self.gg.reshape(-1)[idcs] but faster + np.take(self.gg, idcs, out=self.myy[i], mode="wrap") def force_build_Myy(self, hatIy): """Builds the Myy matrix only including domain points in the input vector (not necessarily a mask) @@ -254,7 +308,7 @@ def check_Myy(self, hatIy): def dot(self, hatIy): """Performs the product with a vector defined on the reduced domain, i.e. inside the limiter. - Returns a vector on the same domain. + Returns a vector on the input domain. Parameters ---------- @@ -268,7 +322,7 @@ def dot(self, hatIy): hatIy_myy_red = hatIy_rect_red[self.myy_mask_red] # perform the dot product - result = np.dot(self.myy, hatIy_myy_red) + result = self._myy_dot(hatIy_myy_red) # bring result back to the reduced plasma domain result_rect_red = self.rebuild_map2d( @@ -277,3 +331,219 @@ def dot(self, hatIy): result_red = result_rect_red[self.mask_inside_limiter_red] return result_red + + def _myy_dot(self, vector): + """Performs the dot product Myy@vector for a vector already projected onto the reduced + rectangular domain in which Myy is defined. + + If myy is cached, np.dot is called directly. Otherwise, the product is batch-computed using + an on-the-fly calculation for blocks of myy (myy calculation is parallelized). + + Returns a vector on the input domain. + + Parameters + ---------- + vector: np.ndarray + 1d vector projected onto the reduced rectangular domain of Myy + + """ + + if self.cache_myy: + return np.dot(self.myy, vector) + + else: + nmask = np.sum(self.myy_mask_red) + + inshape = vector.shape + + if len(inshape) > 1: + outshape = (nmask, *inshape[:-2], inshape[-1]) + else: + outshape = (nmask,) + + dz_idxs = self.idxs_myy_mask_red[1] + r_idxs = self.idxs_myy_mask_red[0] + + d1, d2, d3 = self.gg.shape + d23 = d2 * d3 + + num_slices = 20 # TODO: perhaps instead of fixing the number of slices, fix the block size? + step = (nmask - 1) // num_slices + 1 + + idcs = np.empty((step, nmask), dtype=np.int64) + myy_buff = np.empty(idcs.shape) + result = np.empty(outshape) + + for i in range(num_slices): + + start = i * step + end = start + step + end = min(end, nmask) + + idxs1 = r_idxs[np.newaxis] + idxs2 = r_idxs[start:end, np.newaxis] + idxs3a = dz_idxs[start:end, np.newaxis] + idxs3b = dz_idxs[np.newaxis] + + # idcs is flattened version of (idxs1,idxs2,idxs3) + # TODO: check why there are casting issues with abs() + ne.evaluate( + "idxs1*d23 + idxs2*d3 + abs(idxs3a-idxs3b)", + out=idcs[: end - start], + casting="unsafe", + ) + + # same as self.myy_buff[:end-start] = self.gg.reshape(-1)[idcs_slice] but faster + threaded_take( + self.gg, + idcs[: end - start], + out=myy_buff[: end - start], + mode="wrap", + ) + + np.dot(myy_buff[: end - start], vector, out=result[start:end]) + + return result + + +class FFT_Myy_handler(Myy_handler): + """Object handling all operations which involve the Myy matrix, + i.e. the mututal inductance matrix of all domain grid points. + To reduce memory usage, the Green's function is computed in Fourier + space, which reduces the size scaling from quartic to cubic. + """ + + def __init__(self, limiter_handler): + """Instantiates the object + + Parameters + ---------- + limiter_handler : FreeGSNKE limiter object, i.e. eq.limiter_handler + Sets the properties of the domain grid and those of the limiter + """ + + self.up_project = limiter_handler.up_project + self.down_project = limiter_handler.down_project + + R_1D, Z_1D = self.make_1D(limiter_handler.eqR, limiter_handler.eqZ) + + self.gg = self.grid_greens(R_1D, Z_1D) # in Fourier space, (nR, nR, L_fft) + + def grid_greens(self, R_1D, Z_1D): + """Calculates the Green's function values on the Fourier space of the full geometric + domain defined by the limiter. Uses that the green functions are invariant for vertical + translations. + + Parameters + ---------- + R : np.ndarray + Like eq.R, but only in 1D + i.e. self.make_1D(limiter_handler.eqR,...) + Z : np.ndarray + Like eq.Z, but only in 1D + + Returns + ------- + gg_fft: np.ndarray + Array of shape (nR,nR,L_fft) containing the values of the Green's function in Fourier + space. (L_fft = (2*nZ - 1) // 2 + 1) + """ + + # TODO: worth evaluating whether L_fft is ever different from nZ + + nR = len(R_1D) + nZ = len(Z_1D) + + # Linear convolution length: L = 2*nZ - 1 + L = 2 * nZ - 1 + + # h[i,j,k] = 2π * Greens(R_i, Z_k, R_j, 0) + h_full = np.empty((nR, nR, nZ)) + + num_slices = 10 # fine-tuned to balance memory vs. compute needs + step = nR // num_slices + + for i in range(num_slices): + + start = i * step + end = start + step + end = end if i != num_slices - 1 else nR # last slice gets the remainder + + # Fill up slice of h_full in-place. Applies 2π factor automatically. + Greens( + R_1D[start:end, np.newaxis, np.newaxis], + Z_1D[np.newaxis, np.newaxis, :], + R_1D[np.newaxis, :, np.newaxis], + 0.0, + scale_factor=2.0 * np.pi, + out=h_full[start:end], + ) + + # build symmetric kernel g of length L for *linear* Toeplitz convolution: + # g[..., k] = h[..., k] for k=0..nZ-1 + # g[..., L-k] = h[..., k] for k=1..nZ-1 + g = np.zeros((nR, nR, L), dtype=h_full.dtype) + k = np.arange(1, nZ) + + g[:, :, 0] = h_full[:, :, 0] + g[:, :, k] = h_full[:, :, k] + g[:, :, L - k] = h_full[:, :, k] + + # Perform rFFT of g along Z, and return result + gg_fft = np.fft.rfft(g, axis=2) + + return gg_fft # (nR, nR, L_fft) + + def check_Myy(self, hatIy): + """Dummy, always returns False. Defined only for compatibility with non_linear_solver""" + return False + + def force_build_Myy(self, hatIy): + """Ignored. Defined only for compatibility with non_linear_solver""" + pass + + def dot(self, hatIy): + """Performs the product with a vector defined on the reduced domain, i.e. inside the limiter. + Returns a vector on the input domain. + + Parameters + ---------- + hatIy: np.ndarray + 1d vector on reduced plasma domain, e.g. inside the limiter + """ + up_hatIy = self.up_project(hatIy) + Myy_hatIy = self._myy_dot(up_hatIy) + reduced_prod = self.down_project(Myy_hatIy) + return reduced_prod + + def _myy_dot(self, vector): + """Performs the dot product Myy@vector for a vector that has already been up-projected onto + the full grid domain (where Myy is defined), but NOT yet transformed into Fourier space. + + Returns a vector on the input domain. + + Parameters + ---------- + vector: np.ndarray + 1d vector on the full domain, in real space + """ + + nR, nZ = self.gg.shape[1], self.gg.shape[2] + x = vector.reshape(nR, nZ) + + # Zero-pad vec along Z to length L + L = 2 * nZ - 1 + pad_width = L - nZ # equals nZ - 1 + x_padded = np.pad(x, ((0, 0), (0, pad_width))) # (nR, L) + + # rFFT of padded vec + x_fft = np.fft.rfft(x_padded, axis=1) # (nR, L_fft) + + gg_fft = self.gg + x_fft = x_fft[np.newaxis, :, :] + conv_fft = ne.evaluate("sum(gg_fft*x_fft, axis=1)") + + y_full = np.fft.irfft(conv_fft, n=L, axis=1) + y_full = y_full[:, :nZ] # .reshape(-1) + + return y_full diff --git a/freegsnke/limiter_func.py b/freegsnke/limiter_func.py index d967b1b..eebb112 100644 --- a/freegsnke/limiter_func.py +++ b/freegsnke/limiter_func.py @@ -656,6 +656,16 @@ def hat_Iy_from_jtor(self, jtor): hat_Iy = self.normalize_sum(hat_Iy) return hat_Iy + def up_project(self, vec): + """Projects a (discretized current) vector defined on the reduced plasma domain onto the full plasma domain.""" + new_vec = np.zeros_like(self.map2d) + new_vec[self.mask_inside_limiter] = vec + return new_vec + + def down_project(self, vec): + """Projects a (discretized current) vector defined on the full plasma domain onto the reduced plasma domain.""" + return vec[self.mask_inside_limiter] + def rebuild_map2d(self, reduced_vector, map_dummy, idxs_mask): """Rebuilds 2d map on full domain corresponding to 1d vector reduced_vector on smaller plasma domain diff --git a/freegsnke/machine_config.py b/freegsnke/machine_config.py index 8c2c9ca..b605473 100644 --- a/freegsnke/machine_config.py +++ b/freegsnke/machine_config.py @@ -168,6 +168,7 @@ def _calc_mutual_inductance_entry(tokamak, name_i, name_j): coords_i[1][np.newaxis, :], coords_j[0][:, np.newaxis], coords_j[1][:, np.newaxis], + limit_threading=True, ) if name_j == name_i: diff --git a/freegsnke/nonlinear_solve.py b/freegsnke/nonlinear_solve.py index 6c3e1d7..6e8f393 100644 --- a/freegsnke/nonlinear_solve.py +++ b/freegsnke/nonlinear_solve.py @@ -35,7 +35,7 @@ from .circuit_eq_metal import metal_currents from .GSstaticsolver import NKGSsolver from .linear_solve import linear_solver -from .Myy_builder import Myy_handler +from .Myy_builder import make_Myy_handler from .simplified_solve import simplified_solver_J1 _parallel_linearization_solver = None @@ -73,6 +73,8 @@ class nl_solver: _MAX_STARTING_DI_RATIO = np.sqrt(10.0) _MAX_REUSED_STARTING_DI_RATIO = 4.0 / 3.0 + # TODO: would be wise to make these kwargs kw-only, to prevent users from using + # them as positional, which risks human error and regression bugs def __init__( self, profiles, @@ -103,6 +105,8 @@ def __init__( plasma_descriptor_function=None, mode_selection="coupling", n_linearization_workers=1, + myy_type="reduced", + cache_myy=True, ): """ Initialize the nonlinear solver. @@ -191,6 +195,11 @@ def __init__( Number of worker processes used to build independent ``dIydI`` and ``dIydtheta`` columns during initial and later linearisations. A value of 1 retains the serial calculation. + myy_type: str, default="reduced" + Either "reduced" to define Myy over reduced domain in real space, or "fft" to define Myy + over the full domain in Fourier space. + cache_myy: bool, default=True + Controls whether the Myy matrix is cached at initialization or rebuilt for every product. """ print("-----") @@ -266,7 +275,9 @@ def __init__( self.nIy = np.linalg.norm(self.Iy) # instantiate the Myy_handler object - self.handleMyy = Myy_handler(eq.limiter_handler) + self.handleMyy = make_Myy_handler( + myy_type, eq.limiter_handler, cache_myy=cache_myy + ) # Extract relevant information on the type of profiles function used and on the actual value of associated parameters self.get_profiles_values(profiles) diff --git a/freegsnke/normal_modes.py b/freegsnke/normal_modes.py index fd738de..6b5e72f 100644 --- a/freegsnke/normal_modes.py +++ b/freegsnke/normal_modes.py @@ -137,6 +137,10 @@ def normal_modes_greens(self, eq_vgreen): Can be found at eq._vgreen. np.shape(eq_vgreen)=(n_coils, nx, ny) """ - grid_shape = eq_vgreen.shape[1:] - physical_greens = eq_vgreen.reshape(self.n_coils, -1) - return (self.Pmatrix.T @ physical_greens).reshape(self.n_coils, *grid_shape) + # grid_shape = eq_vgreen.shape[1:] + # physical_greens = eq_vgreen.reshape(self.n_coils, -1) + # dgreen = (self.Pmatrix.T @ physical_greens).reshape(self.n_coils, *grid_shape) + + dgreen = np.tensordot(self.Pmatrix, eq_vgreen, axes=([0], [0])) + + return dgreen diff --git a/freegsnke/tests/test_dynamics.py b/freegsnke/tests/test_dynamics.py index b40fe6b..d433d69 100644 --- a/freegsnke/tests/test_dynamics.py +++ b/freegsnke/tests/test_dynamics.py @@ -53,9 +53,6 @@ def create_machine(): alpha_n=1.2, ) - from freegsnke import GSstaticsolver - - NK = GSstaticsolver.NKGSsolver(eq) currents = np.array( [ 40000, @@ -75,6 +72,27 @@ def create_machine(): keys = list(eq.tokamak.getCurrents().keys()) for i in np.arange(12): eq.tokamak.set_coil_current(keys[i], currents[i]) + + return tokamak, eq, profiles + + +@pytest.fixture() +def create_solver(create_machine, request): + + tokamak, eq, profiles = create_machine + if request.param == "default": + request.param = { + "solver_type": "LUsparse", + "myy_type": "reduced", + "cache_myy": True, + } + + from freegsnke import GSstaticsolver + + NK = GSstaticsolver.NKGSsolver( + eq, + solver_type=request.param["solver_type"], + ) NK.solve(eq, profiles, target_relative_tolerance=1e-8) # Initialize the evolution object @@ -94,12 +112,27 @@ def create_machine(): # modes. Plasma-coupling metrics may calibrate finite-difference steps but # must not change which modes are retained. mode_selection="timescale", + myy_type=request.param["myy_type"], + cache_myy=request.param["cache_myy"], ) - return tokamak, eq, profiles, stepping - - -def test_linearised_growth_rate(create_machine): - tokamak, eq, profiles, stepping = create_machine + return stepping + + +@pytest.mark.parametrize( + "create_solver", + [ + {"solver_type": "LUsparse", "myy_type": "reduced", "cache_myy": True}, + pytest.param( + {"solver_type": "DST", "myy_type": "reduced", "cache_myy": True}, + marks=pytest.mark.xfail, + ), + ], + indirect=True, +) +def test_linearised_growth_rate(create_machine, create_solver): + + tokamak, eq, profiles = create_machine + stepping = create_solver selected_passive_modes = stepping.evol_metal_curr.selected_modes_mask[ stepping.n_active_coils : ] @@ -108,6 +141,7 @@ def test_linearised_growth_rate(create_machine): np.arange(50), ) true_GR = 0.05900 + # check that assert ( abs((stepping.linearised_sol.instability_timescale[0] - true_GR) / true_GR) @@ -115,8 +149,20 @@ def test_linearised_growth_rate(create_machine): ), f"Growth rate deviates { abs((stepping.linearised_sol.growth_rates[0]-true_GR)/true_GR)}% from baseline" -def test_linearised_stepper(create_machine): - tokamak, eq, profiles, stepping = create_machine +@pytest.mark.parametrize( + "create_solver", + [ + {"solver_type": "LUsparse", "myy_type": "reduced", "cache_myy": True}, + {"solver_type": "LUsparse", "myy_type": "fft", "cache_myy": True}, + {"solver_type": "LUsparse", "myy_type": "reduced", "cache_myy": False}, + ], + indirect=True, +) +def test_linearised_stepper(create_machine, create_solver): + + tokamak, eq, profiles = create_machine + stepping = create_solver + U_active = (stepping.vessel_currents_vec * stepping.evol_metal_curr.coil_resist)[ : stepping.evol_metal_curr.n_active_coils ] @@ -207,8 +253,18 @@ def test_linearised_stepper(create_machine): ), "X-point location deviates more than 1/2 of pixel size." -def test_non_linear_stepper(create_machine): - tokamak, eq, profiles, stepping = create_machine +@pytest.mark.parametrize( + "create_solver", + [ + "default", + ], + indirect=True, +) +def test_non_linear_stepper(create_machine, create_solver): + + tokamak, eq, profiles = create_machine + stepping = create_solver + U_active = (stepping.vessel_currents_vec * stepping.evol_metal_curr.coil_resist)[ : stepping.evol_metal_curr.n_active_coils ] diff --git a/freegsnke/tests/test_static_solver.py b/freegsnke/tests/test_static_solver.py index b7b87a3..880a2ee 100644 --- a/freegsnke/tests/test_static_solver.py +++ b/freegsnke/tests/test_static_solver.py @@ -132,7 +132,15 @@ def create_test_files_static_solve(create_machine): test_psi = np.load(STATIC_PSI_BASELINE) -def test_static_solve(create_machine): +@pytest.mark.parametrize( + "solver,order,error,error_msg", + [ + ("xx", 2, ValueError, "Solver type xx"), + ("LUsparse", 3, ValueError, "operator of order 3"), + ("DST", 4, Warning, "ignored"), + ], +) +def test_NKGS_invalid(create_machine, solver, order, error, error_msg): """Tests the implementation of the static solver. Parameters @@ -145,7 +153,60 @@ def test_static_solve(create_machine): from freegsnke import GSstaticsolver - NK = GSstaticsolver.NKGSsolver(eq) + catcher = pytest.warns if issubclass(error, Warning) else pytest.raises + + with catcher(error, match=error_msg): + NK = GSstaticsolver.NKGSsolver( + eq, + solver_type=solver, + gs_operator_order=order, + ) + + +# if issubclass(error,Warning): +# with pytest.warns(error,match=error_msg): +# NK = GSstaticsolver.NKGSsolver( +# eq, +# solver_type=solver, +# gs_operator_order=order, +# ) +# else: +# with pytest.raises(error,match=error_msg): +# NK = GSstaticsolver.NKGSsolver( +# eq, +# solver_type=solver, +# gs_operator_order=order, +# ) + + +@pytest.mark.parametrize( + "cache,solver,order", + [ + (True, "LUsparse", 4), + (False, "DST", None), + (True, "LUsparse", 2), + pytest.param(True, "multigrid", 2, marks=pytest.mark.xfail), + ], +) +def test_static_solve(create_machine, cache, solver, order): + """Tests the implementation of the static solver. + + Parameters + ---------- + create_machine : pytest.fixture + the equilibirum, profiles and constrain object to generate the test set + from. + """ + eq, profiles, constrain = create_machine + + from freegsnke import GSstaticsolver + + NK = GSstaticsolver.NKGSsolver( + eq, + cache_greens=cache, + solver_type=solver, + gs_operator_order=order, + ) # from freegsnke import newtonkrylov # NK = newtonkrylov.NewtonKrylov(eq) @@ -201,16 +262,18 @@ def test_second_order_static_solve(create_machine): assert np.allclose(eq.psi(), reference_psi, atol=tolerance) +@pytest.mark.skip(reason="this should be covered by test_NKGS_invalid above") def test_static_solver_rejects_invalid_operator_order(create_machine): """Only the two finite-difference operators supplied by FreeGS4E are valid.""" eq, _, _ = create_machine from freegsnke import GSstaticsolver - with pytest.raises(ValueError, match="gs_operator_order"): + with pytest.raises(ValueError, match="order"): GSstaticsolver.NKGSsolver(eq, gs_operator_order=3) +@pytest.mark.skip(reason="not currently using the masked/reduced version") def test_limiter_reduced_boundary_green_is_exact(create_machine): """Limiter reduction preserves the boundary flux for confined current.""" eq, _, _ = create_machine diff --git a/requirements.txt b/requirements.txt index 01c3b68..7db724c 100644 --- a/requirements.txt +++ b/requirements.txt @@ -8,3 +8,4 @@ deepdiff>=7.0.1,<9 scikit-image>=0.25.2,<1 notebook>=7.4.2,<8 cvxpy>=1.7.5,<2 +numexpr~=2.14