From 0aad758294e66a1dd8f96816ff19e5e172782b2d Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Fri, 31 Oct 2025 06:39:23 +0100 Subject: [PATCH 01/16] add all files from new-multigrid --- src/struphy/feec/linear_operators.py | 4 +- src/struphy/feec/mass.py | 6 +- src/struphy/feec/tests/multigrid.py | 3247 +++++++++++++++++ .../linear_algebra/multigrid_solver.py | 1070 ++++++ .../linear_algebra/tests/test_multigrid.py | 83 + src/struphy/pic/base.py | 1 - 6 files changed, 4404 insertions(+), 7 deletions(-) create mode 100644 src/struphy/feec/tests/multigrid.py create mode 100644 src/struphy/linear_algebra/multigrid_solver.py create mode 100644 src/struphy/linear_algebra/tests/test_multigrid.py diff --git a/src/struphy/feec/linear_operators.py b/src/struphy/feec/linear_operators.py index 28b4a0805..cb26228bc 100644 --- a/src/struphy/feec/linear_operators.py +++ b/src/struphy/feec/linear_operators.py @@ -448,11 +448,11 @@ def dtype(self): @property def tosparse(self): - raise NotImplementedError() + return self.toarray_struphy(is_sparse=True) @property def toarray(self): - raise NotImplementedError() + return self.toarray_struphy(is_sparse=False) @property def bc(self): diff --git a/src/struphy/feec/mass.py b/src/struphy/feec/mass.py index 5964f5f7c..021637a5c 100644 --- a/src/struphy/feec/mass.py +++ b/src/struphy/feec/mass.py @@ -2513,7 +2513,6 @@ def codomain_femspace(self): def dtype(self): return self._dtype - @property def tosparse(self): if all(op is None for op in (self._W_extraction_op, self._V_extraction_op)): for bl in self._V_boundary_op.bc: @@ -2529,7 +2528,6 @@ def tosparse(self): else: raise NotImplementedError() - @property def toarray(self): if all(op is None for op in (self._W_extraction_op, self._V_extraction_op)): for bl in self._V_boundary_op.bc: @@ -3251,11 +3249,11 @@ def derham(self): @property def tosparse(self): - raise NotImplementedError() + return self.toarray_struphy(is_sparse=True) @property def toarray(self): - raise NotImplementedError() + return self.toarray_struphy(is_sparse=False) def transpose(self, conjugate=False): return StencilMatrixFreeMassOperator( diff --git a/src/struphy/feec/tests/multigrid.py b/src/struphy/feec/tests/multigrid.py new file mode 100644 index 000000000..536d383dd --- /dev/null +++ b/src/struphy/feec/tests/multigrid.py @@ -0,0 +1,3247 @@ +import time + +import matplotlib.pyplot as plt +import numpy as np +from mpi4py import MPI + +from struphy.bsplines.bsplines import basis_funs, find_span +from struphy.bsplines.evaluation_kernels_1d import evaluation_kernel_1d +from struphy.feec.basis_projection_ops import BasisProjectionOperators +from struphy.feec.local_projectors_kernels import fill_matrix_column +from struphy.feec.psydac_derham import Derham +from struphy.feec.utilities_local_projectors import get_one_spline, get_span_and_basis, get_values_and_indices_splines + +from psydac.linalg.solvers import inverse +from struphy.feec import preconditioner +from psydac.linalg.basic import VectorSpace, Vector, LinearOperator +from psydac.fem.projectors import knot_insertion_projection_operator +from struphy.feec.linear_operators import LinOpWithTransp +from psydac.fem.basic import FemSpace +from psydac.fem.tensor import TensorFemSpace +from struphy.feec.mass import WeightedMassOperators +from struphy.geometry.domains import Tokamak, Cuboid, HollowCylinder +from struphy.fields_background.equils import AdhocTorusQPsi +from math import comb, log2 +import random + + +def jacobi(A, b, x_init=None, tol=1e-6, max_iter=1000, verbose = False): + """ + Solves the linear system Ax = b using the Jacobi iterative method. + + Parameters: + A : numpy.ndarray + Coefficient matrix (2D array) + b : numpy.ndarray + Right-hand side vector (1D array) + x_init : numpy.ndarray, optional + Initial guess for the solution (1D array) + tol : float, optional + Convergence tolerance (default: 1e-10) + max_iter : int, optional + Maximum number of iterations (default: 1000) + + Returns: + x : numpy.ndarray + Solution vector + """ + n = A.shape[0] + x = x_init if x_init is not None else np.zeros(n) + x_new = np.zeros(n) + converged = False + for itterations in range(max_iter): + for i in range(n): + sum_except_i = np.dot(A[i, :], x) - A[i, i] * x[i] + x_new[i] = (b[i] - sum_except_i) / A[i, i] + + + error = np.linalg.norm(b - np.dot(A,x_new), ord=np.inf) + if error < tol: + converged = True + if verbose: + print(f'{converged = }') + print(f'{itterations = }') + print(f'{error = }') + return x_new, itterations + + x[:] = x_new # Update x + + + if verbose: + print(f'{converged = }') + print(f'{itterations = }') + print(f'{error = }') + return x_new, itterations + #raise ValueError("Jacobi method did not converge within the maximum number of iterations") + + +def will_jacobi_converge(A, verbose = False): + converges = False + #We get the diagonal, lower triangular and uper tireangular parts of A + Darr = np.diag(np.diag(A)) + Darr_inv = np.diag(1.0 / np.diag(Darr)) + Larr = np.tril(A, k=-1) + Uarr = np.triu(A, k=1) + #Then the Jacobi iteration matrix is + Jarr = np.matmul(Darr_inv,(Larr+Uarr)) + #Now we get its eigenvalues + eigenvalues = np.linalg.eigvals(Jarr) + max_eigen_value = abs(eigenvalues[np.argmax(np.abs(eigenvalues))]) + + if max_eigen_value < 1.0: + converges = True + + if(verbose): + print(f'{max_eigen_value = }') + + return converges + + +def will_gauss_seidel_converge(A, verbose = False): + converges = False + #We get the diagonal, lower triangular and uper triangular parts of A + Darr = np.diag(np.diag(A)) + Larr = np.tril(A, k=-1) + Uarr = np.triu(A, k=1) + Aux = np.linalg.inv(Darr - Larr) + #Then the Gauss-Seidel iteration matrix is + Garr = np.matmul(Aux, Uarr) + #Now we get its eigenvalues + eigenvalues = np.linalg.eigvals(Garr) + max_eigen_value = abs(eigenvalues[np.argmax(np.abs(eigenvalues))]) + + if max_eigen_value < 1.0: + converges = True + + if(verbose): + print(f'{max_eigen_value = }') + + return converges + + +def from_array_to_psydac(x_vector, fem_space): + #fem_space = derham.Vh_fem[sp_key] + symbolic_name = fem_space.symbolic_space.name + + if(symbolic_name == 'H1' or symbolic_name == "L2"): + spaces = [fem_space.spaces] + N = [spaces[0][i].nbasis for i in range(3)] + starts = np.array(fem_space.vector_space.starts) + x = fem_space.vector_space.zeros() + + cont= 0 + for i0 in range(N[0]): + for i1 in range(N[1]): + for i2 in range(N[2]): + x[starts[0]+i0,starts[1]+i1,starts[2]+i2] = x_vector[cont] + cont += 1 + + else: + spaces = [comp.spaces for comp in fem_space.spaces] + N = [[spaces[h][i].nbasis for i in range(3)] for h in range(3)] + starts = np.array([vi.starts for vi in fem_space.vector_space.spaces]) + x = fem_space.vector_space.zeros() + + cont = 0 + for h in range(3): + for i0 in range(N[h][0]): + for i1 in range(N[h][1]): + for i2 in range(N[h][2]): + x[h][starts[h][0]+i0,starts[h][1]+i1,starts[h][2]+i2] = x_vector[cont] + cont += 1 + return x + + +def direct_solver(A_inv,b, fem_space): + # A_inv is already the inverse matrix of A + #fem_space = derham.Vh_fem[sp_key] + symbolic_name = fem_space.symbolic_space.name + + if(symbolic_name == 'H1' or symbolic_name == "L2"): + spaces = [fem_space.spaces] + N = [spaces[0][i].nbasis for i in range(3)] + starts = np.array(fem_space.vector_space.starts) + + b_vector = remove_padding(fem_space, b) + x_vector = np.dot(A_inv, b_vector) + x = fem_space.vector_space.zeros() + + cont= 0 + for i0 in range(N[0]): + for i1 in range(N[1]): + for i2 in range(N[2]): + x[starts[0]+i0,starts[1]+i1,starts[2]+i2] = x_vector[cont] + cont += 1 + + else: + spaces = [comp.spaces for comp in fem_space.spaces] + N = [[spaces[h][i].nbasis for i in range(3)] for h in range(3)] + starts = np.array([vi.starts for vi in fem_space.vector_space.spaces]) + + b_vector = remove_padding(fem_space, b) + x_vector = np.dot(A_inv, b_vector) + x = fem_space.vector_space.zeros() + + cont = 0 + for h in range(3): + for i0 in range(N[h][0]): + for i1 in range(N[h][1]): + for i2 in range(N[h][2]): + x[h][starts[h][0]+i0,starts[h][1]+i1,starts[h][2]+i2] = x_vector[cont] + cont += 1 + return x + + +def get_b_spline_degree(V): + """ + Determines the degree of the B-splines. + + Parameters + ---------- + V : psydac.fem.basic.FemSpace + Finite element spline space (domain, input space). + + Returns + ------- + p : numpy array + numpy array of 3 ints containing the B-spline degrees for each spatial direction. + """ + + assert isinstance(V, FemSpace) + p = np.zeros(3, dtype=int) + + if hasattr(V.symbolic_space, "name"): + V_name = V.symbolic_space.name + + if(V_name == "H1"): + for i, space in enumerate(V.spaces): + p[i] = space.degree + elif(V_name == "L2"): + for i, space in enumerate(V.spaces): + p[i] = space.degree+1 + elif(V_name == "Hcurl"): + V1ds = [comp.spaces for comp in V.spaces] + for i in range(3): + p[i] = V1ds[i][i].degree+1 + elif(V_name == "Hdiv"): + V1ds = [comp.spaces for comp in V.spaces] + for i in range(3): + p[i] = V1ds[i][i].degree + elif(V_name == "H1H1H1"): + V1ds = [comp.spaces for comp in V.spaces] + for i in range(3): + p[i] = V1ds[i][i].degree + else: + raise Exception("Invalid symbolic name.") + + return p + + +def remove_padding(fem_space, v): + + #fem_space = derham.Vh_fem[sp_key] + symbolic_name = fem_space.symbolic_space.name + + if(symbolic_name == 'H1' or symbolic_name == "L2"): + + spaces = [fem_space.spaces] + N = [spaces[0][i].nbasis for i in range(3)] + + starts = np.array(fem_space.vector_space.starts) + + #To make it easier to read I will extract the data out of out disregarding all the padding it come with + v_array = np.zeros(N[0]*N[1]*N[2], dtype=float) + cont = 0 + for i0 in range(N[0]): + for i1 in range(N[1]): + for i2 in range(N[2]): + v_array[cont] = v[starts[0]+i0,starts[1]+i1,starts[2]+i2] + cont += 1 + + else: + spaces = [comp.spaces for comp in fem_space.spaces] + N = [[spaces[h][i].nbasis for i in range(3)] for h in range(3)] + + starts = np.array([vi.starts for vi in fem_space.vector_space.spaces]) + + #To make it easier to read I will extract the data out of out disregarding all the padding it come with + v_array = np.zeros(N[0][0]*N[0][1]*N[0][2]+N[1][0]*N[1][1]*N[1][2]+N[2][0]*N[2][1]*N[2][2], dtype=float) + cont = 0 + for h in range(3): + for i0 in range(N[h][0]): + for i1 in range(N[h][1]): + for i2 in range(N[h][2]): + v_array[cont] = v[h][starts[h][0]+i0,starts[h][1]+i1,starts[h][2]+i2] + cont += 1 + + + return v_array + + +class RestrictionOperator(LinOpWithTransp): + """ + Linear operator which operates between vector spaces of the same kind but different resolutions. + + We assume that the vectors in the domain belong to a De-rham space with n elements, while the codomain + belong to the coarser De-rham space with n/2 elements. + + At the moment we also assume that we are halving only the first spatial direction. + + Parameters + ---------- + V : psydac.fem.basic.FemSpace + Finite element spline space (domain, input space). + + W : psydac.fem.basic.FemSpace + Finite element spline space (codomain, output space). + + """ + def __init__(self, V, W): + + # Check domain and codomain + assert isinstance(V, FemSpace) + assert isinstance(W, FemSpace) + + self._V = V + self._W = W + + # Store info in object + self._domain = V.vector_space + self._codomain = W.vector_space + self._dtype = V.vector_space.dtype + + #Can be "H1", "L2", "Hcurl", "Hdiv", "H1H1H1" + self._V_name = V.symbolic_space.name + self._W_name = W.symbolic_space.name + assert(self._V_name == self._W_name) + + #This list will tell us in which spatial direction we are halving the problem. + self._halving_directions = [False,False,False] + + # input space: 3d StencilVectorSpaces and 1d SplineSpaces of each component + if isinstance(V, TensorFemSpace): + self._V1ds = [V.spaces] + self._VNbasis = np.array([self._V1ds[0][0].nbasis, self._V1ds[0][1].nbasis, self._V1ds[0][2].nbasis]) + + # We get the start and endpoint for each sublist in input + self._in_starts = np.array(V.vector_space.starts) + self._in_ends = np.array(V.vector_space.ends) + else: + self._V1ds = [comp.spaces for comp in V.spaces] + self._VNbasis = np.array( + [ + [self._V1ds[0][0].nbasis, self._V1ds[0][1].nbasis, self._V1ds[0][2].nbasis], + [ + self._V1ds[1][0].nbasis, + self._V1ds[1][1].nbasis, + self._V1ds[1][2].nbasis, + ], + [self._V1ds[2][0].nbasis, self._V1ds[2][1].nbasis, self._V1ds[2][2].nbasis], + ] + ) + + # We get the start and endpoint for each sublist in input + self._in_starts = np.array([vi.starts for vi in V.vector_space.spaces]) + self._in_ends = np.array([vi.ends for vi in V.vector_space.spaces]) + + # output space: 3d StencilVectorSpaces and 1d SplineSpaces of each component + if isinstance(W, TensorFemSpace): + self._W1ds = [W.spaces] + self._WNbasis = np.array([self._W1ds[0][0].nbasis, self._W1ds[0][1].nbasis, self._W1ds[0][2].nbasis]) + + for i in range(3): + if(self._WNbasis[i] < self._VNbasis[i]): + #If this breaks for clamped splines it means .nbasis gives you the number of basis functions, not the number of elements + assert self._VNbasis[i] == self._WNbasis[i]*2 + self._halving_directions[i] = True + + # We get the start and endpoint for each sublist in out + self._out_starts = np.array(W.vector_space.starts) + self._out_ends = np.array(W.vector_space.ends) + + else: + self._W1ds = [comp.spaces for comp in W.spaces] + self._WNbasis = np.array( + [ + [self._W1ds[0][0].nbasis, self._W1ds[0][1].nbasis, self._W1ds[0][2].nbasis], + [ + self._W1ds[1][0].nbasis, + self._W1ds[1][1].nbasis, + self._W1ds[1][2].nbasis, + ], + [self._W1ds[2][0].nbasis, self._W1ds[2][1].nbasis, self._W1ds[2][2].nbasis], + ] + ) + for i in range(3): + if(self._WNbasis[1][i] < self._VNbasis[1][i]): + #If this breaks for clamped splines it means .nbasis gives you the number of basis functions, not the number of elements + assert self._VNbasis[0][i] == self._WNbasis[0][i]*2 and self._VNbasis[1][i] == self._WNbasis[1][i]*2 and self._VNbasis[2][i] == self._WNbasis[2][i]*2 + self._halving_directions[i] = True + + # We get the start and endpoint for each sublist in out + self._out_starts = np.array([vi.starts for vi in W.vector_space.spaces]) + self._out_ends = np.array([vi.ends for vi in W.vector_space.spaces]) + + + + # Degree of the B-spline space, not to be confused with the degrees given by fem_space.spaces.degree since depending on the situation + # it will give the D-spline degree instead + self._p = get_b_spline_degree(V) + + #We also get the D-spline degree + self._pD = self._p - 1 + + #Now we compute the weights that define this linear operator + + #We begin by defining a list that will contain the 3 numpy arrays, each one with the weights for one spatial direction. + #In the case there are direction over which we do not halve the resolution we shall have an array with only one 1.0 + self._all_weights = [] + self._all_weightsD = [] + for i in range(3): + if self._halving_directions[i]: + #Here we store the weights needed for B-splines + weights = np.zeros(self._p[i]+2, dtype=float) + #Here we store the weights needed for D-splines + weightsD = np.zeros(self._pD[i]+2, dtype=float) + for j in range(self._p[i]+2): + weights[j] = 2.0**(-self._p[i])*comb(self._p[i]+1,j) + for j in range(self._pD[i]+2): + weightsD[j] = 2.0**-(self._pD[i]+1)*comb(self._pD[i]+1,j) + self._all_weights.append(weights) + self._all_weightsD.append(weightsD) + else: + self._all_weights.append(np.array([1.0],dtype=float)) + self._all_weightsD.append(np.array([1.0],dtype=float)) + + + + #-------------------------------------- + # Abstract interface + #-------------------------------------- + @property + def domain(self): + return self._domain + + @property + def codomain(self): + return self._codomain + + @property + def dtype(self): + return self._dtype + + def _dot_helper(self, v, out, p, weights, h=None): + """Helper function to perform dot product computation.""" + #First we get the number of weights in each direction + weights_len = [] + for i in range(3): + if self._halving_directions[i]: + weights_len.append(p[i] + 2) + else: + weights_len.append(1) + if h is None: # Scalar case (H1, L2) + for i0 in range(self._out_starts[0], self._out_ends[0] + 1): + for i1 in range(self._out_starts[1], self._out_ends[1] + 1): + for i2 in range(self._out_starts[2], self._out_ends[2] + 1): + for j0 in range(weights_len[0]): + if self._halving_directions[0]: + pos0 = (2 * i0 - p[0] + j0) % self._VNbasis[0] + else: + pos0 = i0 + for j1 in range(weights_len[1]): + if self._halving_directions[1]: + pos1 = (2 * i1 - p[1] + j1) % self._VNbasis[1] + else: + pos1 = i1 + for j2 in range(weights_len[2]): + if self._halving_directions[2]: + pos2 = (2 * i2 - p[2] + j2) % self._VNbasis[2] + else: + pos2 = i2 + out[i0, i1, i2] += weights[0][j0]* weights[1][j1] *weights[2][j2] * v[pos0, pos1, pos2] + else: # Vector case (Hcurl, Hdiv, H1H1H1) + for i0 in range(self._out_starts[h][0], self._out_ends[h][0] + 1): + for i1 in range(self._out_starts[h][1], self._out_ends[h][1] + 1): + for i2 in range(self._out_starts[h][2], self._out_ends[h][2] + 1): + for j0 in range(weights_len[0]): + if self._halving_directions[0]: + pos0 = (2 * i0 - p[0] + j0) % self._VNbasis[h][0] + else: + pos0 = i0 + for j1 in range(weights_len[1]): + if self._halving_directions[1]: + pos1 = (2 * i1 - p[1] + j1) % self._VNbasis[h][1] + else: + pos1 = i1 + for j2 in range(weights_len[2]): + if self._halving_directions[2]: + pos2 = (2 * i2 - p[2] + j2) % self._VNbasis[h][2] + else: + pos2 = i2 + out[h][i0, i1, i2] += weights[0][j0]* weights[1][j1] *weights[2][j2] * v[h][pos0, pos1, pos2] + + return out + + def dot_H1(self, v, out): + return self._dot_helper(v, out, self._p, self._all_weights) + + def dot_L2(self, v, out): + return self._dot_helper(v, out, self._pD, self._all_weightsD) + + def dot_Hcurl(self, v, out): + out = self._dot_helper(v, out, [self._pD[0],self._p[1],self._p[2]], [self._all_weightsD[0], self._all_weights[1], self._all_weights[2]], h = 0) + self._dot_helper(v, out, [self._p[0],self._pD[1],self._p[2]], [self._all_weights[0], self._all_weightsD[1], self._all_weights[2]], h = 1) + return self._dot_helper(v, out, [self._p[0],self._p[1],self._pD[2]], [self._all_weights[0], self._all_weights[1], self._all_weightsD[2]], h = 2) + + def dot_Hdiv(self, v, out): + out = self._dot_helper(v, out, [self._p[0],self._pD[1],self._pD[2]], [self._all_weights[0], self._all_weightsD[1], self._all_weightsD[2]], h = 0) + self._dot_helper(v, out, [self._pD[0],self._p[1],self._pD[2]], [self._all_weightsD[0], self._all_weights[1], self._all_weightsD[2]], h = 1) + return self._dot_helper(v, out, [self._pD[0],self._pD[1],self._p[2]], [self._all_weightsD[0], self._all_weightsD[1], self._all_weights[2]], h = 2) + + def dot_H1H1H1(self, v, out): + out = self._dot_helper(v, out, self._p, self._all_weights, h=0) + self._dot_helper(v, out, self._p, self._all_weights, h=1) + return self._dot_helper(v, out, self._p, self._all_weights, h=2) + + def dot(self, v, out=None): + + assert isinstance(v, Vector) and v.space == self.domain + + if out is None: + out = self.codomain.zeros() + else: + assert isinstance(out, Vector) and out.space == self.codomain + + if self._V_name == 'H1' or self._V_name == 'L2': + for i0 in range(self._out_starts[0], self._out_ends[0]+1): + for i1 in range(self._out_starts[1], self._out_ends[1]+1): + for i2 in range(self._out_starts[2], self._out_ends[2]+1): + out[i0,i1,i2] = 0.0 + else: + for h in range(3): + for i0 in range(self._out_starts[h][0], self._out_ends[h][0]+1): + for i1 in range(self._out_starts[h][1], self._out_ends[h][1]+1): + for i2 in range(self._out_starts[h][2], self._out_ends[h][2]+1): + out[h][i0,i1,i2] = 0.0 + + dot_methods = { + "H1": self.dot_H1, + "L2": self.dot_L2, + "Hcurl": self.dot_Hcurl, + "Hdiv": self.dot_Hdiv, + "H1H1H1": self.dot_H1H1H1, + } + + return dot_methods.get(self._V_name)(v, out) + + def transpose(self, *, out = None): + if out is None: + out = ExtensionOperator(self._W, self._V) + else: + assert isinstance(out, ExtensionOperator) + assert out.domain is self.codomain + assert out.codomain is self.domain + + return out + + +class ExtensionOperator(LinOpWithTransp): + """ + Linear operator which operates between vector spaces of the same kind but different resolutions. + + We assume that the vectors in the domain belong to a De-rham space with n/2 elements, while the codomain + belong to the coarser De-rham space with n elements. + + At the moment we also assume that we are halving only the first spatial direction. + + Parameters + ---------- + V : psydac.fem.basic.FemSpace + Finite element spline space (domain, input space). + + W : psydac.fem.basic.FemSpace + Finite element spline space (codomain, output space). + + """ + def __init__(self, V, W): + + # Check domain and codomain + assert isinstance(V, FemSpace) + assert isinstance(W, FemSpace) + + self._V = V + self._W = W + + # Store info in object + self._domain = V.vector_space + self._codomain = W.vector_space + self._dtype = V.vector_space.dtype + + #Can be "H1", "L2", "Hcurl", "Hdiv", "H1H1H1" + self._V_name = V.symbolic_space.name + self._W_name = W.symbolic_space.name + assert(self._V_name == self._W_name) + + #This list will tell us in which spatial direction we are halving the problem. + self._halving_directions = [False,False,False] + + # input space: 3d StencilVectorSpaces and 1d SplineSpaces of each component + if isinstance(V, TensorFemSpace): + self._V1ds = [V.spaces] + self._VNbasis = np.array([self._V1ds[0][0].nbasis, self._V1ds[0][1].nbasis, self._V1ds[0][2].nbasis]) + + # We get the start and endpoint for each sublist in input + self._in_starts = np.array(V.vector_space.starts) + self._in_ends = np.array(V.vector_space.ends) + else: + self._V1ds = [comp.spaces for comp in V.spaces] + self._VNbasis = np.array( + [ + [self._V1ds[0][0].nbasis, self._V1ds[0][1].nbasis, self._V1ds[0][2].nbasis], + [ + self._V1ds[1][0].nbasis, + self._V1ds[1][1].nbasis, + self._V1ds[1][2].nbasis, + ], + [self._V1ds[2][0].nbasis, self._V1ds[2][1].nbasis, self._V1ds[2][2].nbasis], + ] + ) + + # We get the start and endpoint for each sublist in input + self._in_starts = np.array([vi.starts for vi in V.vector_space.spaces]) + self._in_ends = np.array([vi.ends for vi in V.vector_space.spaces]) + + # output space: 3d StencilVectorSpaces and 1d SplineSpaces of each component + if isinstance(W, TensorFemSpace): + self._W1ds = [W.spaces] + self._WNbasis = np.array([self._W1ds[0][0].nbasis, self._W1ds[0][1].nbasis, self._W1ds[0][2].nbasis]) + + for i in range(3): + if(self._VNbasis[i] < self._WNbasis[i]): + #If this breaks for clamped splines it means .nbasis gives you the number of basis functions, not the number of elements + assert self._WNbasis[i] == self._VNbasis[i]*2 + self._halving_directions[i] = True + + # We get the start and endpoint for each sublist in out + self._out_starts = np.array(W.vector_space.starts) + self._out_ends = np.array(W.vector_space.ends) + + else: + self._W1ds = [comp.spaces for comp in W.spaces] + self._WNbasis = np.array( + [ + [self._W1ds[0][0].nbasis, self._W1ds[0][1].nbasis, self._W1ds[0][2].nbasis], + [ + self._W1ds[1][0].nbasis, + self._W1ds[1][1].nbasis, + self._W1ds[1][2].nbasis, + ], + [self._W1ds[2][0].nbasis, self._W1ds[2][1].nbasis, self._W1ds[2][2].nbasis], + ] + ) + + for i in range(3): + if(self._VNbasis[1][i] < self._WNbasis[1][i]): + assert self._WNbasis[0][i] == self._VNbasis[0][i]*2 and self._WNbasis[1][i] == self._VNbasis[1][i]*2 and self._WNbasis[2][i] == self._VNbasis[2][i]*2 + self._halving_directions[i] = True + + # We get the start and endpoint for each sublist in out + self._out_starts = np.array([vi.starts for vi in W.vector_space.spaces]) + self._out_ends = np.array([vi.ends for vi in W.vector_space.spaces]) + + + + # Degree of the B-spline space, not to be confused with the degrees given by fem_space.spaces.degree since depending on the situation + # it will give the D-spline degree instead + self._p = get_b_spline_degree(V) + #We also get the D-splines degree + self._pD = self._p-1 + + #Now we compute the weights that define this linear operator + + #We begin by defining a list that will contain the 3 numpy arrays, each one with the weights for one spatial direction. + #In the case there are a direction over which we do not halve the resolution we shall have an array with only one 1.0 + self._all_weights_even = [] + self._all_weights_evenD = [] + self._all_weights_odd = [] + self._all_weights_oddD = [] + + #Each list has 3 integers, each one denoting the number of weights in the corresponding weights array. + self._all_size_even = [] + self._all_size_evenD = [] + self._all_size_odd = [] + self._all_size_oddD = [] + + for i in range(3): + if self._halving_directions[i]: + #First for B-splines + if(self._p[i]%2 == 0): + size_even = self._p[i]//2 +1 + size_odd = self._p[i]//2 +1 + weights_even = np.zeros(size_even, dtype=float) + weights_odd = np.zeros(size_odd, dtype=float) + for j in range(size_even): + weights_even[j] = 2.0**(-self._p[i])*comb(self._p[i]+1,2*j) + weights_odd[j] = 2.0**(-self._p[i])*comb(self._p[i]+1,2*j+1) + else: + size_even = (self._p[i]+1)//2 +1 + size_odd = (self._p[i]-1)//2 +1 + weights_even = np.zeros(size_even, dtype=float) + weights_odd = np.zeros(size_odd, dtype=float) + for j in range(size_even): + weights_even[j] = 2.0**(-self._p[i])*comb(self._p[i]+1,2*j) + for j in range(size_odd): + weights_odd[j] = 2.0**(-self._p[i])*comb(self._p[i]+1,2*j+1) + + #Second for D-splines + if(self._pD[i]%2 == 0): + size_evenD = self._pD[i]//2 +1 + size_oddD = self._pD[i]//2 +1 + weights_evenD = np.zeros(size_evenD, dtype=float) + weights_oddD = np.zeros(size_oddD, dtype=float) + for j in range(size_evenD): + weights_evenD[j] = 2.0**-(self._pD[i]+1)*comb(self._pD[i]+1,2*j) + weights_oddD[j] = 2.0**-(self._pD[i]+1)*comb(self._pD[i]+1,2*j+1) + else: + size_evenD = (self._pD[i]+1)//2 +1 + size_oddD = (self._pD[i]-1)//2 +1 + weights_evenD = np.zeros(size_evenD, dtype=float) + weights_oddD = np.zeros(size_oddD, dtype=float) + for j in range(size_evenD): + weights_evenD[j] = 2.0**-(self._pD[i]+1)*comb(self._pD[i]+1,2*j) + for j in range(size_oddD): + weights_oddD[j] = 2.0**-(self._pD[i]+1)*comb(self._pD[i]+1,2*j+1) + + self._all_weights_even.append(weights_even) + self._all_weights_evenD.append(weights_evenD) + self._all_weights_odd.append(weights_odd) + self._all_weights_oddD.append(weights_oddD) + self._all_size_even.append(size_even) + self._all_size_evenD.append(size_evenD) + self._all_size_odd.append(size_odd) + self._all_size_oddD.append(size_oddD) + + else: + self._all_weights_even.append(np.array([1.0],dtype=float)) + self._all_weights_evenD.append(np.array([1.0],dtype=float)) + self._all_weights_odd.append(np.array([1.0],dtype=float)) + self._all_weights_oddD.append(np.array([1.0],dtype=float)) + self._all_size_even.append(1) + self._all_size_evenD.append(1) + self._all_size_odd.append(1) + self._all_size_oddD.append(1) + + + #-------------------------------------- + # Abstract interface + #-------------------------------------- + @property + def domain(self): + return self._domain + + @property + def codomain(self): + return self._codomain + + @property + def dtype(self): + return self._dtype + + def tosparse(self): + pass + + def toarray(self): + pass + + def _dot_helper(self, v, out, p, weights_even, weights_odd, size_even, size_odd, h=None): + """Helper function to perform dot product computation.""" + parity_match = [] + for i in range(3): + parity_match.append(p[i] % 2) + + if h is None: # Scalar case (H1, L2) + for j0 in range(self._out_starts[0], self._out_ends[0] + 1): + parity_j0 = j0 % 2 + weights0, size0, offset0 = ((weights_even[0], size_even[0], 0) if parity_j0 == parity_match[0] else (weights_odd[0], size_odd[0], 1)) + for j1 in range(self._out_starts[1], self._out_ends[1] + 1): + parity_j1 = j1 % 2 + weights1, size1, offset1 = ((weights_even[1], size_even[1], 0) if parity_j1 == parity_match[1] else (weights_odd[1], size_odd[1], 1)) + for j2 in range(self._out_starts[2], self._out_ends[2] + 1): + parity_j2 = j2 % 2 + weights2, size2, offset2 = ((weights_even[2], size_even[2], 0) if parity_j2 == parity_match[2] else (weights_odd[2], size_odd[2], 1)) + for i0 in range(size0): + if self._halving_directions[0]: + pos0 = ((j0 + p[0] - 2 * i0 - offset0) // 2) % self._VNbasis[0] + else: + pos0 = j0 + for i1 in range(size1): + if self._halving_directions[1]: + pos1 = ((j1 + p[1] - 2 * i1 - offset1) // 2) % self._VNbasis[1] + else: + pos1 = j1 + for i2 in range(size2): + if self._halving_directions[2]: + pos2 = ((j2 + p[2] - 2 * i2 - offset2) // 2) % self._VNbasis[2] + else: + pos2 = j2 + out[j0, j1, j2] += weights0[i0] * weights1[i1] * weights2[i2] * v[pos0, pos1, pos2] + + else: # Vector case (Hcurl, Hdiv, H1H1H1) + for j0 in range(self._out_starts[h][0], self._out_ends[h][0] + 1): + parity_j0 = j0 % 2 + weights0, size0, offset0 = ((weights_even[0], size_even[0], 0) if parity_j0 == parity_match[0] else (weights_odd[0], size_odd[0], 1)) + for j1 in range(self._out_starts[h][1], self._out_ends[h][1] + 1): + parity_j1 = j1 % 2 + weights1, size1, offset1 = ((weights_even[1], size_even[1], 0) if parity_j1 == parity_match[1] else (weights_odd[1], size_odd[1], 1)) + for j2 in range(self._out_starts[h][2], self._out_ends[h][2] + 1): + parity_j2 = j2 % 2 + weights2, size2, offset2 = ((weights_even[2], size_even[2], 0) if parity_j2 == parity_match[2] else (weights_odd[2], size_odd[2], 1)) + for i0 in range(size0): + if self._halving_directions[0]: + pos0 = ((j0 + p[0] - 2 * i0 - offset0) // 2) % self._VNbasis[h][0] + else: + pos0 = j0 + for i1 in range(size1): + if self._halving_directions[1]: + pos1 = ((j1 + p[1] - 2 * i1 - offset1) // 2) % self._VNbasis[h][1] + else: + pos1 = j1 + for i2 in range(size2): + if self._halving_directions[2]: + pos2 = ((j2 + p[2] - 2 * i2 - offset2) // 2) % self._VNbasis[h][2] + else: + pos2 = j2 + out[h][j0, j1, j2] += weights0[i0] * weights1[i1] * weights2[i2] * v[h][pos0, pos1, pos2] + + return out + + def dot_H1(self, v, out): + return self._dot_helper(v, out, self._p, self._all_weights_even, self._all_weights_odd, self._all_size_even, self._all_size_odd) + + def dot_L2(self, v, out): + return self._dot_helper(v, out, self._pD, self._all_weights_evenD, self._all_weights_oddD, self._all_size_evenD, self._all_size_oddD) + + def dot_Hcurl(self, v, out): + out = self._dot_helper(v, out, [self._pD[0], self._p[1],self._p[2]], [self._all_weights_evenD[0],self._all_weights_even[1],self._all_weights_even[2]], [self._all_weights_oddD[0],self._all_weights_odd[1],self._all_weights_odd[2]], [self._all_size_evenD[0],self._all_size_even[1],self._all_size_even[2]], [self._all_size_oddD[0],self._all_size_odd[1],self._all_size_odd[2]], h=0) + out = self._dot_helper(v, out, [self._p[0], self._pD[1],self._p[2]], [self._all_weights_even[0],self._all_weights_evenD[1],self._all_weights_even[2]], [self._all_weights_odd[0],self._all_weights_oddD[1],self._all_weights_odd[2]], [self._all_size_even[0],self._all_size_evenD[1],self._all_size_even[2]], [self._all_size_odd[0],self._all_size_oddD[1],self._all_size_odd[2]], h=1) + return self._dot_helper(v, out, [self._p[0], self._p[1],self._pD[2]], [self._all_weights_even[0],self._all_weights_even[1],self._all_weights_evenD[2]], [self._all_weights_odd[0],self._all_weights_odd[1],self._all_weights_oddD[2]], [self._all_size_even[0],self._all_size_even[1],self._all_size_evenD[2]], [self._all_size_odd[0],self._all_size_odd[1],self._all_size_oddD[2]], h=2) + + + def dot_Hdiv(self, v, out): + out = self._dot_helper(v, out, [self._p[0], self._pD[1],self._pD[2]], [self._all_weights_even[0],self._all_weights_evenD[1],self._all_weights_evenD[2]], [self._all_weights_odd[0],self._all_weights_oddD[1],self._all_weights_oddD[2]], [self._all_size_even[0],self._all_size_evenD[1],self._all_size_evenD[2]], [self._all_size_odd[0],self._all_size_oddD[1],self._all_size_oddD[2]], h=0) + out = self._dot_helper(v, out, [self._pD[0], self._p[1],self._pD[2]], [self._all_weights_evenD[0],self._all_weights_even[1],self._all_weights_evenD[2]], [self._all_weights_oddD[0],self._all_weights_odd[1],self._all_weights_oddD[2]], [self._all_size_evenD[0],self._all_size_even[1],self._all_size_evenD[2]], [self._all_size_oddD[0],self._all_size_odd[1],self._all_size_oddD[2]], h=1) + return self._dot_helper(v, out, [self._pD[0], self._pD[1],self._p[2]], [self._all_weights_evenD[0],self._all_weights_evenD[1],self._all_weights_even[2]], [self._all_weights_oddD[0],self._all_weights_oddD[1],self._all_weights_odd[2]], [self._all_size_evenD[0],self._all_size_evenD[1],self._all_size_even[2]], [self._all_size_oddD[0],self._all_size_oddD[1],self._all_size_odd[2]], h=2) + + + def dot_H1H1H1(self, v, out): + out = self._dot_helper(v, out, self._p, self._all_weights_even, self._all_weights_odd, self._all_size_even, self._all_size_odd, h = 0) + self._dot_helper(v, out, self._p, self._all_weights_even, self._all_weights_odd, self._all_size_even, self._all_size_odd, h = 1) + return self._dot_helper(v, out, self._p, self._all_weights_even, self._all_weights_odd, self._all_size_even, self._all_size_odd, h = 2) + + + + def dot(self, v, out=None): + + assert isinstance(v, Vector) and v.space == self.domain + + if out is None: + out = self.codomain.zeros() + else: + assert isinstance(out, Vector) and out.space == self.codomain + + if self._V_name == 'H1' or self._V_name == 'L2': + for i0 in range(self._out_starts[0], self._out_ends[0]+1): + for i1 in range(self._out_starts[1], self._out_ends[1]+1): + for i2 in range(self._out_starts[2], self._out_ends[2]+1): + out[i0,i1,i2] = 0.0 + else: + for h in range(3): + for i0 in range(self._out_starts[h][0], self._out_ends[h][0]+1): + for i1 in range(self._out_starts[h][1], self._out_ends[h][1]+1): + for i2 in range(self._out_starts[h][2], self._out_ends[h][2]+1): + out[h][i0,i1,i2] = 0.0 + + dot_methods = { + "H1": self.dot_H1, + "L2": self.dot_L2, + "Hcurl": self.dot_Hcurl, + "Hdiv": self.dot_Hdiv, + "H1H1H1": self.dot_H1H1H1, + } + + return dot_methods.get(self._V_name)(v, out) + + def transpose(self, *, out = None): + if out is None: + out = RestrictionOperator(self._W, self._V) + else: + assert isinstance(out, RestrictionOperator) + assert out.domain is self.codomain + assert out.codomain is self.domain + + return out + + +def Compute_rate_of_smoothing(Nel, plist, spl_kind): + + from struphy.feec.utilities import create_equal_random_arrays + # get global communicator + comm = MPI.COMM_WORLD + rank = comm.Get_rank() + world_size = comm.Get_size() + + a1= 0.2 + #domain = HollowCylinder(a1 = a1) + domain = Cuboid() + sp_key = '1' + derham = [] + mass_ops = [] + A = [] + + epsilon = 0.0002 + derham.append(Derham([Nel[0],Nel[1],Nel[2]], plist, spl_kind, comm=comm, local_projectors=False)) + mass_ops.append(WeightedMassOperators(derham[0], domain)) + A.append(epsilon*derham[0].curl.T @ mass_ops[0].M2 @ derham[0].curl + mass_ops[0].M1) + + + field_star = derham[0].create_field('fh', 'Hcurl') + field_aprox = derham[0].create_field('fh', 'Hcurl') + + #We are gonna use the method of manufacture solutions to determine the behaviour of our error + #Our exact solution is u_star + u_stararr, u_star = create_equal_random_arrays(derham[0].Vh_fem[sp_key], seed=45) + #We compute the rhs + b = A[0].dot(u_star) + field_star.vector = u_star + + #We store the number of itterations + N_itter = [] + #We define a set of point to evaluate the exact solution and the aproximated one + pointsx = np.linspace(0.0,1.0,100) + pointsy = np.linspace(0.0,1.0,100) + pointsz = np.array([0.0,0.5]) + dx = pointsx[1] - pointsx[0] # Spacing between points + dy = pointsy[1] - pointsy[0] # Spacing between points + dz = pointsz[1] - pointsz[0] # Spacing between points + X, Y, Z = np.meshgrid(pointsx, pointsy, pointsz, indexing="ij") + + + + #We define a list where to store the arrays with the values of the errors for each aproximation and vector component + errorsx = [] + errorsy = [] + errorsz = [] + + #Fourier Transform coefficients of the error function for each vector component + fourier_coeffsx = [] + fourier_coeffs_shiftedx = [] + fourier_coeffsy = [] + fourier_coeffs_shiftedy = [] + fourier_coeffsz = [] + fourier_coeffs_shiftedz = [] + Max_iter = 10 + for i in range(Max_iter): + solver = inverse(A[0],'cg', maxiter = int(i+2)) + u = solver.dot(b) + field_aprox.vector = u + errorsx.append(field_star(X,Y,Z)[0]-field_aprox(X,Y,Z)[0]) + errorsy.append(field_star(X,Y,Z)[1]-field_aprox(X,Y,Z)[1]) + errorsz.append(field_star(X,Y,Z)[2]-field_aprox(X,Y,Z)[2]) + N_itter.append(solver._info['niter']) + + # Compute Fourier Transform coefficients + fourier_coeffsx.append(np.fft.fftn(errorsx[-1])) + fourier_coeffs_shiftedx.append(np.fft.fftshift(fourier_coeffsx[-1])) + fourier_coeffsy.append(np.fft.fftn(errorsy[-1])) + fourier_coeffs_shiftedy.append(np.fft.fftshift(fourier_coeffsy[-1])) + fourier_coeffsz.append(np.fft.fftn(errorsz[-1])) + fourier_coeffs_shiftedz.append(np.fft.fftshift(fourier_coeffsz[-1])) + + # Compute corresponding frequencies + freqsx = np.fft.fftshift(np.fft.fftfreq(len(pointsx), d=dx)) + freqsy = np.fft.fftshift(np.fft.fftfreq(len(pointsy), d=dy)) + freqsz = np.fft.fftshift(np.fft.fftfreq(len(pointsz), d=dz)) + # Create a 3D meshgrid of frequencies + #Fx, Fy, Fz = np.meshgrid(freq_x[-1], freq_y[-1], freq_z[-1], indexing="ij") + + #Now we compute the magnitude of the high frequencies between two interations + def get_smoothing_rate(freqsx,freqsy,freqsz, coeff_old, coeff_new, hx,hy,hz): + smoothing_rate = -1.0 + freq = np.zeros(3,dtype=float) + for ix in range(len(freqsx)): + for iy in range(len(freqsy)): + for iz in range(len(freqsz)): + if((abs(freqsx[ix])> 1.0 / (4.0*hx) or abs(freqsy[iy])> 1.0 / (4.0*hy) or abs(freqsz[iz])> 1.0 / (4.0*hz) ) and abs(coeff_new[ix,iy,iz])>10.0**-6): + smoothing_rate = max(smoothing_rate,abs(coeff_new[ix,iy,iz]/coeff_old[ix,iy,iz])) + freq[0] = freqsx[ix] + freq[1] = freqsx[iy] + freq[2] = freqsx[iz] + + return smoothing_rate, freq + + + for i in range(Max_iter-1): + + smoothing_rate, freq = get_smoothing_rate(freqsx,freqsy,freqsz,fourier_coeffs_shiftedx[i],fourier_coeffs_shiftedx[i+1],1.0/Nel[0],1.0/Nel[1],0.00000001) + + print("#######################") + print(f'{i =}') + print(f'{smoothing_rate =}') + print(f'{freq =}') + print("#######################") + + #for i in range(Max_iter-1): + #smoothing_rate_1d(freqsx, coeff_old, coeff_new, h) + + +def Visualized_high_frequency_dampening(Nel, plist, spl_kind): + + from struphy.feec.utilities import create_equal_random_arrays + # get global communicator + comm = MPI.COMM_WORLD + rank = comm.Get_rank() + world_size = comm.Get_size() + + a1= 0.2 + #domain = HollowCylinder(a1 = a1) + domain = Cuboid() + sp_key = '2' + sp_id = 'Hdiv' + derham = [] + mass_ops = [] + A = [] + + epsilon = 1.0 + derham.append(Derham([Nel[0],Nel[1],Nel[2]], plist, spl_kind, comm=comm, local_projectors=False)) + mass_ops.append(WeightedMassOperators(derham[0], domain)) + #Poisson + #A.append(derham[0].grad.T @ mass_ops[0].M1 @ derham[0].grad) + #Hall-ish + #A.append(epsilon*derham[0].curl.T @ mass_ops[0].M2 @ derham[0].curl + mass_ops[0].M1) + #Hall + #A.append(mass_ops[0].M1 -epsilon*derham[0].curl.T @ mass_ops[0].M2 @ derham[0].curl) + #Shear-Alfven-ish + #A.append(mass_ops[0].M2 -epsilon* mass_ops[0].M2@ derham[0].curl @ mass_ops[0].M1 @ derham[0].curl.T @ mass_ops[0].M2) + #Shear-Alfven + pc_class = getattr(preconditioner,"MassMatrixPreconditioner") + pc = pc_class(mass_ops[0].M1) + M1_inv = inverse( + mass_ops[0].M1, + "pcg", + pc=pc, + maxiter=3000, + verbose=False, + ) + A.append(mass_ops[0].M2 -epsilon* mass_ops[0].M2@ derham[0].curl @ M1_inv @ derham[0].curl.T @ mass_ops[0].M2) + #Shear-Alfven-v2 + #A.append(mass_ops[0].M2 -epsilon* derham[0].curl @ M1_inv @ derham[0].curl.T) + + + + + field_star = derham[0].create_field('fh', sp_id) + field_aprox = derham[0].create_field('fh', sp_id) + + #We are gonna use the method of manufacture solutions to determine the behaviour of our error + #Our exact solution is u_star + u_stararr, u_star = create_equal_random_arrays(derham[0].Vh_fem[sp_key], seed=45) + #We compute the rhs + b = A[0].dot(u_star) + field_star.vector = u_star + + #Turn b into an array + #barr = remove_padding(derham[0].Vh_fem[sp_key],b) + #Turn A[0] into an array + #Aarr = A[0].toarray() + + #Gauss = will_gauss_seidel_converge(Aarr, verbose=True) + #print(f"{Gauss = }") + + + #We store the number of itterations + N_itter = [] + #We define a set of point to evaluate the exact solution and the aproximated one + pointsx = np.linspace(0.0,1.0,Nel[0]) + #pointsx = np.array([0.0,0.5]) + #pointsy = np.linspace(0.0,1.0,Nel[1]) + pointsy = np.array([0.0,0.5]) + pointsz = np.array([0.0,0.5]) + dx = pointsx[1] - pointsx[0] # Spacing between points + dy = pointsy[1] - pointsy[0] # Spacing between points + dz = pointsz[1] - pointsz[0] # Spacing between points + X, Y, Z = np.meshgrid(pointsx, pointsy, pointsz, indexing="ij") + + + + #We define a list where to store the arrays with the values of the errors for each aproximation and vector component + errorsx = [] + errorsy = [] + errorsz = [] + + #Fourier Transform coefficients of the error function for each vector component + fourier_coeffsx = [] + fourier_coeffs_shiftedx = [] + fourier_coeffsy = [] + fourier_coeffs_shiftedy = [] + fourier_coeffsz = [] + fourier_coeffs_shiftedz = [] + max_iter_list = [3] + #max_iter_list = [0,2,3,4,5,6,7,8,9,10] + Number_of_iter = len(max_iter_list) + for i in range(Number_of_iter): + + if(max_iter_list[i]>1): + + solver = inverse(A[0],'cg', maxiter = max_iter_list[i], tol = 10.0**-6) + u = solver.dot(b) + #uarr, itter = jacobi(Aarr,barr,max_iter=max_iter_list[i]) + #u = from_array_to_psydac(uarr, derham[0].Vh_fem[sp_key]) + N_itter.append(solver._info['niter']) + #N_itter.append(itter) + + else: + u = derham[0].Vh[derham[0].space_to_form[sp_id]].zeros() + N_itter.append(0) + + + field_aprox.vector = u + #errorsx.append(field_star(X,Y,Z)-field_aprox(X,Y,Z)) + errorsx.append(field_star(X,Y,Z)[0]-field_aprox(X,Y,Z)[0]) + errorsy.append(field_star(X,Y,Z)[1]-field_aprox(X,Y,Z)[1]) + errorsz.append(field_star(X,Y,Z)[2]-field_aprox(X,Y,Z)[2]) + + + # Compute Fourier Transform coefficients + fourier_coeffsx.append(np.fft.fftn(errorsx[-1])) + fourier_coeffs_shiftedx.append(np.fft.fftshift(fourier_coeffsx[-1])) + fourier_coeffsy.append(np.fft.fftn(errorsy[-1])) + fourier_coeffs_shiftedy.append(np.fft.fftshift(fourier_coeffsy[-1])) + fourier_coeffsz.append(np.fft.fftn(errorsz[-1])) + fourier_coeffs_shiftedz.append(np.fft.fftshift(fourier_coeffsz[-1])) + + # Compute corresponding frequencies + freqsx = np.fft.fftshift(np.fft.fftfreq(len(pointsx), d=dx)) + freqsy = np.fft.fftshift(np.fft.fftfreq(len(pointsy), d=dy)) + freqsz = np.fft.fftshift(np.fft.fftfreq(len(pointsz), d=dz)) + # Create a 3D meshgrid of frequencies + #Fx, Fy, Fz = np.meshgrid(freq_x[-1], freq_y[-1], freq_z[-1], indexing="ij") + + #Now we compute the magnitude of the high frequencies between two interations + def get_magnitude_maximum_nasty_frequency_scalar(freqsx,freqsy,freqsz, coeff_new,hx,hy,hz): + value= 0.0 + freq = np.zeros(3,dtype=float) + for ix in range(len(freqsx)): + for iy in range(len(freqsy)): + for iz in range(len(freqsz)): + if((abs(freqsx[ix])> 1.0/(4.0*hx) or abs(freqsy[iy])> 1.0/(4.0*hy) or abs(freqsz[iz])> 1.0/(4.0*hz) ) and abs(coeff_new[ix,iy,iz])>10.0**-6): + if(abs(coeff_new[ix,iy,iz]) > value ): + value = abs(coeff_new[ix,iy,iz]) + freq[0] = freqsx[ix] + freq[1] = freqsy[iy] + freq[2] = freqsz[iz] + + return value, freq + + + + + def get_magnitude_maximum_nasty_frequency(freqsx,freqsy,freqsz, coeff_newx, coeff_newy, coeff_newz,hx,hy,hz): + value= 0.0 + freq = np.zeros(3,dtype=float) + for ix in range(len(freqsx)): + for iy in range(len(freqsy)): + for iz in range(len(freqsz)): + if((abs(freqsx[ix])> 1.0/(4.0*hx) or abs(freqsy[iy])> 1.0/(4.0*hy) or abs(freqsz[iz])> 1.0/(4.0*hz) ) and (abs(coeff_newx[ix,iy,iz])>10.0**-6 or abs(coeff_newy[ix,iy,iz])>10.0**-6 or abs(coeff_newz[ix,iy,iz])>10.0**-6)): + if(abs(coeff_newx[ix,iy,iz]) > value and abs(coeff_newx[ix,iy,iz]) >= abs(coeff_newy[ix,iy,iz]) and abs(coeff_newx[ix,iy,iz])>= abs(coeff_newz[ix,iy,iz])): + value = abs(coeff_newx[ix,iy,iz]) + freq[0] = freqsx[ix] + freq[1] = freqsy[iy] + freq[2] = freqsz[iz] + + elif(abs(coeff_newy[ix,iy,iz]) > value and abs(coeff_newy[ix,iy,iz]) >= abs(coeff_newx[ix,iy,iz]) and abs(coeff_newy[ix,iy,iz])>= abs(coeff_newz[ix,iy,iz])): + value = abs(coeff_newy[ix,iy,iz]) + freq[0] = freqsx[ix] + freq[1] = freqsy[iy] + freq[2] = freqsz[iz] + + elif(abs(coeff_newz[ix,iy,iz]) > value and abs(coeff_newz[ix,iy,iz]) >= abs(coeff_newx[ix,iy,iz]) and abs(coeff_newz[ix,iy,iz])>= abs(coeff_newy[ix,iy,iz])): + value = abs(coeff_newz[ix,iy,iz]) + freq[0] = freqsx[ix] + freq[1] = freqsy[iy] + freq[2] = freqsz[iz] + + return value, freq + + + magnitudes = [] + bad_frequencies = [] + + for i in range(Number_of_iter): + #magnitude, freq = get_magnitude_maximum_nasty_frequency_scalar(freqsx,freqsy,freqsz,fourier_coeffs_shiftedx[i],1.0/Nel[0],1.0/Nel[1],0.000001) + magnitude, freq = get_magnitude_maximum_nasty_frequency(freqsx,freqsy,freqsz,fourier_coeffs_shiftedx[i],fourier_coeffs_shiftedy[i], fourier_coeffs_shiftedz[i],1.0/Nel[0],0.000001,0.000001) + magnitudes.append(magnitude) + bad_frequencies.append(freq) + + + print(f'{Nel[0] = }') + print(f'{Nel[1] = }') + print("magnitudes") + for i in magnitudes: + print(i) + print("N_itter") + for i in N_itter: + print(i) + print("bad_frequencies") + for i in bad_frequencies: + print(i[0:2]) + + + + plt.figure() + plt.scatter(N_itter,magnitudes) + #plt.yscale("log") + plt.show() + + +def Visualized_all_frequencies_dampening(Nel, plist, spl_kind, Is): + + from struphy.feec.utilities import create_equal_random_arrays + # get global communicator + comm = MPI.COMM_WORLD + rank = comm.Get_rank() + world_size = comm.Get_size() + + a1= 0.2 + #domain = HollowCylinder(a1 = a1) + domain = Cuboid() + model = 'Shear-Alfven' + smoother = 'cg' + derham = [] + mass_ops = [] + A = [] + + epsilon = 10.0**-6.0 + derham.append(Derham([Nel[0],Nel[1],Nel[2]], plist, spl_kind, comm=comm, local_projectors=False)) + mass_ops.append(WeightedMassOperators(derham[0], domain)) + if(model == "Poisson"): + sp_key = '0' + sp_id = 'H1' + #Poisson + A.append(derham[0].grad.T @ mass_ops[0].M1 @ derham[0].grad) + #Hall-ish + #A.append(epsilon*derham[0].curl.T @ mass_ops[0].M2 @ derham[0].curl + mass_ops[0].M1) + elif(model == "Hall"): + sp_key = '1' + sp_id = 'Hcurl' + #Hall + A.append(mass_ops[0].M1 -epsilon*derham[0].curl.T @ mass_ops[0].M2 @ derham[0].curl) + #Shear-Alfven-ish + #A.append(mass_ops[0].M2 -epsilon* mass_ops[0].M2@ derham[0].curl @ mass_ops[0].M1 @ derham[0].curl.T @ mass_ops[0].M2) + elif(model == 'Shear-Alfven' or model == 'Shear-Alfven-v2'): + sp_key = '2' + sp_id = 'Hdiv' + pc_class = getattr(preconditioner,"MassMatrixPreconditioner") + pc = pc_class(mass_ops[0].M1) + M1_inv = inverse( + mass_ops[0].M1, + "pcg", + pc=pc, + maxiter=3000, + verbose=False, + ) + #Shear-Alfven + if (model == "Shear-Alfven"): + A.append(mass_ops[0].M2 -epsilon* mass_ops[0].M2@ derham[0].curl @ M1_inv @ derham[0].curl.T @ mass_ops[0].M2) + else: + #Shear-Alfven-v2 + A.append(mass_ops[0].M2 -epsilon* derham[0].curl @ M1_inv @ derham[0].curl.T) + + field_star = derham[0].create_field('fh', sp_id) + field_aprox = derham[0].create_field('fh', sp_id) + + #We are gonna use the method of manufacture solutions to determine the behaviour of our error + #Our exact solution is u_star + u_stararr, u_star = create_equal_random_arrays(derham[0].Vh_fem[sp_key], seed=45) + #We compute the rhs + b = A[0].dot(u_star) + field_star.vector = u_star + + #Turn b into an array + #barr = remove_padding(derham[0].Vh_fem[sp_key],b) + #Turn A[0] into an array + #Aarr = A[0].toarray() + + #Gauss = will_gauss_seidel_converge(Aarr, verbose=True) + #print(f"{Gauss = }") + + + #We store the number of itterations + N_itter = [] + #We define a set of point to evaluate the exact solution and the aproximated one + pointsx = np.linspace(0.0,1.0,Nel[0]) + #pointsx = np.array([0.0,0.5]) + pointsy = np.linspace(0.0,1.0,Nel[1]) + #pointsy = np.array([0.0,0.5]) + pointsz = np.array([0.0,0.5]) + dx = pointsx[1] - pointsx[0] # Spacing between points + dy = pointsy[1] - pointsy[0] # Spacing between points + dz = pointsz[1] - pointsz[0] # Spacing between points + X, Y, Z = np.meshgrid(pointsx, pointsy, pointsz, indexing="ij") + + + + #We define a list where to store the arrays with the values of the errors for each aproximation and vector component + errorsx = [] + errorsy = [] + errorsz = [] + + #Fourier Transform coefficients of the error function for each vector component + fourier_coeffsx = [] + fourier_coeffs_shiftedx = [] + fourier_coeffsy = [] + fourier_coeffs_shiftedy = [] + fourier_coeffsz = [] + fourier_coeffs_shiftedz = [] + max_iter_list = [0,Is, int(10*Is)] + #max_iter_list = [0,2,3,4,5,6,7,8,9,10] + Number_of_iter = len(max_iter_list) + for i in range(Number_of_iter): + + if(max_iter_list[i]>1): + + solver = inverse(A[0],smoother, maxiter = max_iter_list[i], tol = 10.0**-6) + u = solver.dot(b) + #uarr, itter = jacobi(Aarr,barr,max_iter=max_iter_list[i]) + #u = from_array_to_psydac(uarr, derham[0].Vh_fem[sp_key]) + N_itter.append(solver._info['niter']) + #N_itter.append(itter) + + else: + u = derham[0].Vh[derham[0].space_to_form[sp_id]].zeros() + N_itter.append(0) + + + field_aprox.vector = u + if(model == "Poisson"): + errorsx.append(field_star(X,Y,Z)-field_aprox(X,Y,Z)) + else: + errorsx.append(field_star(X,Y,Z)[0]-field_aprox(X,Y,Z)[0]) + errorsy.append(field_star(X,Y,Z)[1]-field_aprox(X,Y,Z)[1]) + errorsz.append(field_star(X,Y,Z)[2]-field_aprox(X,Y,Z)[2]) + + # Compute Fourier Transform coefficients + fourier_coeffsx.append(np.fft.fftn(errorsx[-1])) + fourier_coeffs_shiftedx.append(np.fft.fftshift(fourier_coeffsx[-1])) + if(model != "Poisson"): + fourier_coeffsy.append(np.fft.fftn(errorsy[-1])) + fourier_coeffs_shiftedy.append(np.fft.fftshift(fourier_coeffsy[-1])) + fourier_coeffsz.append(np.fft.fftn(errorsz[-1])) + fourier_coeffs_shiftedz.append(np.fft.fftshift(fourier_coeffsz[-1])) + + # Compute corresponding frequencies + freqsx = np.fft.fftshift(np.fft.fftfreq(len(pointsx), d=dx)) + freqsy = np.fft.fftshift(np.fft.fftfreq(len(pointsy), d=dy)) + freqsz = np.fft.fftshift(np.fft.fftfreq(len(pointsz), d=dz)) + # Create a 3D meshgrid of frequencies + #Fx, Fy, Fz = np.meshgrid(freq_x[-1], freq_y[-1], freq_z[-1], indexing="ij") + + #i also want to visiualize the maximum freqeuncies the next 3 coarser grids can handle + hx = 1/Nel[0] + next_max_frequenciesx = [1.0/(4.0*hx),1.0/(8.0*hx),1.0/(16.0*hx)] + hy = 1/Nel[1] + next_max_frequenciesy = [1.0/(4.0*hy),1.0/(8.0*hy),1.0/(16.0*hy)] + #I get the amplitude of the maximum initial error wave + max_amplitude = np.max(np.abs(fourier_coeffs_shiftedx[0][:,1,1])) + + if(model == "Poisson"): + plt.figure(figsize=(10,10)) + plt.title("Amplitude of error waves for " + model + " with " +smoother+" solver, with resolution " + str(Nel)+".",size = 10) + plt.scatter(freqsx,np.abs(fourier_coeffs_shiftedx[0][:,1,1]), label = "Initial") + plt.scatter(freqsx,np.abs(fourier_coeffs_shiftedx[1][:,1,1]), label = "After "+str(max_iter_list[1])+" iterations ") + plt.scatter(freqsx,np.abs(fourier_coeffs_shiftedx[2][:,1,1]), label = "After "+str(max_iter_list[2])+" iterations ") + plt.plot(next_max_frequenciesx[0]*np.ones(100),np.linspace(0,max_amplitude,100),linestyle = 'dashed', color ='blue',label = "Max freuency level - 1.") + plt.plot(-next_max_frequenciesx[0]*np.ones(100),np.linspace(0,max_amplitude,100),linestyle = 'dashed', color ='blue') + plt.plot(next_max_frequenciesx[1]*np.ones(100),np.linspace(0,max_amplitude,100),linestyle = 'dashed', color ='purple',label = "Max freuency level - 2.") + plt.plot(-next_max_frequenciesx[1]*np.ones(100),np.linspace(0,max_amplitude,100),linestyle = 'dashed', color ='purple') + plt.plot(next_max_frequenciesx[2]*np.ones(100),np.linspace(0,max_amplitude,100),linestyle = 'dashed', color ='red',label = "Max freuency level - 3.") + plt.plot(-next_max_frequenciesx[2]*np.ones(100),np.linspace(0,max_amplitude,100),linestyle = 'dashed', color ='red') + plt.xlabel("Frequency") + plt.ylabel("Amplitude") + plt.legend() + #plt.ylim(0,30) + plt.yscale("log") + plt.savefig("./Wave-Amplitude-Cuboid-"+model+"-"+smoother+"-"+str(Nel)+".pdf") + plt.close() + + + #We also want to see the relative difference between the initial amplitude and the amplitude after smoothing. + #We compute abs(A_0)-abs(A_i)/abs(A_0), this tell us which percentage of the initial amplitude has been eliminated + #A values of 0 means it all still remains, a value of 1 means we eliminated all the amplitude. And a negative value + #means that the amplitude increased. + eliminated_portion = [] + for i in range(2): + aux = (np.abs(fourier_coeffs_shiftedx[0][:,1,1]) - np.abs(fourier_coeffs_shiftedx[i+1][:,1,1]))/np.abs(fourier_coeffs_shiftedx[0][:,1,1]) + eliminated_portion.append(aux) + + mp = np.min(eliminated_portion[0]) + aux = np.min(eliminated_portion[1]) + mp = min(mp,aux) + + plt.figure(figsize=(10,10)) + plt.title("Eliminated portion of error waves for " + model + " with " +smoother+" solver, with resolution " + str(Nel)+".",size = 10) + plt.scatter(freqsx,eliminated_portion[0], label = "Portion of amplitude eliminated after " + str(max_iter_list[1])) + plt.scatter(freqsx,eliminated_portion[1], label = "Portion of amplitude eliminated after " + str(max_iter_list[2])) + plt.plot(next_max_frequenciesx[0]*np.ones(100),np.linspace(mp,1,100),linestyle = 'dashed', color ='blue',label = "Max freuency level - 1.") + plt.plot(-next_max_frequenciesx[0]*np.ones(100),np.linspace(mp,1,100),linestyle = 'dashed', color ='blue') + plt.plot(next_max_frequenciesx[1]*np.ones(100),np.linspace(mp,1,100),linestyle = 'dashed', color ='purple',label = "Max freuency level - 2.") + plt.plot(-next_max_frequenciesx[1]*np.ones(100),np.linspace(mp,1,100),linestyle = 'dashed', color ='purple') + plt.plot(next_max_frequenciesx[2]*np.ones(100),np.linspace(mp,1,100),linestyle = 'dashed', color ='red',label = "Max freuency level - 3.") + plt.plot(-next_max_frequenciesx[2]*np.ones(100),np.linspace(mp,1,100),linestyle = 'dashed', color ='red') + plt.xlabel("Frequency") + plt.ylabel("Eliminated portion of the amplitude.") + plt.legend() + #plt.ylim(0,30) + #plt.yscale("log") + plt.savefig("./Eliminated-Portion-Cuboid-"+model+"-"+smoother+"-"+str(Nel)+".pdf") + plt.close() + + else: + #I get the amplitude of the maximum initial error wave + max_amplitudey = np.max(np.abs(fourier_coeffs_shiftedy[0][:,1,1])) + #I get the amplitude of the maximum initial error wave + max_amplitudez = np.max(np.abs(fourier_coeffs_shiftedz[0][:,1,1])) + + plt.figure(figsize=(10,10)) + plt.title("X-component. Epsilon 10^-6 Amplitude of error waves for " + model + " with " +smoother+" solver, with resolution " + str(Nel)+".",size = 10) + plt.scatter(freqsx,np.abs(fourier_coeffs_shiftedx[0][:,1,1]), label = "Initial") + plt.scatter(freqsx,np.abs(fourier_coeffs_shiftedx[1][:,1,1]), label = "After "+str(max_iter_list[1])+" iterations ") + plt.scatter(freqsx,np.abs(fourier_coeffs_shiftedx[2][:,1,1]), label = "After "+str(max_iter_list[2])+" iterations ") + plt.plot(next_max_frequenciesx[0]*np.ones(100),np.linspace(0,max_amplitude,100),linestyle = 'dashed', color ='blue',label = "Max freuency level - 1.") + plt.plot(-next_max_frequenciesx[0]*np.ones(100),np.linspace(0,max_amplitude,100),linestyle = 'dashed', color ='blue') + plt.plot(next_max_frequenciesx[1]*np.ones(100),np.linspace(0,max_amplitude,100),linestyle = 'dashed', color ='purple',label = "Max freuency level - 2.") + plt.plot(-next_max_frequenciesx[1]*np.ones(100),np.linspace(0,max_amplitude,100),linestyle = 'dashed', color ='purple') + plt.plot(next_max_frequenciesx[2]*np.ones(100),np.linspace(0,max_amplitude,100),linestyle = 'dashed', color ='red',label = "Max freuency level - 3.") + plt.plot(-next_max_frequenciesx[2]*np.ones(100),np.linspace(0,max_amplitude,100),linestyle = 'dashed', color ='red') + plt.xlabel("Frequency") + plt.ylabel("Amplitude") + plt.legend() + #plt.ylim(0,30) + plt.yscale("log") + plt.savefig("./Wave-Amplitude-Cuboid-"+model+"-"+smoother+"-"+str(Nel)+"X-component-Epsilon-10-6.pdf") + plt.close() + + plt.figure(figsize=(10,10)) + plt.title("Y-component. Epsilon 10^-6 Amplitude of error waves for " + model + " with " +smoother+" solver, with resolution " + str(Nel)+".",size = 10) + plt.scatter(freqsx,np.abs(fourier_coeffs_shiftedy[0][:,1,1]), label = "Initial") + plt.scatter(freqsx,np.abs(fourier_coeffs_shiftedy[1][:,1,1]), label = "After "+str(max_iter_list[1])+" iterations ") + plt.scatter(freqsx,np.abs(fourier_coeffs_shiftedy[2][:,1,1]), label = "After "+str(max_iter_list[2])+" iterations ") + plt.plot(next_max_frequenciesx[0]*np.ones(100),np.linspace(0,max_amplitudey,100),linestyle = 'dashed', color ='blue',label = "Max freuency level - 1.") + plt.plot(-next_max_frequenciesx[0]*np.ones(100),np.linspace(0,max_amplitudey,100),linestyle = 'dashed', color ='blue') + plt.plot(next_max_frequenciesx[1]*np.ones(100),np.linspace(0,max_amplitudey,100),linestyle = 'dashed', color ='purple',label = "Max freuency level - 2.") + plt.plot(-next_max_frequenciesx[1]*np.ones(100),np.linspace(0,max_amplitudey,100),linestyle = 'dashed', color ='purple') + plt.plot(next_max_frequenciesx[2]*np.ones(100),np.linspace(0,max_amplitudey,100),linestyle = 'dashed', color ='red',label = "Max freuency level - 3.") + plt.plot(-next_max_frequenciesx[2]*np.ones(100),np.linspace(0,max_amplitudey,100),linestyle = 'dashed', color ='red') + plt.xlabel("Frequency") + plt.ylabel("Amplitude") + plt.legend() + #plt.ylim(0,30) + plt.yscale("log") + plt.savefig("./Wave-Amplitude-Cuboid-"+model+"-"+smoother+"-"+str(Nel)+"Y-component-Epsilon-10-6.pdf") + plt.close() + + + plt.figure(figsize=(10,10)) + plt.title("Z-component. Epsilon 10^-6 Amplitude of error waves for " + model + " with " +smoother+" solver, with resolution " + str(Nel)+".",size = 10) + plt.scatter(freqsx,np.abs(fourier_coeffs_shiftedz[0][:,1,1]), label = "Initial") + plt.scatter(freqsx,np.abs(fourier_coeffs_shiftedz[1][:,1,1]), label = "After "+str(max_iter_list[1])+" iterations ") + plt.scatter(freqsx,np.abs(fourier_coeffs_shiftedz[2][:,1,1]), label = "After "+str(max_iter_list[2])+" iterations ") + plt.plot(next_max_frequenciesx[0]*np.ones(100),np.linspace(0,max_amplitudez,100),linestyle = 'dashed', color ='blue',label = "Max freuency level - 1.") + plt.plot(-next_max_frequenciesx[0]*np.ones(100),np.linspace(0,max_amplitudez,100),linestyle = 'dashed', color ='blue') + plt.plot(next_max_frequenciesx[1]*np.ones(100),np.linspace(0,max_amplitudez,100),linestyle = 'dashed', color ='purple',label = "Max freuency level - 2.") + plt.plot(-next_max_frequenciesx[1]*np.ones(100),np.linspace(0,max_amplitudez,100),linestyle = 'dashed', color ='purple') + plt.plot(next_max_frequenciesx[2]*np.ones(100),np.linspace(0,max_amplitudez,100),linestyle = 'dashed', color ='red',label = "Max freuency level - 3.") + plt.plot(-next_max_frequenciesx[2]*np.ones(100),np.linspace(0,max_amplitudez,100),linestyle = 'dashed', color ='red') + plt.xlabel("Frequency") + plt.ylabel("Amplitude") + plt.legend() + #plt.ylim(0,30) + plt.yscale("log") + plt.savefig("./Wave-Amplitude-Cuboid-"+model+"-"+smoother+"-"+str(Nel)+"Z-component-Epsilon-10-6.pdf") + plt.close() + + + #We also want to see the relative difference between the initial amplitude and the amplitude after smoothing. + #We compute abs(A_0)-abs(A_i)/abs(A_0), this tell us which percentage of the initial amplitude has been eliminated + #A values of 0 means it all still remains, a value of 1 means we eliminated all the amplitude. And a negative value + #means that the amplitude increased. + eliminated_portionx = [] + for i in range(2): + aux = (np.abs(fourier_coeffs_shiftedx[0][:,1,1]) - np.abs(fourier_coeffs_shiftedx[i+1][:,1,1]))/np.abs(fourier_coeffs_shiftedx[0][:,1,1]) + eliminated_portionx.append(aux) + + mp = np.min(eliminated_portionx[0]) + aux = np.min(eliminated_portionx[1]) + mp = min(mp,aux) + + plt.figure(figsize=(10,10)) + plt.title("X-component. Epsilon 10^-6 Eliminated portion of error waves for " + model + " with " +smoother+" solver, with resolution " + str(Nel)+".",size = 10) + plt.scatter(freqsx,eliminated_portionx[0], label = "Portion of amplitude eliminated after " + str(max_iter_list[1])) + plt.scatter(freqsx,eliminated_portionx[1], label = "Portion of amplitude eliminated after " + str(max_iter_list[2])) + plt.plot(next_max_frequenciesx[0]*np.ones(100),np.linspace(mp,1,100),linestyle = 'dashed', color ='blue',label = "Max freuency level - 1.") + plt.plot(-next_max_frequenciesx[0]*np.ones(100),np.linspace(mp,1,100),linestyle = 'dashed', color ='blue') + plt.plot(next_max_frequenciesx[1]*np.ones(100),np.linspace(mp,1,100),linestyle = 'dashed', color ='purple',label = "Max freuency level - 2.") + plt.plot(-next_max_frequenciesx[1]*np.ones(100),np.linspace(mp,1,100),linestyle = 'dashed', color ='purple') + plt.plot(next_max_frequenciesx[2]*np.ones(100),np.linspace(mp,1,100),linestyle = 'dashed', color ='red',label = "Max freuency level - 3.") + plt.plot(-next_max_frequenciesx[2]*np.ones(100),np.linspace(mp,1,100),linestyle = 'dashed', color ='red') + plt.xlabel("Frequency") + plt.ylabel("Eliminated portion of the amplitude.") + plt.legend() + #plt.ylim(0,30) + #plt.yscale("log") + plt.savefig("./Eliminated-Portion-Cuboid-"+model+"-"+smoother+"-"+str(Nel)+"X-component-Epsilon-10-6.pdf") + plt.close() + + + eliminated_portiony = [] + for i in range(2): + aux = (np.abs(fourier_coeffs_shiftedy[0][:,1,1]) - np.abs(fourier_coeffs_shiftedy[i+1][:,1,1]))/np.abs(fourier_coeffs_shiftedy[0][:,1,1]) + eliminated_portiony.append(aux) + + mp = np.min(eliminated_portiony[0]) + aux = np.min(eliminated_portiony[1]) + mp = min(mp,aux) + + plt.figure(figsize=(10,10)) + plt.title("Y-component. Epsilon 10^-6 Eliminated portion of error waves for " + model + " with " +smoother+" solver, with resolution " + str(Nel)+".",size = 10) + plt.scatter(freqsx,eliminated_portiony[0], label = "Portion of amplitude eliminated after " + str(max_iter_list[1])) + plt.scatter(freqsx,eliminated_portiony[1], label = "Portion of amplitude eliminated after " + str(max_iter_list[2])) + plt.plot(next_max_frequenciesx[0]*np.ones(100),np.linspace(mp,1,100),linestyle = 'dashed', color ='blue',label = "Max freuency level - 1.") + plt.plot(-next_max_frequenciesx[0]*np.ones(100),np.linspace(mp,1,100),linestyle = 'dashed', color ='blue') + plt.plot(next_max_frequenciesx[1]*np.ones(100),np.linspace(mp,1,100),linestyle = 'dashed', color ='purple',label = "Max freuency level - 2.") + plt.plot(-next_max_frequenciesx[1]*np.ones(100),np.linspace(mp,1,100),linestyle = 'dashed', color ='purple') + plt.plot(next_max_frequenciesx[2]*np.ones(100),np.linspace(mp,1,100),linestyle = 'dashed', color ='red',label = "Max freuency level - 3.") + plt.plot(-next_max_frequenciesx[2]*np.ones(100),np.linspace(mp,1,100),linestyle = 'dashed', color ='red') + plt.xlabel("Frequency") + plt.ylabel("Eliminated portion of the amplitude.") + plt.legend() + #plt.ylim(0,30) + #plt.yscale("log") + plt.savefig("./Eliminated-Portion-Cuboid-"+model+"-"+smoother+"-"+str(Nel)+"Y-component-Epsilon-10-6.pdf") + plt.close() + + + eliminated_portionz = [] + for i in range(2): + aux = (np.abs(fourier_coeffs_shiftedz[0][:,1,1]) - np.abs(fourier_coeffs_shiftedz[i+1][:,1,1]))/np.abs(fourier_coeffs_shiftedz[0][:,1,1]) + eliminated_portionz.append(aux) + + mp = np.min(eliminated_portionz[0]) + aux = np.min(eliminated_portionz[1]) + mp = min(mp,aux) + + plt.figure(figsize=(10,10)) + plt.title("Z-component. Epsilon 10^-6 Eliminated portion of error waves for " + model + " with " +smoother+" solver, with resolution " + str(Nel)+".",size = 10) + plt.scatter(freqsx,eliminated_portionz[0], label = "Portion of amplitude eliminated after " + str(max_iter_list[1])) + plt.scatter(freqsx,eliminated_portionz[1], label = "Portion of amplitude eliminated after " + str(max_iter_list[2])) + plt.plot(next_max_frequenciesx[0]*np.ones(100),np.linspace(mp,1,100),linestyle = 'dashed', color ='blue',label = "Max freuency level - 1.") + plt.plot(-next_max_frequenciesx[0]*np.ones(100),np.linspace(mp,1,100),linestyle = 'dashed', color ='blue') + plt.plot(next_max_frequenciesx[1]*np.ones(100),np.linspace(mp,1,100),linestyle = 'dashed', color ='purple',label = "Max freuency level - 2.") + plt.plot(-next_max_frequenciesx[1]*np.ones(100),np.linspace(mp,1,100),linestyle = 'dashed', color ='purple') + plt.plot(next_max_frequenciesx[2]*np.ones(100),np.linspace(mp,1,100),linestyle = 'dashed', color ='red',label = "Max freuency level - 3.") + plt.plot(-next_max_frequenciesx[2]*np.ones(100),np.linspace(mp,1,100),linestyle = 'dashed', color ='red') + plt.xlabel("Frequency") + plt.ylabel("Eliminated portion of the amplitude.") + plt.legend() + #plt.ylim(0,30) + #plt.yscale("log") + plt.savefig("./Eliminated-Portion-Cuboid-"+model+"-"+smoother+"-"+str(Nel)+"Z-component-Epsilon-10-6.pdf") + plt.close() + + +def Visualized_all_frequencies_dampening_2D(Nel, plist, spl_kind, Is): + + from struphy.feec.utilities import create_equal_random_arrays + import plotly.graph_objects as go + + # get global communicator + comm = MPI.COMM_WORLD + rank = comm.Get_rank() + world_size = comm.Get_size() + + a1= 0.2 + #domain = HollowCylinder(a1 = a1) + domain = Cuboid() + model = 'Shear-Alfven' + smoother = 'gmres' + derham = [] + mass_ops = [] + A = [] + + epsilon = 1.0 + derham.append(Derham([Nel[0],Nel[1],Nel[2]], plist, spl_kind, comm=comm, local_projectors=False)) + mass_ops.append(WeightedMassOperators(derham[0], domain)) + if(model == "Poisson"): + sp_key = '0' + sp_id = 'H1' + #Poisson + A.append(derham[0].grad.T @ mass_ops[0].M1 @ derham[0].grad) + #Hall-ish + #A.append(epsilon*derham[0].curl.T @ mass_ops[0].M2 @ derham[0].curl + mass_ops[0].M1) + elif(model == "Hall"): + sp_key = '1' + sp_id = 'Hcurl' + #Hall + A.append(mass_ops[0].M1 -epsilon*derham[0].curl.T @ mass_ops[0].M2 @ derham[0].curl) + #Shear-Alfven-ish + #A.append(mass_ops[0].M2 -epsilon* mass_ops[0].M2@ derham[0].curl @ mass_ops[0].M1 @ derham[0].curl.T @ mass_ops[0].M2) + elif(model == 'Shear-Alfven' or model == 'Shear-Alfven-v2'): + sp_key = '2' + sp_id = 'Hdiv' + pc_class = getattr(preconditioner,"MassMatrixPreconditioner") + pc = pc_class(mass_ops[0].M1) + M1_inv = inverse( + mass_ops[0].M1, + "pcg", + pc=pc, + maxiter=3000, + verbose=False, + ) + #Shear-Alfven + if (model == "Shear-Alfven"): + A.append(mass_ops[0].M2 -epsilon* mass_ops[0].M2@ derham[0].curl @ M1_inv @ derham[0].curl.T @ mass_ops[0].M2) + else: + #Shear-Alfven-v2 + A.append(mass_ops[0].M2 -epsilon* derham[0].curl @ M1_inv @ derham[0].curl.T) + + field_star = derham[0].create_field('fh', sp_id) + field_aprox = derham[0].create_field('fh', sp_id) + + #We are gonna use the method of manufacture solutions to determine the behaviour of our error + #Our exact solution is u_star + u_stararr, u_star = create_equal_random_arrays(derham[0].Vh_fem[sp_key], seed=45) + #We compute the rhs + b = A[0].dot(u_star) + field_star.vector = u_star + + #Turn b into an array + #barr = remove_padding(derham[0].Vh_fem[sp_key],b) + #Turn A[0] into an array + #Aarr = A[0].toarray() + + #Gauss = will_gauss_seidel_converge(Aarr, verbose=True) + #print(f"{Gauss = }") + + + #We store the number of itterations + N_itter = [] + #We define a set of point to evaluate the exact solution and the aproximated one + pointsx = np.linspace(0.0,1.0,Nel[0]) + #pointsx = np.array([0.0,0.5]) + pointsy = np.linspace(0.0,1.0,Nel[1]) + #pointsy = np.array([0.0,0.5]) + pointsz = np.array([0.0,0.5]) + dx = pointsx[1] - pointsx[0] # Spacing between points + dy = pointsy[1] - pointsy[0] # Spacing between points + dz = pointsz[1] - pointsz[0] # Spacing between points + X, Y, Z = np.meshgrid(pointsx, pointsy, pointsz, indexing="ij") + + + + #We define a list where to store the arrays with the values of the errors for each aproximation and vector component + errorsx = [] + errorsy = [] + errorsz = [] + + #Fourier Transform coefficients of the error function for each vector component + fourier_coeffsx = [] + fourier_coeffs_shiftedx = [] + fourier_coeffsy = [] + fourier_coeffs_shiftedy = [] + fourier_coeffsz = [] + fourier_coeffs_shiftedz = [] + max_iter_list = [0,Is, int(10*Is)] + #max_iter_list = [0,2,3,4,5,6,7,8,9,10] + Number_of_iter = len(max_iter_list) + for i in range(Number_of_iter): + + if(max_iter_list[i]>1): + + solver = inverse(A[0],smoother, maxiter = max_iter_list[i], tol = 10.0**-6) + u = solver.dot(b) + #uarr, itter = jacobi(Aarr,barr,max_iter=max_iter_list[i]) + #u = from_array_to_psydac(uarr, derham[0].Vh_fem[sp_key]) + N_itter.append(solver._info['niter']) + #N_itter.append(itter) + + else: + u = derham[0].Vh[derham[0].space_to_form[sp_id]].zeros() + N_itter.append(0) + + + field_aprox.vector = u + if(model == "Poisson"): + errorsx.append(field_star(X,Y,Z)-field_aprox(X,Y,Z)) + else: + errorsx.append(field_star(X,Y,Z)[0]-field_aprox(X,Y,Z)[0]) + errorsy.append(field_star(X,Y,Z)[1]-field_aprox(X,Y,Z)[1]) + errorsz.append(field_star(X,Y,Z)[2]-field_aprox(X,Y,Z)[2]) + + # Compute Fourier Transform coefficients + fourier_coeffsx.append(np.fft.fftn(errorsx[-1])) + fourier_coeffs_shiftedx.append(np.fft.fftshift(fourier_coeffsx[-1])) + if(model != "Poisson"): + fourier_coeffsy.append(np.fft.fftn(errorsy[-1])) + fourier_coeffs_shiftedy.append(np.fft.fftshift(fourier_coeffsy[-1])) + fourier_coeffsz.append(np.fft.fftn(errorsz[-1])) + fourier_coeffs_shiftedz.append(np.fft.fftshift(fourier_coeffsz[-1])) + + # Compute corresponding frequencies + freqsx = np.fft.fftshift(np.fft.fftfreq(len(pointsx), d=dx)) + freqsy = np.fft.fftshift(np.fft.fftfreq(len(pointsy), d=dy)) + freqsz = np.fft.fftshift(np.fft.fftfreq(len(pointsz), d=dz)) + # Create a 3D meshgrid of frequencies + Fx, Fy= np.meshgrid(freqsx, freqsy, indexing="ij") + + #i also want to visiualize the maximum freqeuncies the next 3 coarser grids can handle + hx = 1/Nel[0] + next_max_frequenciesx = [1.0/(4.0*hx),1.0/(8.0*hx),1.0/(16.0*hx)] + hy = 1/Nel[1] + next_max_frequenciesy = [1.0/(4.0*hy),1.0/(8.0*hy),1.0/(16.0*hy)] + #I get the amplitude of the maximum initial error wave + #max_amplitude = np.max(np.abs(fourier_coeffs_shiftedx[0][:,1,1])) + + if(model == "Poisson"): + magnitude_spectrum = [] + log_magnitude_spectrum = [] + #We also want to see the relative difference between the initial amplitude and the amplitude after smoothing. + #We compute abs(A_0)-abs(A_i)/abs(A_0), this tell us which percentage of the initial amplitude has been eliminated + #A values of 0 means it all still remains, a value of 1 means we eliminated all the amplitude. And a negative value + #means that the amplitude increased. + eliminated_portion = [] + for i in range(2): + aux = (np.abs(fourier_coeffs_shiftedx[0][:,:,1]) - np.abs(fourier_coeffs_shiftedx[i+1][:,:,1]))/np.abs(fourier_coeffs_shiftedx[0][:,:,1]) + eliminated_portion.append(aux) + # Create interactive 3D surface plot with Plotly + fig = go.Figure() + + fig.add_trace(go.Surface(z=eliminated_portion[i], x=Fx, y=Fy, colorscale="Viridis")) + + # Set log scale for Z-axis + fig.update_layout( + title="Iteration "+str(max_iter_list[i+1])+". Eliminated portion of error waves for " + model + " with " +smoother+" solver, with resolution " + str(Nel)+".", + scene=dict( + xaxis_title="Frequency X", + yaxis_title="Frequency Y", + zaxis_title="Eliminated portion of initial amplitude", + #zaxis_type="log", # Apply log scale to Z-axis + ) + ) + + # Save as an interactive HTML file + fig.write_html("Eliminated-Portion-Cuboid-"+model+"-"+smoother+"-"+str(Nel)+"-iteration-"+str(max_iter_list[i+1])+".html") + + # Show in browser + #fig.show() + + + for i in range(3): + magnitude_spectrum.append(np.abs(fourier_coeffs_shiftedx[i][:,:,1])) + log_magnitude_spectrum.append(np.log1p(magnitude_spectrum[i])) + + # Plot 3D surface of Fourier magnitude spectrum + #fig = plt.figure(figsize=(10, 7)) + #ax = fig.add_subplot(111, projection="3d") + #ax.plot_surface(Fx, Fy, log_magnitude_spectrum[i], cmap="viridis") + + # Add contour plot at the bottom (Z = 0) + #contour = ax.contourf(Fx, Fy, log_magnitude_spectrum[i], zdir="z", offset=log_magnitude_spectrum[i].min(), cmap="viridis") + + # Labels + #ax.set_xlabel("Frequency X") + #ax.set_ylabel("Frequency Y") + #ax.set_zlabel("Log(Amplitude + 1)") + #ax.set_title("Iteration "+str(max_iter_list[i])+". Amplitude of error waves for " + model + " with " +smoother+" solver, with resolution " + str(Nel)+".") + # Adjust view angle for better visibility + #ax.view_init(elev=30, azim=135) + #plt.show() + #plt.close() + + # Create interactive 3D surface plot with Plotly + fig = go.Figure() + + fig.add_trace(go.Surface(z=log_magnitude_spectrum[i], x=Fx, y=Fy, colorscale="Viridis")) + + # Set log scale for Z-axis + fig.update_layout( + title="Iteration "+str(max_iter_list[i])+". Amplitude of error waves for " + model + " with " +smoother+" solver, with resolution " + str(Nel)+".", + scene=dict( + xaxis_title="Frequency X", + yaxis_title="Frequency Y", + zaxis_title="Log(Amplitude + 1)", + #zaxis_type="log", # Apply log scale to Z-axis + ) + ) + + # Save as an interactive HTML file + fig.write_html("Wave-Amplitude-Cuboid-"+model+"-"+smoother+"-"+str(Nel)+"-iteration-"+str(max_iter_list[i])+".html") + + # Show in browser + #fig.show() + + + + else: + + magnitude_spectrumx = [] + log_magnitude_spectrumx = [] + #We also want to see the relative difference between the initial amplitude and the amplitude after smoothing. + #We compute abs(A_0)-abs(A_i)/abs(A_0), this tell us which percentage of the initial amplitude has been eliminated + #A values of 0 means it all still remains, a value of 1 means we eliminated all the amplitude. And a negative value + #means that the amplitude increased. + eliminated_portionx = [] + for i in range(2): + aux = (np.abs(fourier_coeffs_shiftedx[0][:,:,1]) - np.abs(fourier_coeffs_shiftedx[i+1][:,:,1]))/np.abs(fourier_coeffs_shiftedx[0][:,:,1]) + eliminated_portionx.append(aux) + # Create interactive 3D surface plot with Plotly + fig = go.Figure() + + fig.add_trace(go.Surface(z=eliminated_portionx[i], x=Fx, y=Fy, colorscale="Viridis")) + + # Set log scale for Z-axis + fig.update_layout( + title="X-component. Iteration "+str(max_iter_list[i+1])+". Eliminated portion of error waves for " + model + " with " +smoother+" solver, with resolution " + str(Nel)+".", + scene=dict( + xaxis_title="Frequency X", + yaxis_title="Frequency Y", + zaxis_title="Eliminated portion of initial amplitude", + #zaxis_type="log", # Apply log scale to Z-axis + ) + ) + + # Save as an interactive HTML file + fig.write_html("Eliminated-Portion-Cuboid-"+model+"-"+smoother+"-"+str(Nel)+"-iteration-"+str(max_iter_list[i+1])+"X-component.html") + + # Show in browser + #fig.show() + + + for i in range(3): + magnitude_spectrumx.append(np.abs(fourier_coeffs_shiftedx[i][:,:,1])) + log_magnitude_spectrumx.append(np.log1p(magnitude_spectrumx[i])) + + # Create interactive 3D surface plot with Plotly + fig = go.Figure() + + fig.add_trace(go.Surface(z=log_magnitude_spectrumx[i], x=Fx, y=Fy, colorscale="Viridis")) + + # Set log scale for Z-axis + fig.update_layout( + title="X-component. Iteration "+str(max_iter_list[i])+". Amplitude of error waves for " + model + " with " +smoother+" solver, with resolution " + str(Nel)+".", + scene=dict( + xaxis_title="Frequency X", + yaxis_title="Frequency Y", + zaxis_title="Log(Amplitude + 1)", + #zaxis_type="log", # Apply log scale to Z-axis + ) + ) + + # Save as an interactive HTML file + fig.write_html("Wave-Amplitude-Cuboid-"+model+"-"+smoother+"-"+str(Nel)+"-iteration-"+str(max_iter_list[i])+"X-component.html") + + # Show in browser + #fig.show() + + + magnitude_spectrumy = [] + log_magnitude_spectrumy = [] + #We also want to see the relative difference between the initial amplitude and the amplitude after smoothing. + #We compute abs(A_0)-abs(A_i)/abs(A_0), this tell us which percentage of the initial amplitude has been eliminated + #A values of 0 means it all still remains, a value of 1 means we eliminated all the amplitude. And a negative value + #means that the amplitude increased. + eliminated_portiony = [] + for i in range(2): + aux = (np.abs(fourier_coeffs_shiftedy[0][:,:,1]) - np.abs(fourier_coeffs_shiftedy[i+1][:,:,1]))/np.abs(fourier_coeffs_shiftedy[0][:,:,1]) + eliminated_portiony.append(aux) + # Create interactive 3D surface plot with Plotly + fig = go.Figure() + + fig.add_trace(go.Surface(z=eliminated_portiony[i], x=Fx, y=Fy, colorscale="Viridis")) + + # Set log scale for Z-axis + fig.update_layout( + title="Y-component. Iteration "+str(max_iter_list[i+1])+". Eliminated portion of error waves for " + model + " with " +smoother+" solver, with resolution " + str(Nel)+".", + scene=dict( + xaxis_title="Frequency X", + yaxis_title="Frequency Y", + zaxis_title="Eliminated portion of initial amplitude", + #zaxis_type="log", # Apply log scale to Z-axis + ) + ) + + # Save as an interactive HTML file + fig.write_html("Eliminated-Portion-Cuboid-"+model+"-"+smoother+"-"+str(Nel)+"-iteration-"+str(max_iter_list[i+1])+"Y-component.html") + + # Show in browser + #fig.show() + + + for i in range(3): + magnitude_spectrumy.append(np.abs(fourier_coeffs_shiftedy[i][:,:,1])) + log_magnitude_spectrumy.append(np.log1p(magnitude_spectrumy[i])) + + # Create interactive 3D surface plot with Plotly + fig = go.Figure() + + fig.add_trace(go.Surface(z=log_magnitude_spectrumy[i], x=Fx, y=Fy, colorscale="Viridis")) + + # Set log scale for Z-axis + fig.update_layout( + title="Y-component. Iteration "+str(max_iter_list[i])+". Amplitude of error waves for " + model + " with " +smoother+" solver, with resolution " + str(Nel)+".", + scene=dict( + xaxis_title="Frequency X", + yaxis_title="Frequency Y", + zaxis_title="Log(Amplitude + 1)", + #zaxis_type="log", # Apply log scale to Z-axis + ) + ) + + # Save as an interactive HTML file + fig.write_html("Wave-Amplitude-Cuboid-"+model+"-"+smoother+"-"+str(Nel)+"-iteration-"+str(max_iter_list[i])+"Y-component.html") + + # Show in browser + #fig.show() + + magnitude_spectrumz = [] + log_magnitude_spectrumz = [] + #We also want to see the relative difference between the initial amplitude and the amplitude after smoothing. + #We compute abs(A_0)-abs(A_i)/abs(A_0), this tell us which percentage of the initial amplitude has been eliminated + #A values of 0 means it all still remains, a value of 1 means we eliminated all the amplitude. And a negative value + #means that the amplitude increased. + eliminated_portionz = [] + for i in range(2): + aux = (np.abs(fourier_coeffs_shiftedz[0][:,:,1]) - np.abs(fourier_coeffs_shiftedz[i+1][:,:,1]))/np.abs(fourier_coeffs_shiftedz[0][:,:,1]) + eliminated_portionz.append(aux) + # Create interactive 3D surface plot with Plotly + fig = go.Figure() + + fig.add_trace(go.Surface(z=eliminated_portionz[i], x=Fx, y=Fy, colorscale="Viridis")) + + # Set log scale for Z-axis + fig.update_layout( + title="Z-component. Iteration "+str(max_iter_list[i+1])+". Eliminated portion of error waves for " + model + " with " +smoother+" solver, with resolution " + str(Nel)+".", + scene=dict( + xaxis_title="Frequency X", + yaxis_title="Frequency Y", + zaxis_title="Eliminated portion of initial amplitude", + #zaxis_type="log", # Apply log scale to Z-axis + ) + ) + + # Save as an interactive HTML file + fig.write_html("Eliminated-Portion-Cuboid-"+model+"-"+smoother+"-"+str(Nel)+"-iteration-"+str(max_iter_list[i+1])+"Z-component.html") + + # Show in browser + #fig.show() + + + for i in range(3): + magnitude_spectrumz.append(np.abs(fourier_coeffs_shiftedz[i][:,:,1])) + log_magnitude_spectrumz.append(np.log1p(magnitude_spectrumz[i])) + + # Create interactive 3D surface plot with Plotly + fig = go.Figure() + + fig.add_trace(go.Surface(z=log_magnitude_spectrumz[i], x=Fx, y=Fy, colorscale="Viridis")) + + # Set log scale for Z-axis + fig.update_layout( + title="Z-component. Iteration "+str(max_iter_list[i])+". Amplitude of error waves for " + model + " with " +smoother+" solver, with resolution " + str(Nel)+".", + scene=dict( + xaxis_title="Frequency X", + yaxis_title="Frequency Y", + zaxis_title="Log(Amplitude + 1)", + #zaxis_type="log", # Apply log scale to Z-axis + ) + ) + + # Save as an interactive HTML file + fig.write_html("Wave-Amplitude-Cuboid-"+model+"-"+smoother+"-"+str(Nel)+"-iteration-"+str(max_iter_list[i])+"Z-component.html") + + # Show in browser + #fig.show() + + +def make_plot_smoothing(): + from numpy import array + + ####### + #Example 1 + #Shear-Alfven + #Cuboid + #2D + #CG + ###### + + #Nel = [128,128,1] + #p = [1,1,1] + #Shear-Alfven + #Cuboid + #CG + + #Magnitude of the high frequency (period smaller than 4h) with highest magnitude + + + ############################ + #Nel = [64,64,1] + #p = [1,1,1] + #Shear-Alfven + #Cuboid + #CG + + + ############################ + #Nel = [32,32,1] + #p = [1,1,1] + #Shear-Alfven + #Cuboid + #CG + + + ############################ + #Nel = [16,16,1] + #p = [1,1,1] + #Shear-Alfven + #Cuboid + #CG + + + + ############################ + #Nel = [8,8,1] + #p = [1,1,1] + #Shear-Alfven + #Cuboid + #CG + + + ################################################################ + ################################################################ + ################################################################ + ################################################################ + ####### + #Example 2 + #Shear-Alfven + #Cuboid + #2D + #biCG + ###### + + ############################ + #Nel = [128,128,1] + #p = [1,1,1] + #Shear-Alfven + #Cuboid + #biCG + + + + + ############################ + #Nel = [64,64,1] + #p = [1,1,1] + #Shear-Alfven + #Cuboid + #biCG + + + ############################ + #Nel = [32,32,1] + #p = [1,1,1] + #Shear-Alfven + #Cuboid + #biCG + + + + ############################ + #Nel = [16,16,1] + #p = [1,1,1] + #Shear-Alfven + #Cuboid + #biCG + + + + + ################################################################ + ################################################################ + ################################################################ + ################################################################ + ####### + #Example 3 + #Shear-Alfven + #Cuboid + #2D + #bicgstab + ###### + #Terrible the high frequency errors increase with the number of itterations + + ############################ + #Nel = [128,128,1] + #p = [1,1,1] + #Shear-Alfven + #Cuboid + #bicgstab + + + + + ############################ + #Nel = [64,64,1] + #p = [1,1,1] + #Shear-Alfven + #Cuboid + #bicgstab + + + + + ################################################################ + ################################################################ + ################################################################ + ################################################################ + + + ####### + #Example 4 + #Shear-Alfven + #Cuboid + #2D + #minres + + ############################ + #Nel = [128,128,1] + #p = [1,1,1] + #Hall + #Cuboid + #minres + + magnitudes = [23475.547406074205, 23507.57071761296, 23401.98167905894, 8975.398716277663, 5547.855656375274, 2215.014445147343, 1060.3501137658902, 829.7441809688149, 220.71658497349233, 257.8922866406492] + bad_frequencies = [array([ 44.6484375, -14.8828125, -62.5078125]), array([-44.6484375, 14.8828125, -62.5078125]), array([-44.6484375, 14.8828125, -62.5078125]), array([ 62.5078125, -31.75 , -62.5078125]), array([-63.5 , -43.65625 , -62.5078125]), array([-59.53125 , 62.5078125, -62.5078125]), array([-59.53125 , 62.5078125, -62.5078125]), array([ 59.53125 , -62.5078125, -62.5078125]), array([-59.53125 , 62.5078125, -62.5078125]), array([-59.53125 , 62.5078125, -62.5078125])] + N_itter = [2, 10, 20, 100, 200, 300, 400, 500, 600, 700] + + + ############################ + #Nel = [64,64,1] + #p = [1,1,1] + #Shear-Alfven + #Cuboid + #minres + + magnitudes = [4816.563672246145, 4820.683874727215, 4750.075275409996, 1435.7294516177942, 230.5875341851168, 45.396511954402975, 6.7027103058152555, 6.6240853651862395, 5.862590252251258, 0.753372104197576] + bad_frequencies = [array([-24.609375, -11.8125 , -30.515625]), array([-24.609375, -11.8125 , -30.515625]), array([-24.609375, -11.8125 , -30.515625]), array([-30.515625, 25.59375 , -30.515625]), array([-30.515625, 25.59375 , -30.515625]), array([ 29.53125 , 30.515625, -30.515625]), array([-30.515625, 25.59375 , -30.515625]), array([-29.53125 , -30.515625, -30.515625]), array([-29.53125 , -30.515625, -30.515625]), array([-24.609375, -12.796875, -30.515625])] + N_itter = [2, 10, 20, 100, 200, 300, 400, 500, 600, 700] + + + ############################ + #Nel = [32,32,1] + #p = [1,1,1] + #Shear-Alfven + #Cuboid + #minres + + magnitudes = [1319.382611683582, 1271.7302885073973, 1096.5001351152466, 66.25984056212276, 2.305135582316737, 0.5738656559575724, 0.13428796602415516, 0.0666372580855523, 0.045512532749170435, 0.030227161820733185] + bad_frequencies = [array([-13.5625 , 0. , -14.53125]), array([-13.5625 , 0. , -14.53125]), array([-13.5625 , 0. , -14.53125]), array([-12.59375, -13.5625 , -14.53125]), array([-12.59375, -13.5625 , -14.53125]), array([-13.5625 , -14.53125, -14.53125]), array([-13.5625 , -14.53125, -14.53125]), array([-12.59375, -15.5 , -14.53125]), array([-12.59375, -13.5625 , -14.53125]), array([ 14.53125, 14.53125, -14.53125])] + N_itter = [2, 10, 20, 100, 200, 300, 400, 500, 600, 700] + + + ################################################################ + ################################################################ + ################################################################ + ################################################################ + + ####### + #Example 5 + #Shear-Alfven + #Cuboid + #2D + #lsmr + + ############################ + #Nel = [128,128,1] + #p = [1,1,1] + #Shear-Alfven + #Cuboid + #lsmr + magnitudes = [23625.08035709719, 23546.07810652249, 23555.35959665252, 23553.74183488159, 23553.11191281097, 23552.06535579401, 23550.62355697153, 23548.584106844555, 23546.582190989895, 23543.658419538646] + bad_frequencies = [array([-40.6796875, -0.9921875, -62.5078125]), array([-44.6484375, 14.8828125, -62.5078125]), array([-44.6484375, 14.8828125, -62.5078125]), array([-44.6484375, 14.8828125, -62.5078125]), array([-44.6484375, 14.8828125, -62.5078125]), array([-44.6484375, 14.8828125, -62.5078125]), array([-44.6484375, 14.8828125, -62.5078125]), array([-44.6484375, 14.8828125, -62.5078125]), array([-44.6484375, 14.8828125, -62.5078125]), array([-44.6484375, 14.8828125, -62.5078125])] + N_itter = [2, 10, 20, 100, 200, 300, 400, 500, 600, 700] + + ############################ + #Nel = [64,64,1] + #p = [1,1,1] + #Shear-Alfven + #Cuboid + #lsmr + + magnitudes = [4902.303542373743, 4846.903947432333, 4849.383867406714, 4849.313905985293, 4847.890279297882, 4845.825151419513, 4842.427718581035, 4839.574277539976, 4832.651039101285, 4824.670040396377] + bad_frequencies = [array([-24.609375, -11.8125 , -30.515625]), array([-24.609375, -11.8125 , -30.515625]), array([-24.609375, -11.8125 , -30.515625]), array([ 24.609375, 11.8125 , -30.515625]), array([-24.609375, -11.8125 , -30.515625]), array([-24.609375, -11.8125 , -30.515625]), array([-24.609375, -11.8125 , -30.515625]), array([-24.609375, -11.8125 , -30.515625]), array([-24.609375, -11.8125 , -30.515625]), array([-24.609375, -11.8125 , -30.515625])] + N_itter = [2, 10, 20, 100, 200, 300, 400, 500, 600, 700] + + + ################################################################ + ################################################################ + ################################################################ + ################################################################ + + ####### + #Example 6 + #Shear-Alfven + #Cuboid + #2D + #gmres + + ############################ + #Nel = [128,128,1] + #p = [1,1,1] + #Shear-Alfven + #Cuboid + #gmres + + magnitudes = [23631.325965428965, 23517.072741223346, 23415.519704485465, 9143.805698693457, 5620.759347312477, 2224.5465588280504, 1062.7363141542924, 813.4419688746405, 239.7543262098424, 257.9845124527011] + bad_frequencies = [array([-40.6796875, -0.9921875, -62.5078125]), array([-44.6484375, 14.8828125, -62.5078125]), array([ 44.6484375, -14.8828125, -62.5078125]), array([-62.5078125, 31.75 , -62.5078125]), array([-63.5 , -43.65625 , -62.5078125]), array([ 59.53125 , -62.5078125, -62.5078125]), array([-59.53125 , 62.5078125, -62.5078125]), array([-59.53125 , 62.5078125, -62.5078125]), array([-59.53125 , 62.5078125, -62.5078125]), array([-59.53125 , 62.5078125, -62.5078125])] + N_itter = [2, 10, 20, 100, 200, 300, 400, 500, 600, 700] + + + + + + + + + + + + + ####### + #Example + #Shear-Alfven + #Cuboid + #1D + #CG + ###### + + + ############################ + #Nel = [1024,1,1] + #p = [1,1,1] + #Shear-Alfven + #Cuboid + #CG + + + #As you can see in 1D the CG reduces the magnitude of all problematic high frequencies to 2 e-05. In 1D the CG for the shear-alfven matrix is a good smoother. Thus the Multigrid method works + #with it. But in 2D the CG method is a terrible smoother for this matrix so the multigrid method becomes almost useless with it. + + +def multigrid_Alfven(Nel, plist, spl_kind, u_space): + # get global communicator + comm = MPI.COMM_WORLD + rank = comm.Get_rank() + world_size = comm.Get_size() + derham = Derham(Nel, plist, spl_kind, comm=comm, local_projectors=False) + domain = Tokamak(p = [3,3], psi_shifts = [2.,2.]) + + mhd_equil = AdhocTorusQPsi() + + # must set domain of Cartesian MHD equilibirum + mhd_equil.domain = domain + + + input = derham.Vh[derham.space_to_form[u_space]].zeros() + + + mass_ops = WeightedMassOperators(derham, domain) + basis_ops = BasisProjectionOperators(derham, domain,eq_mhd=mhd_equil) + + id_T = "T" + derham.space_to_form[u_space] + + _T = getattr(basis_ops, id_T) + + _B = -1 / 2 * _T.T @ derham.curl.T @ mass_ops.M2 + _C = 1 / 2 * derham.curl @ _T + + _BC = _B @ _C + + solver = inverse(_BC,'cg') + x = solver.dot(input) + + print(solver._info) + + +def multigrid(Nel, plist, spl_kind, N_levels): + from struphy.feec.utilities import create_equal_random_arrays + # get global communicator + comm = MPI.COMM_WORLD + rank = comm.Get_rank() + world_size = comm.Get_size() + + domain = Cuboid() + #a1 = 0.002 + #domain = HollowCylinder(a1= a1) + sp_key = '0' + + derham = [] + mass_ops = [] + A = [] + + epsilon = 0.0002 + for level in range(N_levels): + derham.append(Derham([Nel[0]//(2**level),Nel[1]//(2**level),Nel[2]], plist, spl_kind, comm=comm, local_projectors=False)) + mass_ops.append(WeightedMassOperators(derham[level], domain)) + A.append(derham[level].grad.T @ mass_ops[level].M1 @ derham[level].grad) + #A.append(epsilon*derham[level].curl.T @ mass_ops[level].M2 @ derham[level].curl + mass_ops[level].M1) + + #We get the inverse of the coarsest system matrix to solve directly the problem in the smaller space + A_inv = np.linalg.inv(A[-1].toarray()) + + R = [] + E = [] + + for level in range(N_levels-1): + R.append(RestrictionOperator(derham[level].Vh_fem[sp_key],derham[level+1].Vh_fem[sp_key])) + E.append(R[level].transpose()) + + method = 'cg' + + #800 + max_iter_list = [14,14,14,18,12,10] + #40 + N_cycles = 2 + + u_stararr, u_star = create_equal_random_arrays(derham[0].Vh_fem[sp_key], seed=45) + + #We compute the rhs + b = A[0].dot(u_star) + + timei = time.time() + + solver_no = inverse(A[0],method, maxiter = 1000000, tol = 10**(-6)) + + u = solver_no.dot(b) + + timef = time.time() + + #We get the total number of itterations + No_Multigrid_itterations = solver_no._info['niter'] + No_Multigrid_error = solver_no._info['res_norm'] + No_Multigrid_time = timef-timei + + #No_Multigrid_itterations = 35088 + #No_Multigrid_error = 9.22E-07 + #No_Multigrid_time = 451.228641271591 + + print("################") + print(f'{No_Multigrid_itterations = }') + print(f'{No_Multigrid_error = }') + print(f'{No_Multigrid_time = }') + print("################") + + + def call_multigrid(max_iter, N_cycles): + u_stararr, u_star = create_equal_random_arrays(derham[0].Vh_fem[sp_key], seed=45) + #We compute the rhs + b = A[0].dot(u_star) + + #We define a list where to store the number of itteration it takes at each multigrid level + #Change N_levels for 1D case + Multigrid_itterations = np.zeros(N_levels, dtype=int) + converged = np.zeros(N_levels, dtype=bool) + + def V_cycle(l, r_l): + #Change for N_levels-1 for 1D case + if (l < N_levels-1): + solver_ini = inverse(A[l],method, maxiter= max_iter[l]) + x_l = solver_ini.dot(r_l) + + #We count the number of itterations + Multigrid_itterations[l] += solver_ini._info['niter'] + + #We determine if the itterative solver converged in the maximum number of itterations + converged[l] = solver_ini._info['success'] + if converged[l] == True: + return x_l + + r_l = r_l - A[l].dot(x_l) + + r_l_plus_1 = R[l].dot(r_l) + x_l_plus_1 = V_cycle(l+1, r_l_plus_1) + #New + #x_l_aux = E[l].dot(x_l_plus_1) + #x_l = x_l + x_l_aux + #r_l = r_l - A[l].dot(x_l_aux) + #solver_end = inverse(A[l].T,method, maxiter= max_iter, x0 =x_l) + #x_l = solver_end.dot(r_l) + #Multigrid_itterations[l] += solver_end._info['niter'] + #### + #old + x_l = x_l + E[l].dot(x_l_plus_1) + ### + + else: + #Solve directly + x_l = direct_solver(A_inv,r_l, derham[l].Vh_fem[sp_key]) + + return x_l + + #N_cycles = 6 + x_0 = derham[0].Vh_fem[sp_key].vector_space.zeros() + + timei = time.time() + for cycle in range(N_cycles): + solver = inverse(A[0],method, maxiter= max_iter[0], x0 = x_0) + x_0 = solver.dot(b) + + Multigrid_itterations[0] += solver._info['niter'] + + #We determine if the itterative solver converged in the maximum number of itterations + converged[0] = solver._info['success'] + if converged[0] == True: + print("Hello") + x = x_0 + break + + r_0 = b - A[0].dot(x_0) + r_1 = R[0].dot(r_0) + + x_0 = x_0 + E[0].dot(V_cycle(1,r_1)) + + if converged[0] == False: + solver = inverse(A[0],method,x0 = x_0, tol = 10**(-6)) + x = solver.dot(b) + Multigrid_itterations[0] += solver._info['niter'] + + timef = time.time() + + + + #We get the final error + Multigrid_error = solver._info['res_norm'] + + Multigrid_time = timef- timei + + speed_up = No_Multigrid_time / Multigrid_time + #speed_up = 27.3452200889587 / Multigrid_time + + print("################") + print("################") + print("################") + #print(f'{a1 = }') + print(f'{max_iter = }') + print(f'{N_cycles = }') + print("################") + print("################") + print(f'{Multigrid_itterations = }') + print(f'{Multigrid_error = }') + print(f'{Multigrid_time = }') + print("################") + print("################") + print(f'{speed_up = }') + + call_multigrid(max_iter_list, N_cycles) + + +def Error_analysis(Nel, plist, spl_kind, N_levels): + + def determine_error(exact,aprox,x): + errors = exact(x,0.0,0.0)-aprox(x,0.0,0.0) + errors = errors.flatten() + return errors + + + + + from struphy.feec.utilities import create_equal_random_arrays + # get global communicator + comm = MPI.COMM_WORLD + rank = comm.Get_rank() + world_size = comm.Get_size() + + domain = Cuboid() + + sp_key = '0' + derham = [] + mass_ops = [] + A = [] + + + for level in range(N_levels): + + derham.append(Derham([Nel[0]//(2**level),Nel[1],Nel[2]], plist, spl_kind, comm=comm, local_projectors=False)) + mass_ops.append(WeightedMassOperators(derham[level], domain)) + A.append(derham[level].grad.T @ mass_ops[level].M1 @ derham[level].grad) + + + field_star = derham[0].create_field('fh', 'H1') + field_aprox = derham[0].create_field('fh', 'H1') + + #We are gonna use the method of manufacture solutions to determine the behaviour of our error + #Our exact solution is u_star + u_stararr, u_star = create_equal_random_arrays(derham[0].Vh_fem[sp_key], seed=45) + #We compute the rhs + b = A[0].dot(u_star) + field_star.vector = u_star + + #We store the number of itterations + N_itter = [] + #We define a set of point to evaluate the exact solution and the aproximated one + points = np.linspace(0.0,1.0,1000) + dx = points[1] - points[0] # Spacing between points + #We define a list where to store the arrays with the values of the errors for each aproximation + errors = [] + + #Fourier Transform coefficients of the error function + fourier_coeffs = [] + #Corresponding frequencies + freqs = [] + for i in range(1, 10): + solver = inverse(A[0],'cg', maxiter = int(i*50) ) + u = solver.dot(b) + field_aprox.vector = u + errors.append(determine_error(field_star,field_aprox,points)) + N_itter.append(solver._info['niter']) + + # Compute Fourier Transform coefficients + fourier_coeffs.append(np.fft.fft(errors[-1])) + # Compute corresponding frequencies + freqs.append(np.fft.fftfreq(len(errors[-1]), d=dx)) + + + + plt.figure() + plt.plot(points,errors[0], label = str(N_itter[0])) + plt.plot(points,errors[-1], label = str(N_itter[-1])) + #plt.scatter(freqs[0],fourier_coeffs[0], label = str(N_itter[0])) + #plt.scatter(freqs[3],fourier_coeffs[3], label = str(N_itter[3])) + #plt.xlim(-50,50) + plt.legend() + plt.show() + plt.close() + + +def Gather_data_V_cycle_parameter_study(Nel, plist, spl_kind, N_levels): + from struphy.feec.utilities import create_equal_random_arrays + # get global communicator + comm = MPI.COMM_WORLD + rank = comm.Get_rank() + world_size = comm.Get_size() + + domain = Cuboid() + a1 = 0.002 + #domain = HollowCylinder(a1= a1) + sp_key = '2' + + derham = [] + mass_ops = [] + A = [] + + epsilon = 1.0 + for level in range(N_levels): + derham.append(Derham([Nel[0]//(2**level),Nel[1]//(2**level),Nel[2]], plist, spl_kind, comm=comm, local_projectors=False)) + mass_ops.append(WeightedMassOperators(derham[level], domain)) + #Poisson + #A.append(derham[level].grad.T @ mass_ops[level].M1 @ derham[level].grad) + #Hall-ish + #A.append(epsilon*derham[level].curl.T @ mass_ops[level].M2 @ derham[level].curl + mass_ops[level].M1) + #Hall + #A.append(mass_ops[level].M1 -epsilon*derham[level].curl.T @ mass_ops[level].M2 @ derham[level].curl) + #Shear-Alfven-ish + #A.append(mass_ops[level].M2 -epsilon* mass_ops[level].M2@ derham[level].curl @ mass_ops[level].M1 @ derham[level].curl.T @ mass_ops[level].M2) + #Shear-Alfven-ish-v2 + #A.append(mass_ops[level].M2 -epsilon* derham[level].curl @ mass_ops[level].M1 @ derham[level].curl.T) + #Shear-Alfven + pc_class = getattr(preconditioner,"MassMatrixPreconditioner") + pc = pc_class(mass_ops[level].M1) + M1_inv = inverse( + mass_ops[level].M1, + "pcg", + pc=pc, + maxiter=3000, + verbose=False, + ) + A.append(mass_ops[level].M2 -epsilon* mass_ops[level].M2@ derham[level].curl @ M1_inv @ derham[level].curl.T @ mass_ops[level].M2) + #Shear-Alfven-v2 + #A.append(mass_ops[level].M2 -epsilon* derham[level].curl @ M1_inv @ derham[level].curl.T) + + #We get the inverse of the coarsest system matrix to solve directly the problem in the smaller space + A_inv = np.linalg.inv(A[-1].toarray()) + + R = [] + E = [] + + for level in range(N_levels-1): + R.append(RestrictionOperator(derham[level].Vh_fem[sp_key],derham[level+1].Vh_fem[sp_key])) + E.append(R[level].transpose()) + + method = 'cg' + + #800 + max_iter_list = [93] + #40 + N_cycles_list = [5] + + u_stararr, u_star = create_equal_random_arrays(derham[0].Vh_fem[sp_key], seed=45) + #We compute the rhs + b = A[0].dot(u_star) + + timei = time.time() + + solver_no = inverse(A[0],method, maxiter = 10000) + + u = solver_no.dot(b) + + timef = time.time() + + #We get the total number of itterations + No_Multigrid_itterations = solver_no._info['niter'] + No_Multigrid_error = solver_no._info['res_norm'] + No_Multigrid_time = timef-timei + + #No_Multigrid_itterations = 9485 + #No_Multigrid_error = 9.72E-07 + #No_Multigrid_time = 867.623549938202 + print("################") + print("################") + print(f'{Nel[0] = }') + print(f'{Nel[1] = }') + print("################") + print("################") + + + print("################") + print(f'{No_Multigrid_itterations = }') + print(f'{No_Multigrid_error = }') + print(f'{No_Multigrid_time = }') + print("################") + + + def call_multigrid(max_iter, N_cycles): + u_stararr, u_star = create_equal_random_arrays(derham[0].Vh_fem[sp_key], seed=45) + #We compute the rhs + b = A[0].dot(u_star) + + #We define a list where to store the number of itteration it takes at each multigrid level + #Change N_levels for 1D case + Multigrid_itterations = np.zeros(N_levels, dtype=int) + converged = np.zeros(N_levels, dtype=bool) + + def V_cycle(l, r_l): + #Change for N_levels-1 for 1D case + if (l < N_levels-1): + solver_ini = inverse(A[l],method, maxiter= max_iter) + x_l = solver_ini.dot(r_l) + + #We count the number of itterations + Multigrid_itterations[l] += solver_ini._info['niter'] + + #We determine if the itterative solver converged in the maximum number of itterations + converged[l] = solver_ini._info['success'] + if converged[l] == True: + return x_l + + r_l = r_l - A[l].dot(x_l) + + r_l_plus_1 = R[l].dot(r_l) + x_l_plus_1 = V_cycle(l+1, r_l_plus_1) + #New + #x_l_aux = E[l].dot(x_l_plus_1) + #x_l = x_l + x_l_aux + #r_l = r_l - A[l].dot(x_l_aux) + #solver_end = inverse(A[l].T,method, maxiter= max_iter, x0 =x_l) + #x_l = solver_end.dot(r_l) + #Multigrid_itterations[l] += solver_end._info['niter'] + #### + #old + x_l = x_l + E[l].dot(x_l_plus_1) + ### + + else: + #Solve directly + x_l = direct_solver(A_inv,r_l, derham[l].Vh_fem[sp_key]) + + return x_l + + #N_cycles = 6 + x_0 = derham[0].Vh_fem[sp_key].vector_space.zeros() + + timei = time.time() + for cycle in range(N_cycles): + solver = inverse(A[0],method, maxiter= max_iter, x0 = x_0) + x_0 = solver.dot(b) + + Multigrid_itterations[0] += solver._info['niter'] + + #We determine if the itterative solver converged in the maximum number of itterations + converged[0] = solver._info['success'] + if converged[0] == True: + print("Hello") + x = x_0 + break + + r_0 = b - A[0].dot(x_0) + r_1 = R[0].dot(r_0) + + x_0 = x_0 + E[0].dot(V_cycle(1,r_1)) + + if converged[0] == False: + solver = inverse(A[0],method, maxiter = 1000,x0 = x_0,) + x = solver.dot(b) + Multigrid_itterations[0] += solver._info['niter'] + + timef = time.time() + + + + #We get the final error + Multigrid_error = solver._info['res_norm'] + + Multigrid_time = timef- timei + + speed_up = No_Multigrid_time / Multigrid_time + + print("################") + print("################") + print("################") + #print(f'{a1 = }') + print(f'{max_iter = }') + print(f'{N_cycles = }') + print("################") + print("################") + print(f'{Multigrid_itterations = }') + print(f'{Multigrid_error = }') + print(f'{Multigrid_time = }') + print("################") + print("################") + print(f'{speed_up = }') + + for max_iter in max_iter_list: + for N_cycles in N_cycles_list: + call_multigrid(max_iter, N_cycles) + + +def Gather_data_V_cycle_scalability(Nellist, plist, spl_kind): + for Nel in Nellist: + #First we compute the number of levels for each Nel + N_levels = int(log2(Nel[0])-1) + Gather_data_V_cycle_parameter_study(Nel, plist, spl_kind, N_levels) + + +def make_plot_scalability(): + + model = 'Hall' + + Nel = [int(2**i * 2**i) for i in range(4,8)] + Multi_time = [0.933451652526856, + 4.96578407287598, + 21.1356129646301, + 121.159014463425 + ] + No_Multi_time = [0.853443145751953, + 6.89243221282959, + 75.1050372123718, + 1334.72232866287 + ] + speed_up = [] + + for i in range(len(Multi_time)): + speed_up.append(No_Multi_time[i]/Multi_time[i]) + + plt.figure() + plt.title('2D ' +model+' run times.') + plt.plot(Nel, Multi_time, label = 'Multigrid.') + plt.scatter(Nel, No_Multi_time, label = 'CG run time.') + plt.xlabel("Nel[0] x Nel[1]") + plt.ylabel('Time (s)') + plt.legend() + plt.savefig("2D-"+model+"-runtime.pdf") + #plt.show() + plt.close() + + plt.figure() + plt.title('2D ' +model+ ' speed up.') + plt.scatter(Nel, speed_up, label = 'Speed_up.') + plt.xlabel("Nel[0] x Nel[1]") + plt.ylabel('Speed_up') + plt.legend() + #plt.show() + plt.savefig("2D-"+model+"-speed-up.pdf") + plt.close() + + +def verify_formula(Nel, plist, spl_kind): + comm = MPI.COMM_WORLD + rank = comm.Get_rank() + derham = Derham(Nel, plist, spl_kind, comm=comm) + + # For B-splines + sp_key = "0" + spaces = derham.Vh_fem[sp_key].spaces + space = spaces[0] + N = space.nbasis + ncells = space.ncells + p = space.degree + T = space.knots + periodic = space.periodic + basis = space.basis + normalize = basis == "M" + + def make_basis_fun(i): + def fun(etas, eta2, eta3): + if isinstance(etas, float) or isinstance(etas, int): + etas = np.array([etas]) + out = np.zeros_like(etas) + for j, eta in enumerate(etas): + span = find_span(T, p, eta) + inds = np.arange(span - p, span + 1) % N + pos = np.argwhere(inds == i) + # print(f'{pos = }') + if pos.size > 0: + pos = pos[0, 0] + out[j] = basis_funs(T, p, eta, span, normalize=normalize)[pos] + else: + out[j] = 0.0 + return out + + return fun + + i = random.randint(0,N-1) + fun = make_basis_fun(i) + points = np.linspace(0.0,1.0,100) + values = fun(points,0.0,0.0) + + + #Now we build the fine grid + derham_fine = Derham([Nel[0]*2,Nel[1],Nel[2]], plist, spl_kind, comm=comm) + + # For B-splines + spaces_fine = derham_fine.Vh_fem[sp_key].spaces + space_fine = spaces_fine[0] + N_fine = space_fine.nbasis + ncells_fine = space_fine.ncells + p_fine = space_fine.degree + T_fine = space_fine.knots + periodic_fine = space_fine.periodic + basis_fine = space_fine.basis + normalize_fine = basis_fine == "M" + + def make_basis_fun_fine(i): + def fun(etas, eta2, eta3): + if isinstance(etas, float) or isinstance(etas, int): + etas = np.array([etas]) + out = np.zeros_like(etas) + for j, eta in enumerate(etas): + span = find_span(T_fine, p_fine, eta) + inds = np.arange(span - p_fine, span + 1) % N_fine + pos = np.argwhere(inds == i) + # print(f'{pos = }') + if pos.size > 0: + pos = pos[0, 0] + out[j] = basis_funs(T_fine, p_fine, eta, span, normalize=normalize_fine)[pos] + else: + out[j] = 0.0 + return out + + return fun + + fun_fine = [] + weights_fine = [] + for j in range(p_fine+2): + fun_fine.append(make_basis_fun_fine((2*i-p_fine+j)%N_fine)) + weights_fine.append(2.0**(-p_fine)*comb(p_fine+1,j)) + + + def fun_combine(etas, eta2, eta3): + if isinstance(etas, float) or isinstance(etas, int): + etas = np.array([etas]) + out = np.zeros_like(etas) + for j in range(p_fine+2): + out += weights_fine[j]*fun_fine[j](etas,0.0,0.0) + return out + + + + values_fine = fun_combine(points,0.0,0.0) + + Equal = True + where = -1 + for j in range(len(values)): + if(abs(values[j] -values_fine[j])>10**-5): + Equal = False + where = j + + print(i) + print(Equal) + print(where) + if Equal == False: + print(values[where]) + print(values_fine[where]) + + +def verify_Restriction_Operator(Nel, plist, spl_kind): + from struphy.feec.utilities import create_equal_random_arrays, compare_arrays + + comm = MPI.COMM_WORLD + rank = comm.Get_rank() + derham = Derham(Nel, plist, spl_kind, comm=comm) + + # For B-splines + sp_key = "0" + spaces = derham.Vh_fem[sp_key].spaces + space = spaces[0] + N = space.nbasis + ncells = space.ncells + p = space.degree + T = space.knots + periodic = space.periodic + basis = space.basis + normalize = basis == "M" + + #Now we build the fine grid + derham_fine = Derham([Nel[0]*2,Nel[1],Nel[2]], plist, spl_kind, comm=comm) + + # For B-splines + spaces_fine = derham_fine.Vh_fem[sp_key].spaces + space_fine = spaces_fine[0] + N_fine = space_fine.nbasis + ncells_fine = space_fine.ncells + p_fine = space_fine.degree + T_fine = space_fine.knots + periodic_fine = space_fine.periodic + basis_fine = space_fine.basis + normalize_fine = basis_fine == "M" + + #We intialize the restriction operator + R = RestrictionOperator(derham_fine.Vh_fem[sp_key],derham.Vh_fem[sp_key]) + + varr, v = create_equal_random_arrays(derham_fine.Vh_fem[sp_key], seed=4568) + varr = varr[0].flatten() + + + out = R.dot(v) + + #To make it easier to read I will extract the data out of out disregarding all the padding it come with + out_array = np.zeros(N, dtype=float) + for i in range(N): + out_array[i] = out[R._out_starts[0]+i,0,0] + + #print(out._data) + + + ##### + #Now we build the Restriction matrix directly to verify our RestrictionOperator is working properly. + R_matrix = np.zeros((N,N_fine),dtype=float) + for i in range(N): + start = 2*i-p + for j in range(p+2): + R_matrix[i,(start+j)%N_fine] = 2.0**(-p)*comb(p+1,j) + + out2 = np.matmul(R_matrix, varr) + + + Equal =True + where = -1 + for i in range(len(out2)): + if(abs(out2[i]-out_array[i])>10.0**(-6)): + Equal = False + where = i + + + + print(f'{Equal = }') + print(f'{out_array = }') + print(f'{out2 = }') + if( Equal == False): + print(f'{where = }') + + +def verify_Extension_Operator(Nel, plist, spl_kind): + from struphy.feec.utilities import create_equal_random_arrays, compare_arrays + + comm = MPI.COMM_WORLD + rank = comm.Get_rank() + derham = Derham(Nel, plist, spl_kind, comm=comm) + + # For B-splines + sp_key = "0" + spaces = derham.Vh_fem[sp_key].spaces + space = spaces[0] + N = space.nbasis + ncells = space.ncells + p = space.degree + T = space.knots + periodic = space.periodic + basis = space.basis + normalize = basis == "M" + + #Now we build the fine grid + derham_fine = Derham([Nel[0]*2,Nel[1],Nel[2]], plist, spl_kind, comm=comm) + + # For B-splines + spaces_fine = derham_fine.Vh_fem[sp_key].spaces + space_fine = spaces_fine[0] + N_fine = space_fine.nbasis + ncells_fine = space_fine.ncells + p_fine = space_fine.degree + T_fine = space_fine.knots + periodic_fine = space_fine.periodic + basis_fine = space_fine.basis + normalize_fine = basis_fine == "M" + + #We intialize the restriction operator + E = ExtensionOperator(derham.Vh_fem[sp_key],derham_fine.Vh_fem[sp_key]) + + v = derham.Vh[sp_key].zeros() + varr, v = create_equal_random_arrays(derham.Vh_fem[sp_key], seed=4568) + varr = varr[0].flatten() + + + out = E.dot(v) + + #To make it easier to read I will extract the data out of out disregarding all the padding it come with + out_array = np.zeros(N_fine, dtype=float) + for i in range(N_fine): + out_array[i] = out[E._out_starts[0]+i,0,0] + + #print(out._data) + + + ##### + #Now we build the Restriction matrix directly to verify our RestrictionOperator is working properly. + R_matrix = np.zeros((N,N_fine),dtype=float) + for i in range(N): + start = 2*i-p + for j in range(p+2): + R_matrix[i,(start+j)%N_fine] = 2.0**(-p)*comb(p+1,j) + + E_matrix = R_matrix.T + + out2 = np.matmul(E_matrix, varr) + + + Equal =True + where = -1 + for i in range(len(out2)): + if(abs(out2[i]-out_array[i])>10.0**(-6)): + Equal = False + where = i + + + + print(f'{Equal = }') + print(f'{out_array = }') + print(f'{out2 = }') + if( Equal == False): + print(f'{where = }') + + +def testing_random_stuff(Nel,plist,spl_kind): + epsilon = 10.0**-6.0 + Nel = 512 + jarray = [] + for j in range(Nel): + jarray.append(float(j)) + + jarray = np.array(jarray) + + compute_directly = True + + #matrix = 'Shear-Alfven' + matrix = 'Hall' + if matrix == 'Shear-Alfven': + + def funl(j): + return 9.0*Nel - epsilon * (243.0/2.0) * (Nel**3.0) * (1.0-np.cos(2.0*np.pi *j / Nel)) /(2.0+np.cos(2.0*np.pi *j / Nel)) + + max_eigen_value_x = 4.0/Nel + + + + + elif( matrix == 'Hall'): + def funl(j): + return 8.0 /(3.0*Nel) - 18.0*epsilon*Nel + (4.0/(3.0*Nel)+18.0*epsilon*Nel)*np.cos(2.0*np.pi*j/Nel) + + max_eigen_value_x = 9.0*Nel + + + + eigen_values_y_z = funl(jarray) + max_eigen_value_y_z=abs(eigen_values_y_z[np.argmax(np.abs(eigen_values_y_z))]) + + + order_of_magnitude_disparity = np.log10(max_eigen_value_y_z/max_eigen_value_x) + print(f'{max_eigen_value_y_z = }') + print(f'{max_eigen_value_x = }') + print(f'{order_of_magnitude_disparity = }') + + plt.figure() + plt.scatter(jarray,eigen_values_y_z) + plt.show() + + + if compute_directly: + # get global communicator + comm = MPI.COMM_WORLD + + domain = Cuboid() + derham = [] + mass_ops = [] + A = [] + + derham.append(Derham([Nel,1,1], [1,1,1], [True,True,True], comm=comm, local_projectors=False)) + mass_ops.append(WeightedMassOperators(derham[0], domain)) + if(matrix == "Poisson"): + sp_key = '0' + sp_id = 'H1' + #Poisson + A.append(derham[0].grad.T @ mass_ops[0].M1 @ derham[0].grad) + #Hall-ish + #A.append(epsilon*derham[0].curl.T @ mass_ops[0].M2 @ derham[0].curl + mass_ops[0].M1) + elif(matrix == "Hall"): + sp_key = '1' + sp_id = 'Hcurl' + #Hall + A.append(mass_ops[0].M1 -epsilon*derham[0].curl.T @ mass_ops[0].M2 @ derham[0].curl) + #Shear-Alfven-ish + #A.append(mass_ops[0].M2 -epsilon* mass_ops[0].M2@ derham[0].curl @ mass_ops[0].M1 @ derham[0].curl.T @ mass_ops[0].M2) + elif(matrix == 'Shear-Alfven' or matrix == 'Shear-Alfven-v2'): + sp_key = '2' + sp_id = 'Hdiv' + pc_class = getattr(preconditioner,"MassMatrixPreconditioner") + pc = pc_class(mass_ops[0].M1) + M1_inv = inverse( + mass_ops[0].M1, + "pcg", + pc=pc, + maxiter=3000, + verbose=False, + ) + #Shear-Alfven + if (matrix == "Shear-Alfven"): + A.append(mass_ops[0].M2 -epsilon* mass_ops[0].M2@ derham[0].curl @ M1_inv @ derham[0].curl.T @ mass_ops[0].M2) + + Aarr = A[0].toarray() + Aarrx = Aarr[0:Nel,0:Nel] + Aarry = Aarr[Nel:2*Nel,Nel:2*Nel] + eigenvaluesx = np.linalg.eigvals(Aarrx) + max_x = abs(eigenvaluesx[np.argmax(np.abs(eigenvaluesx))]) + eigenvaluesy = np.linalg.eigvals(Aarry) + + plt.figure() + plt.plot(jarray,eigenvaluesy) + plt.show() + + max_y = abs(eigenvaluesy[np.argmax(np.abs(eigenvaluesy))]) + print(f'{max_x =}') + + + + + + + + #from struphy.feec.utilities import create_equal_random_arrays + # get global communicator + #comm = MPI.COMM_WORLD + #rank = comm.Get_rank() + #world_size = comm.Get_size() + + #domain = Cuboid() + #a1 = 0.002 + #domain = HollowCylinder(a1= a1) + #sp_key = '0' + + #derham =Derham([Nel[0],Nel[1],Nel[2]], plist, spl_kind, comm=comm, local_projectors=False) + #mass_ops=WeightedMassOperators(derham, domain) + + #u_stararr, u_star = create_equal_random_arrays(derham.Vh_fem[sp_key], seed=45) + #u_stararr = remove_padding(derham.Vh_fem[sp_key], u_star) + + #G = derham.grad + #Garr =G.toarray() + + #GT = derham.grad.T + #GTarr = GT.toarray() + + #C = derham.curl + #Carr = C.toarray() + + #M1 = mass_ops.M1 + #M1arr = M1.toarray() + + #Po = GT @ M1 @ G + #Poarr = Po.toarray() + + + #eigenvalues = np.linalg.eigvals(Poarr) + #max_eigen_value = abs(eigenvalues[np.argmax(np.abs(eigenvalues))]) + #min_eigen_value = abs(eigenvalues[np.argmin(np.abs(eigenvalues))]) + #b = G.dot(u_star) + #barr = np.matmul(Garr, u_stararr) + + #b = remove_padding(derham.Vh_fem['1'], b) + + #same = True + #for i in range(np.size(b)): + #if(barr[i]!= b[i]): + #same = False + #break + + #print(same) + #print(f'{M1arr[32,32:64] =}') + #print(f'{max_eigen_value =}') + #print(f'{min_eigen_value =}') + #for i in range(0,3): + #print(f'{Garr[i,:] =}') + + + + + +if __name__ == '__main__': + from struphy.feec.utilities import create_equal_random_arrays + comm = MPI.COMM_WORLD + derham = Derham([10,1,1], [1,1,1], [True,True,True], comm=comm, local_projectors=False) + u_stararr, u_star = create_equal_random_arrays(derham.Vh_fem['0'], seed=45) + print("################") + print(vars(u_star._space)) + print("################") + #for i in range(7,8): + #Nel = [int(2**i),int(2**i), 1] + #p = [1, 1, 1] + #spl_kind = [True, True, True] + #Gather_data_V_cycle_parameter_study(Nel, p, spl_kind, i) + #Visualized_high_frequency_dampening(Nel, p, spl_kind) + #Visualized_all_frequencies_dampening(Nel, p, spl_kind,7) + #Visualized_all_frequencies_dampening_2D(Nel, p, spl_kind,100) + + #testing_random_stuff([32,1,1],[1,1,1],[True,True,True]) + #128,64,32,16, 8 + # h, 2h,4h,8h,16h + + #p=1, Nel= 8192, level = 12. Coarsest one is 4x4 matrix + #p=2, Nel= 8192, level = 11. Coarsest one is 8x8 matrix + #p=4, Nel= 8192, level = 10. Coarsest one is 16x16 matrix + + #multigrid(Nel, p, spl_kind,6) + #Gather_data_V_cycle_parameter_study(Nel, p, spl_kind, 11) + #Gather_data_V_cycle_scalability([[int(2**i),1,1] for i in range(4,10)], p, spl_kind) + #make_plot_scalability() + #verify_formula(Nel, p, spl_kind) + #verify_Restriction_Operator(Nel, p, spl_kind) + #verify_Extension_Operator(Nel, p, spl_kind) + #Error_analysis(Nel, p, spl_kind, 1) + #trying_New_Restriction(Nel, p, spl_kind, 4) + #Compute_rate_of_smoothing(Nel, p, spl_kind) + #Visualized_high_frequency_dampening(Nel, p, spl_kind) + \ No newline at end of file diff --git a/src/struphy/linear_algebra/multigrid_solver.py b/src/struphy/linear_algebra/multigrid_solver.py new file mode 100644 index 000000000..1f2d886da --- /dev/null +++ b/src/struphy/linear_algebra/multigrid_solver.py @@ -0,0 +1,1070 @@ +import time + +import matplotlib.pyplot as plt +import numpy as np +from mpi4py import MPI +import re +from struphy.feec.psydac_derham import Derham + +from psydac.linalg.solvers import inverse +from psydac.linalg.basic import Vector +from psydac.fem.projectors import knot_insertion_projection_operator +from struphy.feec.linear_operators import LinOpWithTransp +from psydac.fem.basic import FemSpace +from psydac.fem.tensor import TensorFemSpace +from struphy.feec.mass import WeightedMassOperators +from struphy.geometry.domains import Cuboid +from math import comb +from struphy.feec.utilities import create_equal_random_arrays + +def build_operator(expr: str, namespace: dict): + """ + Build a composite LinOpWithTransp (or scalar * LinOpWithTransp expression) + from a string expression. + + Parameters + ---------- + expr : str + String expression defining the operator composition. + Example: + "epsilon * curl.T @ M2 @ curl + mu * M1" + + The expression can contain: + - scalars (e.g., epsilon, mu, ...) + - LinOpWithTransp objects (already supporting +, -, *, @) + - parentheses to control precedence + + namespace : dict + Dictionary mapping symbol names in `expr` to actual Python objects + (scalars, operators, lists, etc.). + Example: + { + "epsilon": 1.3, + "mu": 0.7, + "curl.T": derham.curl.T, + "curl": derham.curl, + "M1": WeightedMassOperators(derham, domain).M1, + "M2": WeightedMassOperators(derham, domain).M2, + } + + Returns + ------- + LinOpWithTransp + A composite operator built according to the expression. The returned + object supports `.dot(v)` and any other functionality of LinOpWithTransp. + + Raises + ------ + ValueError + If some symbols in the expression are not found in `namespace`. + """ + + # Find all potential variable-like tokens (ignore numbers) + tokens = set(re.findall(r"[A-Za-z_][A-Za-z0-9_]*(?:\.[A-Za-z_][A-Za-z0-9_]*)*", expr)) + + missing = tokens - set(namespace.keys()) + + if missing: + raise ValueError( + f"The following symbols are missing from the namespace: {', '.join(sorted(missing))}" + ) + + # Let Python evaluate the expression with operator overloads + return eval(expr, {}, namespace) + +def remove_padding(fem_space, v): + + #fem_space = derham.Vh_fem[sp_key] + symbolic_name = fem_space.symbolic_space + + if(symbolic_name == 'H1' or symbolic_name == "L2"): + + spaces = [fem_space.spaces] + N = [spaces[0][i].nbasis for i in range(3)] + + starts = np.array(fem_space.coeff_space.starts) + + #To make it easier to read I will extract the data out of out disregarding all the padding it come with + v_array = np.zeros(N[0]*N[1]*N[2], dtype=float) + cont = 0 + for i0 in range(N[0]): + for i1 in range(N[1]): + for i2 in range(N[2]): + v_array[cont] = v[starts[0]+i0,starts[1]+i1,starts[2]+i2] + cont += 1 + + else: + spaces = [comp.spaces for comp in fem_space.spaces] + N = [[spaces[h][i].nbasis for i in range(3)] for h in range(3)] + + starts = np.array([vi.starts for vi in fem_space.coeff_space.spaces]) + + #To make it easier to read I will extract the data out of out disregarding all the padding it come with + v_array = np.zeros(N[0][0]*N[0][1]*N[0][2]+N[1][0]*N[1][1]*N[1][2]+N[2][0]*N[2][1]*N[2][2], dtype=float) + cont = 0 + for h in range(3): + for i0 in range(N[h][0]): + for i1 in range(N[h][1]): + for i2 in range(N[h][2]): + v_array[cont] = v[h][starts[h][0]+i0,starts[h][1]+i1,starts[h][2]+i2] + cont += 1 + + return v_array + + +def direct_solver(A_inv,b, fem_space): + # A_inv is already the inverse matrix of A + #fem_space = derham.Vh_fem[sp_key] + symbolic_name = fem_space.symbolic_space + + if(symbolic_name == 'H1' or symbolic_name == "L2"): + spaces = [fem_space.spaces] + N = [spaces[0][i].nbasis for i in range(3)] + starts = np.array(fem_space.coeff_space.starts) + + b_vector = remove_padding(fem_space, b) + x_vector = np.dot(A_inv, b_vector) + x = fem_space.coeff_space.zeros() + + cont= 0 + for i0 in range(N[0]): + for i1 in range(N[1]): + for i2 in range(N[2]): + x[starts[0]+i0,starts[1]+i1,starts[2]+i2] = x_vector[cont] + cont += 1 + + else: + spaces = [comp.spaces for comp in fem_space.spaces] + N = [[spaces[h][i].nbasis for i in range(3)] for h in range(3)] + starts = np.array([vi.starts for vi in fem_space.coeff_space.spaces]) + + b_vector = remove_padding(fem_space, b) + x_vector = np.dot(A_inv, b_vector) + x = fem_space.coeff_space.zeros() + + cont = 0 + for h in range(3): + for i0 in range(N[h][0]): + for i1 in range(N[h][1]): + for i2 in range(N[h][2]): + x[h][starts[h][0]+i0,starts[h][1]+i1,starts[h][2]+i2] = x_vector[cont] + cont += 1 + return x + + +def get_b_spline_degree(V): + """ + Determines the degree of the B-splines. + + Parameters + ---------- + V : psydac.fem.basic.FemSpace + Finite element spline space (domain, input space). + + Returns + ------- + p : numpy array + numpy array of 3 ints containing the B-spline degrees for each spatial direction. + """ + + assert isinstance(V, FemSpace) + p = np.zeros(3, dtype=int) + + if hasattr(V, "symbolic_space"): + V_name = V.symbolic_space + + if(V_name == "H1"): + for i, space in enumerate(V.spaces): + p[i] = space.degree + elif(V_name == "L2"): + for i, space in enumerate(V.spaces): + p[i] = space.degree+1 + elif(V_name == "Hcurl"): + V1ds = [comp.spaces for comp in V.spaces] + for i in range(3): + p[i] = V1ds[i][i].degree+1 + elif(V_name == "Hdiv"): + V1ds = [comp.spaces for comp in V.spaces] + for i in range(3): + p[i] = V1ds[i][i].degree + elif(V_name == "H1H1H1"): + V1ds = [comp.spaces for comp in V.spaces] + for i in range(3): + p[i] = V1ds[i][i].degree + else: + raise Exception("Invalid symbolic name.") + + return p + + + +class MultiGridSolver: + + def __init__(self, derham, sp_key, N_levels, domain, expr: str, namespace: dict, max_iter_list, N_cycles, method='cg'): + """ + + Initialize the multigrid solver with the given parameters. + Parameters + ---------- + derham : Derham + The Derham object representing the finest grid. + sp_key : str + Symbolic space key, e.g., "0", "1", "2", "3". + N_levels : int + Number of multigrid levels. + domain : struphy.geometry.base.Domain + The Geometric domain of the simulation. + + expr : str + String expression defining the operator composition. + Example: + "epsilon * curl.T @ M2 @ curl + mu * M1" + + The expression can contain: + - scalars (e.g., epsilon, mu, ...) + - LinOpWithTransp objects (already supporting +, -, *, @) + - parentheses to control precedence + + namespace : dict + A dictionary that serves as a **string template** or "recipe book" for + building the operators and scalars used in the `expr` string. + + It maps symbol names to **strings of Python code** that describe how to + create the corresponding object. This allows the solver to dynamically + generate the correct operators for each grid level. + + The string recipes can use variables like `derham`, `WeightedMassOperators` and `domain`, which + the solver makes available during the object creation at each level. + + Example: + { + "epsilon": "1.3", + "mu": "0.7", + "curl.T": "derham.curl.T", + "curl": "derham.curl", + "M1": "WeightedMassOperators(derham, domain).M1", + "M2": "WeightedMassOperators(derham, domain).M2", + } + max_iter_list : list + List of maximum iterations for each multigrid level. The first element corresponds to the finest level. + N_cycles : int + Number of multigrid cycles. + method : str + Solution method to use (e.g., 'cg' for conjugate gradient). + + """ + self.sp_key = sp_key + self.Nel = derham.Nel + self.plist = derham.p + self.spl_kind = derham._spl_kind + self.N_levels = N_levels + self.domain = domain + self.expr = expr + self.namespace = namespace + self.derham = derham + self.method = method + self.max_iter_list = max_iter_list + self.N_cycles = N_cycles + + # get global communicator + self.comm = MPI.COMM_WORLD + self.derhamlist = [] + self.A = [] + + for level in range(self.N_levels): + if level == 0: + self.derhamlist.append(self.derham) + else: + self.derhamlist.append(Derham([self.Nel[0]//(2**level),self.Nel[1]//(2**level),self.Nel[2]], self.plist, self.spl_kind, comm=self.comm, local_projectors=False)) + #It might look like we are not using derham but it is being used by updated_namespace + #for rebounding the expressions in self.expr to the correct derham level. + derhamaux = self.derhamlist[level] + + # 1. Create a local context for the eval. It maps string names + # to the actual objects for the CURRENT level. + eval_context = { + "derham": derhamaux, + "domain": self.domain, + "WeightedMassOperators": WeightedMassOperators + } + + # 2. Build the namespace by evaluating each string from the template + # within the context of the current level. + updated_namespace = { + key: eval(recipe_string, {}, eval_context) + for key, recipe_string in self.namespace.items() + } + + self.A.append(build_operator(self.expr, updated_namespace)) + + #We get the inverse of the coarsest system matrix to solve directly the problem in the smaller space + self.A_inv = np.linalg.inv(self.A[-1].toarray()) + + self.R = [] + self.E = [] + + for level in range(self.N_levels-1): + self.R.append(self.RestrictionOperator(self.derhamlist[level].Vh_fem[self.sp_key],self.derhamlist[level+1].Vh_fem[self.sp_key])) + self.E.append(self.R[level].transpose()) + + + class RestrictionOperator(LinOpWithTransp): + """ + Linear operator which operates between vector spaces of the same kind but different resolutions. + + We assume that the vectors in the domain belong to a De-rham space with n elements, while the codomain + belong to the coarser De-rham space with n/2 elements. + + At the moment we also assume that we are halving only the first spatial direction. + + Parameters + ---------- + V : psydac.fem.basic.FemSpace + Finite element spline space (domain, input space). + + W : psydac.fem.basic.FemSpace + Finite element spline space (codomain, output space). + + """ + def __init__(self, V, W): + + # Check domain and codomain + assert isinstance(V, FemSpace) + assert isinstance(W, FemSpace) + + self._V = V + self._W = W + + #print(vars(V)) + + # Store info in object + self._domain = V.coeff_space + self._codomain = W.coeff_space + self._dtype = V.coeff_space.dtype + + #Can be "H1", "L2", "Hcurl", "Hdiv", "H1H1H1" + self._V_name = V.symbolic_space + self._W_name = W.symbolic_space + assert(self._V_name == self._W_name) + + #This list will tell us in which spatial direction we are halving the problem. + self._halving_directions = [False,False,False] + + # input space: 3d StencilVectorSpaces and 1d SplineSpaces of each component + if isinstance(V, TensorFemSpace): + self._V1ds = [V.spaces] + self._VNbasis = np.array([self._V1ds[0][0].nbasis, self._V1ds[0][1].nbasis, self._V1ds[0][2].nbasis]) + + # We get the start and endpoint for each sublist in input + self._in_starts = np.array(V.coeff_space.starts) + self._in_ends = np.array(V.coeff_space.ends) + else: + self._V1ds = [comp.spaces for comp in V.spaces] + self._VNbasis = np.array( + [ + [self._V1ds[0][0].nbasis, self._V1ds[0][1].nbasis, self._V1ds[0][2].nbasis], + [ + self._V1ds[1][0].nbasis, + self._V1ds[1][1].nbasis, + self._V1ds[1][2].nbasis, + ], + [self._V1ds[2][0].nbasis, self._V1ds[2][1].nbasis, self._V1ds[2][2].nbasis], + ] + ) + + # We get the start and endpoint for each sublist in input + self._in_starts = np.array([vi.starts for vi in V.coeff_space.spaces]) + self._in_ends = np.array([vi.ends for vi in V.coeff_space.spaces]) + + # output space: 3d StencilVectorSpaces and 1d SplineSpaces of each component + if isinstance(W, TensorFemSpace): + self._W1ds = [W.spaces] + self._WNbasis = np.array([self._W1ds[0][0].nbasis, self._W1ds[0][1].nbasis, self._W1ds[0][2].nbasis]) + + for i in range(3): + if(self._WNbasis[i] < self._VNbasis[i]): + #If this breaks for clamped splines it means .nbasis gives you the number of basis functions, not the number of elements + assert self._VNbasis[i] == self._WNbasis[i]*2 + self._halving_directions[i] = True + + # We get the start and endpoint for each sublist in out + self._out_starts = np.array(W.coeff_space.starts) + self._out_ends = np.array(W.coeff_space.ends) + + else: + self._W1ds = [comp.spaces for comp in W.spaces] + self._WNbasis = np.array( + [ + [self._W1ds[0][0].nbasis, self._W1ds[0][1].nbasis, self._W1ds[0][2].nbasis], + [ + self._W1ds[1][0].nbasis, + self._W1ds[1][1].nbasis, + self._W1ds[1][2].nbasis, + ], + [self._W1ds[2][0].nbasis, self._W1ds[2][1].nbasis, self._W1ds[2][2].nbasis], + ] + ) + for i in range(3): + if(self._WNbasis[1][i] < self._VNbasis[1][i]): + #If this breaks for clamped splines it means .nbasis gives you the number of basis functions, not the number of elements + assert self._VNbasis[0][i] == self._WNbasis[0][i]*2 and self._VNbasis[1][i] == self._WNbasis[1][i]*2 and self._VNbasis[2][i] == self._WNbasis[2][i]*2 + self._halving_directions[i] = True + + # We get the start and endpoint for each sublist in out + self._out_starts = np.array([vi.starts for vi in W.coeff_space.spaces]) + self._out_ends = np.array([vi.ends for vi in W.coeff_space.spaces]) + + + + # Degree of the B-spline space, not to be confused with the degrees given by fem_space.spaces.degree since depending on the situation + # it will give the D-spline degree instead + self._p = get_b_spline_degree(V) + + #We also get the D-spline degree + self._pD = self._p - 1 + + #Now we compute the weights that define this linear operator + + #We begin by defining a list that will contain the 3 numpy arrays, each one with the weights for one spatial direction. + #In the case there are direction over which we do not halve the resolution we shall have an array with only one 1.0 + self._all_weights = [] + self._all_weightsD = [] + for i in range(3): + if self._halving_directions[i]: + #Here we store the weights needed for B-splines + weights = np.zeros(self._p[i]+2, dtype=float) + #Here we store the weights needed for D-splines + weightsD = np.zeros(self._pD[i]+2, dtype=float) + for j in range(self._p[i]+2): + weights[j] = 2.0**(-self._p[i])*comb(self._p[i]+1,j) + for j in range(self._pD[i]+2): + weightsD[j] = 2.0**-(self._pD[i]+1)*comb(self._pD[i]+1,j) + self._all_weights.append(weights) + self._all_weightsD.append(weightsD) + else: + self._all_weights.append(np.array([1.0],dtype=float)) + self._all_weightsD.append(np.array([1.0],dtype=float)) + + + + #-------------------------------------- + # Abstract interface + #-------------------------------------- + @property + def domain(self): + return self._domain + + @property + def codomain(self): + return self._codomain + + @property + def dtype(self): + return self._dtype + + def _dot_helper(self, v, out, p, weights, h=None): + """Helper function to perform dot product computation.""" + #First we get the number of weights in each direction + weights_len = [] + for i in range(3): + if self._halving_directions[i]: + weights_len.append(p[i] + 2) + else: + weights_len.append(1) + if h is None: # Scalar case (H1, L2) + for i0 in range(self._out_starts[0], self._out_ends[0] + 1): + for i1 in range(self._out_starts[1], self._out_ends[1] + 1): + for i2 in range(self._out_starts[2], self._out_ends[2] + 1): + for j0 in range(weights_len[0]): + if self._halving_directions[0]: + pos0 = (2 * i0 - p[0] + j0) % self._VNbasis[0] + else: + pos0 = i0 + for j1 in range(weights_len[1]): + if self._halving_directions[1]: + pos1 = (2 * i1 - p[1] + j1) % self._VNbasis[1] + else: + pos1 = i1 + for j2 in range(weights_len[2]): + if self._halving_directions[2]: + pos2 = (2 * i2 - p[2] + j2) % self._VNbasis[2] + else: + pos2 = i2 + out[i0, i1, i2] += weights[0][j0]* weights[1][j1] *weights[2][j2] * v[pos0, pos1, pos2] + else: # Vector case (Hcurl, Hdiv, H1H1H1) + for i0 in range(self._out_starts[h][0], self._out_ends[h][0] + 1): + for i1 in range(self._out_starts[h][1], self._out_ends[h][1] + 1): + for i2 in range(self._out_starts[h][2], self._out_ends[h][2] + 1): + for j0 in range(weights_len[0]): + if self._halving_directions[0]: + pos0 = (2 * i0 - p[0] + j0) % self._VNbasis[h][0] + else: + pos0 = i0 + for j1 in range(weights_len[1]): + if self._halving_directions[1]: + pos1 = (2 * i1 - p[1] + j1) % self._VNbasis[h][1] + else: + pos1 = i1 + for j2 in range(weights_len[2]): + if self._halving_directions[2]: + pos2 = (2 * i2 - p[2] + j2) % self._VNbasis[h][2] + else: + pos2 = i2 + out[h][i0, i1, i2] += weights[0][j0]* weights[1][j1] *weights[2][j2] * v[h][pos0, pos1, pos2] + + return out + + def dot_H1(self, v, out): + return self._dot_helper(v, out, self._p, self._all_weights) + + def dot_L2(self, v, out): + return self._dot_helper(v, out, self._pD, self._all_weightsD) + + def dot_Hcurl(self, v, out): + out = self._dot_helper(v, out, [self._pD[0],self._p[1],self._p[2]], [self._all_weightsD[0], self._all_weights[1], self._all_weights[2]], h = 0) + self._dot_helper(v, out, [self._p[0],self._pD[1],self._p[2]], [self._all_weights[0], self._all_weightsD[1], self._all_weights[2]], h = 1) + return self._dot_helper(v, out, [self._p[0],self._p[1],self._pD[2]], [self._all_weights[0], self._all_weights[1], self._all_weightsD[2]], h = 2) + + def dot_Hdiv(self, v, out): + out = self._dot_helper(v, out, [self._p[0],self._pD[1],self._pD[2]], [self._all_weights[0], self._all_weightsD[1], self._all_weightsD[2]], h = 0) + self._dot_helper(v, out, [self._pD[0],self._p[1],self._pD[2]], [self._all_weightsD[0], self._all_weights[1], self._all_weightsD[2]], h = 1) + return self._dot_helper(v, out, [self._pD[0],self._pD[1],self._p[2]], [self._all_weightsD[0], self._all_weightsD[1], self._all_weights[2]], h = 2) + + def dot_H1H1H1(self, v, out): + out = self._dot_helper(v, out, self._p, self._all_weights, h=0) + self._dot_helper(v, out, self._p, self._all_weights, h=1) + return self._dot_helper(v, out, self._p, self._all_weights, h=2) + + def dot(self, v, out=None): + + assert isinstance(v, Vector) and v.space == self.domain + + if out is None: + out = self.codomain.zeros() + else: + assert isinstance(out, Vector) and out.space == self.codomain + + if self._V_name == 'H1' or self._V_name == 'L2': + for i0 in range(self._out_starts[0], self._out_ends[0]+1): + for i1 in range(self._out_starts[1], self._out_ends[1]+1): + for i2 in range(self._out_starts[2], self._out_ends[2]+1): + out[i0,i1,i2] = 0.0 + else: + for h in range(3): + for i0 in range(self._out_starts[h][0], self._out_ends[h][0]+1): + for i1 in range(self._out_starts[h][1], self._out_ends[h][1]+1): + for i2 in range(self._out_starts[h][2], self._out_ends[h][2]+1): + out[h][i0,i1,i2] = 0.0 + + dot_methods = { + "H1": self.dot_H1, + "L2": self.dot_L2, + "Hcurl": self.dot_Hcurl, + "Hdiv": self.dot_Hdiv, + "H1H1H1": self.dot_H1H1H1, + } + + return dot_methods.get(self._V_name)(v, out) + + def transpose(self, *, out = None): + if out is None: + out = MultiGridSolver.ExtensionOperator(self._W, self._V) + else: + assert isinstance(out, MultiGridSolver.ExtensionOperator) + assert out.domain is self.codomain + assert out.codomain is self.domain + + return out + + + class ExtensionOperator(LinOpWithTransp): + """ + Linear operator which operates between vector spaces of the same kind but different resolutions. + + We assume that the vectors in the domain belong to a De-rham space with n/2 elements, while the codomain + belong to the coarser De-rham space with n elements. + + At the moment we also assume that we are halving only the first spatial direction. + + Parameters + ---------- + V : psydac.fem.basic.FemSpace + Finite element spline space (domain, input space). + + W : psydac.fem.basic.FemSpace + Finite element spline space (codomain, output space). + + """ + def __init__(self, V, W): + + # Check domain and codomain + assert isinstance(V, FemSpace) + assert isinstance(W, FemSpace) + + self._V = V + self._W = W + + # Store info in object + self._domain = V.coeff_space + self._codomain = W.coeff_space + self._dtype = V.coeff_space.dtype + + #Can be "H1", "L2", "Hcurl", "Hdiv", "H1H1H1" + self._V_name = V.symbolic_space + self._W_name = W.symbolic_space + assert(self._V_name == self._W_name) + + #This list will tell us in which spatial direction we are halving the problem. + self._halving_directions = [False,False,False] + + # input space: 3d StencilVectorSpaces and 1d SplineSpaces of each component + if isinstance(V, TensorFemSpace): + self._V1ds = [V.spaces] + self._VNbasis = np.array([self._V1ds[0][0].nbasis, self._V1ds[0][1].nbasis, self._V1ds[0][2].nbasis]) + + # We get the start and endpoint for each sublist in input + self._in_starts = np.array(V.coeff_space.starts) + self._in_ends = np.array(V.coeff_space.ends) + else: + self._V1ds = [comp.spaces for comp in V.spaces] + self._VNbasis = np.array( + [ + [self._V1ds[0][0].nbasis, self._V1ds[0][1].nbasis, self._V1ds[0][2].nbasis], + [ + self._V1ds[1][0].nbasis, + self._V1ds[1][1].nbasis, + self._V1ds[1][2].nbasis, + ], + [self._V1ds[2][0].nbasis, self._V1ds[2][1].nbasis, self._V1ds[2][2].nbasis], + ] + ) + + # We get the start and endpoint for each sublist in input + self._in_starts = np.array([vi.starts for vi in V.coeff_space.spaces]) + self._in_ends = np.array([vi.ends for vi in V.coeff_space.spaces]) + + # output space: 3d StencilVectorSpaces and 1d SplineSpaces of each component + if isinstance(W, TensorFemSpace): + self._W1ds = [W.spaces] + self._WNbasis = np.array([self._W1ds[0][0].nbasis, self._W1ds[0][1].nbasis, self._W1ds[0][2].nbasis]) + + for i in range(3): + if(self._VNbasis[i] < self._WNbasis[i]): + #If this breaks for clamped splines it means .nbasis gives you the number of basis functions, not the number of elements + assert self._WNbasis[i] == self._VNbasis[i]*2 + self._halving_directions[i] = True + + # We get the start and endpoint for each sublist in out + self._out_starts = np.array(W.coeff_space.starts) + self._out_ends = np.array(W.coeff_space.ends) + + else: + self._W1ds = [comp.spaces for comp in W.spaces] + self._WNbasis = np.array( + [ + [self._W1ds[0][0].nbasis, self._W1ds[0][1].nbasis, self._W1ds[0][2].nbasis], + [ + self._W1ds[1][0].nbasis, + self._W1ds[1][1].nbasis, + self._W1ds[1][2].nbasis, + ], + [self._W1ds[2][0].nbasis, self._W1ds[2][1].nbasis, self._W1ds[2][2].nbasis], + ] + ) + + for i in range(3): + if(self._VNbasis[1][i] < self._WNbasis[1][i]): + assert self._WNbasis[0][i] == self._VNbasis[0][i]*2 and self._WNbasis[1][i] == self._VNbasis[1][i]*2 and self._WNbasis[2][i] == self._VNbasis[2][i]*2 + self._halving_directions[i] = True + + # We get the start and endpoint for each sublist in out + self._out_starts = np.array([vi.starts for vi in W.coeff_space.spaces]) + self._out_ends = np.array([vi.ends for vi in W.coeff_space.spaces]) + + + + # Degree of the B-spline space, not to be confused with the degrees given by fem_space.spaces.degree since depending on the situation + # it will give the D-spline degree instead + self._p = get_b_spline_degree(V) + #We also get the D-splines degree + self._pD = self._p-1 + + #Now we compute the weights that define this linear operator + + #We begin by defining a list that will contain the 3 numpy arrays, each one with the weights for one spatial direction. + #In the case there are a direction over which we do not halve the resolution we shall have an array with only one 1.0 + self._all_weights_even = [] + self._all_weights_evenD = [] + self._all_weights_odd = [] + self._all_weights_oddD = [] + + #Each list has 3 integers, each one denoting the number of weights in the corresponding weights array. + self._all_size_even = [] + self._all_size_evenD = [] + self._all_size_odd = [] + self._all_size_oddD = [] + + for i in range(3): + if self._halving_directions[i]: + #First for B-splines + if(self._p[i]%2 == 0): + size_even = self._p[i]//2 +1 + size_odd = self._p[i]//2 +1 + weights_even = np.zeros(size_even, dtype=float) + weights_odd = np.zeros(size_odd, dtype=float) + for j in range(size_even): + weights_even[j] = 2.0**(-self._p[i])*comb(self._p[i]+1,2*j) + weights_odd[j] = 2.0**(-self._p[i])*comb(self._p[i]+1,2*j+1) + else: + size_even = (self._p[i]+1)//2 +1 + size_odd = (self._p[i]-1)//2 +1 + weights_even = np.zeros(size_even, dtype=float) + weights_odd = np.zeros(size_odd, dtype=float) + for j in range(size_even): + weights_even[j] = 2.0**(-self._p[i])*comb(self._p[i]+1,2*j) + for j in range(size_odd): + weights_odd[j] = 2.0**(-self._p[i])*comb(self._p[i]+1,2*j+1) + + #Second for D-splines + if(self._pD[i]%2 == 0): + size_evenD = self._pD[i]//2 +1 + size_oddD = self._pD[i]//2 +1 + weights_evenD = np.zeros(size_evenD, dtype=float) + weights_oddD = np.zeros(size_oddD, dtype=float) + for j in range(size_evenD): + weights_evenD[j] = 2.0**-(self._pD[i]+1)*comb(self._pD[i]+1,2*j) + weights_oddD[j] = 2.0**-(self._pD[i]+1)*comb(self._pD[i]+1,2*j+1) + else: + size_evenD = (self._pD[i]+1)//2 +1 + size_oddD = (self._pD[i]-1)//2 +1 + weights_evenD = np.zeros(size_evenD, dtype=float) + weights_oddD = np.zeros(size_oddD, dtype=float) + for j in range(size_evenD): + weights_evenD[j] = 2.0**-(self._pD[i]+1)*comb(self._pD[i]+1,2*j) + for j in range(size_oddD): + weights_oddD[j] = 2.0**-(self._pD[i]+1)*comb(self._pD[i]+1,2*j+1) + + self._all_weights_even.append(weights_even) + self._all_weights_evenD.append(weights_evenD) + self._all_weights_odd.append(weights_odd) + self._all_weights_oddD.append(weights_oddD) + self._all_size_even.append(size_even) + self._all_size_evenD.append(size_evenD) + self._all_size_odd.append(size_odd) + self._all_size_oddD.append(size_oddD) + + else: + self._all_weights_even.append(np.array([1.0],dtype=float)) + self._all_weights_evenD.append(np.array([1.0],dtype=float)) + self._all_weights_odd.append(np.array([1.0],dtype=float)) + self._all_weights_oddD.append(np.array([1.0],dtype=float)) + self._all_size_even.append(1) + self._all_size_evenD.append(1) + self._all_size_odd.append(1) + self._all_size_oddD.append(1) + + + #-------------------------------------- + # Abstract interface + #-------------------------------------- + @property + def domain(self): + return self._domain + + @property + def codomain(self): + return self._codomain + + @property + def dtype(self): + return self._dtype + + def tosparse(self): + pass + + def toarray(self): + pass + + def _dot_helper(self, v, out, p, weights_even, weights_odd, size_even, size_odd, h=None): + """Helper function to perform dot product computation.""" + parity_match = [] + for i in range(3): + parity_match.append(p[i] % 2) + + if h is None: # Scalar case (H1, L2) + for j0 in range(self._out_starts[0], self._out_ends[0] + 1): + parity_j0 = j0 % 2 + weights0, size0, offset0 = ((weights_even[0], size_even[0], 0) if parity_j0 == parity_match[0] else (weights_odd[0], size_odd[0], 1)) + for j1 in range(self._out_starts[1], self._out_ends[1] + 1): + parity_j1 = j1 % 2 + weights1, size1, offset1 = ((weights_even[1], size_even[1], 0) if parity_j1 == parity_match[1] else (weights_odd[1], size_odd[1], 1)) + for j2 in range(self._out_starts[2], self._out_ends[2] + 1): + parity_j2 = j2 % 2 + weights2, size2, offset2 = ((weights_even[2], size_even[2], 0) if parity_j2 == parity_match[2] else (weights_odd[2], size_odd[2], 1)) + for i0 in range(size0): + if self._halving_directions[0]: + pos0 = ((j0 + p[0] - 2 * i0 - offset0) // 2) % self._VNbasis[0] + else: + pos0 = j0 + for i1 in range(size1): + if self._halving_directions[1]: + pos1 = ((j1 + p[1] - 2 * i1 - offset1) // 2) % self._VNbasis[1] + else: + pos1 = j1 + for i2 in range(size2): + if self._halving_directions[2]: + pos2 = ((j2 + p[2] - 2 * i2 - offset2) // 2) % self._VNbasis[2] + else: + pos2 = j2 + out[j0, j1, j2] += weights0[i0] * weights1[i1] * weights2[i2] * v[pos0, pos1, pos2] + + else: # Vector case (Hcurl, Hdiv, H1H1H1) + for j0 in range(self._out_starts[h][0], self._out_ends[h][0] + 1): + parity_j0 = j0 % 2 + weights0, size0, offset0 = ((weights_even[0], size_even[0], 0) if parity_j0 == parity_match[0] else (weights_odd[0], size_odd[0], 1)) + for j1 in range(self._out_starts[h][1], self._out_ends[h][1] + 1): + parity_j1 = j1 % 2 + weights1, size1, offset1 = ((weights_even[1], size_even[1], 0) if parity_j1 == parity_match[1] else (weights_odd[1], size_odd[1], 1)) + for j2 in range(self._out_starts[h][2], self._out_ends[h][2] + 1): + parity_j2 = j2 % 2 + weights2, size2, offset2 = ((weights_even[2], size_even[2], 0) if parity_j2 == parity_match[2] else (weights_odd[2], size_odd[2], 1)) + for i0 in range(size0): + if self._halving_directions[0]: + pos0 = ((j0 + p[0] - 2 * i0 - offset0) // 2) % self._VNbasis[h][0] + else: + pos0 = j0 + for i1 in range(size1): + if self._halving_directions[1]: + pos1 = ((j1 + p[1] - 2 * i1 - offset1) // 2) % self._VNbasis[h][1] + else: + pos1 = j1 + for i2 in range(size2): + if self._halving_directions[2]: + pos2 = ((j2 + p[2] - 2 * i2 - offset2) // 2) % self._VNbasis[h][2] + else: + pos2 = j2 + out[h][j0, j1, j2] += weights0[i0] * weights1[i1] * weights2[i2] * v[h][pos0, pos1, pos2] + + return out + + def dot_H1(self, v, out): + return self._dot_helper(v, out, self._p, self._all_weights_even, self._all_weights_odd, self._all_size_even, self._all_size_odd) + + def dot_L2(self, v, out): + return self._dot_helper(v, out, self._pD, self._all_weights_evenD, self._all_weights_oddD, self._all_size_evenD, self._all_size_oddD) + + def dot_Hcurl(self, v, out): + out = self._dot_helper(v, out, [self._pD[0], self._p[1],self._p[2]], [self._all_weights_evenD[0],self._all_weights_even[1],self._all_weights_even[2]], [self._all_weights_oddD[0],self._all_weights_odd[1],self._all_weights_odd[2]], [self._all_size_evenD[0],self._all_size_even[1],self._all_size_even[2]], [self._all_size_oddD[0],self._all_size_odd[1],self._all_size_odd[2]], h=0) + out = self._dot_helper(v, out, [self._p[0], self._pD[1],self._p[2]], [self._all_weights_even[0],self._all_weights_evenD[1],self._all_weights_even[2]], [self._all_weights_odd[0],self._all_weights_oddD[1],self._all_weights_odd[2]], [self._all_size_even[0],self._all_size_evenD[1],self._all_size_even[2]], [self._all_size_odd[0],self._all_size_oddD[1],self._all_size_odd[2]], h=1) + return self._dot_helper(v, out, [self._p[0], self._p[1],self._pD[2]], [self._all_weights_even[0],self._all_weights_even[1],self._all_weights_evenD[2]], [self._all_weights_odd[0],self._all_weights_odd[1],self._all_weights_oddD[2]], [self._all_size_even[0],self._all_size_even[1],self._all_size_evenD[2]], [self._all_size_odd[0],self._all_size_odd[1],self._all_size_oddD[2]], h=2) + + + def dot_Hdiv(self, v, out): + out = self._dot_helper(v, out, [self._p[0], self._pD[1],self._pD[2]], [self._all_weights_even[0],self._all_weights_evenD[1],self._all_weights_evenD[2]], [self._all_weights_odd[0],self._all_weights_oddD[1],self._all_weights_oddD[2]], [self._all_size_even[0],self._all_size_evenD[1],self._all_size_evenD[2]], [self._all_size_odd[0],self._all_size_oddD[1],self._all_size_oddD[2]], h=0) + out = self._dot_helper(v, out, [self._pD[0], self._p[1],self._pD[2]], [self._all_weights_evenD[0],self._all_weights_even[1],self._all_weights_evenD[2]], [self._all_weights_oddD[0],self._all_weights_odd[1],self._all_weights_oddD[2]], [self._all_size_evenD[0],self._all_size_even[1],self._all_size_evenD[2]], [self._all_size_oddD[0],self._all_size_odd[1],self._all_size_oddD[2]], h=1) + return self._dot_helper(v, out, [self._pD[0], self._pD[1],self._p[2]], [self._all_weights_evenD[0],self._all_weights_evenD[1],self._all_weights_even[2]], [self._all_weights_oddD[0],self._all_weights_oddD[1],self._all_weights_odd[2]], [self._all_size_evenD[0],self._all_size_evenD[1],self._all_size_even[2]], [self._all_size_oddD[0],self._all_size_oddD[1],self._all_size_odd[2]], h=2) + + + def dot_H1H1H1(self, v, out): + out = self._dot_helper(v, out, self._p, self._all_weights_even, self._all_weights_odd, self._all_size_even, self._all_size_odd, h = 0) + self._dot_helper(v, out, self._p, self._all_weights_even, self._all_weights_odd, self._all_size_even, self._all_size_odd, h = 1) + return self._dot_helper(v, out, self._p, self._all_weights_even, self._all_weights_odd, self._all_size_even, self._all_size_odd, h = 2) + + + + def dot(self, v, out=None): + + assert isinstance(v, Vector) and v.space == self.domain + + if out is None: + out = self.codomain.zeros() + else: + assert isinstance(out, Vector) and out.space == self.codomain + + if self._V_name == 'H1' or self._V_name == 'L2': + for i0 in range(self._out_starts[0], self._out_ends[0]+1): + for i1 in range(self._out_starts[1], self._out_ends[1]+1): + for i2 in range(self._out_starts[2], self._out_ends[2]+1): + out[i0,i1,i2] = 0.0 + else: + for h in range(3): + for i0 in range(self._out_starts[h][0], self._out_ends[h][0]+1): + for i1 in range(self._out_starts[h][1], self._out_ends[h][1]+1): + for i2 in range(self._out_starts[h][2], self._out_ends[h][2]+1): + out[h][i0,i1,i2] = 0.0 + + dot_methods = { + "H1": self.dot_H1, + "L2": self.dot_L2, + "Hcurl": self.dot_Hcurl, + "Hdiv": self.dot_Hdiv, + "H1H1H1": self.dot_H1H1H1, + } + + return dot_methods.get(self._V_name)(v, out) + + def transpose(self, *, out = None): + if out is None: + out = MultiGridSolver.RestrictionOperator(self._W, self._V) + else: + assert isinstance(out, MultiGridSolver.RestrictionOperator) + assert out.domain is self.codomain + assert out.codomain is self.domain + + return out + + + def solve(self, b, verbose=False): + max_iter = self.max_iter_list + N_cycles = self.N_cycles + + #We define a list where to store the number of itteration it takes at each multigrid level + #Change N_levels for 1D case + Multigrid_itterations = np.zeros(self.N_levels, dtype=int) + converged = np.ones(self.N_levels, dtype=bool) + + def V_cycle(l, r_l): + #Change for N_levels-1 for 1D case + if (l < self.N_levels-1): + solver_ini = inverse(self.A[l],self.method, maxiter= max_iter[l]) + x_l = solver_ini.dot(r_l) + + #We count the number of itterations + Multigrid_itterations[l] += solver_ini._info['niter'] + + #We determine if the itterative solver converged in the maximum number of itterations + converged[l] = solver_ini._info['success'] + if converged[l] == True: + return x_l + + r_l = r_l - self.A[l].dot(x_l) + + r_l_plus_1 = self.R[l].dot(r_l) + x_l_plus_1 = V_cycle(l+1, r_l_plus_1) + #New + #x_l_aux = self.E[l].dot(x_l_plus_1) + #x_l = x_l + x_l_aux + #r_l = r_l - self.A[l].dot(x_l_aux) + #solver_end = inverse(self.A[l].T,self.method, maxiter= max_iter, x0 =x_l) + #x_l = solver_end.dot(r_l) + #Multigrid_itterations[l] += solver_end._info['niter'] + #### + #old + x_l = x_l + self.E[l].dot(x_l_plus_1) + ### + + else: + #Solve directly + x_l = direct_solver(self.A_inv,r_l, self.derhamlist[l].Vh_fem[self.sp_key]) + return x_l + + timei = time.time() + for cycle in range(N_cycles): + if cycle == 0: + solver = inverse(self.A[0],self.method, maxiter= max_iter[0], tol=1e-8) + else: + solver = inverse(self.A[0],self.method, maxiter= max_iter[0], x0 = x_0, tol=1e-8) + x_0 = solver.dot(b) + + Multigrid_itterations[0] += solver._info['niter'] + + #We determine if the itterative solver converged in the maximum number of itterations + converged[0] = solver._info['success'] + if converged[0] == True: + x = x_0 + break + + r_0 = b - self.A[0].dot(x_0) + r_1 = self.R[0].dot(r_0) + + x_0 = x_0 + self.E[0].dot(V_cycle(1,r_1)) + + if converged[0] == False: + solver = inverse(self.A[0],self.method,x0 = x_0, tol = 10**(-10)) + x = solver.dot(b) + Multigrid_itterations[0] += solver._info['niter'] + + timef = time.time() + + #We get the final error + Multigrid_error = solver._info['res_norm'] + + Multigrid_time = timef- timei + + if verbose: + print("################") + print("################") + print("################") + print(f'{max_iter = }') + print(f'{N_cycles = }') + print("################") + print("################") + print(f'{Multigrid_itterations = }') + print(f'{Multigrid_error = }') + print(f'{Multigrid_time = }') + print(f'{converged = }') + print("################") + print("################") + return x + + +class MultiGridPoissonSolver(MultiGridSolver): + def __init__(self, derham: Derham, sp_key: str, N_levels: int, domain, max_iter_list: list, N_cycles: int, method: str = 'cg'): + expr = " grad.T @ M1 @ grad" + namespace = {"grad.T": "derham.grad.T", + "grad": "derham.grad", + "M1": "WeightedMassOperators(derham, domain).M1"} + super().__init__(derham, sp_key, N_levels, domain, expr, namespace, max_iter_list, N_cycles, method) + + + +if __name__ == "__main__": + comm = MPI.COMM_WORLD + Nel = [16,16,1] + sp_key = '0' + plist = [2,2,1] + spl_kind = [True,True,True] + N_levels = 3 + domain = Cuboid() + derham = Derham([Nel[0],Nel[1],Nel[2]], plist, spl_kind, comm=comm, local_projectors=False) + max_iter = [10, 25, 25] + N_cycles = 5 + method = 'cg' + + expr = " grad.T @ M1 @ grad" + namespace = {"grad.T": "derham.grad.T", + "grad": "derham.grad", + "M1": "WeightedMassOperators(derham, domain).M1"} + + + # 1. Create a local context for the eval. It maps string names + # to the actual objects for the CURRENT level. + eval_context = { + "derham": derham, + "domain": domain, + "WeightedMassOperators": WeightedMassOperators + } + + # 2. Build the namespace by evaluating each string from the template + # within the context of the current level. + updated_namespace = { + key: eval(recipe_string, {}, eval_context) + for key, recipe_string in namespace.items() + } + + A = build_operator(expr, updated_namespace) + + multigrid = MultiGridPoissonSolver(derham, sp_key, N_levels, domain, max_iter, N_cycles, method) + + u_stararr, u_star = create_equal_random_arrays(derham.Vh_fem[sp_key], seed=8765) + #We compute the rhs + b = A.dot(u_star) + b_arr = b.toarray() + u = multigrid.solve(b, verbose=True) + b_ans_arr = A.dot(u).toarray() + if np.allclose(b_ans_arr, b_arr, atol=1e-6): + print("The multigrid solver computed the correct solution.") + else: + print("The multigrid solver did not compute the correct solution.") + print(f"{b_ans_arr = }") + print(f"{b_arr = }") \ No newline at end of file diff --git a/src/struphy/linear_algebra/tests/test_multigrid.py b/src/struphy/linear_algebra/tests/test_multigrid.py new file mode 100644 index 000000000..9db8c7dd2 --- /dev/null +++ b/src/struphy/linear_algebra/tests/test_multigrid.py @@ -0,0 +1,83 @@ +from struphy.linear_algebra.multigrid_solver import MultiGridPoissonSolver, build_operator +from struphy.geometry.domains import Cuboid +from struphy.feec.psydac_derham import Derham +import numpy as np +from mpi4py import MPI +from struphy.feec.mass import WeightedMassOperators +from struphy.feec.utilities import create_equal_random_arrays + +def test_multigrid_poisson_solver(): + """Test for the MultiGridPoissonSolver class solving a Poisson problem.""" + comm = MPI.COMM_WORLD + Nel = [16,16,1] + sp_key = '0' + plist = [2,2,1] + spl_kind = [True,True,True] + N_levels = 3 + domain = Cuboid() + derham = Derham([Nel[0],Nel[1],Nel[2]], plist, spl_kind, comm=comm, local_projectors=False) + max_iter = [10, 25, 25] + N_cycles = 5 + method = 'cg' + + # We build the operator A = grad^T * M1 * grad + A = derham.grad.T @ WeightedMassOperators(derham, domain).M1 @ derham.grad + + # We build the multigrid solver + multigrid = MultiGridPoissonSolver(derham, sp_key, N_levels, domain, max_iter, N_cycles, method) + + # We create a random solution u_star + _, u_star = create_equal_random_arrays(derham.Vh_fem[sp_key], seed=8765) + #We compute the rhs + b = A.dot(u_star) + b_arr = b.toarray() + # We solve the system using the multigrid solver + u = multigrid.solve(b, verbose=True) + + # We check that the solution is correct + b_ans_arr = A.dot(u).toarray() + assert np.allclose(b_ans_arr, b_arr, atol=1e-6) + +def test_build_operator(): + """Test for the build_operator function.""" + comm = MPI.COMM_WORLD + Nel = [8,8,1] + plist = [2,2,1] + spl_kind = [True,True,True] + domain = Cuboid() + derham = Derham([Nel[0],Nel[1],Nel[2]], plist, spl_kind, comm=comm, local_projectors=False) + + expr = " grad.T @ M1 @ grad" + + namespace = {"grad.T": "derham.grad.T", + "grad": "derham.grad", + "M1": "WeightedMassOperators(derham, domain).M1"} + + + # 1. Create a local context for the eval. It maps string names + # to the actual objects for the CURRENT level. + eval_context = { + "derham": derham, + "domain": domain, + "WeightedMassOperators": WeightedMassOperators + } + + # 2. Build the namespace by evaluating each string from the template + # within the context of the current level. + updated_namespace = { + key: eval(recipe_string, {}, eval_context) + for key, recipe_string in namespace.items() + } + + A = build_operator(expr, updated_namespace) + + A_ref = derham.grad.T @ WeightedMassOperators(derham, domain).M1 @ derham.grad + A_arr = A.toarray() + A_ref_arr = A_ref.toarray() + + assert np.allclose(A_arr, A_ref_arr, atol=1e-10) + + +if __name__ == "__main__": + test_build_operator() + test_multigrid_poisson_solver() \ No newline at end of file diff --git a/src/struphy/pic/base.py b/src/struphy/pic/base.py index 84900418d..7deb2571c 100644 --- a/src/struphy/pic/base.py +++ b/src/struphy/pic/base.py @@ -1,4 +1,3 @@ -import copy import os import warnings from abc import ABCMeta, abstractmethod From b6154c05fef6823732fa4f40c6e53bec7acd4ae3 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Fri, 2 Oct 2026 13:53:57 +0200 Subject: [PATCH 02/16] use LinearOperator in tests --- .../tests/{multigrid.py => test_multigrid.py} | 6 +++--- src/struphy/linear_algebra/multigrid_solver.py | 16 ++++++++-------- 2 files changed, 11 insertions(+), 11 deletions(-) rename src/struphy/feec/tests/{multigrid.py => test_multigrid.py} (99%) diff --git a/src/struphy/feec/tests/multigrid.py b/src/struphy/feec/tests/test_multigrid.py similarity index 99% rename from src/struphy/feec/tests/multigrid.py rename to src/struphy/feec/tests/test_multigrid.py index 536d383dd..2b4c3ddf4 100644 --- a/src/struphy/feec/tests/multigrid.py +++ b/src/struphy/feec/tests/test_multigrid.py @@ -15,7 +15,7 @@ from struphy.feec import preconditioner from psydac.linalg.basic import VectorSpace, Vector, LinearOperator from psydac.fem.projectors import knot_insertion_projection_operator -from struphy.feec.linear_operators import LinOpWithTransp +from feectools.linalg.basic import LinearOperator from psydac.fem.basic import FemSpace from psydac.fem.tensor import TensorFemSpace from struphy.feec.mass import WeightedMassOperators @@ -278,7 +278,7 @@ def remove_padding(fem_space, v): return v_array -class RestrictionOperator(LinOpWithTransp): +class RestrictionOperator(LinearOperator): """ Linear operator which operates between vector spaces of the same kind but different resolutions. @@ -545,7 +545,7 @@ def transpose(self, *, out = None): return out -class ExtensionOperator(LinOpWithTransp): +class ExtensionOperator(LinearOperator): """ Linear operator which operates between vector spaces of the same kind but different resolutions. diff --git a/src/struphy/linear_algebra/multigrid_solver.py b/src/struphy/linear_algebra/multigrid_solver.py index 1f2d886da..65feeae85 100644 --- a/src/struphy/linear_algebra/multigrid_solver.py +++ b/src/struphy/linear_algebra/multigrid_solver.py @@ -9,7 +9,7 @@ from psydac.linalg.solvers import inverse from psydac.linalg.basic import Vector from psydac.fem.projectors import knot_insertion_projection_operator -from struphy.feec.linear_operators import LinOpWithTransp +from feectools.linalg.basic import LinearOperator from psydac.fem.basic import FemSpace from psydac.fem.tensor import TensorFemSpace from struphy.feec.mass import WeightedMassOperators @@ -19,7 +19,7 @@ def build_operator(expr: str, namespace: dict): """ - Build a composite LinOpWithTransp (or scalar * LinOpWithTransp expression) + Build a composite LinearOperator (or scalar * LinearOperator expression) from a string expression. Parameters @@ -31,7 +31,7 @@ def build_operator(expr: str, namespace: dict): The expression can contain: - scalars (e.g., epsilon, mu, ...) - - LinOpWithTransp objects (already supporting +, -, *, @) + - LinearOperator objects (already supporting +, -, *, @) - parentheses to control precedence namespace : dict @@ -49,9 +49,9 @@ def build_operator(expr: str, namespace: dict): Returns ------- - LinOpWithTransp + LinearOperator A composite operator built according to the expression. The returned - object supports `.dot(v)` and any other functionality of LinOpWithTransp. + object supports `.dot(v)` and any other functionality of LinearOperator. Raises ------ @@ -222,7 +222,7 @@ def __init__(self, derham, sp_key, N_levels, domain, expr: str, namespace: dict, The expression can contain: - scalars (e.g., epsilon, mu, ...) - - LinOpWithTransp objects (already supporting +, -, *, @) + - LinearOperator objects (already supporting +, -, *, @) - parentheses to control precedence namespace : dict @@ -308,7 +308,7 @@ def __init__(self, derham, sp_key, N_levels, domain, expr: str, namespace: dict, self.E.append(self.R[level].transpose()) - class RestrictionOperator(LinOpWithTransp): + class RestrictionOperator(LinearOperator): """ Linear operator which operates between vector spaces of the same kind but different resolutions. @@ -577,7 +577,7 @@ def transpose(self, *, out = None): return out - class ExtensionOperator(LinOpWithTransp): + class ExtensionOperator(LinearOperator): """ Linear operator which operates between vector spaces of the same kind but different resolutions. From 0bfe1a75261676baa6e15e75d60e34984e76265d Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Sat, 3 Oct 2026 17:35:33 +0200 Subject: [PATCH 03/16] Remove old multigrid solver and research tests The old MultiGridSolver (eval-string operators, serial-only transfer operators, 2D coarsening) and the research script feec/tests/test_multigrid.py are replaced by a new implementation. Co-Authored-By: Claude Opus 5.5 --- src/struphy/feec/tests/test_multigrid.py | 3247 ----------------- .../linear_algebra/multigrid_solver.py | 1070 ------ .../linear_algebra/tests/test_multigrid.py | 83 - 3 files changed, 4400 deletions(-) delete mode 100644 src/struphy/feec/tests/test_multigrid.py delete mode 100644 src/struphy/linear_algebra/multigrid_solver.py delete mode 100644 src/struphy/linear_algebra/tests/test_multigrid.py diff --git a/src/struphy/feec/tests/test_multigrid.py b/src/struphy/feec/tests/test_multigrid.py deleted file mode 100644 index 2b4c3ddf4..000000000 --- a/src/struphy/feec/tests/test_multigrid.py +++ /dev/null @@ -1,3247 +0,0 @@ -import time - -import matplotlib.pyplot as plt -import numpy as np -from mpi4py import MPI - -from struphy.bsplines.bsplines import basis_funs, find_span -from struphy.bsplines.evaluation_kernels_1d import evaluation_kernel_1d -from struphy.feec.basis_projection_ops import BasisProjectionOperators -from struphy.feec.local_projectors_kernels import fill_matrix_column -from struphy.feec.psydac_derham import Derham -from struphy.feec.utilities_local_projectors import get_one_spline, get_span_and_basis, get_values_and_indices_splines - -from psydac.linalg.solvers import inverse -from struphy.feec import preconditioner -from psydac.linalg.basic import VectorSpace, Vector, LinearOperator -from psydac.fem.projectors import knot_insertion_projection_operator -from feectools.linalg.basic import LinearOperator -from psydac.fem.basic import FemSpace -from psydac.fem.tensor import TensorFemSpace -from struphy.feec.mass import WeightedMassOperators -from struphy.geometry.domains import Tokamak, Cuboid, HollowCylinder -from struphy.fields_background.equils import AdhocTorusQPsi -from math import comb, log2 -import random - - -def jacobi(A, b, x_init=None, tol=1e-6, max_iter=1000, verbose = False): - """ - Solves the linear system Ax = b using the Jacobi iterative method. - - Parameters: - A : numpy.ndarray - Coefficient matrix (2D array) - b : numpy.ndarray - Right-hand side vector (1D array) - x_init : numpy.ndarray, optional - Initial guess for the solution (1D array) - tol : float, optional - Convergence tolerance (default: 1e-10) - max_iter : int, optional - Maximum number of iterations (default: 1000) - - Returns: - x : numpy.ndarray - Solution vector - """ - n = A.shape[0] - x = x_init if x_init is not None else np.zeros(n) - x_new = np.zeros(n) - converged = False - for itterations in range(max_iter): - for i in range(n): - sum_except_i = np.dot(A[i, :], x) - A[i, i] * x[i] - x_new[i] = (b[i] - sum_except_i) / A[i, i] - - - error = np.linalg.norm(b - np.dot(A,x_new), ord=np.inf) - if error < tol: - converged = True - if verbose: - print(f'{converged = }') - print(f'{itterations = }') - print(f'{error = }') - return x_new, itterations - - x[:] = x_new # Update x - - - if verbose: - print(f'{converged = }') - print(f'{itterations = }') - print(f'{error = }') - return x_new, itterations - #raise ValueError("Jacobi method did not converge within the maximum number of iterations") - - -def will_jacobi_converge(A, verbose = False): - converges = False - #We get the diagonal, lower triangular and uper tireangular parts of A - Darr = np.diag(np.diag(A)) - Darr_inv = np.diag(1.0 / np.diag(Darr)) - Larr = np.tril(A, k=-1) - Uarr = np.triu(A, k=1) - #Then the Jacobi iteration matrix is - Jarr = np.matmul(Darr_inv,(Larr+Uarr)) - #Now we get its eigenvalues - eigenvalues = np.linalg.eigvals(Jarr) - max_eigen_value = abs(eigenvalues[np.argmax(np.abs(eigenvalues))]) - - if max_eigen_value < 1.0: - converges = True - - if(verbose): - print(f'{max_eigen_value = }') - - return converges - - -def will_gauss_seidel_converge(A, verbose = False): - converges = False - #We get the diagonal, lower triangular and uper triangular parts of A - Darr = np.diag(np.diag(A)) - Larr = np.tril(A, k=-1) - Uarr = np.triu(A, k=1) - Aux = np.linalg.inv(Darr - Larr) - #Then the Gauss-Seidel iteration matrix is - Garr = np.matmul(Aux, Uarr) - #Now we get its eigenvalues - eigenvalues = np.linalg.eigvals(Garr) - max_eigen_value = abs(eigenvalues[np.argmax(np.abs(eigenvalues))]) - - if max_eigen_value < 1.0: - converges = True - - if(verbose): - print(f'{max_eigen_value = }') - - return converges - - -def from_array_to_psydac(x_vector, fem_space): - #fem_space = derham.Vh_fem[sp_key] - symbolic_name = fem_space.symbolic_space.name - - if(symbolic_name == 'H1' or symbolic_name == "L2"): - spaces = [fem_space.spaces] - N = [spaces[0][i].nbasis for i in range(3)] - starts = np.array(fem_space.vector_space.starts) - x = fem_space.vector_space.zeros() - - cont= 0 - for i0 in range(N[0]): - for i1 in range(N[1]): - for i2 in range(N[2]): - x[starts[0]+i0,starts[1]+i1,starts[2]+i2] = x_vector[cont] - cont += 1 - - else: - spaces = [comp.spaces for comp in fem_space.spaces] - N = [[spaces[h][i].nbasis for i in range(3)] for h in range(3)] - starts = np.array([vi.starts for vi in fem_space.vector_space.spaces]) - x = fem_space.vector_space.zeros() - - cont = 0 - for h in range(3): - for i0 in range(N[h][0]): - for i1 in range(N[h][1]): - for i2 in range(N[h][2]): - x[h][starts[h][0]+i0,starts[h][1]+i1,starts[h][2]+i2] = x_vector[cont] - cont += 1 - return x - - -def direct_solver(A_inv,b, fem_space): - # A_inv is already the inverse matrix of A - #fem_space = derham.Vh_fem[sp_key] - symbolic_name = fem_space.symbolic_space.name - - if(symbolic_name == 'H1' or symbolic_name == "L2"): - spaces = [fem_space.spaces] - N = [spaces[0][i].nbasis for i in range(3)] - starts = np.array(fem_space.vector_space.starts) - - b_vector = remove_padding(fem_space, b) - x_vector = np.dot(A_inv, b_vector) - x = fem_space.vector_space.zeros() - - cont= 0 - for i0 in range(N[0]): - for i1 in range(N[1]): - for i2 in range(N[2]): - x[starts[0]+i0,starts[1]+i1,starts[2]+i2] = x_vector[cont] - cont += 1 - - else: - spaces = [comp.spaces for comp in fem_space.spaces] - N = [[spaces[h][i].nbasis for i in range(3)] for h in range(3)] - starts = np.array([vi.starts for vi in fem_space.vector_space.spaces]) - - b_vector = remove_padding(fem_space, b) - x_vector = np.dot(A_inv, b_vector) - x = fem_space.vector_space.zeros() - - cont = 0 - for h in range(3): - for i0 in range(N[h][0]): - for i1 in range(N[h][1]): - for i2 in range(N[h][2]): - x[h][starts[h][0]+i0,starts[h][1]+i1,starts[h][2]+i2] = x_vector[cont] - cont += 1 - return x - - -def get_b_spline_degree(V): - """ - Determines the degree of the B-splines. - - Parameters - ---------- - V : psydac.fem.basic.FemSpace - Finite element spline space (domain, input space). - - Returns - ------- - p : numpy array - numpy array of 3 ints containing the B-spline degrees for each spatial direction. - """ - - assert isinstance(V, FemSpace) - p = np.zeros(3, dtype=int) - - if hasattr(V.symbolic_space, "name"): - V_name = V.symbolic_space.name - - if(V_name == "H1"): - for i, space in enumerate(V.spaces): - p[i] = space.degree - elif(V_name == "L2"): - for i, space in enumerate(V.spaces): - p[i] = space.degree+1 - elif(V_name == "Hcurl"): - V1ds = [comp.spaces for comp in V.spaces] - for i in range(3): - p[i] = V1ds[i][i].degree+1 - elif(V_name == "Hdiv"): - V1ds = [comp.spaces for comp in V.spaces] - for i in range(3): - p[i] = V1ds[i][i].degree - elif(V_name == "H1H1H1"): - V1ds = [comp.spaces for comp in V.spaces] - for i in range(3): - p[i] = V1ds[i][i].degree - else: - raise Exception("Invalid symbolic name.") - - return p - - -def remove_padding(fem_space, v): - - #fem_space = derham.Vh_fem[sp_key] - symbolic_name = fem_space.symbolic_space.name - - if(symbolic_name == 'H1' or symbolic_name == "L2"): - - spaces = [fem_space.spaces] - N = [spaces[0][i].nbasis for i in range(3)] - - starts = np.array(fem_space.vector_space.starts) - - #To make it easier to read I will extract the data out of out disregarding all the padding it come with - v_array = np.zeros(N[0]*N[1]*N[2], dtype=float) - cont = 0 - for i0 in range(N[0]): - for i1 in range(N[1]): - for i2 in range(N[2]): - v_array[cont] = v[starts[0]+i0,starts[1]+i1,starts[2]+i2] - cont += 1 - - else: - spaces = [comp.spaces for comp in fem_space.spaces] - N = [[spaces[h][i].nbasis for i in range(3)] for h in range(3)] - - starts = np.array([vi.starts for vi in fem_space.vector_space.spaces]) - - #To make it easier to read I will extract the data out of out disregarding all the padding it come with - v_array = np.zeros(N[0][0]*N[0][1]*N[0][2]+N[1][0]*N[1][1]*N[1][2]+N[2][0]*N[2][1]*N[2][2], dtype=float) - cont = 0 - for h in range(3): - for i0 in range(N[h][0]): - for i1 in range(N[h][1]): - for i2 in range(N[h][2]): - v_array[cont] = v[h][starts[h][0]+i0,starts[h][1]+i1,starts[h][2]+i2] - cont += 1 - - - return v_array - - -class RestrictionOperator(LinearOperator): - """ - Linear operator which operates between vector spaces of the same kind but different resolutions. - - We assume that the vectors in the domain belong to a De-rham space with n elements, while the codomain - belong to the coarser De-rham space with n/2 elements. - - At the moment we also assume that we are halving only the first spatial direction. - - Parameters - ---------- - V : psydac.fem.basic.FemSpace - Finite element spline space (domain, input space). - - W : psydac.fem.basic.FemSpace - Finite element spline space (codomain, output space). - - """ - def __init__(self, V, W): - - # Check domain and codomain - assert isinstance(V, FemSpace) - assert isinstance(W, FemSpace) - - self._V = V - self._W = W - - # Store info in object - self._domain = V.vector_space - self._codomain = W.vector_space - self._dtype = V.vector_space.dtype - - #Can be "H1", "L2", "Hcurl", "Hdiv", "H1H1H1" - self._V_name = V.symbolic_space.name - self._W_name = W.symbolic_space.name - assert(self._V_name == self._W_name) - - #This list will tell us in which spatial direction we are halving the problem. - self._halving_directions = [False,False,False] - - # input space: 3d StencilVectorSpaces and 1d SplineSpaces of each component - if isinstance(V, TensorFemSpace): - self._V1ds = [V.spaces] - self._VNbasis = np.array([self._V1ds[0][0].nbasis, self._V1ds[0][1].nbasis, self._V1ds[0][2].nbasis]) - - # We get the start and endpoint for each sublist in input - self._in_starts = np.array(V.vector_space.starts) - self._in_ends = np.array(V.vector_space.ends) - else: - self._V1ds = [comp.spaces for comp in V.spaces] - self._VNbasis = np.array( - [ - [self._V1ds[0][0].nbasis, self._V1ds[0][1].nbasis, self._V1ds[0][2].nbasis], - [ - self._V1ds[1][0].nbasis, - self._V1ds[1][1].nbasis, - self._V1ds[1][2].nbasis, - ], - [self._V1ds[2][0].nbasis, self._V1ds[2][1].nbasis, self._V1ds[2][2].nbasis], - ] - ) - - # We get the start and endpoint for each sublist in input - self._in_starts = np.array([vi.starts for vi in V.vector_space.spaces]) - self._in_ends = np.array([vi.ends for vi in V.vector_space.spaces]) - - # output space: 3d StencilVectorSpaces and 1d SplineSpaces of each component - if isinstance(W, TensorFemSpace): - self._W1ds = [W.spaces] - self._WNbasis = np.array([self._W1ds[0][0].nbasis, self._W1ds[0][1].nbasis, self._W1ds[0][2].nbasis]) - - for i in range(3): - if(self._WNbasis[i] < self._VNbasis[i]): - #If this breaks for clamped splines it means .nbasis gives you the number of basis functions, not the number of elements - assert self._VNbasis[i] == self._WNbasis[i]*2 - self._halving_directions[i] = True - - # We get the start and endpoint for each sublist in out - self._out_starts = np.array(W.vector_space.starts) - self._out_ends = np.array(W.vector_space.ends) - - else: - self._W1ds = [comp.spaces for comp in W.spaces] - self._WNbasis = np.array( - [ - [self._W1ds[0][0].nbasis, self._W1ds[0][1].nbasis, self._W1ds[0][2].nbasis], - [ - self._W1ds[1][0].nbasis, - self._W1ds[1][1].nbasis, - self._W1ds[1][2].nbasis, - ], - [self._W1ds[2][0].nbasis, self._W1ds[2][1].nbasis, self._W1ds[2][2].nbasis], - ] - ) - for i in range(3): - if(self._WNbasis[1][i] < self._VNbasis[1][i]): - #If this breaks for clamped splines it means .nbasis gives you the number of basis functions, not the number of elements - assert self._VNbasis[0][i] == self._WNbasis[0][i]*2 and self._VNbasis[1][i] == self._WNbasis[1][i]*2 and self._VNbasis[2][i] == self._WNbasis[2][i]*2 - self._halving_directions[i] = True - - # We get the start and endpoint for each sublist in out - self._out_starts = np.array([vi.starts for vi in W.vector_space.spaces]) - self._out_ends = np.array([vi.ends for vi in W.vector_space.spaces]) - - - - # Degree of the B-spline space, not to be confused with the degrees given by fem_space.spaces.degree since depending on the situation - # it will give the D-spline degree instead - self._p = get_b_spline_degree(V) - - #We also get the D-spline degree - self._pD = self._p - 1 - - #Now we compute the weights that define this linear operator - - #We begin by defining a list that will contain the 3 numpy arrays, each one with the weights for one spatial direction. - #In the case there are direction over which we do not halve the resolution we shall have an array with only one 1.0 - self._all_weights = [] - self._all_weightsD = [] - for i in range(3): - if self._halving_directions[i]: - #Here we store the weights needed for B-splines - weights = np.zeros(self._p[i]+2, dtype=float) - #Here we store the weights needed for D-splines - weightsD = np.zeros(self._pD[i]+2, dtype=float) - for j in range(self._p[i]+2): - weights[j] = 2.0**(-self._p[i])*comb(self._p[i]+1,j) - for j in range(self._pD[i]+2): - weightsD[j] = 2.0**-(self._pD[i]+1)*comb(self._pD[i]+1,j) - self._all_weights.append(weights) - self._all_weightsD.append(weightsD) - else: - self._all_weights.append(np.array([1.0],dtype=float)) - self._all_weightsD.append(np.array([1.0],dtype=float)) - - - - #-------------------------------------- - # Abstract interface - #-------------------------------------- - @property - def domain(self): - return self._domain - - @property - def codomain(self): - return self._codomain - - @property - def dtype(self): - return self._dtype - - def _dot_helper(self, v, out, p, weights, h=None): - """Helper function to perform dot product computation.""" - #First we get the number of weights in each direction - weights_len = [] - for i in range(3): - if self._halving_directions[i]: - weights_len.append(p[i] + 2) - else: - weights_len.append(1) - if h is None: # Scalar case (H1, L2) - for i0 in range(self._out_starts[0], self._out_ends[0] + 1): - for i1 in range(self._out_starts[1], self._out_ends[1] + 1): - for i2 in range(self._out_starts[2], self._out_ends[2] + 1): - for j0 in range(weights_len[0]): - if self._halving_directions[0]: - pos0 = (2 * i0 - p[0] + j0) % self._VNbasis[0] - else: - pos0 = i0 - for j1 in range(weights_len[1]): - if self._halving_directions[1]: - pos1 = (2 * i1 - p[1] + j1) % self._VNbasis[1] - else: - pos1 = i1 - for j2 in range(weights_len[2]): - if self._halving_directions[2]: - pos2 = (2 * i2 - p[2] + j2) % self._VNbasis[2] - else: - pos2 = i2 - out[i0, i1, i2] += weights[0][j0]* weights[1][j1] *weights[2][j2] * v[pos0, pos1, pos2] - else: # Vector case (Hcurl, Hdiv, H1H1H1) - for i0 in range(self._out_starts[h][0], self._out_ends[h][0] + 1): - for i1 in range(self._out_starts[h][1], self._out_ends[h][1] + 1): - for i2 in range(self._out_starts[h][2], self._out_ends[h][2] + 1): - for j0 in range(weights_len[0]): - if self._halving_directions[0]: - pos0 = (2 * i0 - p[0] + j0) % self._VNbasis[h][0] - else: - pos0 = i0 - for j1 in range(weights_len[1]): - if self._halving_directions[1]: - pos1 = (2 * i1 - p[1] + j1) % self._VNbasis[h][1] - else: - pos1 = i1 - for j2 in range(weights_len[2]): - if self._halving_directions[2]: - pos2 = (2 * i2 - p[2] + j2) % self._VNbasis[h][2] - else: - pos2 = i2 - out[h][i0, i1, i2] += weights[0][j0]* weights[1][j1] *weights[2][j2] * v[h][pos0, pos1, pos2] - - return out - - def dot_H1(self, v, out): - return self._dot_helper(v, out, self._p, self._all_weights) - - def dot_L2(self, v, out): - return self._dot_helper(v, out, self._pD, self._all_weightsD) - - def dot_Hcurl(self, v, out): - out = self._dot_helper(v, out, [self._pD[0],self._p[1],self._p[2]], [self._all_weightsD[0], self._all_weights[1], self._all_weights[2]], h = 0) - self._dot_helper(v, out, [self._p[0],self._pD[1],self._p[2]], [self._all_weights[0], self._all_weightsD[1], self._all_weights[2]], h = 1) - return self._dot_helper(v, out, [self._p[0],self._p[1],self._pD[2]], [self._all_weights[0], self._all_weights[1], self._all_weightsD[2]], h = 2) - - def dot_Hdiv(self, v, out): - out = self._dot_helper(v, out, [self._p[0],self._pD[1],self._pD[2]], [self._all_weights[0], self._all_weightsD[1], self._all_weightsD[2]], h = 0) - self._dot_helper(v, out, [self._pD[0],self._p[1],self._pD[2]], [self._all_weightsD[0], self._all_weights[1], self._all_weightsD[2]], h = 1) - return self._dot_helper(v, out, [self._pD[0],self._pD[1],self._p[2]], [self._all_weightsD[0], self._all_weightsD[1], self._all_weights[2]], h = 2) - - def dot_H1H1H1(self, v, out): - out = self._dot_helper(v, out, self._p, self._all_weights, h=0) - self._dot_helper(v, out, self._p, self._all_weights, h=1) - return self._dot_helper(v, out, self._p, self._all_weights, h=2) - - def dot(self, v, out=None): - - assert isinstance(v, Vector) and v.space == self.domain - - if out is None: - out = self.codomain.zeros() - else: - assert isinstance(out, Vector) and out.space == self.codomain - - if self._V_name == 'H1' or self._V_name == 'L2': - for i0 in range(self._out_starts[0], self._out_ends[0]+1): - for i1 in range(self._out_starts[1], self._out_ends[1]+1): - for i2 in range(self._out_starts[2], self._out_ends[2]+1): - out[i0,i1,i2] = 0.0 - else: - for h in range(3): - for i0 in range(self._out_starts[h][0], self._out_ends[h][0]+1): - for i1 in range(self._out_starts[h][1], self._out_ends[h][1]+1): - for i2 in range(self._out_starts[h][2], self._out_ends[h][2]+1): - out[h][i0,i1,i2] = 0.0 - - dot_methods = { - "H1": self.dot_H1, - "L2": self.dot_L2, - "Hcurl": self.dot_Hcurl, - "Hdiv": self.dot_Hdiv, - "H1H1H1": self.dot_H1H1H1, - } - - return dot_methods.get(self._V_name)(v, out) - - def transpose(self, *, out = None): - if out is None: - out = ExtensionOperator(self._W, self._V) - else: - assert isinstance(out, ExtensionOperator) - assert out.domain is self.codomain - assert out.codomain is self.domain - - return out - - -class ExtensionOperator(LinearOperator): - """ - Linear operator which operates between vector spaces of the same kind but different resolutions. - - We assume that the vectors in the domain belong to a De-rham space with n/2 elements, while the codomain - belong to the coarser De-rham space with n elements. - - At the moment we also assume that we are halving only the first spatial direction. - - Parameters - ---------- - V : psydac.fem.basic.FemSpace - Finite element spline space (domain, input space). - - W : psydac.fem.basic.FemSpace - Finite element spline space (codomain, output space). - - """ - def __init__(self, V, W): - - # Check domain and codomain - assert isinstance(V, FemSpace) - assert isinstance(W, FemSpace) - - self._V = V - self._W = W - - # Store info in object - self._domain = V.vector_space - self._codomain = W.vector_space - self._dtype = V.vector_space.dtype - - #Can be "H1", "L2", "Hcurl", "Hdiv", "H1H1H1" - self._V_name = V.symbolic_space.name - self._W_name = W.symbolic_space.name - assert(self._V_name == self._W_name) - - #This list will tell us in which spatial direction we are halving the problem. - self._halving_directions = [False,False,False] - - # input space: 3d StencilVectorSpaces and 1d SplineSpaces of each component - if isinstance(V, TensorFemSpace): - self._V1ds = [V.spaces] - self._VNbasis = np.array([self._V1ds[0][0].nbasis, self._V1ds[0][1].nbasis, self._V1ds[0][2].nbasis]) - - # We get the start and endpoint for each sublist in input - self._in_starts = np.array(V.vector_space.starts) - self._in_ends = np.array(V.vector_space.ends) - else: - self._V1ds = [comp.spaces for comp in V.spaces] - self._VNbasis = np.array( - [ - [self._V1ds[0][0].nbasis, self._V1ds[0][1].nbasis, self._V1ds[0][2].nbasis], - [ - self._V1ds[1][0].nbasis, - self._V1ds[1][1].nbasis, - self._V1ds[1][2].nbasis, - ], - [self._V1ds[2][0].nbasis, self._V1ds[2][1].nbasis, self._V1ds[2][2].nbasis], - ] - ) - - # We get the start and endpoint for each sublist in input - self._in_starts = np.array([vi.starts for vi in V.vector_space.spaces]) - self._in_ends = np.array([vi.ends for vi in V.vector_space.spaces]) - - # output space: 3d StencilVectorSpaces and 1d SplineSpaces of each component - if isinstance(W, TensorFemSpace): - self._W1ds = [W.spaces] - self._WNbasis = np.array([self._W1ds[0][0].nbasis, self._W1ds[0][1].nbasis, self._W1ds[0][2].nbasis]) - - for i in range(3): - if(self._VNbasis[i] < self._WNbasis[i]): - #If this breaks for clamped splines it means .nbasis gives you the number of basis functions, not the number of elements - assert self._WNbasis[i] == self._VNbasis[i]*2 - self._halving_directions[i] = True - - # We get the start and endpoint for each sublist in out - self._out_starts = np.array(W.vector_space.starts) - self._out_ends = np.array(W.vector_space.ends) - - else: - self._W1ds = [comp.spaces for comp in W.spaces] - self._WNbasis = np.array( - [ - [self._W1ds[0][0].nbasis, self._W1ds[0][1].nbasis, self._W1ds[0][2].nbasis], - [ - self._W1ds[1][0].nbasis, - self._W1ds[1][1].nbasis, - self._W1ds[1][2].nbasis, - ], - [self._W1ds[2][0].nbasis, self._W1ds[2][1].nbasis, self._W1ds[2][2].nbasis], - ] - ) - - for i in range(3): - if(self._VNbasis[1][i] < self._WNbasis[1][i]): - assert self._WNbasis[0][i] == self._VNbasis[0][i]*2 and self._WNbasis[1][i] == self._VNbasis[1][i]*2 and self._WNbasis[2][i] == self._VNbasis[2][i]*2 - self._halving_directions[i] = True - - # We get the start and endpoint for each sublist in out - self._out_starts = np.array([vi.starts for vi in W.vector_space.spaces]) - self._out_ends = np.array([vi.ends for vi in W.vector_space.spaces]) - - - - # Degree of the B-spline space, not to be confused with the degrees given by fem_space.spaces.degree since depending on the situation - # it will give the D-spline degree instead - self._p = get_b_spline_degree(V) - #We also get the D-splines degree - self._pD = self._p-1 - - #Now we compute the weights that define this linear operator - - #We begin by defining a list that will contain the 3 numpy arrays, each one with the weights for one spatial direction. - #In the case there are a direction over which we do not halve the resolution we shall have an array with only one 1.0 - self._all_weights_even = [] - self._all_weights_evenD = [] - self._all_weights_odd = [] - self._all_weights_oddD = [] - - #Each list has 3 integers, each one denoting the number of weights in the corresponding weights array. - self._all_size_even = [] - self._all_size_evenD = [] - self._all_size_odd = [] - self._all_size_oddD = [] - - for i in range(3): - if self._halving_directions[i]: - #First for B-splines - if(self._p[i]%2 == 0): - size_even = self._p[i]//2 +1 - size_odd = self._p[i]//2 +1 - weights_even = np.zeros(size_even, dtype=float) - weights_odd = np.zeros(size_odd, dtype=float) - for j in range(size_even): - weights_even[j] = 2.0**(-self._p[i])*comb(self._p[i]+1,2*j) - weights_odd[j] = 2.0**(-self._p[i])*comb(self._p[i]+1,2*j+1) - else: - size_even = (self._p[i]+1)//2 +1 - size_odd = (self._p[i]-1)//2 +1 - weights_even = np.zeros(size_even, dtype=float) - weights_odd = np.zeros(size_odd, dtype=float) - for j in range(size_even): - weights_even[j] = 2.0**(-self._p[i])*comb(self._p[i]+1,2*j) - for j in range(size_odd): - weights_odd[j] = 2.0**(-self._p[i])*comb(self._p[i]+1,2*j+1) - - #Second for D-splines - if(self._pD[i]%2 == 0): - size_evenD = self._pD[i]//2 +1 - size_oddD = self._pD[i]//2 +1 - weights_evenD = np.zeros(size_evenD, dtype=float) - weights_oddD = np.zeros(size_oddD, dtype=float) - for j in range(size_evenD): - weights_evenD[j] = 2.0**-(self._pD[i]+1)*comb(self._pD[i]+1,2*j) - weights_oddD[j] = 2.0**-(self._pD[i]+1)*comb(self._pD[i]+1,2*j+1) - else: - size_evenD = (self._pD[i]+1)//2 +1 - size_oddD = (self._pD[i]-1)//2 +1 - weights_evenD = np.zeros(size_evenD, dtype=float) - weights_oddD = np.zeros(size_oddD, dtype=float) - for j in range(size_evenD): - weights_evenD[j] = 2.0**-(self._pD[i]+1)*comb(self._pD[i]+1,2*j) - for j in range(size_oddD): - weights_oddD[j] = 2.0**-(self._pD[i]+1)*comb(self._pD[i]+1,2*j+1) - - self._all_weights_even.append(weights_even) - self._all_weights_evenD.append(weights_evenD) - self._all_weights_odd.append(weights_odd) - self._all_weights_oddD.append(weights_oddD) - self._all_size_even.append(size_even) - self._all_size_evenD.append(size_evenD) - self._all_size_odd.append(size_odd) - self._all_size_oddD.append(size_oddD) - - else: - self._all_weights_even.append(np.array([1.0],dtype=float)) - self._all_weights_evenD.append(np.array([1.0],dtype=float)) - self._all_weights_odd.append(np.array([1.0],dtype=float)) - self._all_weights_oddD.append(np.array([1.0],dtype=float)) - self._all_size_even.append(1) - self._all_size_evenD.append(1) - self._all_size_odd.append(1) - self._all_size_oddD.append(1) - - - #-------------------------------------- - # Abstract interface - #-------------------------------------- - @property - def domain(self): - return self._domain - - @property - def codomain(self): - return self._codomain - - @property - def dtype(self): - return self._dtype - - def tosparse(self): - pass - - def toarray(self): - pass - - def _dot_helper(self, v, out, p, weights_even, weights_odd, size_even, size_odd, h=None): - """Helper function to perform dot product computation.""" - parity_match = [] - for i in range(3): - parity_match.append(p[i] % 2) - - if h is None: # Scalar case (H1, L2) - for j0 in range(self._out_starts[0], self._out_ends[0] + 1): - parity_j0 = j0 % 2 - weights0, size0, offset0 = ((weights_even[0], size_even[0], 0) if parity_j0 == parity_match[0] else (weights_odd[0], size_odd[0], 1)) - for j1 in range(self._out_starts[1], self._out_ends[1] + 1): - parity_j1 = j1 % 2 - weights1, size1, offset1 = ((weights_even[1], size_even[1], 0) if parity_j1 == parity_match[1] else (weights_odd[1], size_odd[1], 1)) - for j2 in range(self._out_starts[2], self._out_ends[2] + 1): - parity_j2 = j2 % 2 - weights2, size2, offset2 = ((weights_even[2], size_even[2], 0) if parity_j2 == parity_match[2] else (weights_odd[2], size_odd[2], 1)) - for i0 in range(size0): - if self._halving_directions[0]: - pos0 = ((j0 + p[0] - 2 * i0 - offset0) // 2) % self._VNbasis[0] - else: - pos0 = j0 - for i1 in range(size1): - if self._halving_directions[1]: - pos1 = ((j1 + p[1] - 2 * i1 - offset1) // 2) % self._VNbasis[1] - else: - pos1 = j1 - for i2 in range(size2): - if self._halving_directions[2]: - pos2 = ((j2 + p[2] - 2 * i2 - offset2) // 2) % self._VNbasis[2] - else: - pos2 = j2 - out[j0, j1, j2] += weights0[i0] * weights1[i1] * weights2[i2] * v[pos0, pos1, pos2] - - else: # Vector case (Hcurl, Hdiv, H1H1H1) - for j0 in range(self._out_starts[h][0], self._out_ends[h][0] + 1): - parity_j0 = j0 % 2 - weights0, size0, offset0 = ((weights_even[0], size_even[0], 0) if parity_j0 == parity_match[0] else (weights_odd[0], size_odd[0], 1)) - for j1 in range(self._out_starts[h][1], self._out_ends[h][1] + 1): - parity_j1 = j1 % 2 - weights1, size1, offset1 = ((weights_even[1], size_even[1], 0) if parity_j1 == parity_match[1] else (weights_odd[1], size_odd[1], 1)) - for j2 in range(self._out_starts[h][2], self._out_ends[h][2] + 1): - parity_j2 = j2 % 2 - weights2, size2, offset2 = ((weights_even[2], size_even[2], 0) if parity_j2 == parity_match[2] else (weights_odd[2], size_odd[2], 1)) - for i0 in range(size0): - if self._halving_directions[0]: - pos0 = ((j0 + p[0] - 2 * i0 - offset0) // 2) % self._VNbasis[h][0] - else: - pos0 = j0 - for i1 in range(size1): - if self._halving_directions[1]: - pos1 = ((j1 + p[1] - 2 * i1 - offset1) // 2) % self._VNbasis[h][1] - else: - pos1 = j1 - for i2 in range(size2): - if self._halving_directions[2]: - pos2 = ((j2 + p[2] - 2 * i2 - offset2) // 2) % self._VNbasis[h][2] - else: - pos2 = j2 - out[h][j0, j1, j2] += weights0[i0] * weights1[i1] * weights2[i2] * v[h][pos0, pos1, pos2] - - return out - - def dot_H1(self, v, out): - return self._dot_helper(v, out, self._p, self._all_weights_even, self._all_weights_odd, self._all_size_even, self._all_size_odd) - - def dot_L2(self, v, out): - return self._dot_helper(v, out, self._pD, self._all_weights_evenD, self._all_weights_oddD, self._all_size_evenD, self._all_size_oddD) - - def dot_Hcurl(self, v, out): - out = self._dot_helper(v, out, [self._pD[0], self._p[1],self._p[2]], [self._all_weights_evenD[0],self._all_weights_even[1],self._all_weights_even[2]], [self._all_weights_oddD[0],self._all_weights_odd[1],self._all_weights_odd[2]], [self._all_size_evenD[0],self._all_size_even[1],self._all_size_even[2]], [self._all_size_oddD[0],self._all_size_odd[1],self._all_size_odd[2]], h=0) - out = self._dot_helper(v, out, [self._p[0], self._pD[1],self._p[2]], [self._all_weights_even[0],self._all_weights_evenD[1],self._all_weights_even[2]], [self._all_weights_odd[0],self._all_weights_oddD[1],self._all_weights_odd[2]], [self._all_size_even[0],self._all_size_evenD[1],self._all_size_even[2]], [self._all_size_odd[0],self._all_size_oddD[1],self._all_size_odd[2]], h=1) - return self._dot_helper(v, out, [self._p[0], self._p[1],self._pD[2]], [self._all_weights_even[0],self._all_weights_even[1],self._all_weights_evenD[2]], [self._all_weights_odd[0],self._all_weights_odd[1],self._all_weights_oddD[2]], [self._all_size_even[0],self._all_size_even[1],self._all_size_evenD[2]], [self._all_size_odd[0],self._all_size_odd[1],self._all_size_oddD[2]], h=2) - - - def dot_Hdiv(self, v, out): - out = self._dot_helper(v, out, [self._p[0], self._pD[1],self._pD[2]], [self._all_weights_even[0],self._all_weights_evenD[1],self._all_weights_evenD[2]], [self._all_weights_odd[0],self._all_weights_oddD[1],self._all_weights_oddD[2]], [self._all_size_even[0],self._all_size_evenD[1],self._all_size_evenD[2]], [self._all_size_odd[0],self._all_size_oddD[1],self._all_size_oddD[2]], h=0) - out = self._dot_helper(v, out, [self._pD[0], self._p[1],self._pD[2]], [self._all_weights_evenD[0],self._all_weights_even[1],self._all_weights_evenD[2]], [self._all_weights_oddD[0],self._all_weights_odd[1],self._all_weights_oddD[2]], [self._all_size_evenD[0],self._all_size_even[1],self._all_size_evenD[2]], [self._all_size_oddD[0],self._all_size_odd[1],self._all_size_oddD[2]], h=1) - return self._dot_helper(v, out, [self._pD[0], self._pD[1],self._p[2]], [self._all_weights_evenD[0],self._all_weights_evenD[1],self._all_weights_even[2]], [self._all_weights_oddD[0],self._all_weights_oddD[1],self._all_weights_odd[2]], [self._all_size_evenD[0],self._all_size_evenD[1],self._all_size_even[2]], [self._all_size_oddD[0],self._all_size_oddD[1],self._all_size_odd[2]], h=2) - - - def dot_H1H1H1(self, v, out): - out = self._dot_helper(v, out, self._p, self._all_weights_even, self._all_weights_odd, self._all_size_even, self._all_size_odd, h = 0) - self._dot_helper(v, out, self._p, self._all_weights_even, self._all_weights_odd, self._all_size_even, self._all_size_odd, h = 1) - return self._dot_helper(v, out, self._p, self._all_weights_even, self._all_weights_odd, self._all_size_even, self._all_size_odd, h = 2) - - - - def dot(self, v, out=None): - - assert isinstance(v, Vector) and v.space == self.domain - - if out is None: - out = self.codomain.zeros() - else: - assert isinstance(out, Vector) and out.space == self.codomain - - if self._V_name == 'H1' or self._V_name == 'L2': - for i0 in range(self._out_starts[0], self._out_ends[0]+1): - for i1 in range(self._out_starts[1], self._out_ends[1]+1): - for i2 in range(self._out_starts[2], self._out_ends[2]+1): - out[i0,i1,i2] = 0.0 - else: - for h in range(3): - for i0 in range(self._out_starts[h][0], self._out_ends[h][0]+1): - for i1 in range(self._out_starts[h][1], self._out_ends[h][1]+1): - for i2 in range(self._out_starts[h][2], self._out_ends[h][2]+1): - out[h][i0,i1,i2] = 0.0 - - dot_methods = { - "H1": self.dot_H1, - "L2": self.dot_L2, - "Hcurl": self.dot_Hcurl, - "Hdiv": self.dot_Hdiv, - "H1H1H1": self.dot_H1H1H1, - } - - return dot_methods.get(self._V_name)(v, out) - - def transpose(self, *, out = None): - if out is None: - out = RestrictionOperator(self._W, self._V) - else: - assert isinstance(out, RestrictionOperator) - assert out.domain is self.codomain - assert out.codomain is self.domain - - return out - - -def Compute_rate_of_smoothing(Nel, plist, spl_kind): - - from struphy.feec.utilities import create_equal_random_arrays - # get global communicator - comm = MPI.COMM_WORLD - rank = comm.Get_rank() - world_size = comm.Get_size() - - a1= 0.2 - #domain = HollowCylinder(a1 = a1) - domain = Cuboid() - sp_key = '1' - derham = [] - mass_ops = [] - A = [] - - epsilon = 0.0002 - derham.append(Derham([Nel[0],Nel[1],Nel[2]], plist, spl_kind, comm=comm, local_projectors=False)) - mass_ops.append(WeightedMassOperators(derham[0], domain)) - A.append(epsilon*derham[0].curl.T @ mass_ops[0].M2 @ derham[0].curl + mass_ops[0].M1) - - - field_star = derham[0].create_field('fh', 'Hcurl') - field_aprox = derham[0].create_field('fh', 'Hcurl') - - #We are gonna use the method of manufacture solutions to determine the behaviour of our error - #Our exact solution is u_star - u_stararr, u_star = create_equal_random_arrays(derham[0].Vh_fem[sp_key], seed=45) - #We compute the rhs - b = A[0].dot(u_star) - field_star.vector = u_star - - #We store the number of itterations - N_itter = [] - #We define a set of point to evaluate the exact solution and the aproximated one - pointsx = np.linspace(0.0,1.0,100) - pointsy = np.linspace(0.0,1.0,100) - pointsz = np.array([0.0,0.5]) - dx = pointsx[1] - pointsx[0] # Spacing between points - dy = pointsy[1] - pointsy[0] # Spacing between points - dz = pointsz[1] - pointsz[0] # Spacing between points - X, Y, Z = np.meshgrid(pointsx, pointsy, pointsz, indexing="ij") - - - - #We define a list where to store the arrays with the values of the errors for each aproximation and vector component - errorsx = [] - errorsy = [] - errorsz = [] - - #Fourier Transform coefficients of the error function for each vector component - fourier_coeffsx = [] - fourier_coeffs_shiftedx = [] - fourier_coeffsy = [] - fourier_coeffs_shiftedy = [] - fourier_coeffsz = [] - fourier_coeffs_shiftedz = [] - Max_iter = 10 - for i in range(Max_iter): - solver = inverse(A[0],'cg', maxiter = int(i+2)) - u = solver.dot(b) - field_aprox.vector = u - errorsx.append(field_star(X,Y,Z)[0]-field_aprox(X,Y,Z)[0]) - errorsy.append(field_star(X,Y,Z)[1]-field_aprox(X,Y,Z)[1]) - errorsz.append(field_star(X,Y,Z)[2]-field_aprox(X,Y,Z)[2]) - N_itter.append(solver._info['niter']) - - # Compute Fourier Transform coefficients - fourier_coeffsx.append(np.fft.fftn(errorsx[-1])) - fourier_coeffs_shiftedx.append(np.fft.fftshift(fourier_coeffsx[-1])) - fourier_coeffsy.append(np.fft.fftn(errorsy[-1])) - fourier_coeffs_shiftedy.append(np.fft.fftshift(fourier_coeffsy[-1])) - fourier_coeffsz.append(np.fft.fftn(errorsz[-1])) - fourier_coeffs_shiftedz.append(np.fft.fftshift(fourier_coeffsz[-1])) - - # Compute corresponding frequencies - freqsx = np.fft.fftshift(np.fft.fftfreq(len(pointsx), d=dx)) - freqsy = np.fft.fftshift(np.fft.fftfreq(len(pointsy), d=dy)) - freqsz = np.fft.fftshift(np.fft.fftfreq(len(pointsz), d=dz)) - # Create a 3D meshgrid of frequencies - #Fx, Fy, Fz = np.meshgrid(freq_x[-1], freq_y[-1], freq_z[-1], indexing="ij") - - #Now we compute the magnitude of the high frequencies between two interations - def get_smoothing_rate(freqsx,freqsy,freqsz, coeff_old, coeff_new, hx,hy,hz): - smoothing_rate = -1.0 - freq = np.zeros(3,dtype=float) - for ix in range(len(freqsx)): - for iy in range(len(freqsy)): - for iz in range(len(freqsz)): - if((abs(freqsx[ix])> 1.0 / (4.0*hx) or abs(freqsy[iy])> 1.0 / (4.0*hy) or abs(freqsz[iz])> 1.0 / (4.0*hz) ) and abs(coeff_new[ix,iy,iz])>10.0**-6): - smoothing_rate = max(smoothing_rate,abs(coeff_new[ix,iy,iz]/coeff_old[ix,iy,iz])) - freq[0] = freqsx[ix] - freq[1] = freqsx[iy] - freq[2] = freqsx[iz] - - return smoothing_rate, freq - - - for i in range(Max_iter-1): - - smoothing_rate, freq = get_smoothing_rate(freqsx,freqsy,freqsz,fourier_coeffs_shiftedx[i],fourier_coeffs_shiftedx[i+1],1.0/Nel[0],1.0/Nel[1],0.00000001) - - print("#######################") - print(f'{i =}') - print(f'{smoothing_rate =}') - print(f'{freq =}') - print("#######################") - - #for i in range(Max_iter-1): - #smoothing_rate_1d(freqsx, coeff_old, coeff_new, h) - - -def Visualized_high_frequency_dampening(Nel, plist, spl_kind): - - from struphy.feec.utilities import create_equal_random_arrays - # get global communicator - comm = MPI.COMM_WORLD - rank = comm.Get_rank() - world_size = comm.Get_size() - - a1= 0.2 - #domain = HollowCylinder(a1 = a1) - domain = Cuboid() - sp_key = '2' - sp_id = 'Hdiv' - derham = [] - mass_ops = [] - A = [] - - epsilon = 1.0 - derham.append(Derham([Nel[0],Nel[1],Nel[2]], plist, spl_kind, comm=comm, local_projectors=False)) - mass_ops.append(WeightedMassOperators(derham[0], domain)) - #Poisson - #A.append(derham[0].grad.T @ mass_ops[0].M1 @ derham[0].grad) - #Hall-ish - #A.append(epsilon*derham[0].curl.T @ mass_ops[0].M2 @ derham[0].curl + mass_ops[0].M1) - #Hall - #A.append(mass_ops[0].M1 -epsilon*derham[0].curl.T @ mass_ops[0].M2 @ derham[0].curl) - #Shear-Alfven-ish - #A.append(mass_ops[0].M2 -epsilon* mass_ops[0].M2@ derham[0].curl @ mass_ops[0].M1 @ derham[0].curl.T @ mass_ops[0].M2) - #Shear-Alfven - pc_class = getattr(preconditioner,"MassMatrixPreconditioner") - pc = pc_class(mass_ops[0].M1) - M1_inv = inverse( - mass_ops[0].M1, - "pcg", - pc=pc, - maxiter=3000, - verbose=False, - ) - A.append(mass_ops[0].M2 -epsilon* mass_ops[0].M2@ derham[0].curl @ M1_inv @ derham[0].curl.T @ mass_ops[0].M2) - #Shear-Alfven-v2 - #A.append(mass_ops[0].M2 -epsilon* derham[0].curl @ M1_inv @ derham[0].curl.T) - - - - - field_star = derham[0].create_field('fh', sp_id) - field_aprox = derham[0].create_field('fh', sp_id) - - #We are gonna use the method of manufacture solutions to determine the behaviour of our error - #Our exact solution is u_star - u_stararr, u_star = create_equal_random_arrays(derham[0].Vh_fem[sp_key], seed=45) - #We compute the rhs - b = A[0].dot(u_star) - field_star.vector = u_star - - #Turn b into an array - #barr = remove_padding(derham[0].Vh_fem[sp_key],b) - #Turn A[0] into an array - #Aarr = A[0].toarray() - - #Gauss = will_gauss_seidel_converge(Aarr, verbose=True) - #print(f"{Gauss = }") - - - #We store the number of itterations - N_itter = [] - #We define a set of point to evaluate the exact solution and the aproximated one - pointsx = np.linspace(0.0,1.0,Nel[0]) - #pointsx = np.array([0.0,0.5]) - #pointsy = np.linspace(0.0,1.0,Nel[1]) - pointsy = np.array([0.0,0.5]) - pointsz = np.array([0.0,0.5]) - dx = pointsx[1] - pointsx[0] # Spacing between points - dy = pointsy[1] - pointsy[0] # Spacing between points - dz = pointsz[1] - pointsz[0] # Spacing between points - X, Y, Z = np.meshgrid(pointsx, pointsy, pointsz, indexing="ij") - - - - #We define a list where to store the arrays with the values of the errors for each aproximation and vector component - errorsx = [] - errorsy = [] - errorsz = [] - - #Fourier Transform coefficients of the error function for each vector component - fourier_coeffsx = [] - fourier_coeffs_shiftedx = [] - fourier_coeffsy = [] - fourier_coeffs_shiftedy = [] - fourier_coeffsz = [] - fourier_coeffs_shiftedz = [] - max_iter_list = [3] - #max_iter_list = [0,2,3,4,5,6,7,8,9,10] - Number_of_iter = len(max_iter_list) - for i in range(Number_of_iter): - - if(max_iter_list[i]>1): - - solver = inverse(A[0],'cg', maxiter = max_iter_list[i], tol = 10.0**-6) - u = solver.dot(b) - #uarr, itter = jacobi(Aarr,barr,max_iter=max_iter_list[i]) - #u = from_array_to_psydac(uarr, derham[0].Vh_fem[sp_key]) - N_itter.append(solver._info['niter']) - #N_itter.append(itter) - - else: - u = derham[0].Vh[derham[0].space_to_form[sp_id]].zeros() - N_itter.append(0) - - - field_aprox.vector = u - #errorsx.append(field_star(X,Y,Z)-field_aprox(X,Y,Z)) - errorsx.append(field_star(X,Y,Z)[0]-field_aprox(X,Y,Z)[0]) - errorsy.append(field_star(X,Y,Z)[1]-field_aprox(X,Y,Z)[1]) - errorsz.append(field_star(X,Y,Z)[2]-field_aprox(X,Y,Z)[2]) - - - # Compute Fourier Transform coefficients - fourier_coeffsx.append(np.fft.fftn(errorsx[-1])) - fourier_coeffs_shiftedx.append(np.fft.fftshift(fourier_coeffsx[-1])) - fourier_coeffsy.append(np.fft.fftn(errorsy[-1])) - fourier_coeffs_shiftedy.append(np.fft.fftshift(fourier_coeffsy[-1])) - fourier_coeffsz.append(np.fft.fftn(errorsz[-1])) - fourier_coeffs_shiftedz.append(np.fft.fftshift(fourier_coeffsz[-1])) - - # Compute corresponding frequencies - freqsx = np.fft.fftshift(np.fft.fftfreq(len(pointsx), d=dx)) - freqsy = np.fft.fftshift(np.fft.fftfreq(len(pointsy), d=dy)) - freqsz = np.fft.fftshift(np.fft.fftfreq(len(pointsz), d=dz)) - # Create a 3D meshgrid of frequencies - #Fx, Fy, Fz = np.meshgrid(freq_x[-1], freq_y[-1], freq_z[-1], indexing="ij") - - #Now we compute the magnitude of the high frequencies between two interations - def get_magnitude_maximum_nasty_frequency_scalar(freqsx,freqsy,freqsz, coeff_new,hx,hy,hz): - value= 0.0 - freq = np.zeros(3,dtype=float) - for ix in range(len(freqsx)): - for iy in range(len(freqsy)): - for iz in range(len(freqsz)): - if((abs(freqsx[ix])> 1.0/(4.0*hx) or abs(freqsy[iy])> 1.0/(4.0*hy) or abs(freqsz[iz])> 1.0/(4.0*hz) ) and abs(coeff_new[ix,iy,iz])>10.0**-6): - if(abs(coeff_new[ix,iy,iz]) > value ): - value = abs(coeff_new[ix,iy,iz]) - freq[0] = freqsx[ix] - freq[1] = freqsy[iy] - freq[2] = freqsz[iz] - - return value, freq - - - - - def get_magnitude_maximum_nasty_frequency(freqsx,freqsy,freqsz, coeff_newx, coeff_newy, coeff_newz,hx,hy,hz): - value= 0.0 - freq = np.zeros(3,dtype=float) - for ix in range(len(freqsx)): - for iy in range(len(freqsy)): - for iz in range(len(freqsz)): - if((abs(freqsx[ix])> 1.0/(4.0*hx) or abs(freqsy[iy])> 1.0/(4.0*hy) or abs(freqsz[iz])> 1.0/(4.0*hz) ) and (abs(coeff_newx[ix,iy,iz])>10.0**-6 or abs(coeff_newy[ix,iy,iz])>10.0**-6 or abs(coeff_newz[ix,iy,iz])>10.0**-6)): - if(abs(coeff_newx[ix,iy,iz]) > value and abs(coeff_newx[ix,iy,iz]) >= abs(coeff_newy[ix,iy,iz]) and abs(coeff_newx[ix,iy,iz])>= abs(coeff_newz[ix,iy,iz])): - value = abs(coeff_newx[ix,iy,iz]) - freq[0] = freqsx[ix] - freq[1] = freqsy[iy] - freq[2] = freqsz[iz] - - elif(abs(coeff_newy[ix,iy,iz]) > value and abs(coeff_newy[ix,iy,iz]) >= abs(coeff_newx[ix,iy,iz]) and abs(coeff_newy[ix,iy,iz])>= abs(coeff_newz[ix,iy,iz])): - value = abs(coeff_newy[ix,iy,iz]) - freq[0] = freqsx[ix] - freq[1] = freqsy[iy] - freq[2] = freqsz[iz] - - elif(abs(coeff_newz[ix,iy,iz]) > value and abs(coeff_newz[ix,iy,iz]) >= abs(coeff_newx[ix,iy,iz]) and abs(coeff_newz[ix,iy,iz])>= abs(coeff_newy[ix,iy,iz])): - value = abs(coeff_newz[ix,iy,iz]) - freq[0] = freqsx[ix] - freq[1] = freqsy[iy] - freq[2] = freqsz[iz] - - return value, freq - - - magnitudes = [] - bad_frequencies = [] - - for i in range(Number_of_iter): - #magnitude, freq = get_magnitude_maximum_nasty_frequency_scalar(freqsx,freqsy,freqsz,fourier_coeffs_shiftedx[i],1.0/Nel[0],1.0/Nel[1],0.000001) - magnitude, freq = get_magnitude_maximum_nasty_frequency(freqsx,freqsy,freqsz,fourier_coeffs_shiftedx[i],fourier_coeffs_shiftedy[i], fourier_coeffs_shiftedz[i],1.0/Nel[0],0.000001,0.000001) - magnitudes.append(magnitude) - bad_frequencies.append(freq) - - - print(f'{Nel[0] = }') - print(f'{Nel[1] = }') - print("magnitudes") - for i in magnitudes: - print(i) - print("N_itter") - for i in N_itter: - print(i) - print("bad_frequencies") - for i in bad_frequencies: - print(i[0:2]) - - - - plt.figure() - plt.scatter(N_itter,magnitudes) - #plt.yscale("log") - plt.show() - - -def Visualized_all_frequencies_dampening(Nel, plist, spl_kind, Is): - - from struphy.feec.utilities import create_equal_random_arrays - # get global communicator - comm = MPI.COMM_WORLD - rank = comm.Get_rank() - world_size = comm.Get_size() - - a1= 0.2 - #domain = HollowCylinder(a1 = a1) - domain = Cuboid() - model = 'Shear-Alfven' - smoother = 'cg' - derham = [] - mass_ops = [] - A = [] - - epsilon = 10.0**-6.0 - derham.append(Derham([Nel[0],Nel[1],Nel[2]], plist, spl_kind, comm=comm, local_projectors=False)) - mass_ops.append(WeightedMassOperators(derham[0], domain)) - if(model == "Poisson"): - sp_key = '0' - sp_id = 'H1' - #Poisson - A.append(derham[0].grad.T @ mass_ops[0].M1 @ derham[0].grad) - #Hall-ish - #A.append(epsilon*derham[0].curl.T @ mass_ops[0].M2 @ derham[0].curl + mass_ops[0].M1) - elif(model == "Hall"): - sp_key = '1' - sp_id = 'Hcurl' - #Hall - A.append(mass_ops[0].M1 -epsilon*derham[0].curl.T @ mass_ops[0].M2 @ derham[0].curl) - #Shear-Alfven-ish - #A.append(mass_ops[0].M2 -epsilon* mass_ops[0].M2@ derham[0].curl @ mass_ops[0].M1 @ derham[0].curl.T @ mass_ops[0].M2) - elif(model == 'Shear-Alfven' or model == 'Shear-Alfven-v2'): - sp_key = '2' - sp_id = 'Hdiv' - pc_class = getattr(preconditioner,"MassMatrixPreconditioner") - pc = pc_class(mass_ops[0].M1) - M1_inv = inverse( - mass_ops[0].M1, - "pcg", - pc=pc, - maxiter=3000, - verbose=False, - ) - #Shear-Alfven - if (model == "Shear-Alfven"): - A.append(mass_ops[0].M2 -epsilon* mass_ops[0].M2@ derham[0].curl @ M1_inv @ derham[0].curl.T @ mass_ops[0].M2) - else: - #Shear-Alfven-v2 - A.append(mass_ops[0].M2 -epsilon* derham[0].curl @ M1_inv @ derham[0].curl.T) - - field_star = derham[0].create_field('fh', sp_id) - field_aprox = derham[0].create_field('fh', sp_id) - - #We are gonna use the method of manufacture solutions to determine the behaviour of our error - #Our exact solution is u_star - u_stararr, u_star = create_equal_random_arrays(derham[0].Vh_fem[sp_key], seed=45) - #We compute the rhs - b = A[0].dot(u_star) - field_star.vector = u_star - - #Turn b into an array - #barr = remove_padding(derham[0].Vh_fem[sp_key],b) - #Turn A[0] into an array - #Aarr = A[0].toarray() - - #Gauss = will_gauss_seidel_converge(Aarr, verbose=True) - #print(f"{Gauss = }") - - - #We store the number of itterations - N_itter = [] - #We define a set of point to evaluate the exact solution and the aproximated one - pointsx = np.linspace(0.0,1.0,Nel[0]) - #pointsx = np.array([0.0,0.5]) - pointsy = np.linspace(0.0,1.0,Nel[1]) - #pointsy = np.array([0.0,0.5]) - pointsz = np.array([0.0,0.5]) - dx = pointsx[1] - pointsx[0] # Spacing between points - dy = pointsy[1] - pointsy[0] # Spacing between points - dz = pointsz[1] - pointsz[0] # Spacing between points - X, Y, Z = np.meshgrid(pointsx, pointsy, pointsz, indexing="ij") - - - - #We define a list where to store the arrays with the values of the errors for each aproximation and vector component - errorsx = [] - errorsy = [] - errorsz = [] - - #Fourier Transform coefficients of the error function for each vector component - fourier_coeffsx = [] - fourier_coeffs_shiftedx = [] - fourier_coeffsy = [] - fourier_coeffs_shiftedy = [] - fourier_coeffsz = [] - fourier_coeffs_shiftedz = [] - max_iter_list = [0,Is, int(10*Is)] - #max_iter_list = [0,2,3,4,5,6,7,8,9,10] - Number_of_iter = len(max_iter_list) - for i in range(Number_of_iter): - - if(max_iter_list[i]>1): - - solver = inverse(A[0],smoother, maxiter = max_iter_list[i], tol = 10.0**-6) - u = solver.dot(b) - #uarr, itter = jacobi(Aarr,barr,max_iter=max_iter_list[i]) - #u = from_array_to_psydac(uarr, derham[0].Vh_fem[sp_key]) - N_itter.append(solver._info['niter']) - #N_itter.append(itter) - - else: - u = derham[0].Vh[derham[0].space_to_form[sp_id]].zeros() - N_itter.append(0) - - - field_aprox.vector = u - if(model == "Poisson"): - errorsx.append(field_star(X,Y,Z)-field_aprox(X,Y,Z)) - else: - errorsx.append(field_star(X,Y,Z)[0]-field_aprox(X,Y,Z)[0]) - errorsy.append(field_star(X,Y,Z)[1]-field_aprox(X,Y,Z)[1]) - errorsz.append(field_star(X,Y,Z)[2]-field_aprox(X,Y,Z)[2]) - - # Compute Fourier Transform coefficients - fourier_coeffsx.append(np.fft.fftn(errorsx[-1])) - fourier_coeffs_shiftedx.append(np.fft.fftshift(fourier_coeffsx[-1])) - if(model != "Poisson"): - fourier_coeffsy.append(np.fft.fftn(errorsy[-1])) - fourier_coeffs_shiftedy.append(np.fft.fftshift(fourier_coeffsy[-1])) - fourier_coeffsz.append(np.fft.fftn(errorsz[-1])) - fourier_coeffs_shiftedz.append(np.fft.fftshift(fourier_coeffsz[-1])) - - # Compute corresponding frequencies - freqsx = np.fft.fftshift(np.fft.fftfreq(len(pointsx), d=dx)) - freqsy = np.fft.fftshift(np.fft.fftfreq(len(pointsy), d=dy)) - freqsz = np.fft.fftshift(np.fft.fftfreq(len(pointsz), d=dz)) - # Create a 3D meshgrid of frequencies - #Fx, Fy, Fz = np.meshgrid(freq_x[-1], freq_y[-1], freq_z[-1], indexing="ij") - - #i also want to visiualize the maximum freqeuncies the next 3 coarser grids can handle - hx = 1/Nel[0] - next_max_frequenciesx = [1.0/(4.0*hx),1.0/(8.0*hx),1.0/(16.0*hx)] - hy = 1/Nel[1] - next_max_frequenciesy = [1.0/(4.0*hy),1.0/(8.0*hy),1.0/(16.0*hy)] - #I get the amplitude of the maximum initial error wave - max_amplitude = np.max(np.abs(fourier_coeffs_shiftedx[0][:,1,1])) - - if(model == "Poisson"): - plt.figure(figsize=(10,10)) - plt.title("Amplitude of error waves for " + model + " with " +smoother+" solver, with resolution " + str(Nel)+".",size = 10) - plt.scatter(freqsx,np.abs(fourier_coeffs_shiftedx[0][:,1,1]), label = "Initial") - plt.scatter(freqsx,np.abs(fourier_coeffs_shiftedx[1][:,1,1]), label = "After "+str(max_iter_list[1])+" iterations ") - plt.scatter(freqsx,np.abs(fourier_coeffs_shiftedx[2][:,1,1]), label = "After "+str(max_iter_list[2])+" iterations ") - plt.plot(next_max_frequenciesx[0]*np.ones(100),np.linspace(0,max_amplitude,100),linestyle = 'dashed', color ='blue',label = "Max freuency level - 1.") - plt.plot(-next_max_frequenciesx[0]*np.ones(100),np.linspace(0,max_amplitude,100),linestyle = 'dashed', color ='blue') - plt.plot(next_max_frequenciesx[1]*np.ones(100),np.linspace(0,max_amplitude,100),linestyle = 'dashed', color ='purple',label = "Max freuency level - 2.") - plt.plot(-next_max_frequenciesx[1]*np.ones(100),np.linspace(0,max_amplitude,100),linestyle = 'dashed', color ='purple') - plt.plot(next_max_frequenciesx[2]*np.ones(100),np.linspace(0,max_amplitude,100),linestyle = 'dashed', color ='red',label = "Max freuency level - 3.") - plt.plot(-next_max_frequenciesx[2]*np.ones(100),np.linspace(0,max_amplitude,100),linestyle = 'dashed', color ='red') - plt.xlabel("Frequency") - plt.ylabel("Amplitude") - plt.legend() - #plt.ylim(0,30) - plt.yscale("log") - plt.savefig("./Wave-Amplitude-Cuboid-"+model+"-"+smoother+"-"+str(Nel)+".pdf") - plt.close() - - - #We also want to see the relative difference between the initial amplitude and the amplitude after smoothing. - #We compute abs(A_0)-abs(A_i)/abs(A_0), this tell us which percentage of the initial amplitude has been eliminated - #A values of 0 means it all still remains, a value of 1 means we eliminated all the amplitude. And a negative value - #means that the amplitude increased. - eliminated_portion = [] - for i in range(2): - aux = (np.abs(fourier_coeffs_shiftedx[0][:,1,1]) - np.abs(fourier_coeffs_shiftedx[i+1][:,1,1]))/np.abs(fourier_coeffs_shiftedx[0][:,1,1]) - eliminated_portion.append(aux) - - mp = np.min(eliminated_portion[0]) - aux = np.min(eliminated_portion[1]) - mp = min(mp,aux) - - plt.figure(figsize=(10,10)) - plt.title("Eliminated portion of error waves for " + model + " with " +smoother+" solver, with resolution " + str(Nel)+".",size = 10) - plt.scatter(freqsx,eliminated_portion[0], label = "Portion of amplitude eliminated after " + str(max_iter_list[1])) - plt.scatter(freqsx,eliminated_portion[1], label = "Portion of amplitude eliminated after " + str(max_iter_list[2])) - plt.plot(next_max_frequenciesx[0]*np.ones(100),np.linspace(mp,1,100),linestyle = 'dashed', color ='blue',label = "Max freuency level - 1.") - plt.plot(-next_max_frequenciesx[0]*np.ones(100),np.linspace(mp,1,100),linestyle = 'dashed', color ='blue') - plt.plot(next_max_frequenciesx[1]*np.ones(100),np.linspace(mp,1,100),linestyle = 'dashed', color ='purple',label = "Max freuency level - 2.") - plt.plot(-next_max_frequenciesx[1]*np.ones(100),np.linspace(mp,1,100),linestyle = 'dashed', color ='purple') - plt.plot(next_max_frequenciesx[2]*np.ones(100),np.linspace(mp,1,100),linestyle = 'dashed', color ='red',label = "Max freuency level - 3.") - plt.plot(-next_max_frequenciesx[2]*np.ones(100),np.linspace(mp,1,100),linestyle = 'dashed', color ='red') - plt.xlabel("Frequency") - plt.ylabel("Eliminated portion of the amplitude.") - plt.legend() - #plt.ylim(0,30) - #plt.yscale("log") - plt.savefig("./Eliminated-Portion-Cuboid-"+model+"-"+smoother+"-"+str(Nel)+".pdf") - plt.close() - - else: - #I get the amplitude of the maximum initial error wave - max_amplitudey = np.max(np.abs(fourier_coeffs_shiftedy[0][:,1,1])) - #I get the amplitude of the maximum initial error wave - max_amplitudez = np.max(np.abs(fourier_coeffs_shiftedz[0][:,1,1])) - - plt.figure(figsize=(10,10)) - plt.title("X-component. Epsilon 10^-6 Amplitude of error waves for " + model + " with " +smoother+" solver, with resolution " + str(Nel)+".",size = 10) - plt.scatter(freqsx,np.abs(fourier_coeffs_shiftedx[0][:,1,1]), label = "Initial") - plt.scatter(freqsx,np.abs(fourier_coeffs_shiftedx[1][:,1,1]), label = "After "+str(max_iter_list[1])+" iterations ") - plt.scatter(freqsx,np.abs(fourier_coeffs_shiftedx[2][:,1,1]), label = "After "+str(max_iter_list[2])+" iterations ") - plt.plot(next_max_frequenciesx[0]*np.ones(100),np.linspace(0,max_amplitude,100),linestyle = 'dashed', color ='blue',label = "Max freuency level - 1.") - plt.plot(-next_max_frequenciesx[0]*np.ones(100),np.linspace(0,max_amplitude,100),linestyle = 'dashed', color ='blue') - plt.plot(next_max_frequenciesx[1]*np.ones(100),np.linspace(0,max_amplitude,100),linestyle = 'dashed', color ='purple',label = "Max freuency level - 2.") - plt.plot(-next_max_frequenciesx[1]*np.ones(100),np.linspace(0,max_amplitude,100),linestyle = 'dashed', color ='purple') - plt.plot(next_max_frequenciesx[2]*np.ones(100),np.linspace(0,max_amplitude,100),linestyle = 'dashed', color ='red',label = "Max freuency level - 3.") - plt.plot(-next_max_frequenciesx[2]*np.ones(100),np.linspace(0,max_amplitude,100),linestyle = 'dashed', color ='red') - plt.xlabel("Frequency") - plt.ylabel("Amplitude") - plt.legend() - #plt.ylim(0,30) - plt.yscale("log") - plt.savefig("./Wave-Amplitude-Cuboid-"+model+"-"+smoother+"-"+str(Nel)+"X-component-Epsilon-10-6.pdf") - plt.close() - - plt.figure(figsize=(10,10)) - plt.title("Y-component. Epsilon 10^-6 Amplitude of error waves for " + model + " with " +smoother+" solver, with resolution " + str(Nel)+".",size = 10) - plt.scatter(freqsx,np.abs(fourier_coeffs_shiftedy[0][:,1,1]), label = "Initial") - plt.scatter(freqsx,np.abs(fourier_coeffs_shiftedy[1][:,1,1]), label = "After "+str(max_iter_list[1])+" iterations ") - plt.scatter(freqsx,np.abs(fourier_coeffs_shiftedy[2][:,1,1]), label = "After "+str(max_iter_list[2])+" iterations ") - plt.plot(next_max_frequenciesx[0]*np.ones(100),np.linspace(0,max_amplitudey,100),linestyle = 'dashed', color ='blue',label = "Max freuency level - 1.") - plt.plot(-next_max_frequenciesx[0]*np.ones(100),np.linspace(0,max_amplitudey,100),linestyle = 'dashed', color ='blue') - plt.plot(next_max_frequenciesx[1]*np.ones(100),np.linspace(0,max_amplitudey,100),linestyle = 'dashed', color ='purple',label = "Max freuency level - 2.") - plt.plot(-next_max_frequenciesx[1]*np.ones(100),np.linspace(0,max_amplitudey,100),linestyle = 'dashed', color ='purple') - plt.plot(next_max_frequenciesx[2]*np.ones(100),np.linspace(0,max_amplitudey,100),linestyle = 'dashed', color ='red',label = "Max freuency level - 3.") - plt.plot(-next_max_frequenciesx[2]*np.ones(100),np.linspace(0,max_amplitudey,100),linestyle = 'dashed', color ='red') - plt.xlabel("Frequency") - plt.ylabel("Amplitude") - plt.legend() - #plt.ylim(0,30) - plt.yscale("log") - plt.savefig("./Wave-Amplitude-Cuboid-"+model+"-"+smoother+"-"+str(Nel)+"Y-component-Epsilon-10-6.pdf") - plt.close() - - - plt.figure(figsize=(10,10)) - plt.title("Z-component. Epsilon 10^-6 Amplitude of error waves for " + model + " with " +smoother+" solver, with resolution " + str(Nel)+".",size = 10) - plt.scatter(freqsx,np.abs(fourier_coeffs_shiftedz[0][:,1,1]), label = "Initial") - plt.scatter(freqsx,np.abs(fourier_coeffs_shiftedz[1][:,1,1]), label = "After "+str(max_iter_list[1])+" iterations ") - plt.scatter(freqsx,np.abs(fourier_coeffs_shiftedz[2][:,1,1]), label = "After "+str(max_iter_list[2])+" iterations ") - plt.plot(next_max_frequenciesx[0]*np.ones(100),np.linspace(0,max_amplitudez,100),linestyle = 'dashed', color ='blue',label = "Max freuency level - 1.") - plt.plot(-next_max_frequenciesx[0]*np.ones(100),np.linspace(0,max_amplitudez,100),linestyle = 'dashed', color ='blue') - plt.plot(next_max_frequenciesx[1]*np.ones(100),np.linspace(0,max_amplitudez,100),linestyle = 'dashed', color ='purple',label = "Max freuency level - 2.") - plt.plot(-next_max_frequenciesx[1]*np.ones(100),np.linspace(0,max_amplitudez,100),linestyle = 'dashed', color ='purple') - plt.plot(next_max_frequenciesx[2]*np.ones(100),np.linspace(0,max_amplitudez,100),linestyle = 'dashed', color ='red',label = "Max freuency level - 3.") - plt.plot(-next_max_frequenciesx[2]*np.ones(100),np.linspace(0,max_amplitudez,100),linestyle = 'dashed', color ='red') - plt.xlabel("Frequency") - plt.ylabel("Amplitude") - plt.legend() - #plt.ylim(0,30) - plt.yscale("log") - plt.savefig("./Wave-Amplitude-Cuboid-"+model+"-"+smoother+"-"+str(Nel)+"Z-component-Epsilon-10-6.pdf") - plt.close() - - - #We also want to see the relative difference between the initial amplitude and the amplitude after smoothing. - #We compute abs(A_0)-abs(A_i)/abs(A_0), this tell us which percentage of the initial amplitude has been eliminated - #A values of 0 means it all still remains, a value of 1 means we eliminated all the amplitude. And a negative value - #means that the amplitude increased. - eliminated_portionx = [] - for i in range(2): - aux = (np.abs(fourier_coeffs_shiftedx[0][:,1,1]) - np.abs(fourier_coeffs_shiftedx[i+1][:,1,1]))/np.abs(fourier_coeffs_shiftedx[0][:,1,1]) - eliminated_portionx.append(aux) - - mp = np.min(eliminated_portionx[0]) - aux = np.min(eliminated_portionx[1]) - mp = min(mp,aux) - - plt.figure(figsize=(10,10)) - plt.title("X-component. Epsilon 10^-6 Eliminated portion of error waves for " + model + " with " +smoother+" solver, with resolution " + str(Nel)+".",size = 10) - plt.scatter(freqsx,eliminated_portionx[0], label = "Portion of amplitude eliminated after " + str(max_iter_list[1])) - plt.scatter(freqsx,eliminated_portionx[1], label = "Portion of amplitude eliminated after " + str(max_iter_list[2])) - plt.plot(next_max_frequenciesx[0]*np.ones(100),np.linspace(mp,1,100),linestyle = 'dashed', color ='blue',label = "Max freuency level - 1.") - plt.plot(-next_max_frequenciesx[0]*np.ones(100),np.linspace(mp,1,100),linestyle = 'dashed', color ='blue') - plt.plot(next_max_frequenciesx[1]*np.ones(100),np.linspace(mp,1,100),linestyle = 'dashed', color ='purple',label = "Max freuency level - 2.") - plt.plot(-next_max_frequenciesx[1]*np.ones(100),np.linspace(mp,1,100),linestyle = 'dashed', color ='purple') - plt.plot(next_max_frequenciesx[2]*np.ones(100),np.linspace(mp,1,100),linestyle = 'dashed', color ='red',label = "Max freuency level - 3.") - plt.plot(-next_max_frequenciesx[2]*np.ones(100),np.linspace(mp,1,100),linestyle = 'dashed', color ='red') - plt.xlabel("Frequency") - plt.ylabel("Eliminated portion of the amplitude.") - plt.legend() - #plt.ylim(0,30) - #plt.yscale("log") - plt.savefig("./Eliminated-Portion-Cuboid-"+model+"-"+smoother+"-"+str(Nel)+"X-component-Epsilon-10-6.pdf") - plt.close() - - - eliminated_portiony = [] - for i in range(2): - aux = (np.abs(fourier_coeffs_shiftedy[0][:,1,1]) - np.abs(fourier_coeffs_shiftedy[i+1][:,1,1]))/np.abs(fourier_coeffs_shiftedy[0][:,1,1]) - eliminated_portiony.append(aux) - - mp = np.min(eliminated_portiony[0]) - aux = np.min(eliminated_portiony[1]) - mp = min(mp,aux) - - plt.figure(figsize=(10,10)) - plt.title("Y-component. Epsilon 10^-6 Eliminated portion of error waves for " + model + " with " +smoother+" solver, with resolution " + str(Nel)+".",size = 10) - plt.scatter(freqsx,eliminated_portiony[0], label = "Portion of amplitude eliminated after " + str(max_iter_list[1])) - plt.scatter(freqsx,eliminated_portiony[1], label = "Portion of amplitude eliminated after " + str(max_iter_list[2])) - plt.plot(next_max_frequenciesx[0]*np.ones(100),np.linspace(mp,1,100),linestyle = 'dashed', color ='blue',label = "Max freuency level - 1.") - plt.plot(-next_max_frequenciesx[0]*np.ones(100),np.linspace(mp,1,100),linestyle = 'dashed', color ='blue') - plt.plot(next_max_frequenciesx[1]*np.ones(100),np.linspace(mp,1,100),linestyle = 'dashed', color ='purple',label = "Max freuency level - 2.") - plt.plot(-next_max_frequenciesx[1]*np.ones(100),np.linspace(mp,1,100),linestyle = 'dashed', color ='purple') - plt.plot(next_max_frequenciesx[2]*np.ones(100),np.linspace(mp,1,100),linestyle = 'dashed', color ='red',label = "Max freuency level - 3.") - plt.plot(-next_max_frequenciesx[2]*np.ones(100),np.linspace(mp,1,100),linestyle = 'dashed', color ='red') - plt.xlabel("Frequency") - plt.ylabel("Eliminated portion of the amplitude.") - plt.legend() - #plt.ylim(0,30) - #plt.yscale("log") - plt.savefig("./Eliminated-Portion-Cuboid-"+model+"-"+smoother+"-"+str(Nel)+"Y-component-Epsilon-10-6.pdf") - plt.close() - - - eliminated_portionz = [] - for i in range(2): - aux = (np.abs(fourier_coeffs_shiftedz[0][:,1,1]) - np.abs(fourier_coeffs_shiftedz[i+1][:,1,1]))/np.abs(fourier_coeffs_shiftedz[0][:,1,1]) - eliminated_portionz.append(aux) - - mp = np.min(eliminated_portionz[0]) - aux = np.min(eliminated_portionz[1]) - mp = min(mp,aux) - - plt.figure(figsize=(10,10)) - plt.title("Z-component. Epsilon 10^-6 Eliminated portion of error waves for " + model + " with " +smoother+" solver, with resolution " + str(Nel)+".",size = 10) - plt.scatter(freqsx,eliminated_portionz[0], label = "Portion of amplitude eliminated after " + str(max_iter_list[1])) - plt.scatter(freqsx,eliminated_portionz[1], label = "Portion of amplitude eliminated after " + str(max_iter_list[2])) - plt.plot(next_max_frequenciesx[0]*np.ones(100),np.linspace(mp,1,100),linestyle = 'dashed', color ='blue',label = "Max freuency level - 1.") - plt.plot(-next_max_frequenciesx[0]*np.ones(100),np.linspace(mp,1,100),linestyle = 'dashed', color ='blue') - plt.plot(next_max_frequenciesx[1]*np.ones(100),np.linspace(mp,1,100),linestyle = 'dashed', color ='purple',label = "Max freuency level - 2.") - plt.plot(-next_max_frequenciesx[1]*np.ones(100),np.linspace(mp,1,100),linestyle = 'dashed', color ='purple') - plt.plot(next_max_frequenciesx[2]*np.ones(100),np.linspace(mp,1,100),linestyle = 'dashed', color ='red',label = "Max freuency level - 3.") - plt.plot(-next_max_frequenciesx[2]*np.ones(100),np.linspace(mp,1,100),linestyle = 'dashed', color ='red') - plt.xlabel("Frequency") - plt.ylabel("Eliminated portion of the amplitude.") - plt.legend() - #plt.ylim(0,30) - #plt.yscale("log") - plt.savefig("./Eliminated-Portion-Cuboid-"+model+"-"+smoother+"-"+str(Nel)+"Z-component-Epsilon-10-6.pdf") - plt.close() - - -def Visualized_all_frequencies_dampening_2D(Nel, plist, spl_kind, Is): - - from struphy.feec.utilities import create_equal_random_arrays - import plotly.graph_objects as go - - # get global communicator - comm = MPI.COMM_WORLD - rank = comm.Get_rank() - world_size = comm.Get_size() - - a1= 0.2 - #domain = HollowCylinder(a1 = a1) - domain = Cuboid() - model = 'Shear-Alfven' - smoother = 'gmres' - derham = [] - mass_ops = [] - A = [] - - epsilon = 1.0 - derham.append(Derham([Nel[0],Nel[1],Nel[2]], plist, spl_kind, comm=comm, local_projectors=False)) - mass_ops.append(WeightedMassOperators(derham[0], domain)) - if(model == "Poisson"): - sp_key = '0' - sp_id = 'H1' - #Poisson - A.append(derham[0].grad.T @ mass_ops[0].M1 @ derham[0].grad) - #Hall-ish - #A.append(epsilon*derham[0].curl.T @ mass_ops[0].M2 @ derham[0].curl + mass_ops[0].M1) - elif(model == "Hall"): - sp_key = '1' - sp_id = 'Hcurl' - #Hall - A.append(mass_ops[0].M1 -epsilon*derham[0].curl.T @ mass_ops[0].M2 @ derham[0].curl) - #Shear-Alfven-ish - #A.append(mass_ops[0].M2 -epsilon* mass_ops[0].M2@ derham[0].curl @ mass_ops[0].M1 @ derham[0].curl.T @ mass_ops[0].M2) - elif(model == 'Shear-Alfven' or model == 'Shear-Alfven-v2'): - sp_key = '2' - sp_id = 'Hdiv' - pc_class = getattr(preconditioner,"MassMatrixPreconditioner") - pc = pc_class(mass_ops[0].M1) - M1_inv = inverse( - mass_ops[0].M1, - "pcg", - pc=pc, - maxiter=3000, - verbose=False, - ) - #Shear-Alfven - if (model == "Shear-Alfven"): - A.append(mass_ops[0].M2 -epsilon* mass_ops[0].M2@ derham[0].curl @ M1_inv @ derham[0].curl.T @ mass_ops[0].M2) - else: - #Shear-Alfven-v2 - A.append(mass_ops[0].M2 -epsilon* derham[0].curl @ M1_inv @ derham[0].curl.T) - - field_star = derham[0].create_field('fh', sp_id) - field_aprox = derham[0].create_field('fh', sp_id) - - #We are gonna use the method of manufacture solutions to determine the behaviour of our error - #Our exact solution is u_star - u_stararr, u_star = create_equal_random_arrays(derham[0].Vh_fem[sp_key], seed=45) - #We compute the rhs - b = A[0].dot(u_star) - field_star.vector = u_star - - #Turn b into an array - #barr = remove_padding(derham[0].Vh_fem[sp_key],b) - #Turn A[0] into an array - #Aarr = A[0].toarray() - - #Gauss = will_gauss_seidel_converge(Aarr, verbose=True) - #print(f"{Gauss = }") - - - #We store the number of itterations - N_itter = [] - #We define a set of point to evaluate the exact solution and the aproximated one - pointsx = np.linspace(0.0,1.0,Nel[0]) - #pointsx = np.array([0.0,0.5]) - pointsy = np.linspace(0.0,1.0,Nel[1]) - #pointsy = np.array([0.0,0.5]) - pointsz = np.array([0.0,0.5]) - dx = pointsx[1] - pointsx[0] # Spacing between points - dy = pointsy[1] - pointsy[0] # Spacing between points - dz = pointsz[1] - pointsz[0] # Spacing between points - X, Y, Z = np.meshgrid(pointsx, pointsy, pointsz, indexing="ij") - - - - #We define a list where to store the arrays with the values of the errors for each aproximation and vector component - errorsx = [] - errorsy = [] - errorsz = [] - - #Fourier Transform coefficients of the error function for each vector component - fourier_coeffsx = [] - fourier_coeffs_shiftedx = [] - fourier_coeffsy = [] - fourier_coeffs_shiftedy = [] - fourier_coeffsz = [] - fourier_coeffs_shiftedz = [] - max_iter_list = [0,Is, int(10*Is)] - #max_iter_list = [0,2,3,4,5,6,7,8,9,10] - Number_of_iter = len(max_iter_list) - for i in range(Number_of_iter): - - if(max_iter_list[i]>1): - - solver = inverse(A[0],smoother, maxiter = max_iter_list[i], tol = 10.0**-6) - u = solver.dot(b) - #uarr, itter = jacobi(Aarr,barr,max_iter=max_iter_list[i]) - #u = from_array_to_psydac(uarr, derham[0].Vh_fem[sp_key]) - N_itter.append(solver._info['niter']) - #N_itter.append(itter) - - else: - u = derham[0].Vh[derham[0].space_to_form[sp_id]].zeros() - N_itter.append(0) - - - field_aprox.vector = u - if(model == "Poisson"): - errorsx.append(field_star(X,Y,Z)-field_aprox(X,Y,Z)) - else: - errorsx.append(field_star(X,Y,Z)[0]-field_aprox(X,Y,Z)[0]) - errorsy.append(field_star(X,Y,Z)[1]-field_aprox(X,Y,Z)[1]) - errorsz.append(field_star(X,Y,Z)[2]-field_aprox(X,Y,Z)[2]) - - # Compute Fourier Transform coefficients - fourier_coeffsx.append(np.fft.fftn(errorsx[-1])) - fourier_coeffs_shiftedx.append(np.fft.fftshift(fourier_coeffsx[-1])) - if(model != "Poisson"): - fourier_coeffsy.append(np.fft.fftn(errorsy[-1])) - fourier_coeffs_shiftedy.append(np.fft.fftshift(fourier_coeffsy[-1])) - fourier_coeffsz.append(np.fft.fftn(errorsz[-1])) - fourier_coeffs_shiftedz.append(np.fft.fftshift(fourier_coeffsz[-1])) - - # Compute corresponding frequencies - freqsx = np.fft.fftshift(np.fft.fftfreq(len(pointsx), d=dx)) - freqsy = np.fft.fftshift(np.fft.fftfreq(len(pointsy), d=dy)) - freqsz = np.fft.fftshift(np.fft.fftfreq(len(pointsz), d=dz)) - # Create a 3D meshgrid of frequencies - Fx, Fy= np.meshgrid(freqsx, freqsy, indexing="ij") - - #i also want to visiualize the maximum freqeuncies the next 3 coarser grids can handle - hx = 1/Nel[0] - next_max_frequenciesx = [1.0/(4.0*hx),1.0/(8.0*hx),1.0/(16.0*hx)] - hy = 1/Nel[1] - next_max_frequenciesy = [1.0/(4.0*hy),1.0/(8.0*hy),1.0/(16.0*hy)] - #I get the amplitude of the maximum initial error wave - #max_amplitude = np.max(np.abs(fourier_coeffs_shiftedx[0][:,1,1])) - - if(model == "Poisson"): - magnitude_spectrum = [] - log_magnitude_spectrum = [] - #We also want to see the relative difference between the initial amplitude and the amplitude after smoothing. - #We compute abs(A_0)-abs(A_i)/abs(A_0), this tell us which percentage of the initial amplitude has been eliminated - #A values of 0 means it all still remains, a value of 1 means we eliminated all the amplitude. And a negative value - #means that the amplitude increased. - eliminated_portion = [] - for i in range(2): - aux = (np.abs(fourier_coeffs_shiftedx[0][:,:,1]) - np.abs(fourier_coeffs_shiftedx[i+1][:,:,1]))/np.abs(fourier_coeffs_shiftedx[0][:,:,1]) - eliminated_portion.append(aux) - # Create interactive 3D surface plot with Plotly - fig = go.Figure() - - fig.add_trace(go.Surface(z=eliminated_portion[i], x=Fx, y=Fy, colorscale="Viridis")) - - # Set log scale for Z-axis - fig.update_layout( - title="Iteration "+str(max_iter_list[i+1])+". Eliminated portion of error waves for " + model + " with " +smoother+" solver, with resolution " + str(Nel)+".", - scene=dict( - xaxis_title="Frequency X", - yaxis_title="Frequency Y", - zaxis_title="Eliminated portion of initial amplitude", - #zaxis_type="log", # Apply log scale to Z-axis - ) - ) - - # Save as an interactive HTML file - fig.write_html("Eliminated-Portion-Cuboid-"+model+"-"+smoother+"-"+str(Nel)+"-iteration-"+str(max_iter_list[i+1])+".html") - - # Show in browser - #fig.show() - - - for i in range(3): - magnitude_spectrum.append(np.abs(fourier_coeffs_shiftedx[i][:,:,1])) - log_magnitude_spectrum.append(np.log1p(magnitude_spectrum[i])) - - # Plot 3D surface of Fourier magnitude spectrum - #fig = plt.figure(figsize=(10, 7)) - #ax = fig.add_subplot(111, projection="3d") - #ax.plot_surface(Fx, Fy, log_magnitude_spectrum[i], cmap="viridis") - - # Add contour plot at the bottom (Z = 0) - #contour = ax.contourf(Fx, Fy, log_magnitude_spectrum[i], zdir="z", offset=log_magnitude_spectrum[i].min(), cmap="viridis") - - # Labels - #ax.set_xlabel("Frequency X") - #ax.set_ylabel("Frequency Y") - #ax.set_zlabel("Log(Amplitude + 1)") - #ax.set_title("Iteration "+str(max_iter_list[i])+". Amplitude of error waves for " + model + " with " +smoother+" solver, with resolution " + str(Nel)+".") - # Adjust view angle for better visibility - #ax.view_init(elev=30, azim=135) - #plt.show() - #plt.close() - - # Create interactive 3D surface plot with Plotly - fig = go.Figure() - - fig.add_trace(go.Surface(z=log_magnitude_spectrum[i], x=Fx, y=Fy, colorscale="Viridis")) - - # Set log scale for Z-axis - fig.update_layout( - title="Iteration "+str(max_iter_list[i])+". Amplitude of error waves for " + model + " with " +smoother+" solver, with resolution " + str(Nel)+".", - scene=dict( - xaxis_title="Frequency X", - yaxis_title="Frequency Y", - zaxis_title="Log(Amplitude + 1)", - #zaxis_type="log", # Apply log scale to Z-axis - ) - ) - - # Save as an interactive HTML file - fig.write_html("Wave-Amplitude-Cuboid-"+model+"-"+smoother+"-"+str(Nel)+"-iteration-"+str(max_iter_list[i])+".html") - - # Show in browser - #fig.show() - - - - else: - - magnitude_spectrumx = [] - log_magnitude_spectrumx = [] - #We also want to see the relative difference between the initial amplitude and the amplitude after smoothing. - #We compute abs(A_0)-abs(A_i)/abs(A_0), this tell us which percentage of the initial amplitude has been eliminated - #A values of 0 means it all still remains, a value of 1 means we eliminated all the amplitude. And a negative value - #means that the amplitude increased. - eliminated_portionx = [] - for i in range(2): - aux = (np.abs(fourier_coeffs_shiftedx[0][:,:,1]) - np.abs(fourier_coeffs_shiftedx[i+1][:,:,1]))/np.abs(fourier_coeffs_shiftedx[0][:,:,1]) - eliminated_portionx.append(aux) - # Create interactive 3D surface plot with Plotly - fig = go.Figure() - - fig.add_trace(go.Surface(z=eliminated_portionx[i], x=Fx, y=Fy, colorscale="Viridis")) - - # Set log scale for Z-axis - fig.update_layout( - title="X-component. Iteration "+str(max_iter_list[i+1])+". Eliminated portion of error waves for " + model + " with " +smoother+" solver, with resolution " + str(Nel)+".", - scene=dict( - xaxis_title="Frequency X", - yaxis_title="Frequency Y", - zaxis_title="Eliminated portion of initial amplitude", - #zaxis_type="log", # Apply log scale to Z-axis - ) - ) - - # Save as an interactive HTML file - fig.write_html("Eliminated-Portion-Cuboid-"+model+"-"+smoother+"-"+str(Nel)+"-iteration-"+str(max_iter_list[i+1])+"X-component.html") - - # Show in browser - #fig.show() - - - for i in range(3): - magnitude_spectrumx.append(np.abs(fourier_coeffs_shiftedx[i][:,:,1])) - log_magnitude_spectrumx.append(np.log1p(magnitude_spectrumx[i])) - - # Create interactive 3D surface plot with Plotly - fig = go.Figure() - - fig.add_trace(go.Surface(z=log_magnitude_spectrumx[i], x=Fx, y=Fy, colorscale="Viridis")) - - # Set log scale for Z-axis - fig.update_layout( - title="X-component. Iteration "+str(max_iter_list[i])+". Amplitude of error waves for " + model + " with " +smoother+" solver, with resolution " + str(Nel)+".", - scene=dict( - xaxis_title="Frequency X", - yaxis_title="Frequency Y", - zaxis_title="Log(Amplitude + 1)", - #zaxis_type="log", # Apply log scale to Z-axis - ) - ) - - # Save as an interactive HTML file - fig.write_html("Wave-Amplitude-Cuboid-"+model+"-"+smoother+"-"+str(Nel)+"-iteration-"+str(max_iter_list[i])+"X-component.html") - - # Show in browser - #fig.show() - - - magnitude_spectrumy = [] - log_magnitude_spectrumy = [] - #We also want to see the relative difference between the initial amplitude and the amplitude after smoothing. - #We compute abs(A_0)-abs(A_i)/abs(A_0), this tell us which percentage of the initial amplitude has been eliminated - #A values of 0 means it all still remains, a value of 1 means we eliminated all the amplitude. And a negative value - #means that the amplitude increased. - eliminated_portiony = [] - for i in range(2): - aux = (np.abs(fourier_coeffs_shiftedy[0][:,:,1]) - np.abs(fourier_coeffs_shiftedy[i+1][:,:,1]))/np.abs(fourier_coeffs_shiftedy[0][:,:,1]) - eliminated_portiony.append(aux) - # Create interactive 3D surface plot with Plotly - fig = go.Figure() - - fig.add_trace(go.Surface(z=eliminated_portiony[i], x=Fx, y=Fy, colorscale="Viridis")) - - # Set log scale for Z-axis - fig.update_layout( - title="Y-component. Iteration "+str(max_iter_list[i+1])+". Eliminated portion of error waves for " + model + " with " +smoother+" solver, with resolution " + str(Nel)+".", - scene=dict( - xaxis_title="Frequency X", - yaxis_title="Frequency Y", - zaxis_title="Eliminated portion of initial amplitude", - #zaxis_type="log", # Apply log scale to Z-axis - ) - ) - - # Save as an interactive HTML file - fig.write_html("Eliminated-Portion-Cuboid-"+model+"-"+smoother+"-"+str(Nel)+"-iteration-"+str(max_iter_list[i+1])+"Y-component.html") - - # Show in browser - #fig.show() - - - for i in range(3): - magnitude_spectrumy.append(np.abs(fourier_coeffs_shiftedy[i][:,:,1])) - log_magnitude_spectrumy.append(np.log1p(magnitude_spectrumy[i])) - - # Create interactive 3D surface plot with Plotly - fig = go.Figure() - - fig.add_trace(go.Surface(z=log_magnitude_spectrumy[i], x=Fx, y=Fy, colorscale="Viridis")) - - # Set log scale for Z-axis - fig.update_layout( - title="Y-component. Iteration "+str(max_iter_list[i])+". Amplitude of error waves for " + model + " with " +smoother+" solver, with resolution " + str(Nel)+".", - scene=dict( - xaxis_title="Frequency X", - yaxis_title="Frequency Y", - zaxis_title="Log(Amplitude + 1)", - #zaxis_type="log", # Apply log scale to Z-axis - ) - ) - - # Save as an interactive HTML file - fig.write_html("Wave-Amplitude-Cuboid-"+model+"-"+smoother+"-"+str(Nel)+"-iteration-"+str(max_iter_list[i])+"Y-component.html") - - # Show in browser - #fig.show() - - magnitude_spectrumz = [] - log_magnitude_spectrumz = [] - #We also want to see the relative difference between the initial amplitude and the amplitude after smoothing. - #We compute abs(A_0)-abs(A_i)/abs(A_0), this tell us which percentage of the initial amplitude has been eliminated - #A values of 0 means it all still remains, a value of 1 means we eliminated all the amplitude. And a negative value - #means that the amplitude increased. - eliminated_portionz = [] - for i in range(2): - aux = (np.abs(fourier_coeffs_shiftedz[0][:,:,1]) - np.abs(fourier_coeffs_shiftedz[i+1][:,:,1]))/np.abs(fourier_coeffs_shiftedz[0][:,:,1]) - eliminated_portionz.append(aux) - # Create interactive 3D surface plot with Plotly - fig = go.Figure() - - fig.add_trace(go.Surface(z=eliminated_portionz[i], x=Fx, y=Fy, colorscale="Viridis")) - - # Set log scale for Z-axis - fig.update_layout( - title="Z-component. Iteration "+str(max_iter_list[i+1])+". Eliminated portion of error waves for " + model + " with " +smoother+" solver, with resolution " + str(Nel)+".", - scene=dict( - xaxis_title="Frequency X", - yaxis_title="Frequency Y", - zaxis_title="Eliminated portion of initial amplitude", - #zaxis_type="log", # Apply log scale to Z-axis - ) - ) - - # Save as an interactive HTML file - fig.write_html("Eliminated-Portion-Cuboid-"+model+"-"+smoother+"-"+str(Nel)+"-iteration-"+str(max_iter_list[i+1])+"Z-component.html") - - # Show in browser - #fig.show() - - - for i in range(3): - magnitude_spectrumz.append(np.abs(fourier_coeffs_shiftedz[i][:,:,1])) - log_magnitude_spectrumz.append(np.log1p(magnitude_spectrumz[i])) - - # Create interactive 3D surface plot with Plotly - fig = go.Figure() - - fig.add_trace(go.Surface(z=log_magnitude_spectrumz[i], x=Fx, y=Fy, colorscale="Viridis")) - - # Set log scale for Z-axis - fig.update_layout( - title="Z-component. Iteration "+str(max_iter_list[i])+". Amplitude of error waves for " + model + " with " +smoother+" solver, with resolution " + str(Nel)+".", - scene=dict( - xaxis_title="Frequency X", - yaxis_title="Frequency Y", - zaxis_title="Log(Amplitude + 1)", - #zaxis_type="log", # Apply log scale to Z-axis - ) - ) - - # Save as an interactive HTML file - fig.write_html("Wave-Amplitude-Cuboid-"+model+"-"+smoother+"-"+str(Nel)+"-iteration-"+str(max_iter_list[i])+"Z-component.html") - - # Show in browser - #fig.show() - - -def make_plot_smoothing(): - from numpy import array - - ####### - #Example 1 - #Shear-Alfven - #Cuboid - #2D - #CG - ###### - - #Nel = [128,128,1] - #p = [1,1,1] - #Shear-Alfven - #Cuboid - #CG - - #Magnitude of the high frequency (period smaller than 4h) with highest magnitude - - - ############################ - #Nel = [64,64,1] - #p = [1,1,1] - #Shear-Alfven - #Cuboid - #CG - - - ############################ - #Nel = [32,32,1] - #p = [1,1,1] - #Shear-Alfven - #Cuboid - #CG - - - ############################ - #Nel = [16,16,1] - #p = [1,1,1] - #Shear-Alfven - #Cuboid - #CG - - - - ############################ - #Nel = [8,8,1] - #p = [1,1,1] - #Shear-Alfven - #Cuboid - #CG - - - ################################################################ - ################################################################ - ################################################################ - ################################################################ - ####### - #Example 2 - #Shear-Alfven - #Cuboid - #2D - #biCG - ###### - - ############################ - #Nel = [128,128,1] - #p = [1,1,1] - #Shear-Alfven - #Cuboid - #biCG - - - - - ############################ - #Nel = [64,64,1] - #p = [1,1,1] - #Shear-Alfven - #Cuboid - #biCG - - - ############################ - #Nel = [32,32,1] - #p = [1,1,1] - #Shear-Alfven - #Cuboid - #biCG - - - - ############################ - #Nel = [16,16,1] - #p = [1,1,1] - #Shear-Alfven - #Cuboid - #biCG - - - - - ################################################################ - ################################################################ - ################################################################ - ################################################################ - ####### - #Example 3 - #Shear-Alfven - #Cuboid - #2D - #bicgstab - ###### - #Terrible the high frequency errors increase with the number of itterations - - ############################ - #Nel = [128,128,1] - #p = [1,1,1] - #Shear-Alfven - #Cuboid - #bicgstab - - - - - ############################ - #Nel = [64,64,1] - #p = [1,1,1] - #Shear-Alfven - #Cuboid - #bicgstab - - - - - ################################################################ - ################################################################ - ################################################################ - ################################################################ - - - ####### - #Example 4 - #Shear-Alfven - #Cuboid - #2D - #minres - - ############################ - #Nel = [128,128,1] - #p = [1,1,1] - #Hall - #Cuboid - #minres - - magnitudes = [23475.547406074205, 23507.57071761296, 23401.98167905894, 8975.398716277663, 5547.855656375274, 2215.014445147343, 1060.3501137658902, 829.7441809688149, 220.71658497349233, 257.8922866406492] - bad_frequencies = [array([ 44.6484375, -14.8828125, -62.5078125]), array([-44.6484375, 14.8828125, -62.5078125]), array([-44.6484375, 14.8828125, -62.5078125]), array([ 62.5078125, -31.75 , -62.5078125]), array([-63.5 , -43.65625 , -62.5078125]), array([-59.53125 , 62.5078125, -62.5078125]), array([-59.53125 , 62.5078125, -62.5078125]), array([ 59.53125 , -62.5078125, -62.5078125]), array([-59.53125 , 62.5078125, -62.5078125]), array([-59.53125 , 62.5078125, -62.5078125])] - N_itter = [2, 10, 20, 100, 200, 300, 400, 500, 600, 700] - - - ############################ - #Nel = [64,64,1] - #p = [1,1,1] - #Shear-Alfven - #Cuboid - #minres - - magnitudes = [4816.563672246145, 4820.683874727215, 4750.075275409996, 1435.7294516177942, 230.5875341851168, 45.396511954402975, 6.7027103058152555, 6.6240853651862395, 5.862590252251258, 0.753372104197576] - bad_frequencies = [array([-24.609375, -11.8125 , -30.515625]), array([-24.609375, -11.8125 , -30.515625]), array([-24.609375, -11.8125 , -30.515625]), array([-30.515625, 25.59375 , -30.515625]), array([-30.515625, 25.59375 , -30.515625]), array([ 29.53125 , 30.515625, -30.515625]), array([-30.515625, 25.59375 , -30.515625]), array([-29.53125 , -30.515625, -30.515625]), array([-29.53125 , -30.515625, -30.515625]), array([-24.609375, -12.796875, -30.515625])] - N_itter = [2, 10, 20, 100, 200, 300, 400, 500, 600, 700] - - - ############################ - #Nel = [32,32,1] - #p = [1,1,1] - #Shear-Alfven - #Cuboid - #minres - - magnitudes = [1319.382611683582, 1271.7302885073973, 1096.5001351152466, 66.25984056212276, 2.305135582316737, 0.5738656559575724, 0.13428796602415516, 0.0666372580855523, 0.045512532749170435, 0.030227161820733185] - bad_frequencies = [array([-13.5625 , 0. , -14.53125]), array([-13.5625 , 0. , -14.53125]), array([-13.5625 , 0. , -14.53125]), array([-12.59375, -13.5625 , -14.53125]), array([-12.59375, -13.5625 , -14.53125]), array([-13.5625 , -14.53125, -14.53125]), array([-13.5625 , -14.53125, -14.53125]), array([-12.59375, -15.5 , -14.53125]), array([-12.59375, -13.5625 , -14.53125]), array([ 14.53125, 14.53125, -14.53125])] - N_itter = [2, 10, 20, 100, 200, 300, 400, 500, 600, 700] - - - ################################################################ - ################################################################ - ################################################################ - ################################################################ - - ####### - #Example 5 - #Shear-Alfven - #Cuboid - #2D - #lsmr - - ############################ - #Nel = [128,128,1] - #p = [1,1,1] - #Shear-Alfven - #Cuboid - #lsmr - magnitudes = [23625.08035709719, 23546.07810652249, 23555.35959665252, 23553.74183488159, 23553.11191281097, 23552.06535579401, 23550.62355697153, 23548.584106844555, 23546.582190989895, 23543.658419538646] - bad_frequencies = [array([-40.6796875, -0.9921875, -62.5078125]), array([-44.6484375, 14.8828125, -62.5078125]), array([-44.6484375, 14.8828125, -62.5078125]), array([-44.6484375, 14.8828125, -62.5078125]), array([-44.6484375, 14.8828125, -62.5078125]), array([-44.6484375, 14.8828125, -62.5078125]), array([-44.6484375, 14.8828125, -62.5078125]), array([-44.6484375, 14.8828125, -62.5078125]), array([-44.6484375, 14.8828125, -62.5078125]), array([-44.6484375, 14.8828125, -62.5078125])] - N_itter = [2, 10, 20, 100, 200, 300, 400, 500, 600, 700] - - ############################ - #Nel = [64,64,1] - #p = [1,1,1] - #Shear-Alfven - #Cuboid - #lsmr - - magnitudes = [4902.303542373743, 4846.903947432333, 4849.383867406714, 4849.313905985293, 4847.890279297882, 4845.825151419513, 4842.427718581035, 4839.574277539976, 4832.651039101285, 4824.670040396377] - bad_frequencies = [array([-24.609375, -11.8125 , -30.515625]), array([-24.609375, -11.8125 , -30.515625]), array([-24.609375, -11.8125 , -30.515625]), array([ 24.609375, 11.8125 , -30.515625]), array([-24.609375, -11.8125 , -30.515625]), array([-24.609375, -11.8125 , -30.515625]), array([-24.609375, -11.8125 , -30.515625]), array([-24.609375, -11.8125 , -30.515625]), array([-24.609375, -11.8125 , -30.515625]), array([-24.609375, -11.8125 , -30.515625])] - N_itter = [2, 10, 20, 100, 200, 300, 400, 500, 600, 700] - - - ################################################################ - ################################################################ - ################################################################ - ################################################################ - - ####### - #Example 6 - #Shear-Alfven - #Cuboid - #2D - #gmres - - ############################ - #Nel = [128,128,1] - #p = [1,1,1] - #Shear-Alfven - #Cuboid - #gmres - - magnitudes = [23631.325965428965, 23517.072741223346, 23415.519704485465, 9143.805698693457, 5620.759347312477, 2224.5465588280504, 1062.7363141542924, 813.4419688746405, 239.7543262098424, 257.9845124527011] - bad_frequencies = [array([-40.6796875, -0.9921875, -62.5078125]), array([-44.6484375, 14.8828125, -62.5078125]), array([ 44.6484375, -14.8828125, -62.5078125]), array([-62.5078125, 31.75 , -62.5078125]), array([-63.5 , -43.65625 , -62.5078125]), array([ 59.53125 , -62.5078125, -62.5078125]), array([-59.53125 , 62.5078125, -62.5078125]), array([-59.53125 , 62.5078125, -62.5078125]), array([-59.53125 , 62.5078125, -62.5078125]), array([-59.53125 , 62.5078125, -62.5078125])] - N_itter = [2, 10, 20, 100, 200, 300, 400, 500, 600, 700] - - - - - - - - - - - - - ####### - #Example - #Shear-Alfven - #Cuboid - #1D - #CG - ###### - - - ############################ - #Nel = [1024,1,1] - #p = [1,1,1] - #Shear-Alfven - #Cuboid - #CG - - - #As you can see in 1D the CG reduces the magnitude of all problematic high frequencies to 2 e-05. In 1D the CG for the shear-alfven matrix is a good smoother. Thus the Multigrid method works - #with it. But in 2D the CG method is a terrible smoother for this matrix so the multigrid method becomes almost useless with it. - - -def multigrid_Alfven(Nel, plist, spl_kind, u_space): - # get global communicator - comm = MPI.COMM_WORLD - rank = comm.Get_rank() - world_size = comm.Get_size() - derham = Derham(Nel, plist, spl_kind, comm=comm, local_projectors=False) - domain = Tokamak(p = [3,3], psi_shifts = [2.,2.]) - - mhd_equil = AdhocTorusQPsi() - - # must set domain of Cartesian MHD equilibirum - mhd_equil.domain = domain - - - input = derham.Vh[derham.space_to_form[u_space]].zeros() - - - mass_ops = WeightedMassOperators(derham, domain) - basis_ops = BasisProjectionOperators(derham, domain,eq_mhd=mhd_equil) - - id_T = "T" + derham.space_to_form[u_space] - - _T = getattr(basis_ops, id_T) - - _B = -1 / 2 * _T.T @ derham.curl.T @ mass_ops.M2 - _C = 1 / 2 * derham.curl @ _T - - _BC = _B @ _C - - solver = inverse(_BC,'cg') - x = solver.dot(input) - - print(solver._info) - - -def multigrid(Nel, plist, spl_kind, N_levels): - from struphy.feec.utilities import create_equal_random_arrays - # get global communicator - comm = MPI.COMM_WORLD - rank = comm.Get_rank() - world_size = comm.Get_size() - - domain = Cuboid() - #a1 = 0.002 - #domain = HollowCylinder(a1= a1) - sp_key = '0' - - derham = [] - mass_ops = [] - A = [] - - epsilon = 0.0002 - for level in range(N_levels): - derham.append(Derham([Nel[0]//(2**level),Nel[1]//(2**level),Nel[2]], plist, spl_kind, comm=comm, local_projectors=False)) - mass_ops.append(WeightedMassOperators(derham[level], domain)) - A.append(derham[level].grad.T @ mass_ops[level].M1 @ derham[level].grad) - #A.append(epsilon*derham[level].curl.T @ mass_ops[level].M2 @ derham[level].curl + mass_ops[level].M1) - - #We get the inverse of the coarsest system matrix to solve directly the problem in the smaller space - A_inv = np.linalg.inv(A[-1].toarray()) - - R = [] - E = [] - - for level in range(N_levels-1): - R.append(RestrictionOperator(derham[level].Vh_fem[sp_key],derham[level+1].Vh_fem[sp_key])) - E.append(R[level].transpose()) - - method = 'cg' - - #800 - max_iter_list = [14,14,14,18,12,10] - #40 - N_cycles = 2 - - u_stararr, u_star = create_equal_random_arrays(derham[0].Vh_fem[sp_key], seed=45) - - #We compute the rhs - b = A[0].dot(u_star) - - timei = time.time() - - solver_no = inverse(A[0],method, maxiter = 1000000, tol = 10**(-6)) - - u = solver_no.dot(b) - - timef = time.time() - - #We get the total number of itterations - No_Multigrid_itterations = solver_no._info['niter'] - No_Multigrid_error = solver_no._info['res_norm'] - No_Multigrid_time = timef-timei - - #No_Multigrid_itterations = 35088 - #No_Multigrid_error = 9.22E-07 - #No_Multigrid_time = 451.228641271591 - - print("################") - print(f'{No_Multigrid_itterations = }') - print(f'{No_Multigrid_error = }') - print(f'{No_Multigrid_time = }') - print("################") - - - def call_multigrid(max_iter, N_cycles): - u_stararr, u_star = create_equal_random_arrays(derham[0].Vh_fem[sp_key], seed=45) - #We compute the rhs - b = A[0].dot(u_star) - - #We define a list where to store the number of itteration it takes at each multigrid level - #Change N_levels for 1D case - Multigrid_itterations = np.zeros(N_levels, dtype=int) - converged = np.zeros(N_levels, dtype=bool) - - def V_cycle(l, r_l): - #Change for N_levels-1 for 1D case - if (l < N_levels-1): - solver_ini = inverse(A[l],method, maxiter= max_iter[l]) - x_l = solver_ini.dot(r_l) - - #We count the number of itterations - Multigrid_itterations[l] += solver_ini._info['niter'] - - #We determine if the itterative solver converged in the maximum number of itterations - converged[l] = solver_ini._info['success'] - if converged[l] == True: - return x_l - - r_l = r_l - A[l].dot(x_l) - - r_l_plus_1 = R[l].dot(r_l) - x_l_plus_1 = V_cycle(l+1, r_l_plus_1) - #New - #x_l_aux = E[l].dot(x_l_plus_1) - #x_l = x_l + x_l_aux - #r_l = r_l - A[l].dot(x_l_aux) - #solver_end = inverse(A[l].T,method, maxiter= max_iter, x0 =x_l) - #x_l = solver_end.dot(r_l) - #Multigrid_itterations[l] += solver_end._info['niter'] - #### - #old - x_l = x_l + E[l].dot(x_l_plus_1) - ### - - else: - #Solve directly - x_l = direct_solver(A_inv,r_l, derham[l].Vh_fem[sp_key]) - - return x_l - - #N_cycles = 6 - x_0 = derham[0].Vh_fem[sp_key].vector_space.zeros() - - timei = time.time() - for cycle in range(N_cycles): - solver = inverse(A[0],method, maxiter= max_iter[0], x0 = x_0) - x_0 = solver.dot(b) - - Multigrid_itterations[0] += solver._info['niter'] - - #We determine if the itterative solver converged in the maximum number of itterations - converged[0] = solver._info['success'] - if converged[0] == True: - print("Hello") - x = x_0 - break - - r_0 = b - A[0].dot(x_0) - r_1 = R[0].dot(r_0) - - x_0 = x_0 + E[0].dot(V_cycle(1,r_1)) - - if converged[0] == False: - solver = inverse(A[0],method,x0 = x_0, tol = 10**(-6)) - x = solver.dot(b) - Multigrid_itterations[0] += solver._info['niter'] - - timef = time.time() - - - - #We get the final error - Multigrid_error = solver._info['res_norm'] - - Multigrid_time = timef- timei - - speed_up = No_Multigrid_time / Multigrid_time - #speed_up = 27.3452200889587 / Multigrid_time - - print("################") - print("################") - print("################") - #print(f'{a1 = }') - print(f'{max_iter = }') - print(f'{N_cycles = }') - print("################") - print("################") - print(f'{Multigrid_itterations = }') - print(f'{Multigrid_error = }') - print(f'{Multigrid_time = }') - print("################") - print("################") - print(f'{speed_up = }') - - call_multigrid(max_iter_list, N_cycles) - - -def Error_analysis(Nel, plist, spl_kind, N_levels): - - def determine_error(exact,aprox,x): - errors = exact(x,0.0,0.0)-aprox(x,0.0,0.0) - errors = errors.flatten() - return errors - - - - - from struphy.feec.utilities import create_equal_random_arrays - # get global communicator - comm = MPI.COMM_WORLD - rank = comm.Get_rank() - world_size = comm.Get_size() - - domain = Cuboid() - - sp_key = '0' - derham = [] - mass_ops = [] - A = [] - - - for level in range(N_levels): - - derham.append(Derham([Nel[0]//(2**level),Nel[1],Nel[2]], plist, spl_kind, comm=comm, local_projectors=False)) - mass_ops.append(WeightedMassOperators(derham[level], domain)) - A.append(derham[level].grad.T @ mass_ops[level].M1 @ derham[level].grad) - - - field_star = derham[0].create_field('fh', 'H1') - field_aprox = derham[0].create_field('fh', 'H1') - - #We are gonna use the method of manufacture solutions to determine the behaviour of our error - #Our exact solution is u_star - u_stararr, u_star = create_equal_random_arrays(derham[0].Vh_fem[sp_key], seed=45) - #We compute the rhs - b = A[0].dot(u_star) - field_star.vector = u_star - - #We store the number of itterations - N_itter = [] - #We define a set of point to evaluate the exact solution and the aproximated one - points = np.linspace(0.0,1.0,1000) - dx = points[1] - points[0] # Spacing between points - #We define a list where to store the arrays with the values of the errors for each aproximation - errors = [] - - #Fourier Transform coefficients of the error function - fourier_coeffs = [] - #Corresponding frequencies - freqs = [] - for i in range(1, 10): - solver = inverse(A[0],'cg', maxiter = int(i*50) ) - u = solver.dot(b) - field_aprox.vector = u - errors.append(determine_error(field_star,field_aprox,points)) - N_itter.append(solver._info['niter']) - - # Compute Fourier Transform coefficients - fourier_coeffs.append(np.fft.fft(errors[-1])) - # Compute corresponding frequencies - freqs.append(np.fft.fftfreq(len(errors[-1]), d=dx)) - - - - plt.figure() - plt.plot(points,errors[0], label = str(N_itter[0])) - plt.plot(points,errors[-1], label = str(N_itter[-1])) - #plt.scatter(freqs[0],fourier_coeffs[0], label = str(N_itter[0])) - #plt.scatter(freqs[3],fourier_coeffs[3], label = str(N_itter[3])) - #plt.xlim(-50,50) - plt.legend() - plt.show() - plt.close() - - -def Gather_data_V_cycle_parameter_study(Nel, plist, spl_kind, N_levels): - from struphy.feec.utilities import create_equal_random_arrays - # get global communicator - comm = MPI.COMM_WORLD - rank = comm.Get_rank() - world_size = comm.Get_size() - - domain = Cuboid() - a1 = 0.002 - #domain = HollowCylinder(a1= a1) - sp_key = '2' - - derham = [] - mass_ops = [] - A = [] - - epsilon = 1.0 - for level in range(N_levels): - derham.append(Derham([Nel[0]//(2**level),Nel[1]//(2**level),Nel[2]], plist, spl_kind, comm=comm, local_projectors=False)) - mass_ops.append(WeightedMassOperators(derham[level], domain)) - #Poisson - #A.append(derham[level].grad.T @ mass_ops[level].M1 @ derham[level].grad) - #Hall-ish - #A.append(epsilon*derham[level].curl.T @ mass_ops[level].M2 @ derham[level].curl + mass_ops[level].M1) - #Hall - #A.append(mass_ops[level].M1 -epsilon*derham[level].curl.T @ mass_ops[level].M2 @ derham[level].curl) - #Shear-Alfven-ish - #A.append(mass_ops[level].M2 -epsilon* mass_ops[level].M2@ derham[level].curl @ mass_ops[level].M1 @ derham[level].curl.T @ mass_ops[level].M2) - #Shear-Alfven-ish-v2 - #A.append(mass_ops[level].M2 -epsilon* derham[level].curl @ mass_ops[level].M1 @ derham[level].curl.T) - #Shear-Alfven - pc_class = getattr(preconditioner,"MassMatrixPreconditioner") - pc = pc_class(mass_ops[level].M1) - M1_inv = inverse( - mass_ops[level].M1, - "pcg", - pc=pc, - maxiter=3000, - verbose=False, - ) - A.append(mass_ops[level].M2 -epsilon* mass_ops[level].M2@ derham[level].curl @ M1_inv @ derham[level].curl.T @ mass_ops[level].M2) - #Shear-Alfven-v2 - #A.append(mass_ops[level].M2 -epsilon* derham[level].curl @ M1_inv @ derham[level].curl.T) - - #We get the inverse of the coarsest system matrix to solve directly the problem in the smaller space - A_inv = np.linalg.inv(A[-1].toarray()) - - R = [] - E = [] - - for level in range(N_levels-1): - R.append(RestrictionOperator(derham[level].Vh_fem[sp_key],derham[level+1].Vh_fem[sp_key])) - E.append(R[level].transpose()) - - method = 'cg' - - #800 - max_iter_list = [93] - #40 - N_cycles_list = [5] - - u_stararr, u_star = create_equal_random_arrays(derham[0].Vh_fem[sp_key], seed=45) - #We compute the rhs - b = A[0].dot(u_star) - - timei = time.time() - - solver_no = inverse(A[0],method, maxiter = 10000) - - u = solver_no.dot(b) - - timef = time.time() - - #We get the total number of itterations - No_Multigrid_itterations = solver_no._info['niter'] - No_Multigrid_error = solver_no._info['res_norm'] - No_Multigrid_time = timef-timei - - #No_Multigrid_itterations = 9485 - #No_Multigrid_error = 9.72E-07 - #No_Multigrid_time = 867.623549938202 - print("################") - print("################") - print(f'{Nel[0] = }') - print(f'{Nel[1] = }') - print("################") - print("################") - - - print("################") - print(f'{No_Multigrid_itterations = }') - print(f'{No_Multigrid_error = }') - print(f'{No_Multigrid_time = }') - print("################") - - - def call_multigrid(max_iter, N_cycles): - u_stararr, u_star = create_equal_random_arrays(derham[0].Vh_fem[sp_key], seed=45) - #We compute the rhs - b = A[0].dot(u_star) - - #We define a list where to store the number of itteration it takes at each multigrid level - #Change N_levels for 1D case - Multigrid_itterations = np.zeros(N_levels, dtype=int) - converged = np.zeros(N_levels, dtype=bool) - - def V_cycle(l, r_l): - #Change for N_levels-1 for 1D case - if (l < N_levels-1): - solver_ini = inverse(A[l],method, maxiter= max_iter) - x_l = solver_ini.dot(r_l) - - #We count the number of itterations - Multigrid_itterations[l] += solver_ini._info['niter'] - - #We determine if the itterative solver converged in the maximum number of itterations - converged[l] = solver_ini._info['success'] - if converged[l] == True: - return x_l - - r_l = r_l - A[l].dot(x_l) - - r_l_plus_1 = R[l].dot(r_l) - x_l_plus_1 = V_cycle(l+1, r_l_plus_1) - #New - #x_l_aux = E[l].dot(x_l_plus_1) - #x_l = x_l + x_l_aux - #r_l = r_l - A[l].dot(x_l_aux) - #solver_end = inverse(A[l].T,method, maxiter= max_iter, x0 =x_l) - #x_l = solver_end.dot(r_l) - #Multigrid_itterations[l] += solver_end._info['niter'] - #### - #old - x_l = x_l + E[l].dot(x_l_plus_1) - ### - - else: - #Solve directly - x_l = direct_solver(A_inv,r_l, derham[l].Vh_fem[sp_key]) - - return x_l - - #N_cycles = 6 - x_0 = derham[0].Vh_fem[sp_key].vector_space.zeros() - - timei = time.time() - for cycle in range(N_cycles): - solver = inverse(A[0],method, maxiter= max_iter, x0 = x_0) - x_0 = solver.dot(b) - - Multigrid_itterations[0] += solver._info['niter'] - - #We determine if the itterative solver converged in the maximum number of itterations - converged[0] = solver._info['success'] - if converged[0] == True: - print("Hello") - x = x_0 - break - - r_0 = b - A[0].dot(x_0) - r_1 = R[0].dot(r_0) - - x_0 = x_0 + E[0].dot(V_cycle(1,r_1)) - - if converged[0] == False: - solver = inverse(A[0],method, maxiter = 1000,x0 = x_0,) - x = solver.dot(b) - Multigrid_itterations[0] += solver._info['niter'] - - timef = time.time() - - - - #We get the final error - Multigrid_error = solver._info['res_norm'] - - Multigrid_time = timef- timei - - speed_up = No_Multigrid_time / Multigrid_time - - print("################") - print("################") - print("################") - #print(f'{a1 = }') - print(f'{max_iter = }') - print(f'{N_cycles = }') - print("################") - print("################") - print(f'{Multigrid_itterations = }') - print(f'{Multigrid_error = }') - print(f'{Multigrid_time = }') - print("################") - print("################") - print(f'{speed_up = }') - - for max_iter in max_iter_list: - for N_cycles in N_cycles_list: - call_multigrid(max_iter, N_cycles) - - -def Gather_data_V_cycle_scalability(Nellist, plist, spl_kind): - for Nel in Nellist: - #First we compute the number of levels for each Nel - N_levels = int(log2(Nel[0])-1) - Gather_data_V_cycle_parameter_study(Nel, plist, spl_kind, N_levels) - - -def make_plot_scalability(): - - model = 'Hall' - - Nel = [int(2**i * 2**i) for i in range(4,8)] - Multi_time = [0.933451652526856, - 4.96578407287598, - 21.1356129646301, - 121.159014463425 - ] - No_Multi_time = [0.853443145751953, - 6.89243221282959, - 75.1050372123718, - 1334.72232866287 - ] - speed_up = [] - - for i in range(len(Multi_time)): - speed_up.append(No_Multi_time[i]/Multi_time[i]) - - plt.figure() - plt.title('2D ' +model+' run times.') - plt.plot(Nel, Multi_time, label = 'Multigrid.') - plt.scatter(Nel, No_Multi_time, label = 'CG run time.') - plt.xlabel("Nel[0] x Nel[1]") - plt.ylabel('Time (s)') - plt.legend() - plt.savefig("2D-"+model+"-runtime.pdf") - #plt.show() - plt.close() - - plt.figure() - plt.title('2D ' +model+ ' speed up.') - plt.scatter(Nel, speed_up, label = 'Speed_up.') - plt.xlabel("Nel[0] x Nel[1]") - plt.ylabel('Speed_up') - plt.legend() - #plt.show() - plt.savefig("2D-"+model+"-speed-up.pdf") - plt.close() - - -def verify_formula(Nel, plist, spl_kind): - comm = MPI.COMM_WORLD - rank = comm.Get_rank() - derham = Derham(Nel, plist, spl_kind, comm=comm) - - # For B-splines - sp_key = "0" - spaces = derham.Vh_fem[sp_key].spaces - space = spaces[0] - N = space.nbasis - ncells = space.ncells - p = space.degree - T = space.knots - periodic = space.periodic - basis = space.basis - normalize = basis == "M" - - def make_basis_fun(i): - def fun(etas, eta2, eta3): - if isinstance(etas, float) or isinstance(etas, int): - etas = np.array([etas]) - out = np.zeros_like(etas) - for j, eta in enumerate(etas): - span = find_span(T, p, eta) - inds = np.arange(span - p, span + 1) % N - pos = np.argwhere(inds == i) - # print(f'{pos = }') - if pos.size > 0: - pos = pos[0, 0] - out[j] = basis_funs(T, p, eta, span, normalize=normalize)[pos] - else: - out[j] = 0.0 - return out - - return fun - - i = random.randint(0,N-1) - fun = make_basis_fun(i) - points = np.linspace(0.0,1.0,100) - values = fun(points,0.0,0.0) - - - #Now we build the fine grid - derham_fine = Derham([Nel[0]*2,Nel[1],Nel[2]], plist, spl_kind, comm=comm) - - # For B-splines - spaces_fine = derham_fine.Vh_fem[sp_key].spaces - space_fine = spaces_fine[0] - N_fine = space_fine.nbasis - ncells_fine = space_fine.ncells - p_fine = space_fine.degree - T_fine = space_fine.knots - periodic_fine = space_fine.periodic - basis_fine = space_fine.basis - normalize_fine = basis_fine == "M" - - def make_basis_fun_fine(i): - def fun(etas, eta2, eta3): - if isinstance(etas, float) or isinstance(etas, int): - etas = np.array([etas]) - out = np.zeros_like(etas) - for j, eta in enumerate(etas): - span = find_span(T_fine, p_fine, eta) - inds = np.arange(span - p_fine, span + 1) % N_fine - pos = np.argwhere(inds == i) - # print(f'{pos = }') - if pos.size > 0: - pos = pos[0, 0] - out[j] = basis_funs(T_fine, p_fine, eta, span, normalize=normalize_fine)[pos] - else: - out[j] = 0.0 - return out - - return fun - - fun_fine = [] - weights_fine = [] - for j in range(p_fine+2): - fun_fine.append(make_basis_fun_fine((2*i-p_fine+j)%N_fine)) - weights_fine.append(2.0**(-p_fine)*comb(p_fine+1,j)) - - - def fun_combine(etas, eta2, eta3): - if isinstance(etas, float) or isinstance(etas, int): - etas = np.array([etas]) - out = np.zeros_like(etas) - for j in range(p_fine+2): - out += weights_fine[j]*fun_fine[j](etas,0.0,0.0) - return out - - - - values_fine = fun_combine(points,0.0,0.0) - - Equal = True - where = -1 - for j in range(len(values)): - if(abs(values[j] -values_fine[j])>10**-5): - Equal = False - where = j - - print(i) - print(Equal) - print(where) - if Equal == False: - print(values[where]) - print(values_fine[where]) - - -def verify_Restriction_Operator(Nel, plist, spl_kind): - from struphy.feec.utilities import create_equal_random_arrays, compare_arrays - - comm = MPI.COMM_WORLD - rank = comm.Get_rank() - derham = Derham(Nel, plist, spl_kind, comm=comm) - - # For B-splines - sp_key = "0" - spaces = derham.Vh_fem[sp_key].spaces - space = spaces[0] - N = space.nbasis - ncells = space.ncells - p = space.degree - T = space.knots - periodic = space.periodic - basis = space.basis - normalize = basis == "M" - - #Now we build the fine grid - derham_fine = Derham([Nel[0]*2,Nel[1],Nel[2]], plist, spl_kind, comm=comm) - - # For B-splines - spaces_fine = derham_fine.Vh_fem[sp_key].spaces - space_fine = spaces_fine[0] - N_fine = space_fine.nbasis - ncells_fine = space_fine.ncells - p_fine = space_fine.degree - T_fine = space_fine.knots - periodic_fine = space_fine.periodic - basis_fine = space_fine.basis - normalize_fine = basis_fine == "M" - - #We intialize the restriction operator - R = RestrictionOperator(derham_fine.Vh_fem[sp_key],derham.Vh_fem[sp_key]) - - varr, v = create_equal_random_arrays(derham_fine.Vh_fem[sp_key], seed=4568) - varr = varr[0].flatten() - - - out = R.dot(v) - - #To make it easier to read I will extract the data out of out disregarding all the padding it come with - out_array = np.zeros(N, dtype=float) - for i in range(N): - out_array[i] = out[R._out_starts[0]+i,0,0] - - #print(out._data) - - - ##### - #Now we build the Restriction matrix directly to verify our RestrictionOperator is working properly. - R_matrix = np.zeros((N,N_fine),dtype=float) - for i in range(N): - start = 2*i-p - for j in range(p+2): - R_matrix[i,(start+j)%N_fine] = 2.0**(-p)*comb(p+1,j) - - out2 = np.matmul(R_matrix, varr) - - - Equal =True - where = -1 - for i in range(len(out2)): - if(abs(out2[i]-out_array[i])>10.0**(-6)): - Equal = False - where = i - - - - print(f'{Equal = }') - print(f'{out_array = }') - print(f'{out2 = }') - if( Equal == False): - print(f'{where = }') - - -def verify_Extension_Operator(Nel, plist, spl_kind): - from struphy.feec.utilities import create_equal_random_arrays, compare_arrays - - comm = MPI.COMM_WORLD - rank = comm.Get_rank() - derham = Derham(Nel, plist, spl_kind, comm=comm) - - # For B-splines - sp_key = "0" - spaces = derham.Vh_fem[sp_key].spaces - space = spaces[0] - N = space.nbasis - ncells = space.ncells - p = space.degree - T = space.knots - periodic = space.periodic - basis = space.basis - normalize = basis == "M" - - #Now we build the fine grid - derham_fine = Derham([Nel[0]*2,Nel[1],Nel[2]], plist, spl_kind, comm=comm) - - # For B-splines - spaces_fine = derham_fine.Vh_fem[sp_key].spaces - space_fine = spaces_fine[0] - N_fine = space_fine.nbasis - ncells_fine = space_fine.ncells - p_fine = space_fine.degree - T_fine = space_fine.knots - periodic_fine = space_fine.periodic - basis_fine = space_fine.basis - normalize_fine = basis_fine == "M" - - #We intialize the restriction operator - E = ExtensionOperator(derham.Vh_fem[sp_key],derham_fine.Vh_fem[sp_key]) - - v = derham.Vh[sp_key].zeros() - varr, v = create_equal_random_arrays(derham.Vh_fem[sp_key], seed=4568) - varr = varr[0].flatten() - - - out = E.dot(v) - - #To make it easier to read I will extract the data out of out disregarding all the padding it come with - out_array = np.zeros(N_fine, dtype=float) - for i in range(N_fine): - out_array[i] = out[E._out_starts[0]+i,0,0] - - #print(out._data) - - - ##### - #Now we build the Restriction matrix directly to verify our RestrictionOperator is working properly. - R_matrix = np.zeros((N,N_fine),dtype=float) - for i in range(N): - start = 2*i-p - for j in range(p+2): - R_matrix[i,(start+j)%N_fine] = 2.0**(-p)*comb(p+1,j) - - E_matrix = R_matrix.T - - out2 = np.matmul(E_matrix, varr) - - - Equal =True - where = -1 - for i in range(len(out2)): - if(abs(out2[i]-out_array[i])>10.0**(-6)): - Equal = False - where = i - - - - print(f'{Equal = }') - print(f'{out_array = }') - print(f'{out2 = }') - if( Equal == False): - print(f'{where = }') - - -def testing_random_stuff(Nel,plist,spl_kind): - epsilon = 10.0**-6.0 - Nel = 512 - jarray = [] - for j in range(Nel): - jarray.append(float(j)) - - jarray = np.array(jarray) - - compute_directly = True - - #matrix = 'Shear-Alfven' - matrix = 'Hall' - if matrix == 'Shear-Alfven': - - def funl(j): - return 9.0*Nel - epsilon * (243.0/2.0) * (Nel**3.0) * (1.0-np.cos(2.0*np.pi *j / Nel)) /(2.0+np.cos(2.0*np.pi *j / Nel)) - - max_eigen_value_x = 4.0/Nel - - - - - elif( matrix == 'Hall'): - def funl(j): - return 8.0 /(3.0*Nel) - 18.0*epsilon*Nel + (4.0/(3.0*Nel)+18.0*epsilon*Nel)*np.cos(2.0*np.pi*j/Nel) - - max_eigen_value_x = 9.0*Nel - - - - eigen_values_y_z = funl(jarray) - max_eigen_value_y_z=abs(eigen_values_y_z[np.argmax(np.abs(eigen_values_y_z))]) - - - order_of_magnitude_disparity = np.log10(max_eigen_value_y_z/max_eigen_value_x) - print(f'{max_eigen_value_y_z = }') - print(f'{max_eigen_value_x = }') - print(f'{order_of_magnitude_disparity = }') - - plt.figure() - plt.scatter(jarray,eigen_values_y_z) - plt.show() - - - if compute_directly: - # get global communicator - comm = MPI.COMM_WORLD - - domain = Cuboid() - derham = [] - mass_ops = [] - A = [] - - derham.append(Derham([Nel,1,1], [1,1,1], [True,True,True], comm=comm, local_projectors=False)) - mass_ops.append(WeightedMassOperators(derham[0], domain)) - if(matrix == "Poisson"): - sp_key = '0' - sp_id = 'H1' - #Poisson - A.append(derham[0].grad.T @ mass_ops[0].M1 @ derham[0].grad) - #Hall-ish - #A.append(epsilon*derham[0].curl.T @ mass_ops[0].M2 @ derham[0].curl + mass_ops[0].M1) - elif(matrix == "Hall"): - sp_key = '1' - sp_id = 'Hcurl' - #Hall - A.append(mass_ops[0].M1 -epsilon*derham[0].curl.T @ mass_ops[0].M2 @ derham[0].curl) - #Shear-Alfven-ish - #A.append(mass_ops[0].M2 -epsilon* mass_ops[0].M2@ derham[0].curl @ mass_ops[0].M1 @ derham[0].curl.T @ mass_ops[0].M2) - elif(matrix == 'Shear-Alfven' or matrix == 'Shear-Alfven-v2'): - sp_key = '2' - sp_id = 'Hdiv' - pc_class = getattr(preconditioner,"MassMatrixPreconditioner") - pc = pc_class(mass_ops[0].M1) - M1_inv = inverse( - mass_ops[0].M1, - "pcg", - pc=pc, - maxiter=3000, - verbose=False, - ) - #Shear-Alfven - if (matrix == "Shear-Alfven"): - A.append(mass_ops[0].M2 -epsilon* mass_ops[0].M2@ derham[0].curl @ M1_inv @ derham[0].curl.T @ mass_ops[0].M2) - - Aarr = A[0].toarray() - Aarrx = Aarr[0:Nel,0:Nel] - Aarry = Aarr[Nel:2*Nel,Nel:2*Nel] - eigenvaluesx = np.linalg.eigvals(Aarrx) - max_x = abs(eigenvaluesx[np.argmax(np.abs(eigenvaluesx))]) - eigenvaluesy = np.linalg.eigvals(Aarry) - - plt.figure() - plt.plot(jarray,eigenvaluesy) - plt.show() - - max_y = abs(eigenvaluesy[np.argmax(np.abs(eigenvaluesy))]) - print(f'{max_x =}') - - - - - - - - #from struphy.feec.utilities import create_equal_random_arrays - # get global communicator - #comm = MPI.COMM_WORLD - #rank = comm.Get_rank() - #world_size = comm.Get_size() - - #domain = Cuboid() - #a1 = 0.002 - #domain = HollowCylinder(a1= a1) - #sp_key = '0' - - #derham =Derham([Nel[0],Nel[1],Nel[2]], plist, spl_kind, comm=comm, local_projectors=False) - #mass_ops=WeightedMassOperators(derham, domain) - - #u_stararr, u_star = create_equal_random_arrays(derham.Vh_fem[sp_key], seed=45) - #u_stararr = remove_padding(derham.Vh_fem[sp_key], u_star) - - #G = derham.grad - #Garr =G.toarray() - - #GT = derham.grad.T - #GTarr = GT.toarray() - - #C = derham.curl - #Carr = C.toarray() - - #M1 = mass_ops.M1 - #M1arr = M1.toarray() - - #Po = GT @ M1 @ G - #Poarr = Po.toarray() - - - #eigenvalues = np.linalg.eigvals(Poarr) - #max_eigen_value = abs(eigenvalues[np.argmax(np.abs(eigenvalues))]) - #min_eigen_value = abs(eigenvalues[np.argmin(np.abs(eigenvalues))]) - #b = G.dot(u_star) - #barr = np.matmul(Garr, u_stararr) - - #b = remove_padding(derham.Vh_fem['1'], b) - - #same = True - #for i in range(np.size(b)): - #if(barr[i]!= b[i]): - #same = False - #break - - #print(same) - #print(f'{M1arr[32,32:64] =}') - #print(f'{max_eigen_value =}') - #print(f'{min_eigen_value =}') - #for i in range(0,3): - #print(f'{Garr[i,:] =}') - - - - - -if __name__ == '__main__': - from struphy.feec.utilities import create_equal_random_arrays - comm = MPI.COMM_WORLD - derham = Derham([10,1,1], [1,1,1], [True,True,True], comm=comm, local_projectors=False) - u_stararr, u_star = create_equal_random_arrays(derham.Vh_fem['0'], seed=45) - print("################") - print(vars(u_star._space)) - print("################") - #for i in range(7,8): - #Nel = [int(2**i),int(2**i), 1] - #p = [1, 1, 1] - #spl_kind = [True, True, True] - #Gather_data_V_cycle_parameter_study(Nel, p, spl_kind, i) - #Visualized_high_frequency_dampening(Nel, p, spl_kind) - #Visualized_all_frequencies_dampening(Nel, p, spl_kind,7) - #Visualized_all_frequencies_dampening_2D(Nel, p, spl_kind,100) - - #testing_random_stuff([32,1,1],[1,1,1],[True,True,True]) - #128,64,32,16, 8 - # h, 2h,4h,8h,16h - - #p=1, Nel= 8192, level = 12. Coarsest one is 4x4 matrix - #p=2, Nel= 8192, level = 11. Coarsest one is 8x8 matrix - #p=4, Nel= 8192, level = 10. Coarsest one is 16x16 matrix - - #multigrid(Nel, p, spl_kind,6) - #Gather_data_V_cycle_parameter_study(Nel, p, spl_kind, 11) - #Gather_data_V_cycle_scalability([[int(2**i),1,1] for i in range(4,10)], p, spl_kind) - #make_plot_scalability() - #verify_formula(Nel, p, spl_kind) - #verify_Restriction_Operator(Nel, p, spl_kind) - #verify_Extension_Operator(Nel, p, spl_kind) - #Error_analysis(Nel, p, spl_kind, 1) - #trying_New_Restriction(Nel, p, spl_kind, 4) - #Compute_rate_of_smoothing(Nel, p, spl_kind) - #Visualized_high_frequency_dampening(Nel, p, spl_kind) - \ No newline at end of file diff --git a/src/struphy/linear_algebra/multigrid_solver.py b/src/struphy/linear_algebra/multigrid_solver.py deleted file mode 100644 index 65feeae85..000000000 --- a/src/struphy/linear_algebra/multigrid_solver.py +++ /dev/null @@ -1,1070 +0,0 @@ -import time - -import matplotlib.pyplot as plt -import numpy as np -from mpi4py import MPI -import re -from struphy.feec.psydac_derham import Derham - -from psydac.linalg.solvers import inverse -from psydac.linalg.basic import Vector -from psydac.fem.projectors import knot_insertion_projection_operator -from feectools.linalg.basic import LinearOperator -from psydac.fem.basic import FemSpace -from psydac.fem.tensor import TensorFemSpace -from struphy.feec.mass import WeightedMassOperators -from struphy.geometry.domains import Cuboid -from math import comb -from struphy.feec.utilities import create_equal_random_arrays - -def build_operator(expr: str, namespace: dict): - """ - Build a composite LinearOperator (or scalar * LinearOperator expression) - from a string expression. - - Parameters - ---------- - expr : str - String expression defining the operator composition. - Example: - "epsilon * curl.T @ M2 @ curl + mu * M1" - - The expression can contain: - - scalars (e.g., epsilon, mu, ...) - - LinearOperator objects (already supporting +, -, *, @) - - parentheses to control precedence - - namespace : dict - Dictionary mapping symbol names in `expr` to actual Python objects - (scalars, operators, lists, etc.). - Example: - { - "epsilon": 1.3, - "mu": 0.7, - "curl.T": derham.curl.T, - "curl": derham.curl, - "M1": WeightedMassOperators(derham, domain).M1, - "M2": WeightedMassOperators(derham, domain).M2, - } - - Returns - ------- - LinearOperator - A composite operator built according to the expression. The returned - object supports `.dot(v)` and any other functionality of LinearOperator. - - Raises - ------ - ValueError - If some symbols in the expression are not found in `namespace`. - """ - - # Find all potential variable-like tokens (ignore numbers) - tokens = set(re.findall(r"[A-Za-z_][A-Za-z0-9_]*(?:\.[A-Za-z_][A-Za-z0-9_]*)*", expr)) - - missing = tokens - set(namespace.keys()) - - if missing: - raise ValueError( - f"The following symbols are missing from the namespace: {', '.join(sorted(missing))}" - ) - - # Let Python evaluate the expression with operator overloads - return eval(expr, {}, namespace) - -def remove_padding(fem_space, v): - - #fem_space = derham.Vh_fem[sp_key] - symbolic_name = fem_space.symbolic_space - - if(symbolic_name == 'H1' or symbolic_name == "L2"): - - spaces = [fem_space.spaces] - N = [spaces[0][i].nbasis for i in range(3)] - - starts = np.array(fem_space.coeff_space.starts) - - #To make it easier to read I will extract the data out of out disregarding all the padding it come with - v_array = np.zeros(N[0]*N[1]*N[2], dtype=float) - cont = 0 - for i0 in range(N[0]): - for i1 in range(N[1]): - for i2 in range(N[2]): - v_array[cont] = v[starts[0]+i0,starts[1]+i1,starts[2]+i2] - cont += 1 - - else: - spaces = [comp.spaces for comp in fem_space.spaces] - N = [[spaces[h][i].nbasis for i in range(3)] for h in range(3)] - - starts = np.array([vi.starts for vi in fem_space.coeff_space.spaces]) - - #To make it easier to read I will extract the data out of out disregarding all the padding it come with - v_array = np.zeros(N[0][0]*N[0][1]*N[0][2]+N[1][0]*N[1][1]*N[1][2]+N[2][0]*N[2][1]*N[2][2], dtype=float) - cont = 0 - for h in range(3): - for i0 in range(N[h][0]): - for i1 in range(N[h][1]): - for i2 in range(N[h][2]): - v_array[cont] = v[h][starts[h][0]+i0,starts[h][1]+i1,starts[h][2]+i2] - cont += 1 - - return v_array - - -def direct_solver(A_inv,b, fem_space): - # A_inv is already the inverse matrix of A - #fem_space = derham.Vh_fem[sp_key] - symbolic_name = fem_space.symbolic_space - - if(symbolic_name == 'H1' or symbolic_name == "L2"): - spaces = [fem_space.spaces] - N = [spaces[0][i].nbasis for i in range(3)] - starts = np.array(fem_space.coeff_space.starts) - - b_vector = remove_padding(fem_space, b) - x_vector = np.dot(A_inv, b_vector) - x = fem_space.coeff_space.zeros() - - cont= 0 - for i0 in range(N[0]): - for i1 in range(N[1]): - for i2 in range(N[2]): - x[starts[0]+i0,starts[1]+i1,starts[2]+i2] = x_vector[cont] - cont += 1 - - else: - spaces = [comp.spaces for comp in fem_space.spaces] - N = [[spaces[h][i].nbasis for i in range(3)] for h in range(3)] - starts = np.array([vi.starts for vi in fem_space.coeff_space.spaces]) - - b_vector = remove_padding(fem_space, b) - x_vector = np.dot(A_inv, b_vector) - x = fem_space.coeff_space.zeros() - - cont = 0 - for h in range(3): - for i0 in range(N[h][0]): - for i1 in range(N[h][1]): - for i2 in range(N[h][2]): - x[h][starts[h][0]+i0,starts[h][1]+i1,starts[h][2]+i2] = x_vector[cont] - cont += 1 - return x - - -def get_b_spline_degree(V): - """ - Determines the degree of the B-splines. - - Parameters - ---------- - V : psydac.fem.basic.FemSpace - Finite element spline space (domain, input space). - - Returns - ------- - p : numpy array - numpy array of 3 ints containing the B-spline degrees for each spatial direction. - """ - - assert isinstance(V, FemSpace) - p = np.zeros(3, dtype=int) - - if hasattr(V, "symbolic_space"): - V_name = V.symbolic_space - - if(V_name == "H1"): - for i, space in enumerate(V.spaces): - p[i] = space.degree - elif(V_name == "L2"): - for i, space in enumerate(V.spaces): - p[i] = space.degree+1 - elif(V_name == "Hcurl"): - V1ds = [comp.spaces for comp in V.spaces] - for i in range(3): - p[i] = V1ds[i][i].degree+1 - elif(V_name == "Hdiv"): - V1ds = [comp.spaces for comp in V.spaces] - for i in range(3): - p[i] = V1ds[i][i].degree - elif(V_name == "H1H1H1"): - V1ds = [comp.spaces for comp in V.spaces] - for i in range(3): - p[i] = V1ds[i][i].degree - else: - raise Exception("Invalid symbolic name.") - - return p - - - -class MultiGridSolver: - - def __init__(self, derham, sp_key, N_levels, domain, expr: str, namespace: dict, max_iter_list, N_cycles, method='cg'): - """ - - Initialize the multigrid solver with the given parameters. - Parameters - ---------- - derham : Derham - The Derham object representing the finest grid. - sp_key : str - Symbolic space key, e.g., "0", "1", "2", "3". - N_levels : int - Number of multigrid levels. - domain : struphy.geometry.base.Domain - The Geometric domain of the simulation. - - expr : str - String expression defining the operator composition. - Example: - "epsilon * curl.T @ M2 @ curl + mu * M1" - - The expression can contain: - - scalars (e.g., epsilon, mu, ...) - - LinearOperator objects (already supporting +, -, *, @) - - parentheses to control precedence - - namespace : dict - A dictionary that serves as a **string template** or "recipe book" for - building the operators and scalars used in the `expr` string. - - It maps symbol names to **strings of Python code** that describe how to - create the corresponding object. This allows the solver to dynamically - generate the correct operators for each grid level. - - The string recipes can use variables like `derham`, `WeightedMassOperators` and `domain`, which - the solver makes available during the object creation at each level. - - Example: - { - "epsilon": "1.3", - "mu": "0.7", - "curl.T": "derham.curl.T", - "curl": "derham.curl", - "M1": "WeightedMassOperators(derham, domain).M1", - "M2": "WeightedMassOperators(derham, domain).M2", - } - max_iter_list : list - List of maximum iterations for each multigrid level. The first element corresponds to the finest level. - N_cycles : int - Number of multigrid cycles. - method : str - Solution method to use (e.g., 'cg' for conjugate gradient). - - """ - self.sp_key = sp_key - self.Nel = derham.Nel - self.plist = derham.p - self.spl_kind = derham._spl_kind - self.N_levels = N_levels - self.domain = domain - self.expr = expr - self.namespace = namespace - self.derham = derham - self.method = method - self.max_iter_list = max_iter_list - self.N_cycles = N_cycles - - # get global communicator - self.comm = MPI.COMM_WORLD - self.derhamlist = [] - self.A = [] - - for level in range(self.N_levels): - if level == 0: - self.derhamlist.append(self.derham) - else: - self.derhamlist.append(Derham([self.Nel[0]//(2**level),self.Nel[1]//(2**level),self.Nel[2]], self.plist, self.spl_kind, comm=self.comm, local_projectors=False)) - #It might look like we are not using derham but it is being used by updated_namespace - #for rebounding the expressions in self.expr to the correct derham level. - derhamaux = self.derhamlist[level] - - # 1. Create a local context for the eval. It maps string names - # to the actual objects for the CURRENT level. - eval_context = { - "derham": derhamaux, - "domain": self.domain, - "WeightedMassOperators": WeightedMassOperators - } - - # 2. Build the namespace by evaluating each string from the template - # within the context of the current level. - updated_namespace = { - key: eval(recipe_string, {}, eval_context) - for key, recipe_string in self.namespace.items() - } - - self.A.append(build_operator(self.expr, updated_namespace)) - - #We get the inverse of the coarsest system matrix to solve directly the problem in the smaller space - self.A_inv = np.linalg.inv(self.A[-1].toarray()) - - self.R = [] - self.E = [] - - for level in range(self.N_levels-1): - self.R.append(self.RestrictionOperator(self.derhamlist[level].Vh_fem[self.sp_key],self.derhamlist[level+1].Vh_fem[self.sp_key])) - self.E.append(self.R[level].transpose()) - - - class RestrictionOperator(LinearOperator): - """ - Linear operator which operates between vector spaces of the same kind but different resolutions. - - We assume that the vectors in the domain belong to a De-rham space with n elements, while the codomain - belong to the coarser De-rham space with n/2 elements. - - At the moment we also assume that we are halving only the first spatial direction. - - Parameters - ---------- - V : psydac.fem.basic.FemSpace - Finite element spline space (domain, input space). - - W : psydac.fem.basic.FemSpace - Finite element spline space (codomain, output space). - - """ - def __init__(self, V, W): - - # Check domain and codomain - assert isinstance(V, FemSpace) - assert isinstance(W, FemSpace) - - self._V = V - self._W = W - - #print(vars(V)) - - # Store info in object - self._domain = V.coeff_space - self._codomain = W.coeff_space - self._dtype = V.coeff_space.dtype - - #Can be "H1", "L2", "Hcurl", "Hdiv", "H1H1H1" - self._V_name = V.symbolic_space - self._W_name = W.symbolic_space - assert(self._V_name == self._W_name) - - #This list will tell us in which spatial direction we are halving the problem. - self._halving_directions = [False,False,False] - - # input space: 3d StencilVectorSpaces and 1d SplineSpaces of each component - if isinstance(V, TensorFemSpace): - self._V1ds = [V.spaces] - self._VNbasis = np.array([self._V1ds[0][0].nbasis, self._V1ds[0][1].nbasis, self._V1ds[0][2].nbasis]) - - # We get the start and endpoint for each sublist in input - self._in_starts = np.array(V.coeff_space.starts) - self._in_ends = np.array(V.coeff_space.ends) - else: - self._V1ds = [comp.spaces for comp in V.spaces] - self._VNbasis = np.array( - [ - [self._V1ds[0][0].nbasis, self._V1ds[0][1].nbasis, self._V1ds[0][2].nbasis], - [ - self._V1ds[1][0].nbasis, - self._V1ds[1][1].nbasis, - self._V1ds[1][2].nbasis, - ], - [self._V1ds[2][0].nbasis, self._V1ds[2][1].nbasis, self._V1ds[2][2].nbasis], - ] - ) - - # We get the start and endpoint for each sublist in input - self._in_starts = np.array([vi.starts for vi in V.coeff_space.spaces]) - self._in_ends = np.array([vi.ends for vi in V.coeff_space.spaces]) - - # output space: 3d StencilVectorSpaces and 1d SplineSpaces of each component - if isinstance(W, TensorFemSpace): - self._W1ds = [W.spaces] - self._WNbasis = np.array([self._W1ds[0][0].nbasis, self._W1ds[0][1].nbasis, self._W1ds[0][2].nbasis]) - - for i in range(3): - if(self._WNbasis[i] < self._VNbasis[i]): - #If this breaks for clamped splines it means .nbasis gives you the number of basis functions, not the number of elements - assert self._VNbasis[i] == self._WNbasis[i]*2 - self._halving_directions[i] = True - - # We get the start and endpoint for each sublist in out - self._out_starts = np.array(W.coeff_space.starts) - self._out_ends = np.array(W.coeff_space.ends) - - else: - self._W1ds = [comp.spaces for comp in W.spaces] - self._WNbasis = np.array( - [ - [self._W1ds[0][0].nbasis, self._W1ds[0][1].nbasis, self._W1ds[0][2].nbasis], - [ - self._W1ds[1][0].nbasis, - self._W1ds[1][1].nbasis, - self._W1ds[1][2].nbasis, - ], - [self._W1ds[2][0].nbasis, self._W1ds[2][1].nbasis, self._W1ds[2][2].nbasis], - ] - ) - for i in range(3): - if(self._WNbasis[1][i] < self._VNbasis[1][i]): - #If this breaks for clamped splines it means .nbasis gives you the number of basis functions, not the number of elements - assert self._VNbasis[0][i] == self._WNbasis[0][i]*2 and self._VNbasis[1][i] == self._WNbasis[1][i]*2 and self._VNbasis[2][i] == self._WNbasis[2][i]*2 - self._halving_directions[i] = True - - # We get the start and endpoint for each sublist in out - self._out_starts = np.array([vi.starts for vi in W.coeff_space.spaces]) - self._out_ends = np.array([vi.ends for vi in W.coeff_space.spaces]) - - - - # Degree of the B-spline space, not to be confused with the degrees given by fem_space.spaces.degree since depending on the situation - # it will give the D-spline degree instead - self._p = get_b_spline_degree(V) - - #We also get the D-spline degree - self._pD = self._p - 1 - - #Now we compute the weights that define this linear operator - - #We begin by defining a list that will contain the 3 numpy arrays, each one with the weights for one spatial direction. - #In the case there are direction over which we do not halve the resolution we shall have an array with only one 1.0 - self._all_weights = [] - self._all_weightsD = [] - for i in range(3): - if self._halving_directions[i]: - #Here we store the weights needed for B-splines - weights = np.zeros(self._p[i]+2, dtype=float) - #Here we store the weights needed for D-splines - weightsD = np.zeros(self._pD[i]+2, dtype=float) - for j in range(self._p[i]+2): - weights[j] = 2.0**(-self._p[i])*comb(self._p[i]+1,j) - for j in range(self._pD[i]+2): - weightsD[j] = 2.0**-(self._pD[i]+1)*comb(self._pD[i]+1,j) - self._all_weights.append(weights) - self._all_weightsD.append(weightsD) - else: - self._all_weights.append(np.array([1.0],dtype=float)) - self._all_weightsD.append(np.array([1.0],dtype=float)) - - - - #-------------------------------------- - # Abstract interface - #-------------------------------------- - @property - def domain(self): - return self._domain - - @property - def codomain(self): - return self._codomain - - @property - def dtype(self): - return self._dtype - - def _dot_helper(self, v, out, p, weights, h=None): - """Helper function to perform dot product computation.""" - #First we get the number of weights in each direction - weights_len = [] - for i in range(3): - if self._halving_directions[i]: - weights_len.append(p[i] + 2) - else: - weights_len.append(1) - if h is None: # Scalar case (H1, L2) - for i0 in range(self._out_starts[0], self._out_ends[0] + 1): - for i1 in range(self._out_starts[1], self._out_ends[1] + 1): - for i2 in range(self._out_starts[2], self._out_ends[2] + 1): - for j0 in range(weights_len[0]): - if self._halving_directions[0]: - pos0 = (2 * i0 - p[0] + j0) % self._VNbasis[0] - else: - pos0 = i0 - for j1 in range(weights_len[1]): - if self._halving_directions[1]: - pos1 = (2 * i1 - p[1] + j1) % self._VNbasis[1] - else: - pos1 = i1 - for j2 in range(weights_len[2]): - if self._halving_directions[2]: - pos2 = (2 * i2 - p[2] + j2) % self._VNbasis[2] - else: - pos2 = i2 - out[i0, i1, i2] += weights[0][j0]* weights[1][j1] *weights[2][j2] * v[pos0, pos1, pos2] - else: # Vector case (Hcurl, Hdiv, H1H1H1) - for i0 in range(self._out_starts[h][0], self._out_ends[h][0] + 1): - for i1 in range(self._out_starts[h][1], self._out_ends[h][1] + 1): - for i2 in range(self._out_starts[h][2], self._out_ends[h][2] + 1): - for j0 in range(weights_len[0]): - if self._halving_directions[0]: - pos0 = (2 * i0 - p[0] + j0) % self._VNbasis[h][0] - else: - pos0 = i0 - for j1 in range(weights_len[1]): - if self._halving_directions[1]: - pos1 = (2 * i1 - p[1] + j1) % self._VNbasis[h][1] - else: - pos1 = i1 - for j2 in range(weights_len[2]): - if self._halving_directions[2]: - pos2 = (2 * i2 - p[2] + j2) % self._VNbasis[h][2] - else: - pos2 = i2 - out[h][i0, i1, i2] += weights[0][j0]* weights[1][j1] *weights[2][j2] * v[h][pos0, pos1, pos2] - - return out - - def dot_H1(self, v, out): - return self._dot_helper(v, out, self._p, self._all_weights) - - def dot_L2(self, v, out): - return self._dot_helper(v, out, self._pD, self._all_weightsD) - - def dot_Hcurl(self, v, out): - out = self._dot_helper(v, out, [self._pD[0],self._p[1],self._p[2]], [self._all_weightsD[0], self._all_weights[1], self._all_weights[2]], h = 0) - self._dot_helper(v, out, [self._p[0],self._pD[1],self._p[2]], [self._all_weights[0], self._all_weightsD[1], self._all_weights[2]], h = 1) - return self._dot_helper(v, out, [self._p[0],self._p[1],self._pD[2]], [self._all_weights[0], self._all_weights[1], self._all_weightsD[2]], h = 2) - - def dot_Hdiv(self, v, out): - out = self._dot_helper(v, out, [self._p[0],self._pD[1],self._pD[2]], [self._all_weights[0], self._all_weightsD[1], self._all_weightsD[2]], h = 0) - self._dot_helper(v, out, [self._pD[0],self._p[1],self._pD[2]], [self._all_weightsD[0], self._all_weights[1], self._all_weightsD[2]], h = 1) - return self._dot_helper(v, out, [self._pD[0],self._pD[1],self._p[2]], [self._all_weightsD[0], self._all_weightsD[1], self._all_weights[2]], h = 2) - - def dot_H1H1H1(self, v, out): - out = self._dot_helper(v, out, self._p, self._all_weights, h=0) - self._dot_helper(v, out, self._p, self._all_weights, h=1) - return self._dot_helper(v, out, self._p, self._all_weights, h=2) - - def dot(self, v, out=None): - - assert isinstance(v, Vector) and v.space == self.domain - - if out is None: - out = self.codomain.zeros() - else: - assert isinstance(out, Vector) and out.space == self.codomain - - if self._V_name == 'H1' or self._V_name == 'L2': - for i0 in range(self._out_starts[0], self._out_ends[0]+1): - for i1 in range(self._out_starts[1], self._out_ends[1]+1): - for i2 in range(self._out_starts[2], self._out_ends[2]+1): - out[i0,i1,i2] = 0.0 - else: - for h in range(3): - for i0 in range(self._out_starts[h][0], self._out_ends[h][0]+1): - for i1 in range(self._out_starts[h][1], self._out_ends[h][1]+1): - for i2 in range(self._out_starts[h][2], self._out_ends[h][2]+1): - out[h][i0,i1,i2] = 0.0 - - dot_methods = { - "H1": self.dot_H1, - "L2": self.dot_L2, - "Hcurl": self.dot_Hcurl, - "Hdiv": self.dot_Hdiv, - "H1H1H1": self.dot_H1H1H1, - } - - return dot_methods.get(self._V_name)(v, out) - - def transpose(self, *, out = None): - if out is None: - out = MultiGridSolver.ExtensionOperator(self._W, self._V) - else: - assert isinstance(out, MultiGridSolver.ExtensionOperator) - assert out.domain is self.codomain - assert out.codomain is self.domain - - return out - - - class ExtensionOperator(LinearOperator): - """ - Linear operator which operates between vector spaces of the same kind but different resolutions. - - We assume that the vectors in the domain belong to a De-rham space with n/2 elements, while the codomain - belong to the coarser De-rham space with n elements. - - At the moment we also assume that we are halving only the first spatial direction. - - Parameters - ---------- - V : psydac.fem.basic.FemSpace - Finite element spline space (domain, input space). - - W : psydac.fem.basic.FemSpace - Finite element spline space (codomain, output space). - - """ - def __init__(self, V, W): - - # Check domain and codomain - assert isinstance(V, FemSpace) - assert isinstance(W, FemSpace) - - self._V = V - self._W = W - - # Store info in object - self._domain = V.coeff_space - self._codomain = W.coeff_space - self._dtype = V.coeff_space.dtype - - #Can be "H1", "L2", "Hcurl", "Hdiv", "H1H1H1" - self._V_name = V.symbolic_space - self._W_name = W.symbolic_space - assert(self._V_name == self._W_name) - - #This list will tell us in which spatial direction we are halving the problem. - self._halving_directions = [False,False,False] - - # input space: 3d StencilVectorSpaces and 1d SplineSpaces of each component - if isinstance(V, TensorFemSpace): - self._V1ds = [V.spaces] - self._VNbasis = np.array([self._V1ds[0][0].nbasis, self._V1ds[0][1].nbasis, self._V1ds[0][2].nbasis]) - - # We get the start and endpoint for each sublist in input - self._in_starts = np.array(V.coeff_space.starts) - self._in_ends = np.array(V.coeff_space.ends) - else: - self._V1ds = [comp.spaces for comp in V.spaces] - self._VNbasis = np.array( - [ - [self._V1ds[0][0].nbasis, self._V1ds[0][1].nbasis, self._V1ds[0][2].nbasis], - [ - self._V1ds[1][0].nbasis, - self._V1ds[1][1].nbasis, - self._V1ds[1][2].nbasis, - ], - [self._V1ds[2][0].nbasis, self._V1ds[2][1].nbasis, self._V1ds[2][2].nbasis], - ] - ) - - # We get the start and endpoint for each sublist in input - self._in_starts = np.array([vi.starts for vi in V.coeff_space.spaces]) - self._in_ends = np.array([vi.ends for vi in V.coeff_space.spaces]) - - # output space: 3d StencilVectorSpaces and 1d SplineSpaces of each component - if isinstance(W, TensorFemSpace): - self._W1ds = [W.spaces] - self._WNbasis = np.array([self._W1ds[0][0].nbasis, self._W1ds[0][1].nbasis, self._W1ds[0][2].nbasis]) - - for i in range(3): - if(self._VNbasis[i] < self._WNbasis[i]): - #If this breaks for clamped splines it means .nbasis gives you the number of basis functions, not the number of elements - assert self._WNbasis[i] == self._VNbasis[i]*2 - self._halving_directions[i] = True - - # We get the start and endpoint for each sublist in out - self._out_starts = np.array(W.coeff_space.starts) - self._out_ends = np.array(W.coeff_space.ends) - - else: - self._W1ds = [comp.spaces for comp in W.spaces] - self._WNbasis = np.array( - [ - [self._W1ds[0][0].nbasis, self._W1ds[0][1].nbasis, self._W1ds[0][2].nbasis], - [ - self._W1ds[1][0].nbasis, - self._W1ds[1][1].nbasis, - self._W1ds[1][2].nbasis, - ], - [self._W1ds[2][0].nbasis, self._W1ds[2][1].nbasis, self._W1ds[2][2].nbasis], - ] - ) - - for i in range(3): - if(self._VNbasis[1][i] < self._WNbasis[1][i]): - assert self._WNbasis[0][i] == self._VNbasis[0][i]*2 and self._WNbasis[1][i] == self._VNbasis[1][i]*2 and self._WNbasis[2][i] == self._VNbasis[2][i]*2 - self._halving_directions[i] = True - - # We get the start and endpoint for each sublist in out - self._out_starts = np.array([vi.starts for vi in W.coeff_space.spaces]) - self._out_ends = np.array([vi.ends for vi in W.coeff_space.spaces]) - - - - # Degree of the B-spline space, not to be confused with the degrees given by fem_space.spaces.degree since depending on the situation - # it will give the D-spline degree instead - self._p = get_b_spline_degree(V) - #We also get the D-splines degree - self._pD = self._p-1 - - #Now we compute the weights that define this linear operator - - #We begin by defining a list that will contain the 3 numpy arrays, each one with the weights for one spatial direction. - #In the case there are a direction over which we do not halve the resolution we shall have an array with only one 1.0 - self._all_weights_even = [] - self._all_weights_evenD = [] - self._all_weights_odd = [] - self._all_weights_oddD = [] - - #Each list has 3 integers, each one denoting the number of weights in the corresponding weights array. - self._all_size_even = [] - self._all_size_evenD = [] - self._all_size_odd = [] - self._all_size_oddD = [] - - for i in range(3): - if self._halving_directions[i]: - #First for B-splines - if(self._p[i]%2 == 0): - size_even = self._p[i]//2 +1 - size_odd = self._p[i]//2 +1 - weights_even = np.zeros(size_even, dtype=float) - weights_odd = np.zeros(size_odd, dtype=float) - for j in range(size_even): - weights_even[j] = 2.0**(-self._p[i])*comb(self._p[i]+1,2*j) - weights_odd[j] = 2.0**(-self._p[i])*comb(self._p[i]+1,2*j+1) - else: - size_even = (self._p[i]+1)//2 +1 - size_odd = (self._p[i]-1)//2 +1 - weights_even = np.zeros(size_even, dtype=float) - weights_odd = np.zeros(size_odd, dtype=float) - for j in range(size_even): - weights_even[j] = 2.0**(-self._p[i])*comb(self._p[i]+1,2*j) - for j in range(size_odd): - weights_odd[j] = 2.0**(-self._p[i])*comb(self._p[i]+1,2*j+1) - - #Second for D-splines - if(self._pD[i]%2 == 0): - size_evenD = self._pD[i]//2 +1 - size_oddD = self._pD[i]//2 +1 - weights_evenD = np.zeros(size_evenD, dtype=float) - weights_oddD = np.zeros(size_oddD, dtype=float) - for j in range(size_evenD): - weights_evenD[j] = 2.0**-(self._pD[i]+1)*comb(self._pD[i]+1,2*j) - weights_oddD[j] = 2.0**-(self._pD[i]+1)*comb(self._pD[i]+1,2*j+1) - else: - size_evenD = (self._pD[i]+1)//2 +1 - size_oddD = (self._pD[i]-1)//2 +1 - weights_evenD = np.zeros(size_evenD, dtype=float) - weights_oddD = np.zeros(size_oddD, dtype=float) - for j in range(size_evenD): - weights_evenD[j] = 2.0**-(self._pD[i]+1)*comb(self._pD[i]+1,2*j) - for j in range(size_oddD): - weights_oddD[j] = 2.0**-(self._pD[i]+1)*comb(self._pD[i]+1,2*j+1) - - self._all_weights_even.append(weights_even) - self._all_weights_evenD.append(weights_evenD) - self._all_weights_odd.append(weights_odd) - self._all_weights_oddD.append(weights_oddD) - self._all_size_even.append(size_even) - self._all_size_evenD.append(size_evenD) - self._all_size_odd.append(size_odd) - self._all_size_oddD.append(size_oddD) - - else: - self._all_weights_even.append(np.array([1.0],dtype=float)) - self._all_weights_evenD.append(np.array([1.0],dtype=float)) - self._all_weights_odd.append(np.array([1.0],dtype=float)) - self._all_weights_oddD.append(np.array([1.0],dtype=float)) - self._all_size_even.append(1) - self._all_size_evenD.append(1) - self._all_size_odd.append(1) - self._all_size_oddD.append(1) - - - #-------------------------------------- - # Abstract interface - #-------------------------------------- - @property - def domain(self): - return self._domain - - @property - def codomain(self): - return self._codomain - - @property - def dtype(self): - return self._dtype - - def tosparse(self): - pass - - def toarray(self): - pass - - def _dot_helper(self, v, out, p, weights_even, weights_odd, size_even, size_odd, h=None): - """Helper function to perform dot product computation.""" - parity_match = [] - for i in range(3): - parity_match.append(p[i] % 2) - - if h is None: # Scalar case (H1, L2) - for j0 in range(self._out_starts[0], self._out_ends[0] + 1): - parity_j0 = j0 % 2 - weights0, size0, offset0 = ((weights_even[0], size_even[0], 0) if parity_j0 == parity_match[0] else (weights_odd[0], size_odd[0], 1)) - for j1 in range(self._out_starts[1], self._out_ends[1] + 1): - parity_j1 = j1 % 2 - weights1, size1, offset1 = ((weights_even[1], size_even[1], 0) if parity_j1 == parity_match[1] else (weights_odd[1], size_odd[1], 1)) - for j2 in range(self._out_starts[2], self._out_ends[2] + 1): - parity_j2 = j2 % 2 - weights2, size2, offset2 = ((weights_even[2], size_even[2], 0) if parity_j2 == parity_match[2] else (weights_odd[2], size_odd[2], 1)) - for i0 in range(size0): - if self._halving_directions[0]: - pos0 = ((j0 + p[0] - 2 * i0 - offset0) // 2) % self._VNbasis[0] - else: - pos0 = j0 - for i1 in range(size1): - if self._halving_directions[1]: - pos1 = ((j1 + p[1] - 2 * i1 - offset1) // 2) % self._VNbasis[1] - else: - pos1 = j1 - for i2 in range(size2): - if self._halving_directions[2]: - pos2 = ((j2 + p[2] - 2 * i2 - offset2) // 2) % self._VNbasis[2] - else: - pos2 = j2 - out[j0, j1, j2] += weights0[i0] * weights1[i1] * weights2[i2] * v[pos0, pos1, pos2] - - else: # Vector case (Hcurl, Hdiv, H1H1H1) - for j0 in range(self._out_starts[h][0], self._out_ends[h][0] + 1): - parity_j0 = j0 % 2 - weights0, size0, offset0 = ((weights_even[0], size_even[0], 0) if parity_j0 == parity_match[0] else (weights_odd[0], size_odd[0], 1)) - for j1 in range(self._out_starts[h][1], self._out_ends[h][1] + 1): - parity_j1 = j1 % 2 - weights1, size1, offset1 = ((weights_even[1], size_even[1], 0) if parity_j1 == parity_match[1] else (weights_odd[1], size_odd[1], 1)) - for j2 in range(self._out_starts[h][2], self._out_ends[h][2] + 1): - parity_j2 = j2 % 2 - weights2, size2, offset2 = ((weights_even[2], size_even[2], 0) if parity_j2 == parity_match[2] else (weights_odd[2], size_odd[2], 1)) - for i0 in range(size0): - if self._halving_directions[0]: - pos0 = ((j0 + p[0] - 2 * i0 - offset0) // 2) % self._VNbasis[h][0] - else: - pos0 = j0 - for i1 in range(size1): - if self._halving_directions[1]: - pos1 = ((j1 + p[1] - 2 * i1 - offset1) // 2) % self._VNbasis[h][1] - else: - pos1 = j1 - for i2 in range(size2): - if self._halving_directions[2]: - pos2 = ((j2 + p[2] - 2 * i2 - offset2) // 2) % self._VNbasis[h][2] - else: - pos2 = j2 - out[h][j0, j1, j2] += weights0[i0] * weights1[i1] * weights2[i2] * v[h][pos0, pos1, pos2] - - return out - - def dot_H1(self, v, out): - return self._dot_helper(v, out, self._p, self._all_weights_even, self._all_weights_odd, self._all_size_even, self._all_size_odd) - - def dot_L2(self, v, out): - return self._dot_helper(v, out, self._pD, self._all_weights_evenD, self._all_weights_oddD, self._all_size_evenD, self._all_size_oddD) - - def dot_Hcurl(self, v, out): - out = self._dot_helper(v, out, [self._pD[0], self._p[1],self._p[2]], [self._all_weights_evenD[0],self._all_weights_even[1],self._all_weights_even[2]], [self._all_weights_oddD[0],self._all_weights_odd[1],self._all_weights_odd[2]], [self._all_size_evenD[0],self._all_size_even[1],self._all_size_even[2]], [self._all_size_oddD[0],self._all_size_odd[1],self._all_size_odd[2]], h=0) - out = self._dot_helper(v, out, [self._p[0], self._pD[1],self._p[2]], [self._all_weights_even[0],self._all_weights_evenD[1],self._all_weights_even[2]], [self._all_weights_odd[0],self._all_weights_oddD[1],self._all_weights_odd[2]], [self._all_size_even[0],self._all_size_evenD[1],self._all_size_even[2]], [self._all_size_odd[0],self._all_size_oddD[1],self._all_size_odd[2]], h=1) - return self._dot_helper(v, out, [self._p[0], self._p[1],self._pD[2]], [self._all_weights_even[0],self._all_weights_even[1],self._all_weights_evenD[2]], [self._all_weights_odd[0],self._all_weights_odd[1],self._all_weights_oddD[2]], [self._all_size_even[0],self._all_size_even[1],self._all_size_evenD[2]], [self._all_size_odd[0],self._all_size_odd[1],self._all_size_oddD[2]], h=2) - - - def dot_Hdiv(self, v, out): - out = self._dot_helper(v, out, [self._p[0], self._pD[1],self._pD[2]], [self._all_weights_even[0],self._all_weights_evenD[1],self._all_weights_evenD[2]], [self._all_weights_odd[0],self._all_weights_oddD[1],self._all_weights_oddD[2]], [self._all_size_even[0],self._all_size_evenD[1],self._all_size_evenD[2]], [self._all_size_odd[0],self._all_size_oddD[1],self._all_size_oddD[2]], h=0) - out = self._dot_helper(v, out, [self._pD[0], self._p[1],self._pD[2]], [self._all_weights_evenD[0],self._all_weights_even[1],self._all_weights_evenD[2]], [self._all_weights_oddD[0],self._all_weights_odd[1],self._all_weights_oddD[2]], [self._all_size_evenD[0],self._all_size_even[1],self._all_size_evenD[2]], [self._all_size_oddD[0],self._all_size_odd[1],self._all_size_oddD[2]], h=1) - return self._dot_helper(v, out, [self._pD[0], self._pD[1],self._p[2]], [self._all_weights_evenD[0],self._all_weights_evenD[1],self._all_weights_even[2]], [self._all_weights_oddD[0],self._all_weights_oddD[1],self._all_weights_odd[2]], [self._all_size_evenD[0],self._all_size_evenD[1],self._all_size_even[2]], [self._all_size_oddD[0],self._all_size_oddD[1],self._all_size_odd[2]], h=2) - - - def dot_H1H1H1(self, v, out): - out = self._dot_helper(v, out, self._p, self._all_weights_even, self._all_weights_odd, self._all_size_even, self._all_size_odd, h = 0) - self._dot_helper(v, out, self._p, self._all_weights_even, self._all_weights_odd, self._all_size_even, self._all_size_odd, h = 1) - return self._dot_helper(v, out, self._p, self._all_weights_even, self._all_weights_odd, self._all_size_even, self._all_size_odd, h = 2) - - - - def dot(self, v, out=None): - - assert isinstance(v, Vector) and v.space == self.domain - - if out is None: - out = self.codomain.zeros() - else: - assert isinstance(out, Vector) and out.space == self.codomain - - if self._V_name == 'H1' or self._V_name == 'L2': - for i0 in range(self._out_starts[0], self._out_ends[0]+1): - for i1 in range(self._out_starts[1], self._out_ends[1]+1): - for i2 in range(self._out_starts[2], self._out_ends[2]+1): - out[i0,i1,i2] = 0.0 - else: - for h in range(3): - for i0 in range(self._out_starts[h][0], self._out_ends[h][0]+1): - for i1 in range(self._out_starts[h][1], self._out_ends[h][1]+1): - for i2 in range(self._out_starts[h][2], self._out_ends[h][2]+1): - out[h][i0,i1,i2] = 0.0 - - dot_methods = { - "H1": self.dot_H1, - "L2": self.dot_L2, - "Hcurl": self.dot_Hcurl, - "Hdiv": self.dot_Hdiv, - "H1H1H1": self.dot_H1H1H1, - } - - return dot_methods.get(self._V_name)(v, out) - - def transpose(self, *, out = None): - if out is None: - out = MultiGridSolver.RestrictionOperator(self._W, self._V) - else: - assert isinstance(out, MultiGridSolver.RestrictionOperator) - assert out.domain is self.codomain - assert out.codomain is self.domain - - return out - - - def solve(self, b, verbose=False): - max_iter = self.max_iter_list - N_cycles = self.N_cycles - - #We define a list where to store the number of itteration it takes at each multigrid level - #Change N_levels for 1D case - Multigrid_itterations = np.zeros(self.N_levels, dtype=int) - converged = np.ones(self.N_levels, dtype=bool) - - def V_cycle(l, r_l): - #Change for N_levels-1 for 1D case - if (l < self.N_levels-1): - solver_ini = inverse(self.A[l],self.method, maxiter= max_iter[l]) - x_l = solver_ini.dot(r_l) - - #We count the number of itterations - Multigrid_itterations[l] += solver_ini._info['niter'] - - #We determine if the itterative solver converged in the maximum number of itterations - converged[l] = solver_ini._info['success'] - if converged[l] == True: - return x_l - - r_l = r_l - self.A[l].dot(x_l) - - r_l_plus_1 = self.R[l].dot(r_l) - x_l_plus_1 = V_cycle(l+1, r_l_plus_1) - #New - #x_l_aux = self.E[l].dot(x_l_plus_1) - #x_l = x_l + x_l_aux - #r_l = r_l - self.A[l].dot(x_l_aux) - #solver_end = inverse(self.A[l].T,self.method, maxiter= max_iter, x0 =x_l) - #x_l = solver_end.dot(r_l) - #Multigrid_itterations[l] += solver_end._info['niter'] - #### - #old - x_l = x_l + self.E[l].dot(x_l_plus_1) - ### - - else: - #Solve directly - x_l = direct_solver(self.A_inv,r_l, self.derhamlist[l].Vh_fem[self.sp_key]) - return x_l - - timei = time.time() - for cycle in range(N_cycles): - if cycle == 0: - solver = inverse(self.A[0],self.method, maxiter= max_iter[0], tol=1e-8) - else: - solver = inverse(self.A[0],self.method, maxiter= max_iter[0], x0 = x_0, tol=1e-8) - x_0 = solver.dot(b) - - Multigrid_itterations[0] += solver._info['niter'] - - #We determine if the itterative solver converged in the maximum number of itterations - converged[0] = solver._info['success'] - if converged[0] == True: - x = x_0 - break - - r_0 = b - self.A[0].dot(x_0) - r_1 = self.R[0].dot(r_0) - - x_0 = x_0 + self.E[0].dot(V_cycle(1,r_1)) - - if converged[0] == False: - solver = inverse(self.A[0],self.method,x0 = x_0, tol = 10**(-10)) - x = solver.dot(b) - Multigrid_itterations[0] += solver._info['niter'] - - timef = time.time() - - #We get the final error - Multigrid_error = solver._info['res_norm'] - - Multigrid_time = timef- timei - - if verbose: - print("################") - print("################") - print("################") - print(f'{max_iter = }') - print(f'{N_cycles = }') - print("################") - print("################") - print(f'{Multigrid_itterations = }') - print(f'{Multigrid_error = }') - print(f'{Multigrid_time = }') - print(f'{converged = }') - print("################") - print("################") - return x - - -class MultiGridPoissonSolver(MultiGridSolver): - def __init__(self, derham: Derham, sp_key: str, N_levels: int, domain, max_iter_list: list, N_cycles: int, method: str = 'cg'): - expr = " grad.T @ M1 @ grad" - namespace = {"grad.T": "derham.grad.T", - "grad": "derham.grad", - "M1": "WeightedMassOperators(derham, domain).M1"} - super().__init__(derham, sp_key, N_levels, domain, expr, namespace, max_iter_list, N_cycles, method) - - - -if __name__ == "__main__": - comm = MPI.COMM_WORLD - Nel = [16,16,1] - sp_key = '0' - plist = [2,2,1] - spl_kind = [True,True,True] - N_levels = 3 - domain = Cuboid() - derham = Derham([Nel[0],Nel[1],Nel[2]], plist, spl_kind, comm=comm, local_projectors=False) - max_iter = [10, 25, 25] - N_cycles = 5 - method = 'cg' - - expr = " grad.T @ M1 @ grad" - namespace = {"grad.T": "derham.grad.T", - "grad": "derham.grad", - "M1": "WeightedMassOperators(derham, domain).M1"} - - - # 1. Create a local context for the eval. It maps string names - # to the actual objects for the CURRENT level. - eval_context = { - "derham": derham, - "domain": domain, - "WeightedMassOperators": WeightedMassOperators - } - - # 2. Build the namespace by evaluating each string from the template - # within the context of the current level. - updated_namespace = { - key: eval(recipe_string, {}, eval_context) - for key, recipe_string in namespace.items() - } - - A = build_operator(expr, updated_namespace) - - multigrid = MultiGridPoissonSolver(derham, sp_key, N_levels, domain, max_iter, N_cycles, method) - - u_stararr, u_star = create_equal_random_arrays(derham.Vh_fem[sp_key], seed=8765) - #We compute the rhs - b = A.dot(u_star) - b_arr = b.toarray() - u = multigrid.solve(b, verbose=True) - b_ans_arr = A.dot(u).toarray() - if np.allclose(b_ans_arr, b_arr, atol=1e-6): - print("The multigrid solver computed the correct solution.") - else: - print("The multigrid solver did not compute the correct solution.") - print(f"{b_ans_arr = }") - print(f"{b_arr = }") \ No newline at end of file diff --git a/src/struphy/linear_algebra/tests/test_multigrid.py b/src/struphy/linear_algebra/tests/test_multigrid.py deleted file mode 100644 index 9db8c7dd2..000000000 --- a/src/struphy/linear_algebra/tests/test_multigrid.py +++ /dev/null @@ -1,83 +0,0 @@ -from struphy.linear_algebra.multigrid_solver import MultiGridPoissonSolver, build_operator -from struphy.geometry.domains import Cuboid -from struphy.feec.psydac_derham import Derham -import numpy as np -from mpi4py import MPI -from struphy.feec.mass import WeightedMassOperators -from struphy.feec.utilities import create_equal_random_arrays - -def test_multigrid_poisson_solver(): - """Test for the MultiGridPoissonSolver class solving a Poisson problem.""" - comm = MPI.COMM_WORLD - Nel = [16,16,1] - sp_key = '0' - plist = [2,2,1] - spl_kind = [True,True,True] - N_levels = 3 - domain = Cuboid() - derham = Derham([Nel[0],Nel[1],Nel[2]], plist, spl_kind, comm=comm, local_projectors=False) - max_iter = [10, 25, 25] - N_cycles = 5 - method = 'cg' - - # We build the operator A = grad^T * M1 * grad - A = derham.grad.T @ WeightedMassOperators(derham, domain).M1 @ derham.grad - - # We build the multigrid solver - multigrid = MultiGridPoissonSolver(derham, sp_key, N_levels, domain, max_iter, N_cycles, method) - - # We create a random solution u_star - _, u_star = create_equal_random_arrays(derham.Vh_fem[sp_key], seed=8765) - #We compute the rhs - b = A.dot(u_star) - b_arr = b.toarray() - # We solve the system using the multigrid solver - u = multigrid.solve(b, verbose=True) - - # We check that the solution is correct - b_ans_arr = A.dot(u).toarray() - assert np.allclose(b_ans_arr, b_arr, atol=1e-6) - -def test_build_operator(): - """Test for the build_operator function.""" - comm = MPI.COMM_WORLD - Nel = [8,8,1] - plist = [2,2,1] - spl_kind = [True,True,True] - domain = Cuboid() - derham = Derham([Nel[0],Nel[1],Nel[2]], plist, spl_kind, comm=comm, local_projectors=False) - - expr = " grad.T @ M1 @ grad" - - namespace = {"grad.T": "derham.grad.T", - "grad": "derham.grad", - "M1": "WeightedMassOperators(derham, domain).M1"} - - - # 1. Create a local context for the eval. It maps string names - # to the actual objects for the CURRENT level. - eval_context = { - "derham": derham, - "domain": domain, - "WeightedMassOperators": WeightedMassOperators - } - - # 2. Build the namespace by evaluating each string from the template - # within the context of the current level. - updated_namespace = { - key: eval(recipe_string, {}, eval_context) - for key, recipe_string in namespace.items() - } - - A = build_operator(expr, updated_namespace) - - A_ref = derham.grad.T @ WeightedMassOperators(derham, domain).M1 @ derham.grad - A_arr = A.toarray() - A_ref_arr = A_ref.toarray() - - assert np.allclose(A_arr, A_ref_arr, atol=1e-10) - - -if __name__ == "__main__": - test_build_operator() - test_multigrid_poisson_solver() \ No newline at end of file From 0225082cb97561ce8b27d6cfa520dfcb00c9ff93 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Sat, 3 Oct 2026 17:35:33 +0200 Subject: [PATCH 04/16] Derham: optional prescribed domain_decomposition Allows building a coarse Derham whose MPI decomposition is aligned with a finer one (DomainDecomposition.coarsen), as needed for geometric multigrid. Bumps feectools to the ddm-coarsen branch. Co-Authored-By: Claude Opus 5.5 --- feectools | 2 +- src/struphy/feec/psydac_derham.py | 28 +++++++++++++++++++++++++--- 2 files changed, 26 insertions(+), 4 deletions(-) diff --git a/feectools b/feectools index 2e3aa651d..6a11fbd04 160000 --- a/feectools +++ b/feectools @@ -1 +1 @@ -Subproject commit 2e3aa651d6d7fad5072dd39ff8c375fb7ed8ae1d +Subproject commit 6a11fbd0440602920d4a1ae2276742adb1c42f05 diff --git a/src/struphy/feec/psydac_derham.py b/src/struphy/feec/psydac_derham.py index e16af24c5..12f0acc57 100644 --- a/src/struphy/feec/psydac_derham.py +++ b/src/struphy/feec/psydac_derham.py @@ -563,6 +563,11 @@ class Derham: domain : Domain, optional The Struphy domain object for evaluating the mapping F : [0, 1]^3 --> R^3 and the corresponding metric coefficients. + domain_decomposition : DomainDecomposition, optional + Prescribed MPI decomposition of the elements, e.g. ``fine_derham.domain_decomposition.coarsen(...)`` + for an aligned multigrid level. Must match ``grid.num_elements`` and the periodicity from ``options.bcs``, + and be built on ``comm``. If None (default), it is computed from ``comm`` and ``grid.mpi_dims_mask``. + Notes ----- The underlying base sequence is @@ -579,6 +584,8 @@ def __init__( options: DerhamOptions, comm: MPI.Intracomm = None, domain: Domain = None, + *, + domain_decomposition: DomainDecomposition | None = None, ): # inputs @@ -680,6 +687,7 @@ def __init__( comm=self.comm, mpi_dims_mask=mpi_dims_mask, use_feectools=use_feectools, + domain_decomposition=domain_decomposition, ) # FEM spaces @@ -1520,6 +1528,7 @@ def init_derham( comm=None, mpi_dims_mask: tuple[bool, bool, bool] = None, use_feectools: bool = True, + domain_decomposition: DomainDecomposition | None = None, ) -> DiscreteDerham: """Return a discrete Derham complex. Allows for the use of tiny-feectools. @@ -1543,12 +1552,25 @@ def init_derham( use_feectools: bool Use slimmed-down fork `feectools` of Psydac. + + domain_decomposition : DomainDecomposition, optional + Prescribed decomposition of the elements; if None, it is computed from ``comm`` and ``mpi_dims_mask``. """ if use_feectools: - self._domain_decomposition = DomainDecomposition( - num_elements, spl_kind, comm=comm, mpi_dims_mask=mpi_dims_mask - ) + if domain_decomposition is None: + self._domain_decomposition = DomainDecomposition( + num_elements, spl_kind, comm=comm, mpi_dims_mask=mpi_dims_mask + ) + else: + assert tuple(domain_decomposition.ncells) == tuple(num_elements), ( + f"{domain_decomposition.ncells = } does not match {num_elements = }." + ) + assert tuple(domain_decomposition.periods) == tuple(spl_kind), ( + f"{domain_decomposition.periods = } does not match {spl_kind = }." + ) + assert domain_decomposition.comm is comm, "domain_decomposition must be built on the Derham communicator." + self._domain_decomposition = domain_decomposition _derham = self._discretize_derham( num_elements, From b7f250630862492b9504b369043340ca1cc01180 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Sat, 3 Oct 2026 17:42:14 +0200 Subject: [PATCH 05/16] Derham: compare communicators only when the decomposition has one feectools sets DomainDecomposition.comm to None under MockMPI. Co-Authored-By: Claude Opus 5.5 --- src/struphy/feec/psydac_derham.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/struphy/feec/psydac_derham.py b/src/struphy/feec/psydac_derham.py index 12f0acc57..5f8c31503 100644 --- a/src/struphy/feec/psydac_derham.py +++ b/src/struphy/feec/psydac_derham.py @@ -1569,7 +1569,9 @@ def init_derham( assert tuple(domain_decomposition.periods) == tuple(spl_kind), ( f"{domain_decomposition.periods = } does not match {spl_kind = }." ) - assert domain_decomposition.comm is comm, "domain_decomposition must be built on the Derham communicator." + if domain_decomposition.comm is not None and comm is not None: + # (comm is None in the decomposition when feectools runs with MockMPI) + assert domain_decomposition.comm == comm, "domain_decomposition must be built on the Derham communicator." self._domain_decomposition = domain_decomposition _derham = self._discretize_derham( From a93c7c411cd638ffaae5f7c51d1aee321e01385d Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Sat, 3 Oct 2026 17:42:15 +0200 Subject: [PATCH 06/16] Multigrid: Derham hierarchy and MPI-parallel grid transfer MultiGridHierarchy coarsens a Derham in all directions where possible (semi-coarsening otherwise), with MPI decompositions aligned to the finest level. SplineProlongation is the exact embedding of nested spline spaces (1D matrices by collocation, any degree, periodic or clamped, B- and D-splines), applied per direction on the local ghosted arrays; its transpose is the restriction. Homogeneous Dirichlet BCs are handled with the boundary operators. Co-Authored-By: Claude Opus 5.5 --- .../linear_algebra/multigrid/__init__.py | 0 .../linear_algebra/multigrid/hierarchy.py | 106 ++++++++ .../linear_algebra/multigrid/transfer.py | 255 ++++++++++++++++++ .../tests/test_multigrid_transfer.py | 111 ++++++++ 4 files changed, 472 insertions(+) create mode 100644 src/struphy/linear_algebra/multigrid/__init__.py create mode 100644 src/struphy/linear_algebra/multigrid/hierarchy.py create mode 100644 src/struphy/linear_algebra/multigrid/transfer.py create mode 100644 src/struphy/linear_algebra/tests/test_multigrid_transfer.py diff --git a/src/struphy/linear_algebra/multigrid/__init__.py b/src/struphy/linear_algebra/multigrid/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/src/struphy/linear_algebra/multigrid/hierarchy.py b/src/struphy/linear_algebra/multigrid/hierarchy.py new file mode 100644 index 000000000..77bf8c8df --- /dev/null +++ b/src/struphy/linear_algebra/multigrid/hierarchy.py @@ -0,0 +1,106 @@ +"""Hierarchy of nested Derham sequences for geometric multigrid.""" + +import logging + +import numpy as np +from feectools.ddm.cart import DomainDecomposition + +from struphy.feec.psydac_derham import Derham +from struphy.topology.grids import TensorProductGrid + +logger = logging.getLogger("struphy") + + +class MultiGridHierarchy: + r"""Nested Derham sequences :math:`V_0 \supset V_1 \supset \dots \supset V_{L-1}` obtained by + uniform coarsening of the user's (finest) Derham. + + From one level to the next, the number of elements is halved in every direction ``i`` where this + is possible, i.e. where + + * ``num_elements[i]`` is even, + * the coarse grid keeps at least ``max(min_cells, degree[i] + 1)`` elements, + * the MPI decomposition stays aligned: every process owns exactly the coarse elements covering + its fine elements (element starts and ends+1 are even), and owns at least ``degree[i]`` + coarse elements if the direction is split among several processes. + + Directions that cannot be halved are kept (semi-coarsening). Coarsening stops when no direction + can be halved or when ``max_levels`` is reached. All coarse Derhams share the communicator, + the process grid, the options (degree, boundary conditions, quadrature) and the domain of the finest one. + + Parameters + ---------- + derham : Derham + The finest level. + + max_levels : int | None + Maximal number of levels (including the finest); None means as many as possible. + + min_cells : int + Minimal number of elements per direction on the coarsest level. + """ + + def __init__(self, derham: Derham, *, max_levels: int | None = None, min_cells: int = 2): + if derham.polar_splines: + raise NotImplementedError("Multigrid is not yet implemented for polar splines.") + assert max_levels is None or max_levels >= 1 + assert min_cells >= 1 + + self._min_cells = min_cells + self._derhams: list[Derham] = [derham] + self._factors: list[tuple[int, int, int]] = [] + + while max_levels is None or len(self._derhams) < max_levels: + fine = self._derhams[-1] + factors = self._coarsening_factors(fine) + if all(f == 1 for f in factors): + break + + ddm = fine.domain_decomposition.coarsen(factors) + grid = TensorProductGrid( + num_elements=tuple(int(n) for n in ddm.ncells), + mpi_dims_mask=fine.grid.mpi_dims_mask, + ) + coarse = Derham(grid, fine.options, comm=fine.comm, domain=fine.domain, domain_decomposition=ddm) + + self._derhams.append(coarse) + self._factors.append(factors) + + logger.debug(f"Multigrid hierarchy: {[d.num_elements for d in self._derhams]}") + + def _coarsening_factors(self, derham: Derham) -> tuple[int, int, int]: + """Return 2 in every direction that can be halved (see class docstring), else 1.""" + ddm: DomainDecomposition = derham.domain_decomposition + factors = [] + for axis in range(3): + n = derham.num_elements[axis] + p = derham.degree[axis] + starts = np.asarray(ddm.global_element_starts[axis]) + ends = np.asarray(ddm.global_element_ends[axis]) + ok = n % 2 == 0 and n // 2 >= max(self._min_cells, p + 1) + ok = ok and np.all(starts % 2 == 0) and np.all((ends + 1) % 2 == 0) + if ok and ddm.nprocs[axis] > 1: + ok = np.all((ends - starts + 1) // 2 >= p) + factors.append(2 if ok else 1) + return tuple(factors) + + @property + def derhams(self) -> list[Derham]: + """Derham of each level, ``derhams[0]`` is the finest.""" + return self._derhams + + @property + def n_levels(self) -> int: + """Number of levels (including the finest).""" + return len(self._derhams) + + @property + def factors(self) -> list[tuple[int, int, int]]: + """``factors[l]`` is the coarsening factor in each direction from level ``l`` to ``l+1``.""" + return self._factors + + def __getitem__(self, level: int) -> Derham: + return self._derhams[level] + + def __len__(self) -> int: + return self.n_levels diff --git a/src/struphy/linear_algebra/multigrid/transfer.py b/src/struphy/linear_algebra/multigrid/transfer.py new file mode 100644 index 000000000..3bbacd781 --- /dev/null +++ b/src/struphy/linear_algebra/multigrid/transfer.py @@ -0,0 +1,255 @@ +r"""Grid-transfer operators between nested spline spaces. + +On uniformly refined grids the spline spaces are nested, :math:`V_H \subset V_h`, hence every coarse +basis function is a linear combination of fine ones. The prolongation :math:`P: V_H \to V_h` maps the +coefficients of a coarse spline to the coefficients of the *same* function in the fine basis; the +restriction is its transpose :math:`R = P^\top`. Both are tensor products of 1D matrices, applied +component-wise for vector-valued spaces. +""" + +import numpy as np +import scipy.sparse as spa +from feectools.core.bsplines import collocation_matrix +from feectools.fem.splines import SplineSpace +from feectools.fem.tensor import TensorFemSpace +from feectools.linalg.basic import IdentityOperator, LinearOperator, Vector +from feectools.linalg.block import BlockVector +from feectools.linalg.stencil import StencilVector, StencilVectorSpace + +from struphy.feec.psydac_derham import Derham + + +def prolongation_matrix_1d(coarse: SplineSpace, fine: SplineSpace) -> np.ndarray: + r"""Dense 1D prolongation matrix :math:`P \in \mathbb R^{n_h \times n_H}` between nested spline spaces. + + Column ``j`` holds the fine coefficients of the coarse basis function ``j``, i.e. + :math:`\Lambda^H_j = \sum_i P_{ij} \Lambda^h_i`. It is computed by collocation at ``degree + 1`` + points per fine cell (an overdetermined, exactly solvable system). Works for periodic and clamped + splines and for both normalizations (B-splines and D-splines/M-splines). + + Parameters + ---------- + coarse, fine : SplineSpace + 1D spaces of the same degree, kind and normalization; the breaks of ``coarse`` are a subset of those of ``fine``. + """ + assert coarse.degree == fine.degree + assert coarse.periodic == fine.periodic + assert coarse.basis == fine.basis + + if coarse.ncells == fine.ncells: + return np.eye(fine.nbasis) + + breaks = np.asarray(fine.breaks) + nq = fine.degree + 1 + s = (np.arange(nq) + 0.5) / nq + x = (breaks[:-1, None] + np.diff(breaks)[:, None] * s[None, :]).ravel() + + Bh = collocation_matrix(fine.knots, fine.degree, fine.periodic, fine.basis, x) + BH = collocation_matrix(coarse.knots, coarse.degree, coarse.periodic, coarse.basis, x) + P = np.linalg.lstsq(Bh, BH, rcond=None)[0] + + assert np.allclose(Bh @ P, BH, atol=1e-10), "Spline spaces are not nested." + P[np.abs(P) < 1e-13 * np.abs(P).max()] = 0.0 + return P + + +def local_matrix_1d( + A: np.ndarray, + out_start: int, + out_end: int, + in_start: int, + in_end: int, + in_ghost: int, + periodic: bool, +) -> spa.csr_matrix: + r"""Restrict a global 1D matrix to the rows owned by this process and to local (ghosted) input columns. + + Row ``r`` of the result is row ``out_start + r`` of ``A``. Global column ``c`` is mapped to the index + of the local ghosted input array, ``c - in_start + in_ghost``, using periodic images if ``periodic``. + If a column is present several times (owned and as ghost) the owned copy is used. + + Parameters + ---------- + A : numpy.ndarray + Global matrix of shape ``(n_out, n_in)``. + + out_start, out_end : int + Global indices of the first and last output entry owned by this process. + + in_start, in_end : int + Global indices of the first and last input entry owned by this process. + + in_ghost : int + Width of the ghost region on each side of the local input array (``pads * shifts``). + + periodic : bool + Whether the input index is periodic. + """ + n_in = A.shape[1] + n_loc = in_end - in_start + 1 + 2 * in_ghost + + rows, cols, vals = [], [], [] + for r, row in enumerate(range(out_start, out_end + 1)): + for c in np.flatnonzero(A[row]): + images = [c - n_in, c, c + n_in] if periodic else [c] + local = [g - in_start + in_ghost for g in images] + owned = [l for l in local if in_ghost <= l < n_loc - in_ghost] + ghost = [l for l in local if 0 <= l < n_loc] + if owned: + loc = owned[0] + elif ghost: + loc = ghost[0] + else: + raise ValueError(f"Column {c} of row {row} is outside the local ghost region.") + rows.append(r) + cols.append(loc) + vals.append(A[row, c]) + + return spa.csr_matrix((vals, (rows, cols)), shape=(out_end - out_start + 1, n_loc)) + + +class _KronTransfer: + """Tensor product of three local 1D matrices mapping a ghosted input StencilVector to the owned part of the output.""" + + def __init__(self, mats: list[spa.csr_matrix], W: StencilVectorSpace): + self._mats = mats + self._out_slice = tuple(slice(p * m, p * m + e - s + 1) for p, m, s, e in zip(W.pads, W.shifts, W.starts, W.ends)) + + def dot(self, v: StencilVector, out: StencilVector) -> None: + if not v.ghost_regions_in_sync: + v.update_ghost_regions() + x = v._data + for axis, L in enumerate(self._mats): + x = np.moveaxis(x, axis, 0) + shp = x.shape + x = (L @ x.reshape(shp[0], -1)).reshape((L.shape[0],) + shp[1:]) + x = np.moveaxis(x, 0, axis) + out._data[...] = 0.0 + out._data[self._out_slice] = x + out.ghost_regions_in_sync = False + + +def _scalar_spaces(V) -> list[TensorFemSpace]: + """Scalar components of a (vector) FEM space.""" + return [V] if isinstance(V, TensorFemSpace) else list(V.spaces) + + +def _coeff_spaces(W) -> list[StencilVectorSpace]: + """Scalar components of a (block) coefficient space.""" + return [W] if isinstance(W, StencilVectorSpace) else list(W.spaces) + + +class SplineProlongation(LinearOperator): + r"""Prolongation :math:`P: V_H \to V_h` (or, with ``transposed=True``, restriction :math:`R = P^\top`) + between the same space of two nested Derham sequences. + + The MPI decompositions must be aligned (each process owns the coarse elements covering its fine + elements), as produced by :meth:`DomainDecomposition.coarsen`. With homogeneous Dirichlet boundary + conditions, the operator is :math:`\mathbb B_h P \mathbb B_H^\top` (resp. its transpose), which is the + exact embedding of the coarse into the fine space with boundary conditions. + + Parameters + ---------- + coarse, fine : Derham + Coarse and fine level. + + space_id : str + Space key of ``Derham.fem_spaces`` ("0", "1", "2", "3", "v" or "H1", "Hcurl", "Hdiv", "L2", "H1vec"). + + transposed : bool + If True, the restriction :math:`R = P^\top` (fine to coarse) is created. + """ + + def __init__(self, coarse: Derham, fine: Derham, space_id: str, *, transposed: bool = False): + if coarse.polar_splines or fine.polar_splines: + raise NotImplementedError("Grid transfer is not yet implemented for polar splines.") + + self._coarse = coarse + self._fine = fine + self._space_id = space_id + self._transposed = transposed + + form = coarse.space_to_form.get(space_id, space_id) + self._form = form + VH = coarse.fem_spaces[form] + Vh = fine.fem_spaces[form] + + self._Bc = coarse.boundary_ops[form] + self._Bf = fine.boundary_ops[form] + self._apply_bc = not (isinstance(self._Bc, IdentityOperator) and isinstance(self._Bf, IdentityOperator)) + + if transposed: + self._domain, self._codomain = fine.coeff_spaces[form], coarse.coeff_spaces[form] + V_in, V_out = Vh, VH + else: + self._domain, self._codomain = coarse.coeff_spaces[form], fine.coeff_spaces[form] + V_in, V_out = VH, Vh + + self._kron = [] + for cH, ch, Win, Wout in zip( + _scalar_spaces(VH), + _scalar_spaces(Vh), + _coeff_spaces(self._domain), + _coeff_spaces(self._codomain), + ): + mats = [] + for axis, (sH, sh) in enumerate(zip(cH.spaces, ch.spaces)): + P = prolongation_matrix_1d(sH, sh) + A = P.T if transposed else P + mats.append( + local_matrix_1d( + A, + int(Wout.starts[axis]), + int(Wout.ends[axis]), + int(Win.starts[axis]), + int(Win.ends[axis]), + int(Win.pads[axis] * Win.shifts[axis]), + sh.periodic, + ) + ) + self._kron.append(_KronTransfer(mats, Wout)) + + @property + def domain(self): + return self._domain + + @property + def codomain(self): + return self._codomain + + @property + def dtype(self): + return self._domain.dtype + + @property + def space_id(self) -> str: + return self._space_id + + @property + def transposed(self) -> bool: + return self._transposed + + def dot(self, v: Vector, out: Vector | None = None) -> Vector: + """Apply the operator, ``out = P v`` (or ``out = R v`` if transposed).""" + assert isinstance(v, Vector) and v.space == self.domain + if out is None: + out = self.codomain.zeros() + else: + assert isinstance(out, Vector) and out.space == self.codomain + + B_in, B_out = (self._Bf, self._Bc) if self._transposed else (self._Bc, self._Bf) + if self._apply_bc: + v = B_in.dot(v) + + if isinstance(v, BlockVector): + for k, vk, ok in zip(self._kron, v.blocks, out.blocks): + k.dot(vk, ok) + else: + self._kron[0].dot(v, out) + + if self._apply_bc: + B_out.dot(out, out=out) + return out + + def transpose(self, conjugate: bool = False) -> "SplineProlongation": + return SplineProlongation(self._coarse, self._fine, self._space_id, transposed=not self._transposed) diff --git a/src/struphy/linear_algebra/tests/test_multigrid_transfer.py b/src/struphy/linear_algebra/tests/test_multigrid_transfer.py new file mode 100644 index 000000000..de01ee5e0 --- /dev/null +++ b/src/struphy/linear_algebra/tests/test_multigrid_transfer.py @@ -0,0 +1,111 @@ +import numpy as np +import pytest +from feectools.ddm.mpi import mpi as MPI +from mpi4py import MPI as MPI4PY + +from struphy.feec.mass import WeightedMassOperators +from struphy.feec.psydac_derham import Derham +from struphy.feec.utilities import create_equal_random_arrays +from struphy.geometry.domains import Cuboid +from struphy.io.options import DerhamOptions +from struphy.linear_algebra.multigrid.hierarchy import MultiGridHierarchy +from struphy.linear_algebra.multigrid.transfer import SplineProlongation, prolongation_matrix_1d +from struphy.topology.grids import TensorProductGrid + +BCS = [ + (None, None, None), + (("dirichlet", "dirichlet"), None, ("free", "dirichlet")), +] + + +def _hierarchy(num_elements, degree, bcs, max_levels=None): + domain = Cuboid(l1=0.0, r1=2.0, l2=0.0, r2=1.0, l3=0.0, r3=3.0) + derham = Derham( + TensorProductGrid(num_elements=num_elements), + DerhamOptions(degree=degree, bcs=bcs), + comm=MPI.COMM_WORLD, + domain=domain, + ) + return MultiGridHierarchy(derham, max_levels=max_levels), domain + + +@pytest.mark.mpi_skip +@pytest.mark.parametrize("degree", [1, 2, 3, 4]) +@pytest.mark.parametrize("periodic", [True, False]) +@pytest.mark.parametrize("basis", ["B", "M"]) +def test_prolongation_matrix_1d(degree, periodic, basis): + """The coarse basis is reproduced exactly by the prolongated coefficients; partition of unity is kept.""" + from feectools.fem.splines import SplineSpace + + def space(n): + return SplineSpace(degree, grid=np.linspace(0.0, 1.0, n + 1), periodic=periodic, basis=basis) + + coarse, fine = space(8), space(16) + P = prolongation_matrix_1d(coarse, fine) + assert P.shape == (fine.nbasis, coarse.nbasis) + + if basis == "B": + # partition of unity: the constant function has coefficients 1 on both grids + assert np.allclose(P @ np.ones(coarse.nbasis), 1.0) + + # each fine row couples to at most ceil((p+2)/2) coarse functions + assert np.max(np.count_nonzero(P, axis=1)) <= (degree + 3) // 2 + + +@pytest.mark.parametrize("num_elements, degree", [((16, 8, 8), (3, 2, 1)), ((8, 16, 1), (2, 3, 1))]) +@pytest.mark.parametrize("bcs", BCS) +def test_hierarchy(num_elements, degree, bcs): + """Coarse levels have aligned decompositions and at least degree+1 cells per coarsened direction.""" + h, _ = _hierarchy(num_elements, degree, bcs) + assert h.n_levels >= 2 + assert len(h.factors) == h.n_levels - 1 + for l, f in enumerate(h.factors): + fine, coarse = h[l], h[l + 1] + for axis in range(3): + assert fine.num_elements[axis] == f[axis] * coarse.num_elements[axis] + assert fine.domain_decomposition.starts[axis] == f[axis] * coarse.domain_decomposition.starts[axis] + if f[axis] == 2: + assert coarse.num_elements[axis] >= degree[axis] + 1 + assert coarse.options is fine.options + # the coarsest level cannot be coarsened further + assert h._coarsening_factors(h[-1]) == (1, 1, 1) + + +@pytest.mark.parametrize("bcs", BCS) +@pytest.mark.parametrize("space_id", ["H1", "Hcurl", "Hdiv", "L2", "H1vec"]) +def test_transfer(bcs, space_id): + r"""Restriction is the transpose of the prolongation, and P is the exact embedding: R M_h P = M_H.""" + h, domain = _hierarchy((16, 8, 8), (3, 2, 1), bcs, max_levels=3) + comm = MPI4PY.COMM_WORLD + form = h[0].space_to_form[space_id] + + for l in range(h.n_levels - 1): + P = SplineProlongation(h[l + 1], h[l], space_id) + R = P.T + assert R.domain is P.codomain and R.codomain is P.domain + + _, u = create_equal_random_arrays(h[l + 1].fem_spaces[form], seed=1) + _, w = create_equal_random_arrays(h[l].fem_spaces[form], seed=2) + assert np.isclose(R.dot(w).inner(u), w.inner(P.dot(u)), rtol=1e-12) + + Mh = getattr(WeightedMassOperators(h[l], domain), "M" + form) + MH = getattr(WeightedMassOperators(h[l + 1], domain), "M" + form) + a = R.dot(Mh.dot(P.dot(u))) + b = MH.dot(u) + err = comm.allreduce(np.max(np.abs((a - b).toarray())), op=MPI4PY.MAX) + ref = comm.allreduce(np.max(np.abs(b.toarray())), op=MPI4PY.MAX) + assert err < 1e-12 * ref + + +@pytest.mark.parametrize("bcs", BCS) +def test_prolongation_of_spline(bcs): + """The prolongated coefficients represent the same function (point evaluation).""" + h, _ = _hierarchy((8, 8, 4), (2, 3, 1), bcs, max_levels=2) + P = SplineProlongation(h[1], h[0], "H1") + _, u = create_equal_random_arrays(h[1].fem_spaces["0"], seed=4) + u = h[1].boundary_ops["0"].dot(u) + + fH = h[1].create_spline_function("fH", "H1", coeffs=u) + fh = h[0].create_spline_function("fh", "H1", coeffs=P.dot(u)) + e = np.linspace(0.0, 1.0, 7) + assert np.allclose(fH(e, e, e), fh(e, e, e), atol=1e-12) From 6e6dc51cc265d716b8b4ef597410d514ef4204cb Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Sat, 3 Oct 2026 18:05:20 +0200 Subject: [PATCH 07/16] Bump feectools: ghost-region sync fixes in axpy and essential BCs Co-Authored-By: Claude Opus 5.5 --- feectools | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/feectools b/feectools index 6a11fbd04..eb6693cbf 160000 --- a/feectools +++ b/feectools @@ -1 +1 @@ -Subproject commit 6a11fbd0440602920d4a1ae2276742adb1c42f05 +Subproject commit eb6693cbfb9fcb6a1d65ba6130b5775daf7f84c6 From 635a61c40f91f48e254b2ca0258d47b079841407 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Sat, 3 Oct 2026 18:05:20 +0200 Subject: [PATCH 08/16] Mass and basis projection operators: to_dict/from_dict WeightedMassOperator records how it was created by WeightedMassOperators.create_weighted_mass (spaces, name, weights as given, transposition), so it can be re-created on another Derham with from_dict. The recipe is dropped when the data is modified afterwards (assemble with new weights, in-place arithmetic) or the weights are bound to the grid (quadrature values, spline functions). BasisProjectionOperator gets the same, for callable weights. Co-Authored-By: Claude Opus 5.5 --- src/struphy/feec/basis_projection_ops.py | 50 ++++++++++++++ src/struphy/feec/mass.py | 85 ++++++++++++++++++++++++ 2 files changed, 135 insertions(+) diff --git a/src/struphy/feec/basis_projection_ops.py b/src/struphy/feec/basis_projection_ops.py index ff5ecf806..43d3d4c86 100644 --- a/src/struphy/feec/basis_projection_ops.py +++ b/src/struphy/feec/basis_projection_ops.py @@ -1851,6 +1851,56 @@ def dot(self, v, out=None, tol=1e-14, maxiter=1000): return out + @property + def is_reconstructible(self) -> bool: + """Whether the operator can be re-created from :meth:`to_dict` (e.g. on another Derham). + + False if a weight is given as values at the projection points (an array bound to the current Derham). + """ + return not any(isinstance(w, xp.ndarray) for row in self._weights for w in row) + + def to_dict(self) -> dict: + """Recipe for re-creating the operator with :meth:`from_dict` (on any Derham). + + Weights are stored as given (callables are kept as objects, hence the dictionary is in general not JSON serializable). + """ + if not self.is_reconstructible: + raise ValueError("BasisProjectionOperator with weights given as arrays cannot be serialized.") + V_id, W_id = ( + (self._codomain_symbolic_name, self._domain_symbolic_name) + if self._transposed + else (self._domain_symbolic_name, self._codomain_symbolic_name) + ) + return { + "type": self.__class__.__name__, + "params": { + "V_id": V_id, + "W_id": W_id, + "weights": [list(row) for row in self._weights], + "transposed": self._transposed, + "polar_shift": self._polar_shift, + "use_cache": self._use_cache, + }, + } + + @classmethod + def from_dict(cls, dct: dict, derham: Derham) -> "BasisProjectionOperator": + """Re-create a :class:`BasisProjectionOperator` from :meth:`to_dict` on the given Derham, + with the Derham's (global) commuting projector, extraction and boundary operators.""" + assert dct["type"] == cls.__name__ + params = dct["params"] + V_id, W_id = params["V_id"], params["W_id"] + return cls( + derham.projectors[W_id], + derham.fem_spaces[V_id], + [list(row) for row in params["weights"]], + V_extraction_op=derham.extraction_ops[V_id], + V_boundary_op=derham.boundary_ops[V_id], + transposed=params["transposed"], + polar_shift=params["polar_shift"], + use_cache=params["use_cache"], + ) + def transpose(self, conjugate=False): """ Returns the transposed operator. diff --git a/src/struphy/feec/mass.py b/src/struphy/feec/mass.py index 963c123ca..bf40afbad 100644 --- a/src/struphy/feec/mass.py +++ b/src/struphy/feec/mass.py @@ -1254,6 +1254,19 @@ def f_call_matrix(e1, e2, e3): dry_run=dry_run, ) + # weights given at quadrature points or as spline functions are bound to this Derham + grid_bound = len(spline_functions) > 0 or ( + isinstance(weights, list) and any(isinstance(w, xp.ndarray) for row in weights for w in row) + ) + out._creation_info = None if grid_bound else { + "V_id": V_id, + "W_id": W_id, + "name": name, + "weights": weights, + "transposed": transposed, + "is_transpose": False, + } + if assemble and not dry_run: out.assemble() @@ -1465,6 +1478,9 @@ def __init__( self._name = name self._dry_run = dry_run + # recipe for re-creating the operator with WeightedMassOperators.create_weighted_mass, see to_dict() + self._creation_info: dict | None = None + assert not (dry_run and transposed), "dry_run=True is not supported for transposed operators." # spline functions that are used as weights in the operator, to be evaluated at quadrature points @@ -2066,6 +2082,9 @@ def transpose(self, conjugate=False): # weights of M in its own (transposed) block order M._weights = [[self._weights[n][m] for n in range(len(self._weights))] for m in range(len(self._weights[0]))] + if self._creation_info is not None: + M._creation_info = dict(self._creation_info, is_transpose=not self._creation_info["is_transpose"]) + if self._matrix_free: if self._symmetry is not None: M.assemble(weights=M._weights) @@ -2107,6 +2126,9 @@ def assemble(self, weights=None, clear=True): assert not self._dry_run, ( "A dry-run operator has no matrix data and cannot be assembled (memory estimation only)." ) + if weights is not None or not clear: + # the data no longer stems from the creation recipe + self._creation_info = None if self._matrix_free: if weights is not None: @@ -2325,6 +2347,64 @@ def assemble(self, weights=None, clear=True): logger.debug("Done.") + @property + def is_reconstructible(self) -> bool: + """Whether the operator can be re-created from :meth:`to_dict` (e.g. on another Derham). + + True for operators created by :meth:`WeightedMassOperators.create_weighted_mass` (and their transposes) + whose data has not been modified afterwards (by ``assemble(weights=...)``, in-place arithmetic, ...). + """ + return self._creation_info is not None + + def to_dict(self) -> dict: + """Recipe for re-creating the operator with :meth:`WeightedMassOperators.create_weighted_mass`. + + The weights are stored as given at creation. The dictionary is JSON serializable if they are + strings (``'Ginv'``, ``'sqrt_g'``, ...) or nested lists of numbers; callables are kept as objects. + Re-create the operator (on any Derham) with :meth:`from_dict`. + """ + if self._creation_info is None: + raise ValueError( + f"WeightedMassOperator {self.name!r} cannot be serialized: it was not created by " + "WeightedMassOperators.create_weighted_mass or its data was modified afterwards." + ) + params = dict(self._creation_info) + if isinstance(params["weights"], tuple): + params["weights"] = list(params["weights"]) + return { + "type": self.__class__.__name__, + "params": params, + } + + @classmethod + def from_dict(cls, dct: dict, mass_ops: "WeightedMassOperators") -> "WeightedMassOperator": + """Re-create a :class:`WeightedMassOperator` from :meth:`to_dict` with the given collection. + + Parameters + ---------- + dct : dict + Output of :meth:`to_dict`. + + mass_ops : WeightedMassOperators + Collection providing the Derham, domain and matrix_free option of the new operator. + """ + assert dct["type"] == cls.__name__ + params = dct["params"] + name = params["name"] + weights = params["weights"] + if isinstance(weights, list) and not (len(weights) > 0 and isinstance(weights[0], list)): + weights = tuple(weights) + + out = mass_ops.create_weighted_mass( + params["V_id"], + params["W_id"], + name=name, + weights=weights, + assemble=True, + transposed=params["transposed"], + ) + return out.T if params["is_transpose"] else out + def copy(self, out=None): """Create a copy of self, that can potentially be stored in a given WeightedMassOperator. @@ -2361,6 +2441,7 @@ def copy(self, out=None): def __imul__(self, a): self._mat *= a + self._creation_info = None return self def __iadd__(self, M): @@ -2368,10 +2449,12 @@ def __iadd__(self, M): assert M.codomain is self.codomain if isinstance(M, WeightedMassOperator): + self._creation_info = None self._mat += M._mat return self elif isinstance(M, LinearOperator): + self._creation_info = None self._mat += M return self @@ -2383,10 +2466,12 @@ def __isub__(self, M): assert M.codomain is self.codomain if isinstance(M, WeightedMassOperator): + self._creation_info = None self._mat -= M._mat return self elif isinstance(M, LinearOperator): + self._creation_info = None self._mat -= M return self From ef8fc1a580db84b916c71696d81d02be1602addf Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Sat, 3 Oct 2026 18:05:21 +0200 Subject: [PATCH 09/16] Multigrid: re-discretize operators on coarse levels by tree walk OperatorCoarsener rebuilds sums, compositions, scalings, powers and block operators from coarsened children; derivative, boundary, identity and zero operators on the coarse spaces; mass and basis projection operators from to_dict(). Leaves are cached, so changed scalars do not trigger re-assembly. Tests check the Galerkin property R A P = A_H. Co-Authored-By: Claude Opus 5.5 --- .../linear_algebra/multigrid/coarsen.py | 175 ++++++++++++++++++ .../tests/test_multigrid_coarsen.py | 143 ++++++++++++++ 2 files changed, 318 insertions(+) create mode 100644 src/struphy/linear_algebra/multigrid/coarsen.py create mode 100644 src/struphy/linear_algebra/tests/test_multigrid_coarsen.py diff --git a/src/struphy/linear_algebra/multigrid/coarsen.py b/src/struphy/linear_algebra/multigrid/coarsen.py new file mode 100644 index 000000000..b33bce84e --- /dev/null +++ b/src/struphy/linear_algebra/multigrid/coarsen.py @@ -0,0 +1,175 @@ +r"""Re-discretization of a linear operator on a coarser Derham. + +A (composite) operator on the fine level is walked as an expression tree. Composite nodes (sums, +compositions, scalings, powers, block operators) are rebuilt from their coarsened children, scalars are +kept. Leaves are re-created on the coarse Derham: + +* :class:`IdentityOperator`, :class:`ZeroOperator`, :class:`BoundaryOperator`: on the corresponding coarse spaces, +* :class:`DirectionalDerivativeOperator` (blocks of ``derham.grad``, ``curl``, ``div`` and their transposes), +* :class:`WeightedMassOperator` and :class:`BasisProjectionOperator`: from their ``to_dict()`` recipe. + +Further leaf types can be supported with :func:`register_coarsening`. +""" + +from collections.abc import Callable + +from feectools.feec.derivatives import DirectionalDerivativeOperator +from feectools.linalg.basic import ( + ComposedLinearOperator, + IdentityOperator, + LinearOperator, + PowerLinearOperator, + ScaledLinearOperator, + SumLinearOperator, + VectorSpace, + ZeroOperator, +) +from feectools.linalg.block import BlockLinearOperator, BlockVectorSpace + +from struphy.feec.basis_projection_ops import BasisProjectionOperator +from struphy.feec.linear_operators import BoundaryOperator +from struphy.feec.mass import WeightedMassOperator, WeightedMassOperators +from struphy.feec.psydac_derham import Derham +from struphy.geometry.base import Domain + +_REGISTRY: dict[type, Callable[[LinearOperator, "OperatorCoarsener"], LinearOperator]] = {} + + +def register_coarsening(cls: type): + """Decorator registering ``fun(op, coarsener) -> LinearOperator`` as the coarsening rule for leaves of type ``cls``.""" + + def decorator(fun): + _REGISTRY[cls] = fun + return fun + + return decorator + + +class OperatorCoarsener: + r"""Maps linear operators on the coefficient spaces of ``fine`` to the corresponding operators on ``coarse``. + + Coarse leaves are cached (keyed by the fine leaf object), hence calling the coarsener again on an + operator that differs only in its scalars or composition (e.g. ``sigma * M0 + grad.T @ M1 @ grad`` + with a new ``sigma``) re-assembles nothing. + + Parameters + ---------- + fine, coarse : Derham + Fine and coarse level (same options, nested grids). + + domain : Domain + Mapping used for re-assembling mass matrices on the coarse level. + + matrix_free : bool + Whether coarse mass matrices are matrix-free. + """ + + def __init__(self, fine: Derham, coarse: Derham, domain: Domain, *, matrix_free: bool = False): + self._fine = fine + self._coarse = coarse + self._mass_ops = WeightedMassOperators(coarse, domain, matrix_free=matrix_free) + + # fine -> coarse coefficient spaces (also components of block spaces) + self._spaces: dict[int, tuple[VectorSpace, VectorSpace]] = {} + for form in ("0", "1", "2", "3", "v"): + Vf, Vc = fine.coeff_spaces[form], coarse.coeff_spaces[form] + self._spaces[id(Vf)] = (Vf, Vc) + if isinstance(Vf, BlockVectorSpace): + for vf, vc in zip(Vf.spaces, Vc.spaces): + self._spaces.setdefault(id(vf), (vf, vc)) + + self._cache: dict[int, tuple[LinearOperator, LinearOperator]] = {} + + @property + def fine(self) -> Derham: + return self._fine + + @property + def coarse(self) -> Derham: + return self._coarse + + @property + def mass_ops(self) -> WeightedMassOperators: + """Mass operators of the coarse level.""" + return self._mass_ops + + def space(self, V: VectorSpace) -> VectorSpace: + """Coarse counterpart of the fine coefficient space ``V``.""" + try: + return self._spaces[id(V)][1] + except KeyError: + raise ValueError(f"{V} is not a coefficient space of the fine Derham.") from None + + def __call__(self, A: LinearOperator) -> LinearOperator: + """Return the coarse-level version of the fine-level operator ``A``.""" + if isinstance(A, ScaledLinearOperator): + B = self(A.operator) + return ScaledLinearOperator(B.domain, B.codomain, A.scalar, B) + + if isinstance(A, SumLinearOperator): + addends = [self(a) for a in A.addends] + return SumLinearOperator(self.space(A.domain), self.space(A.codomain), *addends) + + if isinstance(A, ComposedLinearOperator): + factors = [self(a) for a in A.multiplicants] + return ComposedLinearOperator(self.space(A.domain), self.space(A.codomain), *factors) + + if isinstance(A, PowerLinearOperator): + return PowerLinearOperator(self.space(A.domain), self.space(A.codomain), self(A.operator), A.factorial) + + if isinstance(A, BlockLinearOperator): + blocks = {ij: self(A[ij]) for ij in A.nonzero_block_indices} + return BlockLinearOperator(self.space(A.domain), self.space(A.codomain), blocks=blocks) + + if isinstance(A, IdentityOperator): + return IdentityOperator(self.space(A.domain), self.space(A.codomain)) + + if isinstance(A, ZeroOperator): + return ZeroOperator(self.space(A.domain), self.space(A.codomain)) + + # leaves: cached + cached = self._cache.get(id(A)) + if cached is not None and cached[0] is A: + return cached[1] + + for cls in type(A).__mro__: + if cls in _REGISTRY: + B = _REGISTRY[cls](A, self) + break + else: + raise NotImplementedError( + f"Cannot re-discretize an operator of type {type(A).__name__} on a coarse grid; " + "use struphy operators or register a rule with register_coarsening." + ) + + assert B.domain is self.space(A.domain) and B.codomain is self.space(A.codomain), ( + f"Coarsening of {type(A).__name__} gave wrong (co)domain." + ) + self._cache[id(A)] = (A, B) + return B + + +@register_coarsening(DirectionalDerivativeOperator) +def _coarsen_derivative(A: DirectionalDerivativeOperator, c: OperatorCoarsener) -> LinearOperator: + return DirectionalDerivativeOperator( + c.space(A._spaceV), + c.space(A._spaceW), + A._diffdir, + negative=A._negative, + transposed=A._transposed, + ) + + +@register_coarsening(BoundaryOperator) +def _coarsen_boundary(A: BoundaryOperator, c: OperatorCoarsener) -> LinearOperator: + return BoundaryOperator(c.space(A.domain), A._space_id, A.bc) + + +@register_coarsening(WeightedMassOperator) +def _coarsen_mass(A: WeightedMassOperator, c: OperatorCoarsener) -> LinearOperator: + return WeightedMassOperator.from_dict(A.to_dict(), c.mass_ops) + + +@register_coarsening(BasisProjectionOperator) +def _coarsen_basis_projection(A: BasisProjectionOperator, c: OperatorCoarsener) -> LinearOperator: + return BasisProjectionOperator.from_dict(A.to_dict(), c.coarse) diff --git a/src/struphy/linear_algebra/tests/test_multigrid_coarsen.py b/src/struphy/linear_algebra/tests/test_multigrid_coarsen.py new file mode 100644 index 000000000..c612e553a --- /dev/null +++ b/src/struphy/linear_algebra/tests/test_multigrid_coarsen.py @@ -0,0 +1,143 @@ +import json + +import numpy as np +import pytest +from feectools.ddm.mpi import mpi as MPI +from mpi4py import MPI as MPI4PY + +from struphy.feec.mass import WeightedMassOperator, WeightedMassOperators +from struphy.feec.psydac_derham import Derham +from struphy.feec.utilities import create_equal_random_arrays +from struphy.geometry.domains import Cuboid +from struphy.io.options import DerhamOptions +from struphy.linear_algebra.multigrid.coarsen import OperatorCoarsener +from struphy.linear_algebra.multigrid.hierarchy import MultiGridHierarchy +from struphy.linear_algebra.multigrid.transfer import SplineProlongation +from struphy.topology.grids import TensorProductGrid + +BCS = [ + (None, None, None), + (("dirichlet", "dirichlet"), None, ("free", "dirichlet")), +] + + +def _derham(num_elements, degree, bcs): + domain = Cuboid(l1=0.0, r1=2.0, l2=0.0, r2=1.0, l3=0.0, r3=3.0) + derham = Derham( + TensorProductGrid(num_elements=num_elements), + DerhamOptions(degree=degree, bcs=bcs), + comm=MPI.COMM_WORLD, + domain=domain, + ) + return derham, domain + + +def _max_diff(a, b): + comm = MPI4PY.COMM_WORLD + err = comm.allreduce(np.max(np.abs((a - b).toarray())), op=MPI4PY.MAX) + ref = comm.allreduce(np.max(np.abs(b.toarray())), op=MPI4PY.MAX) + return err / ref + + +@pytest.mark.parametrize("bcs", BCS) +def test_mass_to_dict(bcs): + derham, domain = _derham((8, 6, 4), (2, 2, 1), bcs) + mass_ops = WeightedMassOperators(derham, domain) + _, u = create_equal_random_arrays(derham.fem_spaces["1"], seed=1) + + # predefined operator with string weights: JSON serializable + M1 = mass_ops.M1 + dct = M1.to_dict() + assert dct["type"] == "WeightedMassOperator" + assert dct["params"]["weights"] == ["Ginv", "sqrt_g"] + json.dumps(dct) + assert _max_diff(WeightedMassOperator.from_dict(dct, mass_ops).dot(u), M1.dot(u)) < 1e-14 + + # callable weights and transposes + Mc = mass_ops.create_weighted_mass( + "Hcurl", "Hdiv", weights=("Ginv", lambda e1, e2, e3: 1.0 + e1 * e2), name="Mc", assemble=True + ) + McT = Mc.T + assert McT.to_dict()["params"]["is_transpose"] + _, w = create_equal_random_arrays(derham.fem_spaces["2"], seed=2) + assert _max_diff(WeightedMassOperator.from_dict(McT.to_dict(), mass_ops).dot(w), McT.dot(w)) < 1e-14 + + # modified data cannot be re-created + M = mass_ops.create_weighted_mass("H1", "H1", weights=("sqrt_g",), assemble=True) + assert M.is_reconstructible + M *= 2.0 + assert not M.is_reconstructible + with pytest.raises(ValueError): + M.to_dict() + + +def test_basis_projection_to_dict(): + from struphy.feec.basis_projection_ops import BasisProjectionOperator + + derham, domain = _derham((8, 6, 4), (2, 2, 1), BCS[0]) + fun = [[lambda e1, e2, e3: 1.0 + e1 * e3]] + K = BasisProjectionOperator( + derham.projectors["L2"], + derham.fem_spaces["H1"], + fun, + V_extraction_op=derham.extraction_ops["H1"], + V_boundary_op=derham.boundary_ops["H1"], + ) + for op, V in [(K, "0"), (K.T, "3")]: + _, u = create_equal_random_arrays(derham.fem_spaces[V], seed=3) + op2 = BasisProjectionOperator.from_dict(op.to_dict(), derham) + assert _max_diff(op2.dot(u), op.dot(u)) < 1e-14 + + +@pytest.mark.parametrize("bcs", BCS) +@pytest.mark.parametrize("degree", [(2, 3, 1), (3, 1, 2)]) +def test_coarsen_poisson(bcs, degree): + r"""Re-discretization of :math:`\sigma M_0 + G^\top M_1 G` equals the Galerkin product :math:`R A P` (Cuboid, exact quadrature).""" + derham, domain = _derham((8, 8, 4), degree, bcs) + h = MultiGridHierarchy(derham, max_levels=2) + mass_ops = WeightedMassOperators(derham, domain) + A = 0.7 * mass_ops.M0 + derham.grad.T @ mass_ops.M1 @ derham.grad + + C = OperatorCoarsener(h[0], h[1], domain) + Ac = C(A) + assert Ac.domain is h[1].coeff_spaces["0"] and Ac.codomain is h[1].coeff_spaces["0"] + + # same as building it directly on the coarse level + mass_c = WeightedMassOperators(h[1], domain) + Ad = 0.7 * mass_c.M0 + h[1].grad.T @ mass_c.M1 @ h[1].grad + _, u = create_equal_random_arrays(h[1].fem_spaces["0"], seed=5) + assert _max_diff(Ac.dot(u), Ad.dot(u)) < 1e-13 + + # Galerkin property + P = SplineProlongation(h[1], h[0], "H1") + assert _max_diff(P.T.dot(A.dot(P.dot(u))), Ac.dot(u)) < 1e-12 + + # leaves are cached: a new scalar re-assembles nothing + A2 = 2.0 * mass_ops.M0 + derham.grad.T @ mass_ops.M1 @ derham.grad + Ac2 = C(A2) + leaves = lambda op: [a for a in op.addends] + assert leaves(Ac2)[0].operator is leaves(Ac)[0].operator + + +@pytest.mark.parametrize("bcs", BCS) +def test_coarsen_curl_curl(bcs): + r"""Re-discretization of :math:`C^\top M_2 C + M_1` equals the Galerkin product on 1-forms.""" + derham, domain = _derham((8, 8, 4), (2, 2, 1), bcs) + h = MultiGridHierarchy(derham, max_levels=2) + mass_ops = WeightedMassOperators(derham, domain) + A = derham.curl.T @ mass_ops.M2 @ derham.curl + mass_ops.M1 + + Ac = OperatorCoarsener(h[0], h[1], domain)(A) + P = SplineProlongation(h[1], h[0], "Hcurl") + _, u = create_equal_random_arrays(h[1].fem_spaces["1"], seed=6) + assert _max_diff(P.T.dot(A.dot(P.dot(u))), Ac.dot(u)) < 1e-12 + + +def test_coarsen_unknown_leaf(): + from feectools.linalg.stencil import StencilMatrix + + derham, domain = _derham((8, 8, 4), (2, 2, 1), BCS[0]) + h = MultiGridHierarchy(derham, max_levels=2) + S = StencilMatrix(derham.coeff_spaces["0"], derham.coeff_spaces["0"]) + with pytest.raises(NotImplementedError): + OperatorCoarsener(h[0], h[1], domain)(S) From 31858b2321b7932ced4710d625b208d994556c40 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Sat, 3 Oct 2026 18:05:21 +0200 Subject: [PATCH 10/16] Multigrid: smoothers, V-cycle preconditioner and MG-PCG solver Chebyshev smoother (default, inner preconditioner: approximate mass inverse, Jacobi or identity; eigenvalue estimate by Lanczos), damped Jacobi and fixed-step PCG smoothers. Exact diagonals of composite operators by colored probing. MultiGridPreconditioner applies a symmetric V-cycle with a replicated direct (or CG) coarse solve and an optional constant null space; MultiGridSolver wraps it in PCG with a relative tolerance. Options in the MultiGridOptions dataclass. Co-Authored-By: Claude Opus 5.5 --- .../multigrid/preconditioner.py | 483 ++++++++++++++++++ .../linear_algebra/multigrid/smoothers.py | 382 ++++++++++++++ .../tests/test_multigrid_solver.py | 166 ++++++ 3 files changed, 1031 insertions(+) create mode 100644 src/struphy/linear_algebra/multigrid/preconditioner.py create mode 100644 src/struphy/linear_algebra/multigrid/smoothers.py create mode 100644 src/struphy/linear_algebra/tests/test_multigrid_solver.py diff --git a/src/struphy/linear_algebra/multigrid/preconditioner.py b/src/struphy/linear_algebra/multigrid/preconditioner.py new file mode 100644 index 000000000..9764cae70 --- /dev/null +++ b/src/struphy/linear_algebra/multigrid/preconditioner.py @@ -0,0 +1,483 @@ +r"""Geometric multigrid V-cycle as a preconditioner, and a multigrid-preconditioned CG solver. + +Given a symmetric positive (semi-)definite operator :math:`A` on one space of a :class:`Derham`, +for example the Poisson operator :math:`\sigma \mathbb M^0 + \mathbb G^\top \mathbb M^1 \mathbb G`: + +1. A hierarchy of coarser Derhams is built (:class:`MultiGridHierarchy`). +2. :math:`A` is re-discretized on every level from its expression tree (:class:`OperatorCoarsener`). +3. One application of the preconditioner is a V-cycle with zero initial guess:: + + x_l = 0 + smooth(A_l, b_l, x_l) # n_pre times + x_l += P_l V_cycle(l+1, R_l (b_l - A_l x_l)) + smooth(A_l, b_l, x_l) # n_post times + + with an exact solve on the coarsest level. With the same linear, :math:`A`-symmetric smoother before and + after the coarse-grid correction, the V-cycle is a symmetric positive definite preconditioner for CG. +""" + +import logging +from dataclasses import dataclass +from typing import Literal + +import numpy as np +import scipy.linalg as sla +from feectools.ddm.mpi import mpi as MPI +from feectools.linalg.basic import IdentityOperator, LinearOperator, Vector +from feectools.linalg.block import BlockVector +from feectools.linalg.solvers import inverse + +from struphy.feec.mass import WeightedMassOperators +from struphy.feec.preconditioner import MassMatrixPreconditioner +from struphy.feec.psydac_derham import Derham +from struphy.geometry.base import Domain +from struphy.io.options import OptionsBase +from struphy.linear_algebra.multigrid.coarsen import OperatorCoarsener +from struphy.linear_algebra.multigrid.hierarchy import MultiGridHierarchy +from struphy.linear_algebra.multigrid.smoothers import ( + ChebyshevSmoother, + DiagonalComputer, + JacobiSmoother, + KrylovSmoother, + Smoother, + _owned_slice, + _stencil_blocks, + inverse_diagonal, +) +from struphy.linear_algebra.multigrid.transfer import SplineProlongation +from struphy.utils.utils import check_option + +logger = logging.getLogger("struphy") + +OptsSmoother = Literal["chebyshev", "jacobi", "cg"] +OptsSmootherPrecond = Literal["mass", "jacobi", "identity"] +OptsCoarseSolver = Literal["direct", "cg"] +OptsNullspace = Literal["constants"] + + +@dataclass +class MultiGridOptions(OptionsBase): + r"""Options of :class:`MultiGridPreconditioner`. + + Parameters + ---------- + smoother : str + "chebyshev" (default), "jacobi" (damped) or "cg" (fixed number of PCG steps, non-linear). + + smoother_precond : str + Preconditioner inside the Chebyshev and CG smoothers: "mass" (Kronecker-product approximation of + the inverse mass matrix of the space, default), "jacobi" (inverse diagonal of the operator) or "identity". + + smoother_degree : int + Polynomial degree of the Chebyshev smoother, number of sweeps of the Jacobi smoother, + or number of iterations of the CG smoother. + + n_pre, n_post : int + Number of smoother applications before and after the coarse-grid correction. + + jacobi_omega : float + Damping factor of the Jacobi smoother. + + chebyshev_bounds : tuple[float, float] + Smoothing interval of the Chebyshev smoother relative to the estimated largest eigenvalue. + + eig_iter : int + Number of Lanczos steps for the eigenvalue estimate of the Chebyshev smoother. + + max_levels : int | None + Maximal number of levels (None: coarsen as long as possible). + + min_cells : int + Minimal number of elements per coarsened direction on the coarsest level. + + coarse_solver : str + "direct" (LU factorization of the assembled coarsest operator, replicated on all processes) + or "cg" (PCG with the smoother preconditioner, to tolerance ``coarse_tol``). + + coarse_tol : float + Relative tolerance of the "cg" coarse solver. + + nullspace : str | None + "constants" if the operator is singular with the constant functions in its kernel (e.g. the + Poisson operator on 0-forms with periodic or Neumann boundary conditions). + + matrix_free_mass : bool + Whether re-discretized mass matrices on coarse levels are matrix-free. + """ + + smoother: OptsSmoother = "chebyshev" + smoother_precond: OptsSmootherPrecond = "mass" + smoother_degree: int = 3 + n_pre: int = 1 + n_post: int = 1 + jacobi_omega: float = 2.0 / 3.0 + chebyshev_bounds: tuple[float, float] = (0.1, 1.1) + eig_iter: int = 15 + max_levels: int | None = None + min_cells: int = 2 + coarse_solver: OptsCoarseSolver = "direct" + coarse_tol: float = 1e-10 + nullspace: OptsNullspace | None = None + matrix_free_mass: bool = False + + def __post_init__(self): + check_option(self.smoother, OptsSmoother) + check_option(self.smoother_precond, OptsSmootherPrecond) + check_option(self.coarse_solver, OptsCoarseSolver) + if self.nullspace is not None: + check_option(self.nullspace, OptsNullspace) + assert self.smoother_degree >= 1 + assert self.n_pre >= 0 and self.n_post >= 0 and self.n_pre + self.n_post >= 1 + + +class MultiGridPreconditioner(LinearOperator): + r"""One geometric multigrid V-cycle as an approximate inverse of ``A`` (see module docstring). + + Parameters + ---------- + A : LinearOperator + Symmetric positive (semi-)definite operator on ``derham.coeff_spaces[form]``, built from struphy + operators (Derham derivatives, boundary operators, :class:`WeightedMassOperator`, ...) with + ``+``, ``-``, ``*`` and ``@``. + + derham : Derham + The finest level. + + domain : Domain + Mapping, used to re-discretize mass matrices on the coarse levels. + + options : MultiGridOptions | None + Options (default: ``MultiGridOptions()``). + + mass_ops : WeightedMassOperators | None + Mass operators of the finest level (only used by the "mass" smoother preconditioner). + """ + + def __init__( + self, + A: LinearOperator, + derham: Derham, + domain: Domain, + options: MultiGridOptions | None = None, + *, + mass_ops: WeightedMassOperators | None = None, + ): + self._options = MultiGridOptions() if options is None else options + opts = self._options + self._derham = derham + self._domain_map = domain + + self._form = _find_form(derham, A.domain) + assert A.codomain is A.domain, "MultiGridPreconditioner requires a square operator." + if opts.nullspace == "constants": + assert self._form == "0", "nullspace='constants' is only implemented for 0-forms." + + self._hierarchy = MultiGridHierarchy(derham, max_levels=opts.max_levels, min_cells=opts.min_cells) + L = self._hierarchy.n_levels + logger.info(f"Multigrid levels: {[d.num_elements for d in self._hierarchy.derhams]}") + + self._P = [SplineProlongation(self._hierarchy[l + 1], self._hierarchy[l], self._form) for l in range(L - 1)] + self._R = [P.T for P in self._P] + self._coarseners = [ + OperatorCoarsener(self._hierarchy[l], self._hierarchy[l + 1], domain, matrix_free=opts.matrix_free_mass) + for l in range(L - 1) + ] + self._diag = [DiagonalComputer(d.degree) for d in self._hierarchy.derhams] + + if opts.smoother_precond == "mass": + fine_mass = WeightedMassOperators(derham, domain) if mass_ops is None else mass_ops + mass = [fine_mass] + [c.mass_ops for c in self._coarseners] + self._mass_pc = [MassMatrixPreconditioner(getattr(m, "M" + self._form)) for m in mass] + + # work vectors per level + self._b = [d.coeff_spaces[self._form].zeros() for d in self._hierarchy.derhams] + self._x = [d.coeff_spaces[self._form].zeros() for d in self._hierarchy.derhams] + self._r = [d.coeff_spaces[self._form].zeros() for d in self._hierarchy.derhams] + self._e = [d.coeff_spaces[self._form].zeros() for d in self._hierarchy.derhams] + + self._A: list[LinearOperator] = [] + self.update(A) + + # ------------------------------------------------------------------ + @property + def domain(self): + return self._A[0].domain + + @property + def codomain(self): + return self._A[0].codomain + + @property + def dtype(self): + return self._A[0].dtype + + @property + def options(self) -> MultiGridOptions: + return self._options + + @property + def hierarchy(self) -> MultiGridHierarchy: + return self._hierarchy + + @property + def operators(self) -> list[LinearOperator]: + """System operator on each level, ``operators[0]`` is the given one.""" + return self._A + + @property + def smoothers(self) -> list[Smoother]: + return self._smoothers + + def transpose(self, conjugate: bool = False) -> "MultiGridPreconditioner": + assert all(s.is_symmetric for s in self._smoothers), "Only a symmetric V-cycle can be transposed." + assert self._options.n_pre == self._options.n_post + return self + + # ------------------------------------------------------------------ + def update(self, A: LinearOperator) -> None: + """Set a new fine-level operator (e.g. with changed scalars) and update all levels. + + Mass matrices and derivative operators of coarse levels are re-used if ``A`` is built from the same objects. + """ + assert A.domain is self._derham.coeff_spaces[self._form] + self._A = [A] + for c in self._coarseners: + self._A.append(c(self._A[-1])) + self._smoothers = [self._make_smoother(l) for l in range(self._hierarchy.n_levels - 1)] + self._setup_coarse_solver() + + def _smoother_precond(self, l: int) -> LinearOperator: + opts = self._options + if opts.smoother_precond == "mass": + return self._mass_pc[l] + if opts.smoother_precond == "jacobi": + return inverse_diagonal(self._diag[l](self._A[l])) + return IdentityOperator(self._A[l].domain) + + def _make_smoother(self, l: int) -> Smoother: + opts = self._options + A = self._A[l] + if opts.smoother == "chebyshev": + return ChebyshevSmoother( + A, + self._smoother_precond(l), + degree=opts.smoother_degree, + bounds=opts.chebyshev_bounds, + eig_iter=opts.eig_iter, + ) + if opts.smoother == "jacobi": + D_inv = inverse_diagonal(self._diag[l](A)) + return JacobiSmoother(A, D_inv, omega=opts.jacobi_omega, sweeps=opts.smoother_degree) + return KrylovSmoother(A, self._smoother_precond(l), iterations=opts.smoother_degree) + + def _setup_coarse_solver(self) -> None: + opts = self._options + A = self._A[-1] + if opts.coarse_solver == "cg": + pc = self._smoother_precond(len(self._A) - 1) if opts.smoother_precond != "identity" else None + self._coarse_cg = inverse(A, "pcg", pc=pc, tol=1e-300, maxiter=1000, recycle=False) + return + + Ad = _assemble_dense(A) + if opts.nullspace == "constants": + # A + s 1 1^T is regular; for b orthogonal to 1 its solution is the zero-mean solution of A x = b + n = Ad.shape[0] + Ad = Ad + np.mean(np.abs(np.diag(Ad))) / n * np.ones((n, n)) + else: + # rows/cols of Dirichlet dofs are zero: put ones on the diagonal + zero = np.flatnonzero(np.all(Ad == 0.0, axis=1)) + Ad[zero, zero] = 1.0 + self._coarse_lu = sla.lu_factor(Ad) + + # ------------------------------------------------------------------ + def dot(self, b: Vector, out: Vector | None = None) -> Vector: + """Apply one V-cycle (zero initial guess) to ``b``.""" + assert b.space is self.domain + if out is None: + out = self.domain.zeros() + b.copy(out=self._b[0]) + if self._options.nullspace == "constants": + _remove_mean(self._b[0]) + self._vcycle(0) + self._x[0].copy(out=out) + if self._options.nullspace == "constants": + _remove_mean(out) + return out + + def _vcycle(self, l: int) -> None: + b, x = self._b[l], self._x[l] + if l == len(self._A) - 1: + self._coarse_solve(b, x) + return + + x *= 0.0 + S = self._smoothers[l] + for _ in range(self._options.n_pre): + S.smooth(b, x) + + r = S.residual(b, x, self._r[l]) + self._R[l].dot(r, out=self._b[l + 1]) + self._vcycle(l + 1) + self._P[l].dot(self._x[l + 1], out=self._e[l]) + x += self._e[l] + + for _ in range(self._options.n_post): + S.smooth(b, x) + + def _coarse_solve(self, b: Vector, x: Vector) -> None: + if self._options.coarse_solver == "cg": + nb = np.sqrt(b.inner(b)) + if nb == 0.0: + x *= 0.0 + return + self._coarse_cg._options["tol"] = self._options.coarse_tol * nb + self._coarse_cg.dot(b, out=x) + return + + bg = _gather(b) + if self._options.nullspace == "constants": + bg -= bg.mean() + _scatter(sla.lu_solve(self._coarse_lu, bg), x) + + +class MultiGridSolver(LinearOperator): + r"""Conjugate gradient method preconditioned with :class:`MultiGridPreconditioner`. + + Parameters + ---------- + A, derham, domain, options, mass_ops : + See :class:`MultiGridPreconditioner`. + + tol : float + Relative tolerance, the iteration stops when :math:`\|b - A x\|_2 \leq \mathrm{tol}\, \|b\|_2`. + + maxiter : int + Maximal number of CG iterations. + + verbose : bool + Print the residual in every iteration. + """ + + def __init__( + self, + A: LinearOperator, + derham: Derham, + domain: Domain, + options: MultiGridOptions | None = None, + *, + mass_ops: WeightedMassOperators | None = None, + tol: float = 1e-8, + maxiter: int = 100, + verbose: bool = False, + ): + self._pc = MultiGridPreconditioner(A, derham, domain, options, mass_ops=mass_ops) + self._tol = tol + self._solver = inverse(A, "pcg", pc=self._pc, tol=tol, maxiter=maxiter, verbose=verbose, recycle=False) + + @property + def domain(self): + return self._pc.domain + + @property + def codomain(self): + return self._pc.codomain + + @property + def dtype(self): + return self._pc.dtype + + @property + def preconditioner(self) -> MultiGridPreconditioner: + return self._pc + + @property + def info(self) -> dict: + """Information of the last solve: ``niter``, ``success``, ``res_norm``.""" + return self._solver._info + + def update(self, A: LinearOperator) -> None: + """Set a new operator (see :meth:`MultiGridPreconditioner.update`).""" + self._pc.update(A) + self._solver.linop = A + + def transpose(self, conjugate: bool = False) -> "MultiGridSolver": + return self + + def dot(self, b: Vector, out: Vector | None = None, x0: Vector | None = None) -> Vector: + """Solve ``A x = b`` (initial guess ``x0``, zero by default).""" + nb = np.sqrt(b.inner(b)) + self._solver._options["tol"] = self._tol * nb if nb > 0.0 else self._tol + self._solver._options["x0"] = x0 if x0 is not None else self.domain.zeros() + return self._solver.dot(b, out=out) + + +# ---------------------------------------------------------------------------------------------------- +def _find_form(derham: Derham, V) -> str: + for form in ("0", "1", "2", "3", "v"): + if derham.coeff_spaces[form] is V: + return form + raise ValueError("The operator does not act on a coefficient space of the given Derham.") + + +def _comm(v: Vector): + blk = _stencil_blocks(v)[0] + return blk.space.cart.comm if blk.space.parallel else None + + +def _gather(v: Vector) -> np.ndarray: + """Global coefficient array of ``v`` (same on all processes).""" + a = v.toarray() + comm = _comm(v) + if comm is not None and comm.size > 1: + comm.Allreduce(MPI.IN_PLACE, a, op=MPI.SUM) + return a + + +def _scatter(a: np.ndarray, v: Vector) -> None: + """Write the owned part of the global array ``a`` into ``v``.""" + offset = 0 + for blk in _stencil_blocks(v): + V = blk.space + n = int(np.prod(V.npts)) + glob = a[offset : offset + n].reshape(tuple(int(m) for m in V.npts)) + blk._data[_owned_slice(blk)] = glob[tuple(slice(s, e + 1) for s, e in zip(V.starts, V.ends))] + blk.ghost_regions_in_sync = False + offset += n + + +def _assemble_dense(A: LinearOperator) -> np.ndarray: + """Global dense matrix of ``A`` (same on all processes), by applying it to all unit vectors.""" + e = A.domain.zeros() + blocks = _stencil_blocks(e) + sizes = [int(np.prod(b.space.npts)) for b in blocks] + N = sum(sizes) + out = np.zeros((N, N)) + y = A.codomain.zeros() + col = 0 + for blk, n in zip(blocks, sizes): + V = blk.space + for flat in range(n): + gidx = np.unravel_index(flat, tuple(int(m) for m in V.npts)) + for b in blocks: + b._data[...] = 0.0 + b.ghost_regions_in_sync = False + if all(s <= i <= e_ for i, s, e_ in zip(gidx, V.starts, V.ends)): + loc = tuple(int(i - s + p * m) for i, s, p, m in zip(gidx, V.starts, V.pads, V.shifts)) + blk._data[loc] = 1.0 + A.dot(e, out=y) + out[:, col] = _gather(y) + col += 1 + return out + + +def _remove_mean(v: Vector) -> None: + """Subtract the mean of all coefficients (the projection orthogonal to the constant vector).""" + comm = _comm(v) + s = sum(float(np.sum(b._data[_owned_slice(b)])) for b in _stencil_blocks(v)) + N = sum(int(np.prod(b.space.npts)) for b in _stencil_blocks(v)) + if comm is not None and comm.size > 1: + s = comm.allreduce(s, op=MPI.SUM) + mean = s / N + for b in _stencil_blocks(v): + b._data[_owned_slice(b)] -= mean + b.ghost_regions_in_sync = False diff --git a/src/struphy/linear_algebra/multigrid/smoothers.py b/src/struphy/linear_algebra/multigrid/smoothers.py new file mode 100644 index 000000000..5ebf23ef9 --- /dev/null +++ b/src/struphy/linear_algebra/multigrid/smoothers.py @@ -0,0 +1,382 @@ +r"""Smoothers for geometric multigrid. + +A smoother approximately solves :math:`A x = b` by the update :math:`x \leftarrow x + S (b - A x)`. +:class:`ChebyshevSmoother` and :class:`JacobiSmoother` are linear and :math:`A`-symmetric, hence a V-cycle +with the same smoother before and after the coarse-grid correction is a symmetric preconditioner, suitable +for the conjugate gradient method. :class:`KrylovSmoother` is non-linear (use with care). +""" + +from abc import ABC, abstractmethod + +import numpy as np +from feectools.linalg.basic import ( + IdentityOperator, + LinearOperator, + ScaledLinearOperator, + SumLinearOperator, + Vector, + ZeroOperator, +) +from feectools.linalg.block import BlockVector +from feectools.linalg.solvers import inverse +from feectools.linalg.stencil import StencilVector + + +class Smoother(ABC): + """Base class of multigrid smoothers for the operator ``A``.""" + + def __init__(self, A: LinearOperator): + assert A.domain is A.codomain + self._A = A + self._r = A.codomain.zeros() + + @property + def A(self) -> LinearOperator: + return self._A + + @property + def is_symmetric(self) -> bool: + """Whether the smoother is linear and A-symmetric (needed for a symmetric V-cycle).""" + return True + + def residual(self, b: Vector, x: Vector, out: Vector) -> Vector: + """``out = b - A x``.""" + self._A.dot(x, out=out) + out *= -1.0 + out += b + return out + + @abstractmethod + def smooth(self, b: Vector, x: Vector) -> None: + """Improve the approximate solution ``x`` of ``A x = b`` in place.""" + + +class JacobiSmoother(Smoother): + r"""Damped Jacobi, :math:`x \leftarrow x + \omega D^{-1}(b - A x)`, repeated ``sweeps`` times. + + Parameters + ---------- + A : LinearOperator + System operator. + + diag_inv : LinearOperator + Inverse diagonal :math:`D^{-1}` of ``A`` (e.g. from :func:`inverse_diagonal`). + + omega : float + Damping factor. + + sweeps : int + Number of iterations per call. + """ + + def __init__(self, A: LinearOperator, diag_inv: LinearOperator, *, omega: float = 2.0 / 3.0, sweeps: int = 1): + super().__init__(A) + self._D_inv = diag_inv + self._omega = omega + self._sweeps = sweeps + self._z = A.domain.zeros() + + def smooth(self, b: Vector, x: Vector) -> None: + for _ in range(self._sweeps): + self.residual(b, x, self._r) + self._D_inv.dot(self._r, out=self._z) + x.mul_iadd(self._omega, self._z) + + +class ChebyshevSmoother(Smoother): + r"""Chebyshev polynomial smoother of degree ``degree`` for :math:`M^{-1} A`. + + The polynomial damps the eigenmodes of :math:`M^{-1} A` in the interval + :math:`[\alpha \lambda_\max, \beta \lambda_\max]` (``bounds = (alpha, beta)``), where + :math:`\lambda_\max` is estimated with a few Lanczos (PCG) steps. Each call costs ``degree`` + applications of ``A`` and of ``M_inv``. No inner products are computed. + + Parameters + ---------- + A : LinearOperator + Symmetric positive (semi-)definite system operator. + + M_inv : LinearOperator + Symmetric positive definite preconditioner, e.g. an approximate mass-matrix inverse or + an inverse diagonal of ``A``. + + degree : int + Polynomial degree. + + bounds : tuple[float, float] + Lower and upper end of the smoothing interval relative to the estimated :math:`\lambda_\max`. + + eig_iter : int + Number of Lanczos steps for estimating :math:`\lambda_\max`. + + lambda_max : float | None + Largest eigenvalue of :math:`M^{-1} A`, estimated if None. + """ + + def __init__( + self, + A: LinearOperator, + M_inv: LinearOperator, + *, + degree: int = 3, + bounds: tuple[float, float] = (0.1, 1.1), + eig_iter: int = 15, + lambda_max: float | None = None, + ): + super().__init__(A) + assert degree >= 1 + assert 0.0 < bounds[0] < bounds[1] + self._M_inv = M_inv + self._degree = degree + + if lambda_max is None: + lambda_max = estimate_lambda_max(A, M_inv, n_iter=eig_iter) + self._lambda_max = lambda_max + a, b = bounds[0] * lambda_max, bounds[1] * lambda_max + self._theta = 0.5 * (b + a) + self._delta = 0.5 * (b - a) + + self._d = A.domain.zeros() + self._z = A.domain.zeros() + + @property + def lambda_max(self) -> float: + return self._lambda_max + + def smooth(self, b: Vector, x: Vector) -> None: + # Saad, Iterative Methods for Sparse Linear Systems, Alg. 12.1 (preconditioned) + theta, delta = self._theta, self._delta + sigma = theta / delta + rho = 1.0 / sigma + r, d, z = self._r, self._d, self._z + + self.residual(b, x, r) + self._M_inv.dot(r, out=d) + d *= 1.0 / theta + for k in range(self._degree): + x += d + if k == self._degree - 1: + break + self._A.dot(d, out=z) + r -= z + rho_new = 1.0 / (2.0 * sigma - rho) + self._M_inv.dot(r, out=z) + d *= rho_new * rho + d.mul_iadd(2.0 * rho_new / delta, z) + rho = rho_new + + +class KrylovSmoother(Smoother): + """A fixed number of (preconditioned) conjugate gradient iterations, warm-started from ``x``. + + Note that this smoother is non-linear; a V-cycle using it is not a fixed linear preconditioner. + """ + + def __init__(self, A: LinearOperator, M_inv: LinearOperator | None = None, *, iterations: int = 3): + super().__init__(A) + self._solver = inverse(A, "pcg", pc=M_inv, maxiter=iterations, tol=1e-300, recycle=False) + + @property + def is_symmetric(self) -> bool: + return False + + def smooth(self, b: Vector, x: Vector) -> None: + self._solver._options["x0"] = x.copy() + self._solver.dot(b, out=x) + + +def estimate_lambda_max(A: LinearOperator, M_inv: LinearOperator, *, n_iter: int = 15, seed: int = 1234) -> float: + r"""Estimate the largest eigenvalue of :math:`M^{-1} A` with ``n_iter`` Lanczos steps (via PCG coefficients). + + The estimate is a lower bound that converges quickly to :math:`\lambda_\max`. + """ + b = A.domain.zeros() + _fill_random(b, seed) + + x_r = b.copy() + z = M_inv.dot(x_r) + p = z.copy() + q = A.domain.zeros() + rz = x_r.inner(z) + alphas, betas = [], [] + for _ in range(n_iter): + A.dot(p, out=q) + pq = p.inner(q) + if pq <= 0.0 or rz <= 0.0: + break + alpha = rz / pq + x_r.mul_iadd(-alpha, q) + M_inv.dot(x_r, out=z) + rz_new = x_r.inner(z) + beta = rz_new / rz + alphas.append(alpha) + betas.append(beta) + if rz_new <= 1e-30 * rz: + break + p *= beta + p += z + rz = rz_new + + k = len(alphas) + assert k > 0, "Lanczos breakdown in the eigenvalue estimate (is A positive semi-definite?)." + T = np.zeros((k, k)) + for j in range(k): + T[j, j] = 1.0 / alphas[j] + (betas[j - 1] / alphas[j - 1] if j > 0 else 0.0) + if j + 1 < k: + T[j, j + 1] = T[j + 1, j] = np.sqrt(betas[j]) / alphas[j] + return float(np.linalg.eigvalsh(T).max()) + + +def _stencil_blocks(v: Vector) -> list[StencilVector]: + return list(v.blocks) if isinstance(v, BlockVector) else [v] + + +def _fill_random(v: Vector, seed: int) -> None: + """Fill the owned entries of ``v`` with uniform random numbers in [-1, 1] (different on each process).""" + for n, blk in enumerate(_stencil_blocks(v)): + V = blk.space + rng = np.random.default_rng([seed, n] + [int(s) for s in V.starts]) + idx = _owned_slice(blk) + blk._data[idx] = rng.uniform(-1.0, 1.0, size=blk._data[idx].shape) + blk.ghost_regions_in_sync = False + + +def _owned_slice(v: StencilVector) -> tuple[slice, ...]: + V = v.space + return tuple(slice(p * m, p * m + e - s + 1) for p, m, s, e in zip(V.pads, V.shifts, V.starts, V.ends)) + + +# ---------------------------------------------------------------------------------------------------- +# diagonal of composite operators +# ---------------------------------------------------------------------------------------------------- +class DiagonalOperator(LinearOperator): + """Pointwise multiplication by the entries of a vector.""" + + def __init__(self, d: Vector): + self._d = d + self._space = d.space + + @property + def domain(self): + return self._space + + @property + def codomain(self): + return self._space + + @property + def dtype(self): + return self._space.dtype + + @property + def vector(self) -> Vector: + return self._d + + def dot(self, v: Vector, out: Vector | None = None) -> Vector: + if out is None: + out = self._space.zeros() + for vb, db, ob in zip(_stencil_blocks(v), _stencil_blocks(self._d), _stencil_blocks(out)): + idx = _owned_slice(ob) + ob._data[idx] = db._data[idx] * vb._data[idx] + ob.ghost_regions_in_sync = False + return out + + def transpose(self, conjugate: bool = False) -> "DiagonalOperator": + return self + + +class DiagonalComputer: + r"""Exact diagonal of (composite) operators, cached per operator object. + + Sums and scalings are combined from the diagonals of their parts; identity and zero operators are + trivial; every other operator (a leaf, or a product such as :math:`G^\top M G`) is probed with + colored unit vectors: dofs with equal index modulo :math:`c_d \geq 2 w_d + 1` in every direction do + not couple, so one application of the operator per color gives the diagonal entries of all dofs of + that color. This costs :math:`\prod_d c_d` operator applications (times the number of components). + + Parameters + ---------- + widths : tuple[int, int, int] + Upper bound for the coupling distance (in index units) of the operators in each direction, + e.g. the spline degrees for mass and stiffness matrices. + """ + + def __init__(self, widths: tuple[int, int, int]): + self._widths = tuple(widths) + self._cache: dict[int, tuple[LinearOperator, Vector]] = {} + + def __call__(self, A: LinearOperator) -> Vector: + """Diagonal of ``A`` as a vector in ``A.domain``.""" + assert A.domain is A.codomain + if isinstance(A, SumLinearOperator): + d = A.domain.zeros() + for a in A.addends: + d += self(a) + return d + if isinstance(A, ScaledLinearOperator): + return self(A.operator) * A.scalar + if isinstance(A, IdentityOperator): + d = A.domain.zeros() + for blk in _stencil_blocks(d): + blk._data[...] = 1.0 + return d + if isinstance(A, ZeroOperator): + return A.domain.zeros() + + cached = self._cache.get(id(A)) + if cached is not None and cached[0] is A: + return cached[1] + d = self._probe(A) + self._cache[id(A)] = (A, d) + return d + + def _probe(self, A: LinearOperator) -> Vector: + d = A.domain.zeros() + e = A.domain.zeros() + y = A.domain.zeros() + d_blocks, e_blocks = _stencil_blocks(d), _stencil_blocks(e) + + for n, (db, eb) in enumerate(zip(d_blocks, e_blocks)): + V = eb.space + colors = [_n_colors(int(npts), 2 * w + 1, bool(per)) for npts, w, per in zip(V.npts, self._widths, V.periods)] + glob = [np.arange(s, e + 1) for s, e in zip(V.starts, V.ends)] + idx = _owned_slice(eb) + for color in np.ndindex(*colors): + mask = np.ones([len(g) for g in glob], dtype=bool) + for axis, (g, c, k) in enumerate(zip(glob, colors, color)): + shape = [1, 1, 1] + shape[axis] = len(g) + mask = mask & (g % c == k).reshape(shape) + for blk in e_blocks: + blk._data[...] = 0.0 + blk.ghost_regions_in_sync = False + eb._data[idx] = mask.astype(float) + A.dot(e, out=y) + yb = _stencil_blocks(y)[n] + db._data[idx][mask] = yb._data[idx][mask] + return d + + +def _n_colors(n: int, c: int, periodic: bool) -> int: + """Number of colors in one direction with ``n`` dofs such that equal colors are at least ``c`` apart.""" + if c >= n: + return n + if not periodic: + return c + # periodic: c must divide n to avoid coupling across the periodic boundary + for k in range(c, n + 1): + if n % k == 0: + return k + return n + + +def inverse_diagonal(diag: Vector) -> DiagonalOperator: + """Inverse of a diagonal (entries equal to zero, e.g. Dirichlet dofs, are mapped to zero).""" + inv = diag.copy() + for blk in _stencil_blocks(inv): + idx = _owned_slice(blk) + data = blk._data[idx] + out = np.zeros_like(data) + np.divide(1.0, data, out=out, where=data != 0.0) + blk._data[idx] = out + return DiagonalOperator(inv) diff --git a/src/struphy/linear_algebra/tests/test_multigrid_solver.py b/src/struphy/linear_algebra/tests/test_multigrid_solver.py new file mode 100644 index 000000000..7b0359686 --- /dev/null +++ b/src/struphy/linear_algebra/tests/test_multigrid_solver.py @@ -0,0 +1,166 @@ +import numpy as np +import pytest +from feectools.ddm.mpi import mpi as MPI +from feectools.linalg.basic import LinearOperator + +from struphy.feec.mass import WeightedMassOperators +from struphy.feec.psydac_derham import Derham +from struphy.feec.utilities import create_equal_random_arrays +from struphy.geometry.domains import Cuboid +from struphy.io.options import DerhamOptions +from struphy.linear_algebra.multigrid.preconditioner import ( + MultiGridOptions, + MultiGridPreconditioner, + MultiGridSolver, + _assemble_dense, +) +from struphy.linear_algebra.multigrid.smoothers import ( + ChebyshevSmoother, + DiagonalComputer, + JacobiSmoother, + inverse_diagonal, +) +from struphy.topology.grids import TensorProductGrid + +DIRICHLET = (("dirichlet", "dirichlet"), ("dirichlet", "dirichlet"), None) +PERIODIC = (None, None, None) + + +def _poisson(n, p, bcs, sigma=0.0): + domain = Cuboid(l1=0.0, r1=1.0, l2=0.0, r2=2.0, l3=0.0, r3=1.0) + derham = Derham( + TensorProductGrid(num_elements=(n, n, 1)), + DerhamOptions(degree=(p, p, 1), bcs=bcs), + comm=MPI.COMM_WORLD, + domain=domain, + ) + mass_ops = WeightedMassOperators(derham, domain) + A = derham.grad.T @ mass_ops.M1 @ derham.grad + if sigma != 0.0: + A = sigma * mass_ops.M0 + A + return derham, domain, mass_ops, A + + +class _SmootherAsOperator(LinearOperator): + """x = S b (one smoother call from zero initial guess).""" + + def __init__(self, S): + self._S = S + + domain = property(lambda self: self._S.A.domain) + codomain = property(lambda self: self._S.A.domain) + dtype = property(lambda self: float) + + def transpose(self, conjugate=False): + return self + + def dot(self, b, out=None): + x = self.domain.zeros() + self._S.smooth(b, x) + if out is None: + return x + x.copy(out=out) + return out + + +@pytest.mark.mpi_skip +@pytest.mark.parametrize("bcs", [DIRICHLET, PERIODIC]) +def test_diagonal(bcs): + """Probed diagonal equals the diagonal of the assembled operator.""" + derham, _, mass_ops, A = _poisson(8, 3, bcs, sigma=0.3) + d = DiagonalComputer(derham.degree)(A) + assert np.allclose(d.toarray(), np.diag(_assemble_dense(A)), atol=1e-14) + + +@pytest.mark.mpi_skip +@pytest.mark.parametrize("kind", ["chebyshev_jacobi", "chebyshev_mass", "jacobi"]) +def test_smoother_symmetric(kind): + """The linear smoothers are symmetric (as matrices from right-hand side to iterate).""" + from struphy.feec.preconditioner import MassMatrixPreconditioner + + derham, _, mass_ops, A = _poisson(8, 2, PERIODIC, sigma=0.5) + D_inv = inverse_diagonal(DiagonalComputer(derham.degree)(A)) + if kind == "chebyshev_jacobi": + S = ChebyshevSmoother(A, D_inv, degree=3) + elif kind == "chebyshev_mass": + S = ChebyshevSmoother(A, MassMatrixPreconditioner(mass_ops.M0), degree=3) + else: + S = JacobiSmoother(A, D_inv, sweeps=3) + Sd = _assemble_dense(_SmootherAsOperator(S)) + assert np.abs(Sd - Sd.T).max() < 1e-12 * np.abs(Sd).max() + + +@pytest.mark.mpi_skip +@pytest.mark.parametrize("bcs, nullspace", [(DIRICHLET, None), (PERIODIC, "constants")]) +@pytest.mark.parametrize("smoother_precond", ["mass", "jacobi"]) +def test_vcycle_spd(bcs, nullspace, smoother_precond): + """The V-cycle is symmetric positive definite (on the complement of the null space) and contracts.""" + derham, domain, mass_ops, A = _poisson(16, 2, bcs) + pc = MultiGridPreconditioner( + A, derham, domain, MultiGridOptions(smoother_precond=smoother_precond, nullspace=nullspace), mass_ops=mass_ops + ) + B = _assemble_dense(pc) + Ad = _assemble_dense(A) + assert np.abs(B - B.T).max() < 1e-12 * np.abs(B).max() + + # restrict to the dofs/modes that matter: interior dofs (Dirichlet) or zero-mean vectors (periodic) + N = Ad.shape[0] + if nullspace == "constants": + Q = np.linalg.qr(np.eye(N) - np.ones((N, N)) / N)[0][:, : N - 1] + else: + Q = np.eye(N)[:, np.flatnonzero(np.diag(Ad) != 0.0)] + assert np.linalg.eigvalsh(Q.T @ B @ Q).min() > 0.0 + E = Q.T @ (np.eye(N) - B @ Ad) @ Q + assert np.abs(np.linalg.eigvals(E)).max() < 0.5 + + +@pytest.mark.parametrize("bcs, nullspace", [(DIRICHLET, None), (PERIODIC, "constants")]) +@pytest.mark.parametrize("p", [2, 3]) +@pytest.mark.parametrize( + "smoother, smoother_precond", + [("chebyshev", "mass"), ("chebyshev", "jacobi"), ("jacobi", "identity"), ("cg", "mass")], +) +def test_poisson_h_independent(bcs, nullspace, p, smoother, smoother_precond): + """MG-preconditioned CG converges in a small, mesh-independent number of iterations.""" + niter = [] + for n in (16, 32): + derham, domain, mass_ops, A = _poisson(n, p, bcs) + _, u = create_equal_random_arrays(derham.fem_spaces["0"], seed=3) + b = A.dot(u) + opts = MultiGridOptions(smoother=smoother, smoother_precond=smoother_precond, nullspace=nullspace) + solver = MultiGridSolver(A, derham, domain, opts, mass_ops=mass_ops, tol=1e-8) + x = solver.dot(b) + r = b - A.dot(x) + assert solver.info["success"] + assert np.sqrt(r.inner(r)) <= 1e-8 * np.sqrt(b.inner(b)) + niter.append(solver.info["niter"]) + assert max(niter) <= 15 + assert niter[1] <= niter[0] + 2 + + +def test_update(): + """Changing a scalar of the operator re-uses the coarse operators and still converges.""" + derham, domain, mass_ops, A = _poisson(16, 2, DIRICHLET, sigma=2.0) + solver = MultiGridSolver(A, derham, domain, mass_ops=mass_ops, tol=1e-10) + coarse_M0 = solver.preconditioner.operators[1].addends[0].operator + + A2 = 100.0 * mass_ops.M0 + derham.grad.T @ mass_ops.M1 @ derham.grad + solver.update(A2) + assert solver.preconditioner.operators[1].addends[0].operator is coarse_M0 + + _, u = create_equal_random_arrays(derham.fem_spaces["0"], seed=4) + b = A2.dot(u) + x = solver.dot(b) + r = b - A2.dot(x) + assert np.sqrt(r.inner(r)) <= 1e-10 * np.sqrt(b.inner(b)) + assert solver.info["niter"] <= 15 + + +@pytest.mark.mpi_skip +def test_options(): + opts = MultiGridOptions(smoother="jacobi", max_levels=3) + assert MultiGridOptions.from_dict(opts.to_dict()) == opts + with pytest.raises(AssertionError): + MultiGridOptions(smoother="gauss-seidel") + with pytest.raises(AssertionError): + MultiGridOptions(n_pre=0, n_post=0) From d8d87d48a237a6bde19ebbe032cbcaf317f5ff65 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Sat, 3 Oct 2026 18:23:47 +0200 Subject: [PATCH 11/16] ImplicitDiffusion/PoissonSolve: multigrid preconditioner option precond="MultiGrid" uses MultiGridPreconditioner (options in the new multigrid field). The preconditioner is updated only when sigma_1 (e.g. sigma_1/dt) changes. Other precond values keep the previous behavior. Co-Authored-By: Claude Opus 5.5 --- src/struphy/io/options.py | 1 + src/struphy/propagators/implicit_diffusion.py | 41 ++++++-- src/struphy/propagators/poisson_solve.py | 22 +++-- src/struphy/propagators/tests/test_poisson.py | 97 +++++++++++++++++++ 4 files changed, 147 insertions(+), 14 deletions(-) diff --git a/src/struphy/io/options.py b/src/struphy/io/options.py index 72299fbbd..289404c52 100644 --- a/src/struphy/io/options.py +++ b/src/struphy/io/options.py @@ -76,6 +76,7 @@ class LiteralOptions: OptsSymmSolver = Literal["pcg", "cg"] OptsGenSolver = Literal["pbicgstab", "bicgstab", "gmres"] OptsMassPrecond = Literal["MassMatrixPreconditioner", "MassMatrixDiagonalPreconditioner", None] + OptsDiffusionPrecond = Literal["MultiGrid", "MassMatrixPreconditioner", "MassMatrixDiagonalPreconditioner", None] OptsSaddlePointSolver = Literal["uzawa"] OptsDirectSolver = Literal["SparseSolver", "ScipySparse", "InexactNPInverse", "DirectNPInverse"] OptsNonlinearSolver = Literal["Picard", "Newton"] diff --git a/src/struphy/propagators/implicit_diffusion.py b/src/struphy/propagators/implicit_diffusion.py index e211e127b..2af8fcd03 100644 --- a/src/struphy/propagators/implicit_diffusion.py +++ b/src/struphy/propagators/implicit_diffusion.py @@ -11,6 +11,7 @@ from struphy.feec.mass import L2Projector, WeightedMassOperator from struphy.io.options import LiteralOptions, OptionsBase +from struphy.linear_algebra.multigrid.preconditioner import MultiGridOptions, MultiGridPreconditioner from struphy.linear_algebra.solver import SolverParameters from struphy.models.variables import FEECVariable, PICVariable, SPHVariable from struphy.pic.accumulation.filter import FilterParameters @@ -182,10 +183,16 @@ class Options(OptionsBase): Name of the symmetric iterative solver passed to :func:`psydac.linalg.solvers.inverse`. - precond : LiteralOptions.OptsMassPrecond, default="MassMatrixPreconditioner" + precond : LiteralOptions.OptsDiffusionPrecond, default="MassMatrixPreconditioner" Name of the preconditioner configuration. - Currently this class sets ``pc=None`` internally, so this option is - reserved for compatibility and future extensions. + ``"MultiGrid"`` uses a geometric multigrid V-cycle + (:class:`~struphy.linear_algebra.multigrid.preconditioner.MultiGridPreconditioner`, + requires ``solver="pcg"``). The other values currently result in ``pc=None``. + + multigrid : MultiGridOptions, default=None + Options of the multigrid preconditioner (if ``precond="MultiGrid"``). + If ``None``, defaults to ``MultiGridOptions()``. Set ``nullspace="constants"`` + for (nearly) singular problems, e.g. a periodic Poisson problem with tiny ``sigma_1``. solver_params : SolverParameters, default=None Iterative-solver controls (for example ``tol``, ``maxiter``, @@ -216,7 +223,8 @@ class Options(OptionsBase): diffusion_mat: OptsDiffusionMat = "M1" x0: StencilVector = None solver: LiteralOptions.OptsSymmSolver = "pcg" - precond: LiteralOptions.OptsMassPrecond = "MassMatrixPreconditioner" + precond: LiteralOptions.OptsDiffusionPrecond = "MassMatrixPreconditioner" + multigrid: MultiGridOptions = None solver_params: SolverParameters = None filter_params: dict[PICVariable | SPHVariable, FilterParameters] = None @@ -225,11 +233,15 @@ def __post_init__(self): check_option(self.stab_mat, self.OptsStabMat) check_option(self.diffusion_mat, self.OptsDiffusionMat) check_option(self.solver, LiteralOptions.OptsSymmSolver) - check_option(self.precond, LiteralOptions.OptsMassPrecond) + check_option(self.precond, LiteralOptions.OptsDiffusionPrecond) + if self.precond == "MultiGrid": + assert self.solver == "pcg", "precond='MultiGrid' requires solver='pcg'." # defaults if self.solver_params is None: self.solver_params = SolverParameters() + if self.multigrid is None: + self.multigrid = MultiGridOptions() @property def options(self) -> Options: @@ -345,10 +357,20 @@ def verify_rhs(rho) -> StencilVector | FEECVariable | AccumulatorVector: self._diffusion_op = self.derham.grad.T @ diffusion_mat @ self.derham.grad # preconditioner and solver for Ax=b - if self.options.precond is None: - pc = None + self._mg = None + if self.options.precond == "MultiGrid": + # the operator is updated in __call__ if sigma_1 changes (e.g. with dt) + self._mg_sig_1 = self._sigma_1 + self._mg = MultiGridPreconditioner( + self._sigma_1 * stab_mat + self._diffusion_op, + self.derham, + self.domain, + self.options.multigrid, + mass_ops=self.mass_ops, + ) + pc = self._mg else: - # TODO: waiting for multigrid preconditioner + # TODO: mass-matrix preconditioners are not effective for this operator pc = None # solver just with A_2, but will be set during call with dt @@ -459,6 +481,9 @@ def __call__(self, dt): # compute lhs self._solver.linop = sig_1 * self._stab_mat + self._diffusion_op + if self._mg is not None and sig_1 != self._mg_sig_1: + self._mg.update(self._solver.linop) + self._mg_sig_1 = sig_1 # solve with ProfileManager.profile_region(self._solve_region, functions=[self._solver.solve]): diff --git a/src/struphy/propagators/poisson_solve.py b/src/struphy/propagators/poisson_solve.py index cb7a1c3be..9e71cb73c 100644 --- a/src/struphy/propagators/poisson_solve.py +++ b/src/struphy/propagators/poisson_solve.py @@ -5,6 +5,7 @@ from feectools.linalg.stencil import StencilVector from struphy.io.options import LiteralOptions, OptionsBase +from struphy.linear_algebra.multigrid.preconditioner import MultiGridOptions from struphy.linear_algebra.solver import SolverParameters from struphy.models.variables import FEECVariable, PICVariable, SPHVariable from struphy.pic.accumulation.filter import FilterParameters @@ -66,10 +67,14 @@ class Options(OptionsBase): Name of the symmetric iterative solver passed to :func:`psydac.linalg.solvers.inverse`. - precond : LiteralOptions.OptsMassPrecond, default="MassMatrixPreconditioner" - Name of the preconditioner configuration. - Currently this class inherits the same behavior as - :class:`ImplicitDiffusion`, where ``pc=None`` is used internally. + precond : LiteralOptions.OptsDiffusionPrecond, default="MassMatrixPreconditioner" + Name of the preconditioner configuration, see :class:`ImplicitDiffusion` + (``"MultiGrid"`` for geometric multigrid). + + multigrid : MultiGridOptions, default=None + Options of the multigrid preconditioner (if ``precond="MultiGrid"``). + For periodic or Neumann boundary conditions with ``stab_eps = 0``, use + ``MultiGridOptions(nullspace="constants")``. solver_params : SolverParameters, default=None Iterative-solver controls (for example ``tol``, ``maxiter``, @@ -96,7 +101,8 @@ class Options(OptionsBase): diffusion_mat: OptsDiffusionMat = "M1" x0: StencilVector = None solver: LiteralOptions.OptsSymmSolver = "pcg" - precond: LiteralOptions.OptsMassPrecond = "MassMatrixPreconditioner" + precond: LiteralOptions.OptsDiffusionPrecond = "MassMatrixPreconditioner" + multigrid: MultiGridOptions = None solver_params: SolverParameters = None filter_params: dict[PICVariable | SPHVariable, FilterParameters] = None @@ -105,11 +111,15 @@ def __post_init__(self): check_option(self.stab_mat, self.OptsStabMat) check_option(self.diffusion_mat, self.OptsDiffusionMat) check_option(self.solver, LiteralOptions.OptsSymmSolver) - check_option(self.precond, LiteralOptions.OptsMassPrecond) + check_option(self.precond, LiteralOptions.OptsDiffusionPrecond) + if self.precond == "MultiGrid": + assert self.solver == "pcg", "precond='MultiGrid' requires solver='pcg'." # defaults if self.solver_params is None: self.solver_params = SolverParameters() + if self.multigrid is None: + self.multigrid = MultiGridOptions() # Poisson solve (-> set some params of parent class) self.sigma_1 = self.stab_eps diff --git a/src/struphy/propagators/tests/test_poisson.py b/src/struphy/propagators/tests/test_poisson.py index dafd9bb48..3f9447071 100644 --- a/src/struphy/propagators/tests/test_poisson.py +++ b/src/struphy/propagators/tests/test_poisson.py @@ -786,6 +786,103 @@ def rho2_pulled(e1, e2, e3): assert error2 < err_lim + +@pytest.mark.parametrize("degree", [[2, 2, 1], [3, 3, 1]]) +@pytest.mark.parametrize("bc_type", ["periodic", "dirichlet", "neumann"]) +def test_poisson_2d_multigrid(degree, bc_type): + """PoissonSolve with precond="MultiGrid" agrees with the unpreconditioned solve, in few iterations.""" + from mpi4py import MPI as MPI4PY + + from struphy.linear_algebra.multigrid.preconditioner import MultiGridOptions + + domain = domains.Colella(Lx=4.0, Ly=2.0, alpha=0.1, Lz=1.0) + bcs = { + "periodic": (None, None, None), + "dirichlet": (("dirichlet", "dirichlet"), None, None), + "neumann": (("free", "free"), None, None), + }[bc_type] + derham = Derham(TensorProductGrid(num_elements=[32, 32, 1]), DerhamOptions(degree=degree, bcs=bcs), comm=comm) + mass_ops = WeightedMassOperators(derham, domain) + Propagator.derham = derham + Propagator.domain = domain + Propagator.mass_ops = mass_ops + + def rho(e1, e2, e3): + return xp.cos(2 * xp.pi * e1) * xp.sin(2 * xp.pi * e2) + 0.3 * xp.sin(4 * xp.pi * e2) + + phis = [] + infos = [] + for precond in ["MassMatrixPreconditioner", "MultiGrid"]: + phi = FEECVariable(space="H1") + phi.allocate(derham=derham, domain=domain) + solver = PoissonSolve(rho=rho) + solver.variables.phi = phi + solver.options = solver.Options( + stab_eps=1e-8, + solver="pcg", + precond=precond, + # the Jacobi smoother is robust w.r.t. the mapping (the default mass smoother is robust w.r.t. the degree); + # no null space: the system is regularized by stab_eps + multigrid=MultiGridOptions(smoother_precond="jacobi"), + solver_params=SolverParameters(tol=1e-11, maxiter=3000, recycle=False), + ) + solver.allocate() + solver(1.0) + phis.append(phi.spline.vector.toarray()) + infos.append(solver._solver._info) + + # global coefficient arrays (toarray only fills the local part) + phis = [MPI4PY.COMM_WORLD.allreduce(p, op=MPI4PY.SUM) for p in phis] + if bc_type != "dirichlet": + # solutions are defined up to a constant (the stabilization is tiny) + phis = [p - xp.mean(p) for p in phis] + assert xp.max(xp.abs(phis[0] - phis[1])) < 1e-6 * xp.max(xp.abs(phis[0])) + assert infos[1]["niter"] <= 25 + assert infos[1]["niter"] < infos[0]["niter"] + + +def test_implicit_diffusion_multigrid_dt(): + """With divide_by_dt, the multigrid preconditioner follows changes of dt.""" + from struphy.linear_algebra.multigrid.preconditioner import MultiGridOptions + from struphy.propagators.implicit_diffusion import ImplicitDiffusion + + domain = domains.Cuboid(l1=0.0, r1=2.0, l2=0.0, r2=1.0, l3=0.0, r3=1.0) + derham = Derham( + TensorProductGrid(num_elements=[32, 16, 1]), + DerhamOptions(degree=[2, 2, 1], bcs=(("dirichlet", "dirichlet"), None, None)), + comm=comm, + ) + mass_ops = WeightedMassOperators(derham, domain) + Propagator.derham = derham + Propagator.domain = domain + Propagator.mass_ops = mass_ops + + phi = FEECVariable(space="H1") + phi.allocate(derham=derham, domain=domain) + phi.spline.vector = derham.P0(lambda e1, e2, e3: xp.sin(xp.pi * e1) * xp.cos(2 * xp.pi * e2)) + + prop = ImplicitDiffusion() + prop.variables.phi = phi + prop.options = prop.Options( + sigma_1=1.0, + sigma_2=1.0, + sigma_3=0.0, + divide_by_dt=True, + precond="MultiGrid", + multigrid=MultiGridOptions(), + solver_params=SolverParameters(tol=1e-12, maxiter=100, recycle=False), + ) + prop.allocate() + + for dt in [0.1, 0.1, 0.01]: + rhs = (1.0 / dt) * mass_ops.M0.dot(phi.spline.vector) + prop(dt) + A = (1.0 / dt) * mass_ops.M0 + derham.grad.T @ mass_ops.M1 @ derham.grad + r = rhs - A.dot(phi.spline.vector) + assert xp.sqrt(r.inner(r)) < 1e-9 * xp.sqrt(rhs.inner(rhs)) + assert prop._solver._info["niter"] <= 15 + + if __name__ == "__main__": # direction = 0 # bc_type = "dirichlet" From 4bdb0b7f86b787eb8bdeedd4589b9243b265bff3 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Sat, 3 Oct 2026 18:24:31 +0200 Subject: [PATCH 12/16] Format with ruff Co-Authored-By: Claude Opus 5.5 --- src/struphy/feec/mass.py | 20 +++++++++++-------- src/struphy/feec/psydac_derham.py | 4 +++- .../linear_algebra/multigrid/smoothers.py | 4 +++- .../linear_algebra/multigrid/transfer.py | 4 +++- src/struphy/propagators/tests/test_poisson.py | 1 - 5 files changed, 21 insertions(+), 12 deletions(-) diff --git a/src/struphy/feec/mass.py b/src/struphy/feec/mass.py index bf40afbad..8fdbae777 100644 --- a/src/struphy/feec/mass.py +++ b/src/struphy/feec/mass.py @@ -1258,14 +1258,18 @@ def f_call_matrix(e1, e2, e3): grid_bound = len(spline_functions) > 0 or ( isinstance(weights, list) and any(isinstance(w, xp.ndarray) for row in weights for w in row) ) - out._creation_info = None if grid_bound else { - "V_id": V_id, - "W_id": W_id, - "name": name, - "weights": weights, - "transposed": transposed, - "is_transpose": False, - } + out._creation_info = ( + None + if grid_bound + else { + "V_id": V_id, + "W_id": W_id, + "name": name, + "weights": weights, + "transposed": transposed, + "is_transpose": False, + } + ) if assemble and not dry_run: out.assemble() diff --git a/src/struphy/feec/psydac_derham.py b/src/struphy/feec/psydac_derham.py index 5f8c31503..2b696bda1 100644 --- a/src/struphy/feec/psydac_derham.py +++ b/src/struphy/feec/psydac_derham.py @@ -1571,7 +1571,9 @@ def init_derham( ) if domain_decomposition.comm is not None and comm is not None: # (comm is None in the decomposition when feectools runs with MockMPI) - assert domain_decomposition.comm == comm, "domain_decomposition must be built on the Derham communicator." + assert domain_decomposition.comm == comm, ( + "domain_decomposition must be built on the Derham communicator." + ) self._domain_decomposition = domain_decomposition _derham = self._discretize_derham( diff --git a/src/struphy/linear_algebra/multigrid/smoothers.py b/src/struphy/linear_algebra/multigrid/smoothers.py index 5ebf23ef9..c209e2f79 100644 --- a/src/struphy/linear_algebra/multigrid/smoothers.py +++ b/src/struphy/linear_algebra/multigrid/smoothers.py @@ -338,7 +338,9 @@ def _probe(self, A: LinearOperator) -> Vector: for n, (db, eb) in enumerate(zip(d_blocks, e_blocks)): V = eb.space - colors = [_n_colors(int(npts), 2 * w + 1, bool(per)) for npts, w, per in zip(V.npts, self._widths, V.periods)] + colors = [ + _n_colors(int(npts), 2 * w + 1, bool(per)) for npts, w, per in zip(V.npts, self._widths, V.periods) + ] glob = [np.arange(s, e + 1) for s, e in zip(V.starts, V.ends)] idx = _owned_slice(eb) for color in np.ndindex(*colors): diff --git a/src/struphy/linear_algebra/multigrid/transfer.py b/src/struphy/linear_algebra/multigrid/transfer.py index 3bbacd781..68299a116 100644 --- a/src/struphy/linear_algebra/multigrid/transfer.py +++ b/src/struphy/linear_algebra/multigrid/transfer.py @@ -113,7 +113,9 @@ class _KronTransfer: def __init__(self, mats: list[spa.csr_matrix], W: StencilVectorSpace): self._mats = mats - self._out_slice = tuple(slice(p * m, p * m + e - s + 1) for p, m, s, e in zip(W.pads, W.shifts, W.starts, W.ends)) + self._out_slice = tuple( + slice(p * m, p * m + e - s + 1) for p, m, s, e in zip(W.pads, W.shifts, W.starts, W.ends) + ) def dot(self, v: StencilVector, out: StencilVector) -> None: if not v.ghost_regions_in_sync: diff --git a/src/struphy/propagators/tests/test_poisson.py b/src/struphy/propagators/tests/test_poisson.py index 3f9447071..d0f98a398 100644 --- a/src/struphy/propagators/tests/test_poisson.py +++ b/src/struphy/propagators/tests/test_poisson.py @@ -786,7 +786,6 @@ def rho2_pulled(e1, e2, e3): assert error2 < err_lim - @pytest.mark.parametrize("degree", [[2, 2, 1], [3, 3, 1]]) @pytest.mark.parametrize("bc_type", ["periodic", "dirichlet", "neumann"]) def test_poisson_2d_multigrid(degree, bc_type): From 91a5ed5a7374053ba95c3ebf88b96b11e430a98c Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Sun, 4 Oct 2026 08:57:56 +0200 Subject: [PATCH 13/16] Tests: import MPI from feectools.ddm.mpi, not mpi4py Works without mpi4py; reductions only when comm.Get_size() > 1 since the MockComm returns None. Co-Authored-By: Claude Opus 5.5 --- .../linear_algebra/tests/test_multigrid_coarsen.py | 10 ++++++---- .../linear_algebra/tests/test_multigrid_transfer.py | 10 ++++++---- src/struphy/propagators/tests/test_poisson.py | 5 ++--- 3 files changed, 14 insertions(+), 11 deletions(-) diff --git a/src/struphy/linear_algebra/tests/test_multigrid_coarsen.py b/src/struphy/linear_algebra/tests/test_multigrid_coarsen.py index c612e553a..455617072 100644 --- a/src/struphy/linear_algebra/tests/test_multigrid_coarsen.py +++ b/src/struphy/linear_algebra/tests/test_multigrid_coarsen.py @@ -3,7 +3,6 @@ import numpy as np import pytest from feectools.ddm.mpi import mpi as MPI -from mpi4py import MPI as MPI4PY from struphy.feec.mass import WeightedMassOperator, WeightedMassOperators from struphy.feec.psydac_derham import Derham @@ -33,9 +32,12 @@ def _derham(num_elements, degree, bcs): def _max_diff(a, b): - comm = MPI4PY.COMM_WORLD - err = comm.allreduce(np.max(np.abs((a - b).toarray())), op=MPI4PY.MAX) - ref = comm.allreduce(np.max(np.abs(b.toarray())), op=MPI4PY.MAX) + comm = MPI.COMM_WORLD + err = np.max(np.abs((a - b).toarray())) + ref = np.max(np.abs(b.toarray())) + if comm.Get_size() > 1: + err = comm.allreduce(err, op=MPI.MAX) + ref = comm.allreduce(ref, op=MPI.MAX) return err / ref diff --git a/src/struphy/linear_algebra/tests/test_multigrid_transfer.py b/src/struphy/linear_algebra/tests/test_multigrid_transfer.py index de01ee5e0..2584f75a4 100644 --- a/src/struphy/linear_algebra/tests/test_multigrid_transfer.py +++ b/src/struphy/linear_algebra/tests/test_multigrid_transfer.py @@ -1,7 +1,6 @@ import numpy as np import pytest from feectools.ddm.mpi import mpi as MPI -from mpi4py import MPI as MPI4PY from struphy.feec.mass import WeightedMassOperators from struphy.feec.psydac_derham import Derham @@ -76,7 +75,7 @@ def test_hierarchy(num_elements, degree, bcs): def test_transfer(bcs, space_id): r"""Restriction is the transpose of the prolongation, and P is the exact embedding: R M_h P = M_H.""" h, domain = _hierarchy((16, 8, 8), (3, 2, 1), bcs, max_levels=3) - comm = MPI4PY.COMM_WORLD + comm = MPI.COMM_WORLD form = h[0].space_to_form[space_id] for l in range(h.n_levels - 1): @@ -92,8 +91,11 @@ def test_transfer(bcs, space_id): MH = getattr(WeightedMassOperators(h[l + 1], domain), "M" + form) a = R.dot(Mh.dot(P.dot(u))) b = MH.dot(u) - err = comm.allreduce(np.max(np.abs((a - b).toarray())), op=MPI4PY.MAX) - ref = comm.allreduce(np.max(np.abs(b.toarray())), op=MPI4PY.MAX) + err = np.max(np.abs((a - b).toarray())) + ref = np.max(np.abs(b.toarray())) + if comm.Get_size() > 1: + err = comm.allreduce(err, op=MPI.MAX) + ref = comm.allreduce(ref, op=MPI.MAX) assert err < 1e-12 * ref diff --git a/src/struphy/propagators/tests/test_poisson.py b/src/struphy/propagators/tests/test_poisson.py index d0f98a398..d1c70e48e 100644 --- a/src/struphy/propagators/tests/test_poisson.py +++ b/src/struphy/propagators/tests/test_poisson.py @@ -790,8 +790,6 @@ def rho2_pulled(e1, e2, e3): @pytest.mark.parametrize("bc_type", ["periodic", "dirichlet", "neumann"]) def test_poisson_2d_multigrid(degree, bc_type): """PoissonSolve with precond="MultiGrid" agrees with the unpreconditioned solve, in few iterations.""" - from mpi4py import MPI as MPI4PY - from struphy.linear_algebra.multigrid.preconditioner import MultiGridOptions domain = domains.Colella(Lx=4.0, Ly=2.0, alpha=0.1, Lz=1.0) @@ -831,7 +829,8 @@ def rho(e1, e2, e3): infos.append(solver._solver._info) # global coefficient arrays (toarray only fills the local part) - phis = [MPI4PY.COMM_WORLD.allreduce(p, op=MPI4PY.SUM) for p in phis] + if comm.Get_size() > 1: + phis = [comm.allreduce(p, op=MPI.SUM) for p in phis] if bc_type != "dirichlet": # solutions are defined up to a constant (the stabilization is tiny) phis = [p - xp.mean(p) for p in phis] From 71f223f4f77d5b2b6096c98021c227c17a36c36e Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Sun, 4 Oct 2026 08:57:56 +0200 Subject: [PATCH 14/16] Update feectools to latest devel-tiny (0.3.0) Co-Authored-By: Claude Opus 5.5 --- feectools | 2 +- pyproject.toml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/feectools b/feectools index eb6693cbf..a991b7182 160000 --- a/feectools +++ b/feectools @@ -1 +1 @@ -Subproject commit eb6693cbfb9fcb6a1d65ba6130b5775daf7f84c6 +Subproject commit a991b7182b5efdeafc2b968ac016b1dc87de335d diff --git a/pyproject.toml b/pyproject.toml index 6538f59f7..d7bec3f8a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -26,7 +26,7 @@ dependencies = [ "numpy<=2.5.0", "cunumpy>=0.2.0, <=0.3.0", "pyccel>=2.2.0, <=2.2.3", - "feectools>=0.1.11, <=0.2.0", + "feectools>=0.3.0, <=0.3.0", "scipy<=1.18.0", "h5py<=3.16.0", "h5netcdf<=1.8.1", From 29856ffa96cb5efc68147f95b5c7259811ed4f0f Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Mon, 5 Oct 2026 09:53:17 +0200 Subject: [PATCH 15/16] copy _creation_info too --- src/struphy/feec/mass.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/src/struphy/feec/mass.py b/src/struphy/feec/mass.py index 8fdbae777..c617c300c 100644 --- a/src/struphy/feec/mass.py +++ b/src/struphy/feec/mass.py @@ -2441,6 +2441,12 @@ def copy(self, out=None): out._weights = [list(row) for row in self._weights] self._mat.copy(out=out._mat) + + if self._creation_info is None: + out._creation_info = None + else: + out._creation_info = dict(self._creation_info) # to create a separate dictionary + return out def __imul__(self, a): From e120979aad6fb94ecf01b8a1d8652b5de1397699 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Mon, 5 Oct 2026 10:07:07 +0200 Subject: [PATCH 16/16] Multigrid: fix KrylovSmoother breakdown/iteration count, tuple weights round-trip - KrylovSmoother: own fixed-step PCG loop instead of feectools' PCG, which produced NaNs for a zero residual (tol**2 underflow), crashed for maxiter=1 and performed one iteration fewer than requested. - WeightedMassOperator.to_dict/from_dict: record whether weights were a tuple instead of guessing, so tuples starting with a constant 3x3 matrix are no longer mistaken for block weights. - Regression tests for both. Co-Authored-By: Claude Opus 5.5 --- src/struphy/feec/mass.py | 13 +++-- .../linear_algebra/multigrid/smoothers.py | 51 +++++++++++++++++-- .../tests/test_multigrid_coarsen.py | 7 +++ .../tests/test_multigrid_solver.py | 36 +++++++++++++ 4 files changed, 97 insertions(+), 10 deletions(-) diff --git a/src/struphy/feec/mass.py b/src/struphy/feec/mass.py index c617c300c..b230a64c0 100644 --- a/src/struphy/feec/mass.py +++ b/src/struphy/feec/mass.py @@ -2373,7 +2373,10 @@ def to_dict(self) -> dict: "WeightedMassOperators.create_weighted_mass or its data was modified afterwards." ) params = dict(self._creation_info) - if isinstance(params["weights"], tuple): + # tuple (1D product of weights) and 2D list (block weights) are different formats; store the + # tuple as a list (JSON) and record its type, since its entries may themselves be (3x3) lists + params["weights_is_tuple"] = isinstance(params["weights"], tuple) + if params["weights_is_tuple"]: params["weights"] = list(params["weights"]) return { "type": self.__class__.__name__, @@ -2396,7 +2399,7 @@ def from_dict(cls, dct: dict, mass_ops: "WeightedMassOperators") -> "WeightedMas params = dct["params"] name = params["name"] weights = params["weights"] - if isinstance(weights, list) and not (len(weights) > 0 and isinstance(weights[0], list)): + if params["weights_is_tuple"]: weights = tuple(weights) out = mass_ops.create_weighted_mass( @@ -2441,12 +2444,12 @@ def copy(self, out=None): out._weights = [list(row) for row in self._weights] self._mat.copy(out=out._mat) - + if self._creation_info is None: out._creation_info = None else: - out._creation_info = dict(self._creation_info) # to create a separate dictionary - + out._creation_info = dict(self._creation_info) # to create a separate dictionary + return out def __imul__(self, a): diff --git a/src/struphy/linear_algebra/multigrid/smoothers.py b/src/struphy/linear_algebra/multigrid/smoothers.py index c209e2f79..67a684332 100644 --- a/src/struphy/linear_algebra/multigrid/smoothers.py +++ b/src/struphy/linear_algebra/multigrid/smoothers.py @@ -18,7 +18,6 @@ ZeroOperator, ) from feectools.linalg.block import BlockVector -from feectools.linalg.solvers import inverse from feectools.linalg.stencil import StencilVector @@ -167,22 +166,64 @@ def smooth(self, b: Vector, x: Vector) -> None: class KrylovSmoother(Smoother): - """A fixed number of (preconditioned) conjugate gradient iterations, warm-started from ``x``. + r"""A fixed number of (preconditioned) conjugate gradient iterations, warm-started from ``x``. + + Each call to :meth:`smooth` performs exactly ``iterations`` PCG steps for :math:`A x = b`, + starting from the current ``x``. No convergence tolerance is used; the iteration stops early only on + breakdown, i.e. when :math:`r^\top M^{-1} r = 0` (e.g. zero residual) or :math:`p^\top A p = 0`. Note that this smoother is non-linear; a V-cycle using it is not a fixed linear preconditioner. + + Parameters + ---------- + A : LinearOperator + System operator (symmetric positive definite). + + M_inv : LinearOperator | None + Preconditioner :math:`M^{-1}` (symmetric positive definite). If None, the identity is used. + + iterations : int + Number of PCG steps per call (at least 1). """ def __init__(self, A: LinearOperator, M_inv: LinearOperator | None = None, *, iterations: int = 3): super().__init__(A) - self._solver = inverse(A, "pcg", pc=M_inv, maxiter=iterations, tol=1e-300, recycle=False) + assert iterations >= 1, f"KrylovSmoother needs at least one iteration, got {iterations}." + self._M_inv = IdentityOperator(A.domain) if M_inv is None else M_inv + self._iterations = iterations + self._z = A.domain.zeros() + self._p = A.domain.zeros() + self._q = A.domain.zeros() @property def is_symmetric(self) -> bool: return False def smooth(self, b: Vector, x: Vector) -> None: - self._solver._options["x0"] = x.copy() - self._solver.dot(b, out=x) + """Perform ``iterations`` PCG steps for ``A x = b``, updating ``x`` in place.""" + r, z, p, q = self._r, self._z, self._p, self._q + self.residual(b, x, r) + self._M_inv.dot(r, out=z) + z.copy(out=p) + rz = r.inner(z) + for k in range(self._iterations): + # inner products are global reductions, hence all ranks break consistently + if rz == 0.0: + break + self._A.dot(p, out=q) + pq = p.inner(q) + if pq == 0.0: + break + alpha = rz / pq + x.mul_iadd(alpha, p) + if k == self._iterations - 1: + break + r.mul_iadd(-alpha, q) + self._M_inv.dot(r, out=z) + rz_new = r.inner(z) + p *= rz_new / rz + p += z + rz = rz_new def estimate_lambda_max(A: LinearOperator, M_inv: LinearOperator, *, n_iter: int = 15, seed: int = 1234) -> float: diff --git a/src/struphy/linear_algebra/tests/test_multigrid_coarsen.py b/src/struphy/linear_algebra/tests/test_multigrid_coarsen.py index 455617072..2b6ad256e 100644 --- a/src/struphy/linear_algebra/tests/test_multigrid_coarsen.py +++ b/src/struphy/linear_algebra/tests/test_multigrid_coarsen.py @@ -64,6 +64,13 @@ def test_mass_to_dict(bcs): _, w = create_equal_random_arrays(derham.fem_spaces["2"], seed=2) assert _max_diff(WeightedMassOperator.from_dict(McT.to_dict(), mass_ops).dot(w), McT.dot(w)) < 1e-14 + # constant 3x3 matrix as first tuple entry (must not be mistaken for block weights) + Mm = mass_ops.create_weighted_mass( + "Hcurl", "Hcurl", weights=([[1.0, 0.0, 0.0], [0.0, 2.0, 0.0], [0.0, 0.0, 3.0]], "sqrt_g"), assemble=True + ) + dct = json.loads(json.dumps(Mm.to_dict())) + assert _max_diff(WeightedMassOperator.from_dict(dct, mass_ops).dot(u), Mm.dot(u)) < 1e-14 + # modified data cannot be re-created M = mass_ops.create_weighted_mass("H1", "H1", weights=("sqrt_g",), assemble=True) assert M.is_reconstructible diff --git a/src/struphy/linear_algebra/tests/test_multigrid_solver.py b/src/struphy/linear_algebra/tests/test_multigrid_solver.py index 7b0359686..3cd809231 100644 --- a/src/struphy/linear_algebra/tests/test_multigrid_solver.py +++ b/src/struphy/linear_algebra/tests/test_multigrid_solver.py @@ -18,6 +18,7 @@ ChebyshevSmoother, DiagonalComputer, JacobiSmoother, + KrylovSmoother, inverse_diagonal, ) from struphy.topology.grids import TensorProductGrid @@ -90,6 +91,41 @@ def test_smoother_symmetric(kind): assert np.abs(Sd - Sd.T).max() < 1e-12 * np.abs(Sd).max() +@pytest.mark.mpi_skip +@pytest.mark.parametrize("iterations", [1, 2, 4]) +def test_krylov_smoother(iterations): + """KrylovSmoother performs exactly ``iterations`` CG steps and is safe for a zero residual.""" + derham, _, mass_ops, A = _poisson(8, 2, PERIODIC, sigma=0.5) + _, b = create_equal_random_arrays(derham.fem_spaces["0"], seed=3) + + # zero right-hand side with zero initial guess: x stays zero (no NaN) + x = A.domain.zeros() + KrylovSmoother(A, iterations=iterations).smooth(A.domain.zeros(), x) + assert np.all(x.toarray() == 0.0) + + # the k-th CG iterate minimizes the A-norm error over the k-th Krylov space, so the error + # decreases strictly with the number of iterations; compare with one call of k - 1 iterations + x_ref = A.domain.zeros() + if iterations > 1: + KrylovSmoother(A, iterations=iterations - 1).smooth(b, x_ref) + x = A.domain.zeros() + KrylovSmoother(A, iterations=iterations).smooth(b, x) + + Ad = _assemble_dense(A) + x_ex = np.linalg.solve(Ad, b.toarray()) + + def err(y): + e = y.toarray() - x_ex + return e @ Ad @ e + + assert err(x) < err(x_ref) + + # one step from zero equals the steepest descent step alpha * b + if iterations == 1: + alpha = b.inner(b) / b.inner(A.dot(b)) + assert np.allclose(x.toarray(), alpha * b.toarray(), rtol=1e-12, atol=1e-14) + + @pytest.mark.mpi_skip @pytest.mark.parametrize("bcs, nullspace", [(DIRICHLET, None), (PERIODIC, "constants")]) @pytest.mark.parametrize("smoother_precond", ["mass", "jacobi"])