diff --git a/psydac/feec/polar/__init__.py b/psydac/feec/polar/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/psydac/feec/polar/conga_projections.py b/psydac/feec/polar/conga_projections.py new file mode 100644 index 000000000..715ea5e6c --- /dev/null +++ b/psydac/feec/polar/conga_projections.py @@ -0,0 +1,1245 @@ +import numpy as np +from numpy import pi + +from scipy.linalg import toeplitz, matmul_toeplitz +from scipy.sparse import coo_matrix, lil_matrix, spmatrix, eye as sp_eye +from scipy.sparse.linalg import inv as sp_inv + +from psydac.fem.vector import VectorFemSpace +from psydac.linalg.basic import LinearOperator, Vector +from psydac.linalg.block import BlockLinearOperator +from psydac.linalg.stencil import StencilVector +from psydac.fem.tensor import TensorFemSpace +from psydac.linalg.utilities import array_to_psydac + + +class C0PolarProjection_V0(LinearOperator): + """ + CONGA Projector P0 from the full spline space S^{p1, p2} on logical domain + to V0 the pre-polar 0-forms splines. The associate matrix is square as in + the CONGA approach we keep using the tensor B-spline basis, instead of the + polar basis of Toshniwal. P0 enforces coefficient relations to be in V0. + + Parameters: + ----------- + + W0 : TensorFemSpace + The full tensor product spline space S^{p1,p2} + + transposed : Boolean + switch between P0 and P0 transposed (defalut is False) + + hbc : Boolean + switch on and off the imposition of homogeneous Dirichlet boundary + conditions (default is False) + """ + + def __init__(self, W0, *, transposed=False, hbc=False): + assert isinstance(W0, TensorFemSpace) + + self.W0 = W0 + self.transposed = transposed + self.hbc = hbc + + @property + def domain(self): + return self.W0.coeff_space + + @property + def codomain(self): + return self.W0.coeff_space + + @property + def dtype(self): + return float + + def dot(self, x, out=None): + assert isinstance(x, StencilVector) + if not x.ghost_regions_in_sync: + x.update_ghost_regions() + + [s1, s2] = self.W0.coeff_space.starts + [e1, e2] = self.W0.coeff_space.ends + [n1, n2] = self.W0.coeff_space.npts + rank_at_polar_edge = (s1 == 0) + rank_at_outer_edge = (e1 == n1 - 1) + + if out is None: + y = self.W0.coeff_space.zeros() + else: + assert isinstance(out, StencilVector) + assert out.space is self.W0.coeff_space + out *= 0 + y = out + + if rank_at_polar_edge: + # compute sum of points with s = 0 for the process + local_sum = np.sum(x[0, s2:e2 + 1]) + if x.space.parallel: + from mpi4py import MPI + angle_comm = x.space.cart.subcomm[1] + local_sum = angle_comm.allreduce(local_sum, op=MPI.SUM) + #compute average of all points with s = 0 + y[0, s2:e2 + 1] = local_sum / n2 + y[1:e1 + 1, s2:e2 + 1] = x[1:e1 + 1, s2:e2 + 1] + else: + y[s1:e1 + 1, s2:e2 + 1] = x[s1:e1 + 1, s2:e2 + 1] + + + if self.hbc: + if rank_at_outer_edge: + y[e1, :] = 0. + + y.update_ghost_regions() + return y + + def transpose(self, conjugate=False): + return C0PolarProjection_V0(self.W0, transposed=not self.transposed, hbc=self.hbc) + + def tosparse(self): + + # Matrix size is n1*n2. Rows with global indices i*n2 + j (C-style flattening) with + # s1 <= i <= e1, s2 <= j <= e2 are filled for the current process + + [n1, n2] = self.W0.coeff_space.npts + [s1, s2] = self.W0.coeff_space.starts + [e1, e2] = self.W0.coeff_space.ends + rank_at_polar_edge = (s1 == 0) + rank_at_outer_edge = (e1 == n1 - 1) + data, cols, rows = [], [], [] + + # Assemble the matrix entries which contribute to the polar edge (s1 = 0) + len_theta = e2 - s2 + 1 + if rank_at_polar_edge: + data = np.tile((1 / n2) * np.ones(n2), len_theta) + rows = np.repeat(np.arange(s2, e2 + 1), n2) + cols = np.tile(np.arange(n2), len_theta) + + # Assemble the rest of the matrix (identity block) + # We do not need entries with s1 = 0 (already accounted for) + start_s = s1 + 1 if rank_at_polar_edge else s1 + end_s = e1 if (rank_at_outer_edge and self.hbc) else e1 + 1 + + i = np.arange(start_s, end_s)[:, None] + j = np.arange(s2, e2 + 1)[None, :] + # rows of identity owned by the process + local_rows = (i * n2 + j).ravel() + + data = np.concatenate((data, np.ones(len_theta * (end_s - start_s)))) + cols = np.concatenate((cols, local_rows)) + rows = np.concatenate((rows, local_rows)) + + rows = np.asarray(rows, dtype=np.int64) + cols = np.asarray(cols, dtype=np.int64) + + P = coo_matrix((data, (rows, cols)), shape=[n1 * n2, n1 * n2], dtype=self.W0.coeff_space.dtype) + P.eliminate_zeros() + + # P is symmetric so we always return P no matter self.transposed + return P + + def toarray(self): + return self.tosparse().toarray() + + +# --------- 1-FORMS CONGA PROJECTOR P1 ----------# +# It is a BlockLinearOperator with 4 blocks Upper-Left (0, 0), Upper-Right (0, 1) +# Lower-Left (1, 0) and Lower-Right (1, 1). +# +# ______________________ +# | | | +# | (0,0) | (0,1) | +# |___________|__________| +# | | | +# | (1,0) | (1,1) | +# |___________|__________| + +class C0PolarProjection_V1_00(LinearOperator): + """ + Upper Left block of P1. + + Parameters: + ---------- + + W1 : VectorFemSpace + Full tensor product spline space of the 1-forms S^{p1-1, p2} x S^{p1, p2-1} + + transposed : Boolean + Switch between P1 and P1 transposed (default is False) + """ + + def __init__(self, W1): + # assert isinstance(W1, ProductFemSpace) + assert isinstance(W1, VectorFemSpace) + + self.W1 = W1 + + @property + def domain(self): + return self.W1.coeff_space[0] + + @property + def codomain(self): + return self.W1.coeff_space[0] + + @property + def dtype(self): + return float + + # The upper-left block of P1 is an identity block + def dot(self, x, out=None): + assert isinstance(x, StencilVector) + + if not x.ghost_regions_in_sync: + x.update_ghost_regions() + + [s1, s2] = self.W1.coeff_space[0].starts + [e1, e2] = self.W1.coeff_space[0].ends + + if out is None: + y = self.W1.coeff_space[0].zeros() + else: + assert isinstance(out, StencilVector) + assert out.space is self.W1.coeff_space[0] + out *= 0 + y = out + + y[s1:e1+1, s2:e2+1] = x[s1:e1+1, s2:e2+1] + + y.update_ghost_regions() + return y + + def transpose(self, conjugate=False): + return self + + @property + def T(self): + return self.transpose() + + def tosparse(self): + + [n01, n02] = self.domain.npts + [s1, s2] = self.domain.starts + [e1, e2] = self.domain.ends + + i = np.arange(s1, e1 + 1)[:, None] + j = np.arange(s2, e2 + 1)[None, :] + local_rows = (i * n02 + j).ravel() + data = np.ones((e1 - s1 + 1) * (e2 - s2 + 1)) + + local_rows = np.asarray(local_rows, dtype=np.int64) + + P = coo_matrix((data, (local_rows, local_rows)), shape=[n01 * n02, n01 * n02], dtype=self.domain.dtype) + P.eliminate_zeros() + + return P + + def toarray(self): + return self.tosparse().toarray() + + + +class C0PolarProjection_V1_10(LinearOperator): + """ + Lower left block of P1. + + Parameters: + ----------- + + W1 : VectorFemSpace (former ProductFemSpace) + Full tensor product spline space of the 1-forms S^{p1-1, p2} x S^{p1, p2-1} + + transposed : Boolean + Switch between P1 and P1 transposed (default is False) + """ + + def __init__(self, W1, transposed=False): + # assert isinstance(W1, ProductFemSpace) + assert isinstance(W1, VectorFemSpace) + + self.W1 = W1 + self.transposed = transposed + + @property + def domain(self): + idx = 1 if self.transposed else 0 + return self.W1.coeff_space[idx] + + @property + def codomain(self): + idx = 0 if self.transposed else 1 + return self.W1.coeff_space[idx] + + @property + def dtype(self): + return float + + def dot(self, x, out=None): + assert isinstance(x, StencilVector) + + if not x.ghost_regions_in_sync: + x.update_ghost_regions() + + # The number of radial basis functions is one less than the number on the angular basis functions along dir x1 + [s1, s2] = self.domain.starts + [e1, e2] = self.domain.ends + rank_at_polar_edge = (s1 == 0) + + if out is None: + y = self.codomain.zeros() + else: + assert isinstance(out, StencilVector) + assert out.space is self.codomain + out *= 0 + y = out + + y[s1:e1 + 1, s2:e2 + 1] = 0 + + if self.transposed: + if rank_at_polar_edge: + y[0, s2:e2 + 1] = x[1, s2 - 1:e2] - x[1, s2:e2 + 1] + else: + if rank_at_polar_edge: + # the cells are periodic in angular dim, so we use ghost regions + y[1, s2:e2 + 1] = x[0, s2 + 1:e2 + 2] - x[0, s2:e2 + 1] + + y.update_ghost_regions() + return y + + def transpose(self, conjugate=False): + return C0PolarProjection_V1_10(self.W1, transposed=not self.transposed) + + @property + def T(self): + return self.transpose() + + def tosparse(self): + + [n01, n02] = self.domain.npts + [n11, n12] = self.codomain.npts + s1, s2 = self.domain.starts + e1, e2 = self.codomain.ends + rank_at_polar_edge = (s1 == 0) + + data, cols, rows = [], [], [] + + # matrix d (derivatives of periodic splines) + # for example: n_theta = 3, single rank at polar edge, not transposed. Then: + # data = [-1 1 -1 1 -1 1], rows = [3 3 4 4 5 5], cols = [0 1 1 2 2 0] + if rank_at_polar_edge: + len_theta = e2 - s2 + 1 + data = np.tile([-1, 1], len_theta) + k = np.arange(s2, e2 + 1) + if self.transposed: + rows = np.repeat(np.arange(s2, e2 + 1), 2) + cols = np.column_stack((k, (k - 1) % n12)).ravel() + n12 + else: + rows = np.repeat(np.arange(s2, e2 + 1), 2) + n12 + cols = np.column_stack((k, (k + 1) % n12 )).ravel() + + dtype = self.domain.dtype + rows = np.asarray(rows, dtype=np.int64) + cols = np.asarray(cols, dtype=np.int64) + + P = coo_matrix((data, (rows, cols)), shape=[n11 * n12, n01 * n02], dtype=dtype) + P.eliminate_zeros() + return P + + def toarray(self): + return self.tosparse().toarray() + + +class C0PolarProjection_V1_11(LinearOperator): + """ + Lower right block of P1. + + Parameters: + ----------- + + W1 : VectorFemSpace (former ProductFemSpace) + Full tensor product spline space of the 1-forms S^{p1-1, p2} x S^{p1, p2-1} + + transposed : Boolean + Switch between P1 and P1 transposed (default is False) + + hbc : Boolean + Switch on and off the imposition of homogeneous Dirichlet boundary + conditions on the tangential (angular) direction (default is False) + """ + + def __init__(self, W1, transposed=False, hbc=False): + # assert isinstance(W1, ProductFemSpace) + assert isinstance(W1, VectorFemSpace) + + self.W1 = W1 + self.transposed = transposed + self.hbc = hbc + + @property + def domain(self): + return self.W1.coeff_space[1] + + @property + def codomain(self): + return self.W1.coeff_space[1] + + @property + def dtype(self): + return float + + # It is diagonal block with a zero block of size (2n2)x(2n2) and an identity block + def dot(self, x, out=None): + assert isinstance(x, StencilVector) + + if not x.ghost_regions_in_sync: + x.update_ghost_regions() + + [s1, s2] = self.codomain.starts + [e1, e2] = self.codomain.ends + [n1, n2] = self.codomain.npts + rank_at_polar_edge = (s1 == 0) + rank_at_outer_edge = (e1 == n1 - 1) + + if out is None: + y = self.W1.coeff_space[1].zeros() + else: + assert isinstance(out, StencilVector) + assert out.space is self.codomain + out *= 0 + y = out + if rank_at_polar_edge: + y[0, s2:e2 + 1] = 0 + y[1, s2:e2 + 1] = 0 + y[2:e1 + 1, s2:e2 + 1] = x[2:e1 + 1, s2:e2 + 1] # Identity block + else: + y[s1:e1 + 1, s2:e2 + 1] = x[s1:e1 + 1, s2:e2 + 1] + + if self.hbc and rank_at_outer_edge: + y[e1, :] = 0. + + y.update_ghost_regions() + return y + + def transpose(self, conjugate=False): + return C0PolarProjection_V1_11(self.W1, transposed=not self.transposed, hbc=self.hbc) + + @property + def T(self): + return self.transpose() + + def tosparse(self): + + [s1, s2] = self.codomain.starts + [e1, e2] = self.codomain.ends + [n01, n02] = self.domain.npts + [n11, n12] = self.codomain.npts + rank_at_outer_edge = (e1 == n01 - 1) + dtype = self.domain.dtype + + # identity block of size n12(n12 - 2) + start_s = max(2, s1) + end_s = e1 if (rank_at_outer_edge and self.hbc) else e1 + 1 + len_theta = e2 - s2 + 1 + data = np.ones(len_theta * (end_s - start_s)) + + i = np.arange(start_s, end_s)[:, None] + j = np.arange(s2, e2 + 1)[None, :] + local_rows = (i * n02 + j).ravel() + local_rows = np.asarray(local_rows, dtype=np.int64) + + P = coo_matrix((data, (local_rows, local_rows)), shape=[n11 * n12, n01 * n02], dtype=dtype) + P.eliminate_zeros() + return P + + + def toarray(self): + return self.tosparse().toarray() + + +class C0PolarProjection_V1(BlockLinearOperator): + """ + CONGA Projector P1 from the full spline space S^{p1-1, p2} x S^{p1, p2-1} + on logical domain to V1 the pre-polar 1-forms splines. The associate matrix + is square as in the CONGA approach we keep using the tensor B-spline basis, + instead of the polar basis of Toshniwal. P1 enforces coefficient relations + to be in V1. + + Parameters: + ----------- + + W1 : VectorFemSpace (former ProductFemSpace) + Full tensor product spline space of the 1-forms S^{p1-1, p2} x S^{p1, p2-1} + + transposed : Boolean + Switch between P1 and P1 transposed (default is False) + + hbc : Boolean + Switch on and off the imposition of homogeneous Dirichlet boundary + conditions on the tangential (angular) direction (default is False) + """ + + def __init__(self, W1, transposed=False, hbc=False): + assert isinstance(W1, VectorFemSpace) + assert W1.symbolic_space.kind.name == 'hcurl' + + T1 = W1.coeff_space + + super().__init__(T1, T1) + + self[0, 0] = C0PolarProjection_V1_00(W1) + # self[0, 1] = C0CongaProjector1_01(W1, transposed = transposed) + self[1, 0] = C0PolarProjection_V1_10(W1, transposed=transposed) + self[1, 1] = C0PolarProjection_V1_11(W1, transposed=transposed, hbc=hbc) + + self.W1 = W1 + self.transposed = transposed + self.hbc = hbc + + + @property + def T(self): + return self.transpose() + + +# ---------------- 2-FORMS CONGA PROJECTOR P2 ----------------# +class C0PolarProjection_V2(LinearOperator): + """ + CONGA Projector P2 from the full spline space S^{p1-1, p2-1} on logical + domain to V2, the pre-polar 2-forms splines. The associate matrix + is square, as in the CONGA approach we keep using the tensor B-spline basis, + instead of the polar basis of Toshniwal. P2 enforces coefficient relations + to be in V2. + + Parameters: + ----------- + + W2 : TensorFemSpace + Full tensor product spline space of the 2-forms S^{p1-1, p2-1} + + transposed : Boolean + Switch between P2 and P2 transposed (default is False) + """ + + def __init__(self, W2, transposed=False): + assert isinstance(W2, TensorFemSpace) + + self.W2 = W2 + self.transposed = transposed + + @property + def domain(self): + return self.W2.coeff_space + + @property + def codomain(self): + return self.W2.coeff_space + + @property + def dtype(self): + return float + + def dot(self, x, out=None): + assert isinstance(x, StencilVector) + + if not x.ghost_regions_in_sync: + x.update_ghost_regions() + + [s1, s2] = self.W2.coeff_space.starts + [e1, e2] = self.W2.coeff_space.ends + [n1, n2] = self.W2.coeff_space.npts + rank_at_polar_edge = (s1 == 0) + + if out is None: + y = self.W2.coeff_space.zeros() + else: + assert isinstance(out, StencilVector) + assert out.space is self.W2.coeff_space + out *= 0 + y = out + + if self.transposed: + if rank_at_polar_edge: + y[0, s2:e2 + 1] = x[1, s2:e2 + 1] + y[1, s2:e2 + 1] = x[1, s2:e2 + 1] + y[2:e1 + 1, s2:e2 + 1] = x[2:e1 + 1, s2:e2 + 1] + else: + y[s1:e1 + 1, s2:e2 + 1] = x[s1:e1 + 1, s2:e2 + 1] + + else: + if rank_at_polar_edge: + y[0, s2:e2 + 1] = 0 + y[1, s2:e2 + 1] = x[0, s2:e2 + 1] + x[1, s2:e2 + 1] + y[2:e1 + 1, s2:e2 + 1] = x[2:e1 + 1, s2:e2 + 1] + else: + y[s1:e1 + 1, s2:e2 + 1] = x[s1:e1 + 1, s2:e2 + 1] + + + y.update_ghost_regions() + return y + + def transpose(self, conjugate=False): + return C0PolarProjection_V2(self.W2, transposed=not self.transposed) + + @property + def T(self): + return self.transpose() + + def tosparse(self): + + [s1, s2] = self.W2.coeff_space.starts + [e1, e2] = self.W2.coeff_space.ends + [n1, n2] = self.W2.coeff_space.npts + + data = np.ones((e1 - s1 + 1) * (e2 - s2 + 1)) + + i = np.arange(s1, e1 + 1)[:, None] + j = np.arange(s2, e2 + 1)[None, :] + local_rows = (i * n2 + j).ravel() + if self.transposed: + rows = local_rows + cols = local_rows[local_rows < n2] + n2 + else: + rows_to_repeat = local_rows[(local_rows >= n2) & (local_rows < 2 * n2)] + rows = np.tile(rows_to_repeat, 2) + cols = rows_to_repeat - n2 + rows = np.concatenate((rows, local_rows[local_rows >= 2 * n2])) + cols = np.concatenate((cols, local_rows[local_rows >= n2])) + + rows = np.asarray(rows, dtype=np.int64) + cols = np.asarray(cols, dtype=np.int64) + + P = coo_matrix((data, (rows, cols)), shape=(n1 * n2, n1 * n2), dtype=self.W2.coeff_space.dtype) + P.eliminate_zeros() + + return P + + def toarray(self): + return self.tosparse().toarray() + + + + +# =============================================================================# +# C1 POLAR SPLINE COMPLEX (C1P) # +# =============================================================================# + +def cos_sin_avg(theta, x, angle_comm, s2, e2, n2, i): + """ + Compute weighted averages in parallel: + sum (2/n2) * cos(theta_k) * x_ik + sum (2/n2) * sin(theta_k) * x_ik + over all k belonging to the process, for fixed i (0 or 1) + then sum over all processes + """ + from mpi4py import MPI + sum_cos = (2 / n2) * np.dot(np.cos(theta), x[i, s2:e2 + 1]) + sum_sin = (2 / n2) * np.dot(np.sin(theta), x[i, s2:e2 + 1]) + buf = np.array([sum_cos, sum_sin]) + + if angle_comm is not None: + buf = angle_comm.allreduce(buf, op=MPI.SUM) + return buf + + +# --------- 0-FORMS CONGA PROJECTOR P0 ----------# + +# matrix p in the paper notation +def toeplitz_columns_sym(t, s2, e2, n2): + + i = np.arange(n2)[:, None] # (n2, 1) + j = np.arange(s2, e2 + 1)[None, :] # (1, m) + idx = np.abs(i - j).ravel('F') # (n2, m) + return np.cos(t[idx]) + + +class C1PolarProjection_V0(LinearOperator): + """ + CONGA Projector P0 from the full spline space S^{p1, p2} on logical domain + to U0 the pre-polar 0-forms splines. The associate matrix is square as in + the CONGA approach we keep using the tensor B-spline basis, instead of the + polar basis of Toshniwal. P0 enforces coefficient relations to be in U0. + + Parameters: + ----------- + + W0 : TensorFemSpace + The full tensor product spline space S^{p1,p2} + + gamma : float + free parameter in the entries of P0. Any values provides a valid CONGA + prjector in U0. However, in order to have the commuting property + grad P0 u = P1 grad u + for u in Im(Pi0) and Pi0 the geometric projector on W0, we should set + gamma = 1 (default) + + transposed : Boolean + switch between P0 and P0 transposed (defalut is False) + + hbc : Boolean + switch on and off the imposition of homogeneous Dirichlet boundary + conditions (default is False) + """ + + def __init__(self, W0, *, gamma=1, transposed=False, hbc=False): + assert isinstance(W0, TensorFemSpace) + + self.gamma = gamma + self.W0 = W0 + self.transposed = transposed + self.hbc = hbc + + @property + def domain(self): + return self.W0.coeff_space + + @property + def codomain(self): + return self.W0.coeff_space + + @property + def dtype(self): + return float + + def dot(self, x, out=None): + assert isinstance(x, StencilVector) + + if not x.ghost_regions_in_sync: + x.update_ghost_regions() + + [s1, s2] = self.W0.coeff_space.starts + [e1, e2] = self.W0.coeff_space.ends + [n1, n2] = self.W0.coeff_space.npts + theta_local = np.linspace(s2 * 2 * pi / n2, e2 * 2 * pi / n2, e2 - s2 + 1) + rank_at_polar_edge = (s1 == 0) + rank_at_outer_edge = (e1 == n1 - 1) + + if out is None: + y = self.W0.coeff_space.zeros() + else: + assert isinstance(out, StencilVector) + assert out.space is self.W0.coeff_space + out *= 0 + y = out + + if rank_at_polar_edge: + # compute sums of points with s = 0, 1 for the process + x0 = np.sum(x[0, s2:e2 + 1]) / n2 + x1 = np.sum(x[1, s2:e2 + 1]) / n2 + + angle_comm = x.space.cart.subcomm[1] if x.space.parallel else None + sum_cos, sum_sin = cos_sin_avg(theta_local, x, angle_comm, s2, e2, n2, 1) + + if x.space.parallel: + + from mpi4py import MPI + # global averages + x0 = angle_comm.allreduce(x0, op=MPI.SUM) + x1 = angle_comm.allreduce(x1, op=MPI.SUM) + + if self.transposed: + y[0, s2:e2 + 1] = self.gamma * x0 + self.gamma * x1 + y[1, s2:e2 + 1] = (1 - self.gamma) * x0 + (1 - self.gamma) * x1 + \ + np.cos(theta_local) * sum_cos + np.sin(theta_local) * sum_sin + else: + y[0, s2:e2 + 1] = self.gamma * x0 + (1 - self.gamma) * x1 + y[1, s2:e2 + 1] = self.gamma * x0 + (1 - self.gamma) * x1 + \ + np.cos(theta_local) * sum_cos + np.sin(theta_local) * sum_sin + y[2:e1 + 1, s2:e2 + 1] = x[2:e1 + 1, s2:e2 + 1] + else: + y[s1:e1 + 1, s2:e2 + 1] = x[s1:e1 + 1, s2:e2 + 1] + + if self.hbc and rank_at_outer_edge: + y[e1, :] = 0. + + y.update_ghost_regions() + return y + + def transpose(self, conjugate=False): + return C1PolarProjection_V0(self.W0, gamma=self.gamma, transposed=not self.transposed, hbc=self.hbc) + + @property + def T(self): + return self.transpose() + + def tosparse(self): + + [n1, n2] = self.W0.coeff_space.npts + [s1, s2] = self.W0.coeff_space.starts + [e1, e2] = self.W0.coeff_space.ends + theta = np.linspace(0, 2 * pi, n2, endpoint=False) + + rank_at_polar_edge = (s1 == 0) + rank_at_outer_edge = (e1 == n1 - 1) + data, cols, rows = [], [], [] + + if rank_at_polar_edge: + # matrix of size n2*n2 with all entries equal to gamma / n2 + data = (self.gamma / n2) * np.ones(n2 * (e2 - s2 + 1)) + # matrix of size n2*n2 with all entries equal to (1 - gamma) / n2 + data = np.concatenate((data, (1 - self.gamma) / n2 * np.ones((e2 - s2 + 1) * n2))) + rows = np.tile(np.repeat(np.arange(s2, e2 + 1), n2), 2) + cols = np.tile(np.arange(n2), e2 - s2 + 1) + cols = np.concatenate((cols, np.tile(np.arange(n2, 2 * n2), e2 - s2 + 1))) + if e1 > 1 > s1: + # matrix of size n2*n2 with all entries equal to gamma / n2 + d_block2 = (self.gamma / n2) * np.ones((e2 - s2 + 1) * n2) + data = np.concatenate((data, d_block2)) + rows = np.concatenate((rows, np.repeat(np.arange(n2 + s2, n2 + e2 + 1), n2))) + cols = np.concatenate((cols, np.tile(np.arange(n2), e2 - s2 + 1))) + + # matrix p + d_block2 = (1 - self.gamma) / n2 * np.ones((e2 - s2 + 1) * n2) + \ + 2 / n2 * toeplitz_columns_sym(theta, s2, e2, n2) + data = np.concatenate((data, d_block2)) + rows = np.concatenate((rows, np.repeat(np.arange(n2 + s2, n2 + e2 + 1), n2))) + cols = np.concatenate((cols, np.tile(np.arange(n2, 2 * n2), e2 - s2 + 1))) + + # Assemble the rest of the matrix (identity block) + start_s = max(s1, 2) + end_s = e1 if (rank_at_outer_edge and self.hbc) else e1 + 1 + i = np.arange(start_s, end_s)[:, None] + j = np.arange(s2, e2 + 1)[None, :] + local_rows = (i * n2 + j).ravel() + + data = np.concatenate((data, np.ones((e2 - s2 + 1) * (end_s - start_s)))) + cols = np.concatenate((cols, local_rows)) + rows = np.concatenate((rows, local_rows)) + + rows = np.asarray(rows, dtype=np.int64) + cols = np.asarray(cols, dtype=np.int64) + + P = coo_matrix((data, (rows, cols)), shape=[n1 * n2, n1 * n2], dtype=self.W0.coeff_space.dtype) + P.eliminate_zeros() + + return P.T if self.transposed else P + + def toarray(self): + return self.tosparse().toarray() + + +# --------- 1-FORMS CONGA PROJECTOR P1 ----------# +# It is a BlockLinearOperator with 4 blocks Upper-Left (0, 0), Upper-Right (0, 1) +# Lower-Left (1, 0) and Lower-Right (1, 1). +# +# ______________________ +# | | | +# | (0,0) | (0,1) | +# |___________|__________| +# | | | +# | (1,0) | (1,1) | +# |___________|__________| + +class C1PolarProjection_V1_00(LinearOperator): + """ + Upper Left block of P1. + + Parameters: + ---------- + + W1 : VectorFemSpace (former ProductFemSpace) + Full tensor product spline space of the 1-forms S^{p1-1, p2} x S^{p1, p2-1} + + transposed : Boolean + Switch between P1 and P1 transposed (default is False) + """ + + def __init__(self, W1, transposed=False): + assert isinstance(W1, VectorFemSpace) + + self.W1 = W1 + self.transposed = transposed + + @property + def domain(self): + return self.W1.coeff_space[0] + + @property + def codomain(self): + return self.W1.coeff_space[0] + + @property + def dtype(self): + return float + + def dot(self, x, out=None): + assert isinstance(x, StencilVector) + + [s1, s2] = self.domain.starts + [e1, e2] = self.domain.ends + [n1, n2] = self.domain.npts + theta_local = np.linspace(s2 * 2 * pi / n2, e2 * 2 * pi / n2, e2 - s2 + 1) + rank_at_polar_edge = (s1 == 0) + + if out is None: + y = self.codomain.zeros() + else: + assert isinstance(out, StencilVector) + assert out.space is self.codomain + out *= 0 + y = out + + if rank_at_polar_edge: + angle_comm = x.space.cart.subcomm[1] if x.space.parallel else None + sum_cos0, sum_sin0 = cos_sin_avg(theta_local, x, angle_comm, s2, e2, n2, 0) + + if self.transposed: + sum_cos1, sum_sin1 = cos_sin_avg(theta_local, x, angle_comm, s2, e2, n2, 1) + y[0, s2:e2 + 1] = sum_cos0 * np.cos(theta_local) + sum_sin0 * np.sin(theta_local) +\ + x[1, s2:e2 + 1] - sum_cos1 * np.cos(theta_local) - sum_sin1 * np.sin(theta_local) + y[1, s2:e2 + 1] = x[1, s2:e2 + 1] + else: + y[0, s2:e2 + 1] = sum_cos0 * np.cos(theta_local) + sum_sin0 * np.sin(theta_local) + y[1, s2:e2 + 1] = x[0, s2:e2 + 1] + x[1, s2:e2 + 1] - y[0, s2:e2 + 1] + + y[2:e1 + 1, s2:e2 + 1] = x[2:e1 + 1, s2:e2 + 1] + else: + y[s1:e1 + 1, s2:e2 + 1] = x[s1:e1 + 1, s2:e2 + 1] + + y.update_ghost_regions() + return y + + def transpose(self, conjugate=False): + return C1PolarProjection_V1_00(self.W1, transposed=not self.transposed) + + @property + def T(self): + return self.transpose() + + def tosparse(self): + + [n01, n02] = self.domain.npts + theta = np.linspace(0, 2 * pi, n02, endpoint=False) # Warning not mpi! + [s1, s2] = self.domain.starts + [e1, e2] = self.domain.ends + + data, cols, rows = [], [], [] + rank_at_polar_edge = (s1 == 0) + + if rank_at_polar_edge: + #block p + data = (2 / n02) * toeplitz_columns_sym(theta, s2, e2, n02) + # block I - p + i_minus_p = -1 * data.copy() + diag_js = np.arange(s2, e2 + 1) + diag_pos = (diag_js - s2) * n02 + diag_js + i_minus_p[diag_pos] += 1.0 + + data = np.concatenate((data, i_minus_p)) + + cols_theta = np.tile(np.arange(n02), e2 - s2 + 1) + rows_theta = np.repeat(np.arange(s2, e2 + 1), n02) + rows = np.concatenate((rows, rows_theta)) + cols = np.concatenate((cols, cols_theta)) + rows = np.concatenate((rows, rows_theta + 1 * n02)) + cols = np.concatenate((cols, cols_theta)) + + # Assemble the rest of the matrix (identity block) + start_s = max(s1, 1) + end_s = e1 + 1 + i = np.arange(start_s, end_s)[:, None] + j = np.arange(s2, e2 + 1)[None, :] + local_rows = (i * n02 + j).ravel() + + data = np.concatenate((data, np.ones((e2 - s2 + 1) * (end_s - start_s)))) + cols = np.concatenate((cols, local_rows)) + rows = np.concatenate((rows, local_rows)) + + rows = np.asarray(rows, dtype=np.int64) + cols = np.asarray(cols, dtype=np.int64) + + P = coo_matrix((data, (rows, cols)), shape=[n01 * n02, n01 * n02], dtype=self.domain.dtype) + P.eliminate_zeros() + + return P.T if self.transposed else P + + def toarray(self): + return self.tosparse().toarray() + + +class C1PolarProjection_V1_10(LinearOperator): + """ + Lower left block of P1. + + Parameters: + ----------- + + W1 : VectorFemSpace (former ProductFemSpace) + Full tensor product spline space of the 1-forms S^{p1-1, p2} x S^{p1, p2-1} + + transposed : Boolean + Switch between P1 and P1 transposed (default is False) + """ + + def __init__(self, W1, transposed=False): + assert isinstance(W1, VectorFemSpace) + + self.W1 = W1 + self.transposed = transposed + + @property + def domain(self): + if self.transposed: + return self.W1.coeff_space[1] + return self.W1.coeff_space[0] + + @property + def codomain(self): + if self.transposed: + return self.W1.coeff_space[0] + return self.W1.coeff_space[1] + + @property + def dtype(self): + return float + + def forward_theta_diff_local(self, z, angle_comm): + + # Compute result[j] = z[j+1] - z[j] for distributed array z (periodic) + result = np.empty_like(z) + + if angle_comm is None or angle_comm.size == 1: + result[:] = np.roll(z, -1) - z + return result + + result[:-1] = z[1:] - z[:-1] + + rank = angle_comm.rank + size = angle_comm.size + right_rank = (rank + 1) % size + left_rank = (rank - 1) % size + + recvbuf = np.empty(1, dtype=z.dtype) + angle_comm.Sendrecv(sendbuf=z[0], dest=left_rank, recvbuf=recvbuf, source=right_rank) + + result[-1] = recvbuf[0] - z[-1] + return result + + def dot(self, x, out=None): + assert isinstance(x, StencilVector) + if not x.ghost_regions_in_sync: + x.update_ghost_regions() + + # The number of radial basis functions is one less than + # the number on the angular basis functions + [s1, s2] = self.domain.starts + [e1, e2] = self.domain.ends + [n1, n2] = self.domain.npts + theta_local = np.linspace(s2 * 2 * pi / n2, e2 * 2 * pi / n2, e2 - s2 + 1) + rank_at_polar_edge = (s1 == 0) + + if out is None: + y = self.codomain.zeros() + else: + assert isinstance(out, StencilVector) + assert out.space is self.codomain + out *= 0 + y = out + + if rank_at_polar_edge: + angle_comm = x.space.cart.subcomm[1] if x.space.parallel else None + + if self.transposed: + y[0, s2:e2 + 1] = x[1, s2 - 1:e2] - x[1, s2:e2 + 1] + sum_cos, sum_sin = cos_sin_avg(theta_local, y, angle_comm, s2, e2, n2, 0) + y[0, s2:e2 + 1] = np.cos(theta_local) * sum_cos + np.sin(theta_local) * sum_sin + y[1:e1 + 1, s2:e2 + 1] = 0 + else: + y[0, s2:e2 + 1] = 0 + sum_cos, sum_sin = cos_sin_avg(theta_local, x, angle_comm, s2, e2, n2, 0) + row = np.cos(theta_local) * sum_cos + np.sin(theta_local) * sum_sin + y[1, s2:e2 + 1] = self.forward_theta_diff_local(row, angle_comm) + y[2:e1 + 1, s2:e2 + 1] = 0 + else: + y[s1:e1 + 1, s2:e2 + 1] = 0 + + + y.update_ghost_regions() + return y + + def transpose(self, conjugate=False): + return C1PolarProjection_V1_10(self.W1, transposed=not self.transposed) + + @property + def T(self): + return self.transpose() + + def tosparse(self): + + if self.transposed: + domain_P1_10 = self.codomain + codomain_P1_10 = self.domain + else: + domain_P1_10 = self.domain + codomain_P1_10 = self.codomain + + [n01, n02] = domain_P1_10.npts # radial grid + [n11, n12] = codomain_P1_10.npts # angular grid + [s1, s2] = domain_P1_10.starts + [e1, e2] = domain_P1_10.ends + + rank_at_polar_edge = (s1 == 0) + data, cols, rows = [], [], [] + if rank_at_polar_edge: + local_j = np.arange(s2, e2 + 1) # rows + all_k = np.arange(n02) # cols + + rows = np.repeat(local_j + n02, n02) + cols = np.tile(all_k, len(local_j)) + + theta = 2.0 * pi * np.arange(n02) / n02 + + j = local_j[:, None] + i = all_k[None, :] + + p = np.cos(theta[j - i]) + p_next = np.cos(theta[(j + 1) % n02 - i]) + q = (2 / n02) * (p_next - p) + data = q.ravel(order='C') + + rows = np.asarray(rows, dtype=np.int64) + cols = np.asarray(cols, dtype=np.int64) + + P = coo_matrix((data, (rows, cols)), shape=[n11 * n12, n01 * n02], dtype=self.domain.dtype) + + return P.T if self.transposed else P + + def toarray(self): + return self.tosparse().toarray() + + + +class C1PolarProjection_V1(BlockLinearOperator): + """ + CONGA Projector P1 from the full spline space S^{p1-1, p2} x S^{p1, p2-1} + on logical domain to U1 the pre-polar 1-forms splines. The associate matrix + is square as in the CONGA approach we keep using the tensor B-spline basis, + instead of the polar basis of Toshniwal. P1 enforces coefficient relations + to be in U1. + + Parameters: + ----------- + + W1 : VectorFemSpace (ProductFemSpace) + Full tensor product spline space of the 1-forms S^{p1-1, p2} x S^{p1, p2-1} + + transposed : Boolean + Switch between P1 and P1 transposed (default is False) + + hbc : Boolean + Switch on and off the imposition of homogeneous Dirichlet boundary + conditions on the tangential (angular) direction (default is False) + """ + + def __init__(self, W1, transposed=False, hbc=False): + assert isinstance(W1, VectorFemSpace) + assert W1.symbolic_space.kind.name == 'hcurl' + + T1 = W1.coeff_space + + super().__init__(T1, T1) + + self[0, 0] = C1PolarProjection_V1_00(W1, transposed=transposed) + self[1, 0] = C1PolarProjection_V1_10(W1, transposed=transposed) + # self[0, 1] is 0 block + # The lower right block is the same for C0 and C1 projections + self[1, 1] = C0PolarProjection_V1_11(W1, transposed=transposed, hbc=hbc) + + self.W1 = W1 + self.transposed = transposed + self.hbc = hbc + + +# -------------- 2-FORMS CONGA PROJECTOR P2 ----------------# +# Equals C0PolarProjection_V2 + +# ------------ strong and weak curl ----------------- + +from scipy.sparse.linalg import spilu, cg +from scipy.sparse.linalg import LinearOperator as SparseLinearOperator + + +class SparseCurlAsOperator(LinearOperator): + + def __init__(self, W1, W2, strong_curl_sp, M1=None, M2=None, strong=False, store_M1inv=False): + assert isinstance(W1, VectorFemSpace) + assert isinstance(W2, TensorFemSpace) + assert isinstance(strong_curl_sp, spmatrix) + + self.W1 = W1 + self.W2 = W2 + + self.strong = strong + self._store_M1inv = store_M1inv + + if strong: + self.curl_sp = strong_curl_sp + + else: + self.curl_sp = strong_curl_sp.T # dual curl (in the dual bases) + assert isinstance(M1, LinearOperator) + assert isinstance(M2, LinearOperator) + self.M1_sp = M1.tosparse() + self.M2_sp = M2.tosparse() + if store_M1inv: + self._M1inv_sp = sp_inv(self.M1_sp) + else: + self._M1_spilu = spilu(self.M1_sp) + self._precond_M = SparseLinearOperator(self.M1_sp.shape, self._M1_spilu.solve) + + @property + def domain(self): + if self.strong: + return self.W1.coeff_space + else: + return self.W2.coeff_space + + @property + def codomain(self): + if self.strong: + return self.W2.coeff_space + else: + return self.W1.coeff_space + + @property + def dtype(self): + return float + + def toarray(self): + # return self + raise NotImplementedError('toarray() is not defined for this class.') + + def tosparse(self): + # return self + raise NotImplementedError('tosparse() is not defined for this class.') + + def transpose(self, conjugate=False): + raise NotImplementedError('transpose() is not defined for this class.') + + # Warning: this dot method has to be revised for mpi! + # the toeplitz multiplication requires all processes along the theta-dir to communicate. + def dot(self, x, out=None): + assert isinstance(x, Vector) + if self.strong: + Cx_arr = self.curl_sp @ x.toarray() + else: + tx_arr = self.M2_sp @ x.toarray() + tCx_arr = self.curl_sp @ tx_arr + if self._store_M1inv: + Cx_arr = self._M1inv_sp.dot(tCx_arr) + else: + # Cx_arr = spsolve(self.M1_sp, tCx_arr) + Cx_arr, exit_code = cg(self.M1_sp, tCx_arr, M=self._precond_M, rtol=1e-7) + # Cx_arr = self._M1_spilu.solve(tCx_arr) + if out is None: + y = self.codomain.zeros() + else: + assert isinstance(out, Vector) + assert out.space is self.codomain + y = out + + y1 = array_to_psydac(Cx_arr, self.codomain) + y1.copy(out=y) + return y + + + diff --git a/psydac/feec/polar/examples/__init__.py b/psydac/feec/polar/examples/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/psydac/feec/polar/examples/analyticalTE.py b/psydac/feec/polar/examples/analyticalTE.py new file mode 100644 index 000000000..03be4702f --- /dev/null +++ b/psydac/feec/polar/examples/analyticalTE.py @@ -0,0 +1,289 @@ +from sympde.topology import Square +from sympde.topology.analytical_mapping import PolarMapping + +from psydac.api.tests.test_api_feec_2d import add_colorbar +from psydac.feec.pull_push import push_2d_hcurl, push_2d_l2 +from numpy import pi +import numpy as np +import matplotlib.pyplot as plt + +class CircularCavitySolution: + """ + Time-harmonic solution of Maxwell's equations in a disk-like domain with + perfectly conducting walls. This is a "transverse electric" solution, with + E = (Ex, Ey) and B = Bz. Domain is [0, R] x [0, 2pi]. + + Parameters + ---------- + R : float + domain radius + + c: float + Speed of light in arbitrary units + + (m, n): int + Mode number. Warning: m > 0, n >= 0 + + D: float + shift of logical center (in "Target" mapping with c0=D*R2, c1=0, k=0, D=D) + + scale: float + Rescaling the values by a real factor. Default = 1 + + """ + + def __init__(self, R, c, m, n, D=0, scale=1): + from numpy import pi + from scipy.special import jnp_zeros + pnm = jnp_zeros(n, m)[-1] + kc = pnm / R + omega = c * kc + + phase = pi / 4 + + self.c = c + self.scale = scale + self.n = n + self.kc = kc + self.omega = omega + self.phase = phase + self._R = R + assert 0 <= D < .5 + self._D = D + + # Exact solutions for electric and magnetic field with polar parametrization of disk domain + def Es_ex(self, t, s, theta): + from numpy import sin # , cos, sqrt, arctan2 + from scipy.special import jv + + scale = self.scale + n = self.n + kc = self.kc + omega = self.omega + phase = self.phase + c = self.c + + return - scale * c * n / (s * kc + 1e-10) * sin(n * theta) * jv(n, kc * s) * sin(omega * t + phase) + + def Et_ex(self, t, s, theta, s_factor=True): + """ + if s_factor: multiply by s (as in logical field) + """ + from numpy import sin, cos + from scipy.special import jvp + + scale = self.scale + n = self.n + kc = self.kc + omega = self.omega + phase = self.phase + c = self.c + + val = - scale * c * cos(n * theta) * jvp(n, kc * s) * sin(omega * t + phase) + if s_factor: + val *= s + return val + + def B_ex(self, t, s, theta, s_factor=True): + """ + if s_factor: multiply by s (as in logical field) + """ + from numpy import cos + from scipy.special import jv + + scale = self.scale + n = self.n + kc = self.kc + omega = self.omega + phase = self.phase + + val = scale * cos(n * theta) * jv(n, kc * s) * cos(omega * t + phase) + if s_factor: + val *= s + return val + # The magnitude of B is approximately equal to scale / 3 + + def dB_dt_ex(self, t, s, theta): + ''' + = dB/dt + ''' + from numpy import cos, sin + from scipy.special import jv + + scale = self.scale + n = self.n + kc = self.kc + omega = self.omega + phase = self.phase + + return scale * omega * s * cos(n * theta) * jv(n, kc * s) * sin(omega * t + phase) + + # physical field + + def get_radius_angle(self, x, y): + from numpy import sqrt, arctan2 # , sin, cos + r = sqrt(x * x + y * y) + alpha = arctan2(y, x) + return r, alpha + + def Ex_ex(self, t, x, y): + from numpy import cos, sin + + r, alpha = self.get_radius_angle(x, y) + return cos(alpha) * self.Es_ex(t, r, alpha) - sin(alpha) * self.Et_ex(t, r, alpha, s_factor=False) + + def Ey_ex(self, t, x, y): + from numpy import cos, sin + + r, alpha = self.get_radius_angle(x, y) + return sin(alpha) * self.Es_ex(t, r, alpha) + cos(alpha) * self.Et_ex(t, r, alpha, s_factor=False) + + def Bz_ex(self, t, x, y): + + r, alpha = self.get_radius_angle(x, y) + return self.B_ex(t, r, alpha, s_factor=False) + + +def main(): + """ + This function is not currently used. It is kept for possible future development. + """ + + # Logical domain is rectangle [0, R] x [0, 2pi] + R = 2.0 + + # Speed of light equal c and scaling of the fields by a scale factor + c = 1 + scale = 1 + + # Mode number + # (m, n) = (1, 0) + # (m, n) = (2, 1) + (m, n) = (2, 3) + + # Exact solution + exact_solution = CircularCavitySolution(R=R, c=c, m=m, n=n, scale=scale) + + # Exact fields, as callable functions of (t, s, theta) + Es_ex = exact_solution.Es_ex + Et_ex = exact_solution.Et_ex + B_ex = exact_solution.B_ex + dB_dt_ex = exact_solution.dB_dt_ex + + # Logical domain: [0, R] x [0, 2pi] + logical_domain = Square('Omega', bounds1=[0, R], bounds2=[0, 2 * pi]) + + # Physical domain: disk of radius R obtained as image of the logical_domain + # with the analytical mapping of a circle + mapping = PolarMapping('F', c1=0, c2=0, rmin=0, rmax=1) + domain = mapping(logical_domain) + F = mapping.get_callable_mapping() + + # Set time + t = 0 + + Es = lambda x, y: Es_ex(t, x, y) + Et = lambda x, y: Et_ex(t, x, y) + B = lambda x, y: B_ex(t, x, y) + dB_dt = lambda x, y: dB_dt_ex(t, x, y) + + # Plot of fields + N = 100 + + rho = np.linspace(1e-20, R, N) + theta = np.linspace(0, 2 * pi, N) + rho, theta = np.meshgrid(rho, theta, indexing='ij') + x, y = F(rho, theta) + + Ex_values = np.empty_like(rho) + Ey_values = np.empty_like(rho) + B_values = np.empty_like(rho) + + for i, x1i in enumerate(rho[:, 0]): + for j, x2j in enumerate(theta[0, :]): + Ex_values[i, j], Ey_values[i, j] = \ + push_2d_hcurl(Es, Et, x1i, x2j, F) + + B_values[i, j] = push_2d_l2(B, x1i, x2j, F) + + fig, axs = plt.subplots(2, 2, figsize=(12, 12)) + im0 = axs[0, 0].contourf(x, y, Ex_values, 50) + im1 = axs[0, 1].contourf(x, y, Ey_values, 50) + im2 = axs[1, 0].contourf(x, y, np.sqrt(Ex_values ** 2 + Ey_values ** 2), 50) + im3 = axs[1, 1].contourf(x, y, B_values, 50) + axs[0, 0].set_title(r'$E_x$') + axs[0, 1].set_title(r'$E_y$') + axs[1, 0].set_title(r'$||\mathbf{E}||$') + axs[1, 1].set_title('$B$') + add_colorbar(im0, axs[0, 0]) + add_colorbar(im1, axs[0, 1]) + add_colorbar(im2, axs[1, 0]) + add_colorbar(im3, axs[1, 1]) + + # Test: curl E = - d_t B_z with finite Differences + # on a tensor grid in the square inscribed the disk + N = 160 + l = R / np.sqrt(2) # edge length + + x, dx = np.linspace(-l, l, N, retstep=True) + y, dy = np.linspace(-l, l, N, retstep=True) + x, y = np.meshgrid(x, y, indexing='ij') + rho = np.sqrt(x ** 2 + y ** 2) + theta = np.arctan2(y, x) % (2 * pi) + + Ex_values = np.empty_like(rho) + Ey_values = np.empty_like(rho) + Bt_values = np.empty_like(rho) + + ni, nj = rho.shape + for i in range(ni): + for j in range(nj): + x1_ij = rho[i, j] + x2_ij = theta[i, j] + Ex_values[i, j], Ey_values[i, j] = \ + push_2d_hcurl(Es, Et, x1_ij, x2_ij, F) + + Bt_values[i, j] = push_2d_l2(dB_dt, x1_ij, x2_ij, F) + + fig, axs = plt.subplots(2, 3, figsize=(15, 15)) + im1 = axs[0, 0].contourf(x, y, np.sqrt(Ex_values ** 2 + Ey_values ** 2)) + im2 = axs[0, 1].contourf(x, y, -Bt_values) + axs[0, 0].set_title(r'$||\mathbf{E}||$') + axs[0, 1].set_title(r'$-\partial_t B$') + for axi in axs.flat: + axi.set_aspect('equal') + add_colorbar(im1, axs[0, 0]) + add_colorbar(im2, axs[0, 1]) + + curlE_values = np.zeros_like(rho) + valerr = 0 + for i in range(1, ni - 1): + for j in range(1, nj - 1): + curlE_values[i, j] = (Ey_values[i + 1, j] - Ey_values[i - 1, j]) / (2 * dx) \ + - (Ex_values[i, j + 1] - Ex_values[i, j - 1]) / (2 * dy) + + d = np.abs(Bt_values[i, j] + curlE_values[i, j]) + + if valerr < d: + valerr = d + + skip = (slice(None, None, int(N / 20)), slice(None, None, int(N / 20))) + im3 = axs[1, 0].contourf(x, y, curlE_values) + axs[1, 0].quiver(x[skip], y[skip], Ex_values[skip], Ey_values[skip]) + add_colorbar(im3, axs[1, 0]) + axs[1, 0].set_title(r'curl $\mathbf{E}$') + im4 = axs[1, 1].contourf(x, y, -Bt_values) + axs[1, 1].quiver(x[skip], y[skip], Ex_values[skip], Ey_values[skip]) + add_colorbar(im4, axs[1, 1]) + im5 = axs[0, 2].contourf(x, y, Ex_values) + im6 = axs[1, 2].contourf(x, y, Ey_values) + add_colorbar(im5, axs[0, 2]) + add_colorbar(im6, axs[1, 2]) + print('|curlE + d_t B| <= ', valerr) + + +if __name__ == "__main__": + main() + plt.show() + + diff --git a/psydac/feec/polar/examples/maxwell_2d.py b/psydac/feec/polar/examples/maxwell_2d.py new file mode 100644 index 000000000..ef252fa57 --- /dev/null +++ b/psydac/feec/polar/examples/maxwell_2d.py @@ -0,0 +1,1019 @@ +""" +Solve the Transverse Electric Time dependent Maxwell Problem on an analytical disk domain. + +Example of run: +python maxwell_2d.py -S -n 16 32 -d 3 3 -T 1 -D 0.2 -s 1 +""" + +import os +import numpy as np +from mpi4py import MPI + +import matplotlib.pyplot as plt + +from psydac.fem.basic import FemField +from utils_congapol import check_regular_ring_map, add_colorbar, create_tensor_spline_space + + +# ====================== TIME DISCRETIZATION ==================================# + +def compute_stable_dt(cfl, C_m, dC_m, V, tau=None, light_c=1): + """ + compute stable time step for a leap-frog Maxwell solver, + given the (discrete) primal and dual curl operators + + cfl: stability factor (1 to choose the maximum stable dt) + C_m: (primal) curl V -> dV + dC_m: (dual) curl dV -> V + V: FEM space + tau: (optional) time to solve with an integer number of time steps + light_c: Maxwell parameter + + """ + print(f" .. compute stable dt by estimating the operator norm of ") + print(f" .. dual_curl_h @ curl_h: V_h -> V_h ") + print(f" .. with dim(V_h) = {V.coeff_space.dimension} ... ") + + def vect_norm_2(vv): + return np.sqrt(vv.inner(vv)) + + vv = V.coeff_space.zeros() + print(f'type(V.coeff_space) = {type(V.coeff_space)}') + print(f'V.coeff_space.shape = {V.coeff_space.shape}, V.coeff_space.dimension = {V.coeff_space.dimension}') + vv[:] = np.random.random(size=V.coeff_space.shape) + norm_vv = vect_norm_2(vv) + max_ncfl = 500 + ncfl = 0 + spectral_rho = 1 + conv = False + CC_m = dC_m @ C_m + while not (conv or ncfl > max_ncfl): + ncfl += 1 + vv *= (1. / norm_vv) + vv = CC_m @ vv + norm_vv = vect_norm_2(vv) + old_spectral_rho = spectral_rho + spectral_rho = norm_vv.copy() # copy ?? + conv = abs((spectral_rho - old_spectral_rho) / spectral_rho) < 0.001 + print(f" ... spectral radius iteration: spectral_rho( dC_m @ C_m ) ~= {spectral_rho}") + norm_op = np.sqrt(spectral_rho) + c_dt_max = 2. / norm_op + dt = cfl * c_dt_max / light_c + if tau is not None: + Nt_per_tau = int(np.ceil(tau / dt)) + assert Nt_per_tau >= 1 + dt = tau / Nt_per_tau + else: + Nt_per_tau = 1 + assert light_c * dt <= cfl * c_dt_max + print(f" Time step dt computed for Maxwell solver:") + print(f" with cfl = repr({cfl} we found dt = {dt} -- that is Nt_per_tau = {Nt_per_tau} on tau = {tau}.") + print( + f" -- note that c*Dt = {light_c * dt} and c_dt_max = {c_dt_max} thus c * dt / c_dt_max = {light_c * dt / c_dt_max}") + print(f" -- and spectral_radius((c*dt)**2* dC_m @ C_m ) = {(light_c * dt * norm_op) ** 2} (should be < 4)") + return Nt_per_tau, dt, norm_op + + +# =============================================================================# + +# =========================== VISUALIZATION ===================================# + +def plot_field_and_error(name, t, x, y, field_h, field_ex, *gridlines, only_field=True): + if only_field: + fig, ax0 = plt.subplots(1, 1, figsize=(7, 6)) + axes = [ax0] + else: + fig, (ax0, ax1) = plt.subplots(1, 2, figsize=(15, 6)) + axes = [ax0, ax1] + im0 = ax0.contourf(x, y, field_h, 50) + ax0.set_title(r'${0}_h$'.format(name)) + if not only_field: + im1 = ax1.contourf(x, y, field_ex - field_h, 50) + ax1.set_title(r'${0} - {0}_h$'.format(name)) + for ax in axes: + ax.plot(*gridlines[0], color='k') + ax.plot(*gridlines[1], color='k') + ax.set_xlabel('x', fontsize=14) + ax.set_ylabel('y', fontsize=14, rotation='horizontal') + ax.set_aspect('equal') + add_colorbar(im0, ax0) + if not only_field: + add_colorbar(im1, ax1) + fig.suptitle('Time t = {:10.3e}'.format(t)) + fig.tight_layout() + return fig + + +def update_plot(fig, t, x, y, field_h, field_ex): + ax0, ax1, cax0, cax1 = fig.axes + im0 = ax0.contourf(x, y, field_h, 50) + im1 = ax1.contourf(x, y, field_ex - field_h, 50) + fig.colorbar(im0, cax=cax0) + fig.colorbar(im1, cax=cax1) + fig.suptitle('Time t = {:10.3e}'.format(t)) + fig.canvas.draw() + + +def plot_curve_along_s(name, s_str, time_str, theta0, + s, curve_h, curve_ref=None, + left_s=None, left_curve_h=None, left_curve_ref=None, ): + + fig, ax = plt.subplots() + ax.plot(s, curve_h, label=f'{name}_h') + ax.plot(s, curve_ref, label=f'{name}_ref') + + if left_s is not None: + ax.plot(left_s, left_curve_h, label=f'{name}_h (left)') + if left_curve_ref is not None: + ax.plot(left_s, left_curve_ref, label=f'{name}_ref (left)') + + ax.set(xlabel=s_str, title=f'field {name} along {s_str} for theta={theta0}, at {time_str}') + ax.legend() + return fig + + +# =============================================================================# + +def run_maxwell_2d_TE(*, ncells, smooth, degree, nsteps, tend, + splitting_order, shift_D, use_spline_mapping, tol, + cfl=0.9, show_figs=True, study='maxwell_bessel', use_scipy=True, verbose=False): + from numpy import pi + + from sympy import cos, sin, Tuple, sqrt + + from sympde.topology import Square, Domain + from sympde.topology.analytical_mapping import PolarMapping, TargetMapping + from sympde.topology import Derham + from sympde.topology import elements_of, element_of + from sympde.topology.mapping import Mapping + from sympde.calculus import dot + from sympde.expr import integral + from sympde.expr import BilinearForm, LinearForm + + from psydac.api.discretization import discretize + from psydac.api.settings import PSYDAC_BACKENDS + from psydac.feec.pull_push import push_2d_hcurl, push_2d_l2 + from psydac.utilities.utils import refine_array_1d + from psydac.linalg.solvers import inverse + from psydac.linalg.basic import IdentityOperator + from psydac.mapping.discrete import SplineMapping + from psydac.cad.geometry import Geometry + + assert splitting_order in [2, 4] + + from psydac.feec.polar.conga_projections import C0PolarProjection_V1, C0PolarProjection_V2, C1PolarProjection_V1, SparseCurlAsOperator + + from analyticalTE import CircularCavitySolution # , constant_field + from waveTE import GaussianSolution + + backend = PSYDAC_BACKENDS['pyccel-gcc'] + + # Radius physical domain + R = 1. # 2.0 + + # Speed of light and scaling + c = 1.0 + scale = 1.0 + + # Mode number + (m, n) = (2, 3) + + # Exact/initial solution + assert study in ['L2_proj', 'maxwell_bessel', 'maxwell_wave'] + study_L2_proj = (study == 'L2_proj') + + visdir = f'plots_{study}' + os.makedirs(visdir, exist_ok=True) + + if study == 'maxwell_wave': + # This is just the initial solution + exact_solution = GaussianSolution(sigma=1e-1, x0=0, y0=0, scale=scale) + else: + exact_solution = CircularCavitySolution(R=R, c=c, m=m, n=n, scale=scale) + + Ex_ex_t = exact_solution.Ex_ex + Ey_ex_t = exact_solution.Ey_ex + Bz_ex_t = exact_solution.Bz_ex + + # Logical domain: [0, R] x [0, 2pi] + logical_bounds = [[0, R], [0, 2 * pi]] + logical_domain = Square('Omega', bounds1=logical_bounds[0], bounds2=logical_bounds[1]) + + # Physical domain: disk of radius R obtained as image of the logical_domain + # with the analytical mapping of a circle + polar_mapping = False + if polar_mapping: + mapping = PolarMapping('PM', c1=0, c2=0, rmin=0, rmax=1) + else: + # domain = ((0, 1), (0, 2 * np.pi)) + mapping = TargetMapping('TM', c1=shift_D * R * R, c2=0, k=0, D=shift_D) + + # run parameters string + rp_str = f'{ncells[0]}_{ncells[1]}_p{degree[0]}_D{shift_D}_s{smooth}' + if use_spline_mapping: + rp_str += '_sm' + else: + rp_str += '_pm' # WARNING: check that polar_mapping == True ? + + # Communicator, size, rank + mpi_comm = MPI.COMM_WORLD + mpi_size = mpi_comm.Get_size() + mpi_rank = mpi_comm.Get_rank() + + if use_spline_mapping: + + # ==================== SPLINE SPACE FOR SPLINE MAPPINGS =======================# + + V = create_tensor_spline_space(ncells, degree, [False, True], + (logical_bounds[0], logical_bounds[1]), mpi_comm) + + # ==================== MAPPING & PHYSICAL DOMAIN ==============================# + + # Create spline mapping by interpolation of analytical mapping + map_analytic = mapping.get_callable_mapping() + map_discrete = SplineMapping.from_mapping(V, map_analytic) + + check_regular_ring_map(map_discrete) + + # Create symbolic mapping with callable mapping as spline + mapping = Mapping('M', dim=2) + mapping.set_callable_mapping(map_discrete) + # In order to create a sympde.Domain object from this mapping we have + # to create first a HDF5 file and then load as sympde.Domain.fromfile + geometry = Geometry.from_discrete_mapping(map_discrete, comm=mpi_comm) + geometry.export('geo.h5') + domain = Domain.from_file('geo.h5') + + else: + # Only symbolic mapping is necessary + domain = mapping(logical_domain) + + # DeRham sequence + derham = Derham(domain, sequence=['h1', 'hcurl', 'l2']) + + # Trial and test functions + u1, v1 = elements_of(derham.V1, names='u1, v1') # electric field E = (Ex, Ey) + u2, v2 = elements_of(derham.V2, names='u2, v2') # magnetic field Bz + + # Bilinear forms that correspond to mass matrices for spaces V1 and V2 + a1 = BilinearForm((u1, v1), integral(domain, dot(u1, v1))) + a2 = BilinearForm((u2, v2), integral(domain, u2 * v2)) + + # Discrete physical domain and discrete DeRham sequence + if use_spline_mapping: + domain_h = discretize(domain, filename='geo.h5', comm=mpi_comm) + derham_h = discretize(derham, domain_h) + F = list(domain_h.mappings.values()).pop() + else: + domain_h = discretize(domain, ncells=ncells, periodic=[False, True], comm=mpi_comm) + derham_h = discretize(derham, domain_h, degree=degree) + F = mapping.get_callable_mapping() + + def l2_norm_of(f_log): + """ + Compute the l2 norm of f over the physical domain. + Interface needed since FEM space `integral` function is over the symbolic domain... + """ + f2_with_det = lambda eta1, eta2: f_log(eta1, eta2) ** 2 * np.sqrt(F.metric_det(eta1, eta2)) + return np.sqrt(derham_h.V0.integral(f2_with_det)) + + def build_eval_context(): + """ + Build all serial objects needed for plotting fields. + + This is used by: + - study='L2_proj' + - initial Maxwell plots + - final Maxwell plots + """ + if mpi_rank != 0: + return None + + if use_spline_mapping: + domain_h_serial = discretize(domain, filename='geo.h5') + F_serial = [*domain_h_serial.mappings.values()].pop() + else: + domain_h_serial = discretize(domain, ncells=ncells, periodic=[False, True]) + F_serial = F + + derham_h_serial = discretize(derham, domain_h_serial, degree=degree) + V0_s, V1_s, V2_s = derham_h_serial.spaces + V1_sx, V1_sy = V1_s.spaces + + grid_x1 = V0_s.breaks[0].copy() + grid_x2 = V0_s.breaks[1].copy() + + # Fix division by zero without taking care of the limit as s --> 0 + grid_x1[0] = 1e-20 + + N = 5 + x1 = refine_array_1d(grid_x1, N) + x2 = refine_array_1d(grid_x2, N) + + x1, x2 = np.meshgrid(x1, x2, indexing='ij') + + x = np.empty_like(x1) + y = np.empty_like(x1) + + for i in range(x1.shape[0]): + for j in range(x1.shape[1]): + x[i, j], y[i, j] = F_serial(x1[i, j], x2[i, j]) + + gridlines_x1 = (x[:, ::N], y[:, ::N]) + gridlines_x2 = (x[::N, :].T, y[::N, :].T) + gridlines = (gridlines_x1, gridlines_x2) + + return { + 'domain_h_serial': domain_h_serial, + 'derham_h_serial': derham_h_serial, + 'V0_s': V0_s, + 'V1_s': V1_s, + 'V2_s': V2_s, + 'V1_sx': V1_sx, + 'V1_sy': V1_sy, + 'F_serial': F_serial, + 'x1': x1, + 'x2': x2, + 'x': x, + 'y': y, + 'gridlines': gridlines, + } + + def run_study_L2_proj(): + """ + Study the H(curl) L2 projection on the mapped disk. Compares the standard and + filtered projected fields against a manufactured physical vector field, + prints L2 projection errors and plots the projected components and their errors. + Executed when --study is 'L2_proj'. + Only for serial runs + """ + omega = 4 + print(f'studying L2 proj of f in H_0(curl) .. with omega = {omega}') + xs, ys = domain.coordinates + r = sqrt(xs * xs + ys * ys) + + # f in H_0(curl;Omega) + f_x = -xs + sin(omega * (ys + 2 * xs * xs)) * (r - R) + f_y = -ys + cos(omega * (2 * xs - ys * ys)) * (r - R) + f_phys = Tuple(f_x, f_y) + + from sympy import lambdify + fx_call = lambdify([xs, ys], f_x) + fy_call = lambdify([xs, ys], f_y) + + v = element_of(derham.V1, name='u') + + l = LinearForm(v, integral(domain, dot(f_phys, v))) + lh = discretize(l, domain_h, V1_h, backend=backend) + tilde_f = lh.assemble() + + # create fields and point to coeffs + + fh = Pi1((fx_call, fy_call)) + fh_filter = Pi1((fx_call, fy_call)) + + M1_inv = inverse(M1, 'cg', verbose=verbose, tol=tol) + print("using standard L2 projection") + fh_c = M1_inv @ tilde_f + fh_c = P1 @ fh_c + print('fh_c:') + print(fh_c.toarray()[:]) + + print("using filtered L2 projection") + PTtilde_f = P1.T @ tilde_f + fh_filter_c = M1_inv @ PTtilde_f + fh_filter_c = P1 @ fh_filter_c + + # compute error and exit + + errx = lambda x1, x2: push_2d_hcurl(fh.fields[0], fh.fields[1], x1, x2, F)[0] - fx_call(F(x1, x2)[0], + F(x1, x2)[1]) + erry = lambda x1, x2: push_2d_hcurl(fh.fields[0], fh.fields[1], x1, x2, F)[1] - fy_call(F(x1, x2)[0], + F(x1, x2)[1]) + errx_filter = lambda x1, x2: push_2d_hcurl(fh_filter.fields[0], fh_filter.fields[1], x1, x2, F)[ + 0] - fx_call(F(x1, x2)[0], F(x1, x2)[1]) + erry_filter = lambda x1, x2: push_2d_hcurl(fh_filter.fields[0], fh_filter.fields[1], x1, x2, F)[ + 1] - fy_call(F(x1, x2)[0], F(x1, x2)[1]) + + error_l2_fx = l2_norm_of(errx) + error_l2_fy = l2_norm_of(erry) + error_l2_fx_filter = l2_norm_of(errx_filter) + error_l2_fy_filter = l2_norm_of(erry_filter) + print('L2 norm of projection error on fx: {:.2e}'.format(error_l2_fx)) + print('L2 norm of projection error on fy: {:.2e}'.format(error_l2_fy)) + print('L2 norm of projection error on fx (filtered): {:.2e}'.format(error_l2_fx_filter)) + print('L2 norm of projection error on fy (filtered): {:.2e}'.format(error_l2_fy_filter)) + + cst_wo_det = lambda x1, x2: 1 + cst_wi_det = lambda x1, x2: 1 * np.sqrt(F.metric_det(x1,x2)) + int_wo_det = derham_h.V0.integral(cst_wo_det) + int_wi_det = derham_h.V0.integral(cst_wi_det) + print('V0 - integral of 1 (no det): {:.2e}'.format(int_wo_det)) + print('V0 - integral of 1 (with det): {:.2e}'.format(int_wi_det)) + print('radius of disk: {:.2e}'.format(R)) + print('area of log domain: {:.2e}'.format(2*pi*R)) + print('area of disk: {:.2e}'.format(pi*R*R)) + + int_wo_det = derham_h.V1.spaces[0].integral(cst_wo_det) + int_wi_det = derham_h.V1.spaces[0].integral(cst_wi_det) + print('V1.x - integral of 1 (no det): {:.2e}'.format(int_wo_det)) + print('V1.x - integral of 1 (with det): {:.2e}'.format(int_wi_det)) + + if not show_figs: + return locals() + + plot_data = build_eval_context() + if plot_data is None: + return locals() + + x1 = plot_data['x1'] + x2 = plot_data['x2'] + x = plot_data['x'] + y = plot_data['y'] + gridlines = plot_data['gridlines'] + + # plot + + fx_values = np.empty_like(x1) + fy_values = np.empty_like(x1) + fx_filter_values = np.empty_like(x1) + fy_filter_values = np.empty_like(x1) + + fx_ex_values = np.empty_like(x1) + fy_ex_values = np.empty_like(x1) + + for i, x1i in enumerate(x1[:, 0]): + for j, x2j in enumerate(x2[0, :]): + + xij, yij = F(x1i, x2j) + fx_values[i, j], fy_values[i, j] = \ + push_2d_hcurl(fh.fields[0], fh.fields[1], x1i, x2j, F) + fx_filter_values[i, j], fy_filter_values[i, j] = \ + push_2d_hcurl(fh_filter.fields[0], fh_filter.fields[1], x1i, x2j, F) + fx_ex_values[i, j], fy_ex_values[i, j] = \ + fx_call(xij, yij), fy_call(xij, yij) + + fig1 = plot_field_and_error(r'f^x', 0, x, y, fx_values, fx_ex_values, *gridlines) + #fig2.savefig(f'{visdir}/fx_{rp_str}.png') + + fig2 = plot_field_and_error(r'f^y', 0, x, y, fy_values, fy_ex_values, *gridlines) + #fig3.savefig(f'{visdir}/fy_{rp_str}.png') + + print('done: showing fh') + + fig3 = plot_field_and_error(r'f^x filter', 0, x, y, fx_filter_values, fx_ex_values, *gridlines) + #fig2.savefig(f'{visdir}/fx_filter_{rp_str}.png') + + fig4 = plot_field_and_error(r'f^y filter', 0, x, y, fy_filter_values, fy_ex_values, *gridlines) + #fig3.savefig(f'{visdir}/fy_filter_{rp_str}.png') + + print('done: showing fh_filter') + return locals() + + # ============================================================================== + # DISCRETIZATION + # ============================================================================== + + # Differential operators + D0, D1 = derham_h.derivatives(kind='linop') + D1_T = D1.T + + # Extract spaces + V0_h, V1_h, V2_h = derham_h.spaces + + I1 = IdentityOperator(V1_h.coeff_space) + I2 = IdentityOperator(V2_h.coeff_space) + + # Conga projectors + if smooth == 0: + P1 = C0PolarProjection_V1(V1_h, hbc=True) + P2 = C0PolarProjection_V2(V2_h) + else: + P1 = C1PolarProjection_V1(V1_h, hbc=True) + P2 = C0PolarProjection_V2(V2_h) + P1_T = P1.T + P2_T = P2.T + + # Discrete bilinear forms + a1_h = discretize(a1, domain_h, (derham_h.V1, derham_h.V1), backend=backend) + a2_h = discretize(a2, domain_h, (derham_h.V2, derham_h.V2), backend=backend) + + hs = R / ncells[0] + htheta = 2 * pi / ncells[1] + + # Mass matrices (StencilMatrix objects) + + # raw Mass matrices in W1 and W2 have unbounded values, this seems to be fine for the assembly but inverting these matrices leads to virtual nonsense... + # so they need to be regularized below + M1_raw = a1_h.assemble() + M2_raw = a2_h.assemble() + + # regularization of M1 and M2 so that they are bounded and invertible: + M1 = (htheta * hs) * (I1 - P1.T) @ (I1 - P1) + P1.T @ M1_raw @ P1 + M2 = (htheta * hs) * (I2 - P2.T) @ (I2 - P2) + P2.T @ M2_raw @ P2 + + Pi0, Pi1, Pi2 = derham_h.projectors(nquads=[degree[0] + 10, degree[1] + 10]) + + if study_L2_proj: + run_study_L2_proj() + return locals() + + # Geometric Projectors + + # Time integration setup + # -------------------------------------------------------------------------- + + t = 0 + + # Callable exact fields + Ex_ex = lambda t: (lambda x, y, t0=t: Ex_ex_t(t0, x, y)) + Ey_ex = lambda t: (lambda x, y, t0=t: Ey_ex_t(t0, x, y)) + Bz_ex = lambda t: (lambda x, y, t0=t: Bz_ex_t(t0, x, y)) + + # Initial conditions, discrete fields -- here with a pull-back in the projections + E_log = Pi1((Ex_ex(t), Ey_ex(t))) + E_log.coeffs.update_ghost_regions() + B_log = Pi2(Bz_ex(t)) + B_log.coeffs.update_ghost_regions() + + # Initial conditions, spline coefficients + e = E_log.coeffs + b = B_log.coeffs + + if study == 'maxwell_wave': + b = D1 @ e + + # Conga Projection + e = P1 @ e + b = P2 @ b + + V1_s, V1_theta = V1_h.spaces + Ex_field = FemField(V1_s, coeffs=e[0]) + Ey_field = FemField(V1_theta, coeffs=e[1]) + B_field = FemField(V2_h, coeffs=b) + V1_s.export_fields('Ex.h5', Ex_field=Ex_field) + V1_theta.export_fields('Ey.h5', Ey_field=Ey_field) + V2_h.export_fields('B.h5', B_field=B_field) + + if use_scipy: + + print(" -------------- SCIPY operators ------------ ") + conga_curl_sp = (D1 @ P1).tosparse() + step_faraday_2d = SparseCurlAsOperator(W1=V1_h, W2=V2_h, strong_curl_sp=conga_curl_sp, strong=True, + store_M1inv=False) + step_ampere_2d = SparseCurlAsOperator(W1=V1_h, W2=V2_h, strong_curl_sp=conga_curl_sp, M1=M1, M2=M2, + strong=False) + + else: + M1_inv = inverse(M1, 'cg', verbose=verbose, tol=tol) + step_ampere_2d = M1_inv @ P1_T @ D1_T @ M2 + step_faraday_2d = D1 @ P1 + + Nt, dt, norm_curlh = compute_stable_dt(cfl, C_m=step_ampere_2d, dC_m=step_faraday_2d, V=V2_h, tau=tend, light_c=1) + + # If final time is given, recompute number of time steps + if tend is None: + tend = nsteps * dt + print(f'final time (re)computed: {tend}') + else: + nsteps = Nt + print(f'nsteps recomputed: {nsteps}') + + # ============================================================================== + # VISUALIZATION SETUP + # ============================================================================== + + def plot_fields_along_s(tstr): + + j0 = 0 + j1 = x2.shape[1] // 2 + theta0 = x2[0, j0] + theta1 = x2[0, j1] + name = 'Ex' + fig_line = plot_curve_along_s(name, 's', tstr, f'{theta0} and {theta1}', + x1[:, j0], Ex_values[:, j0], Ex_ex_values[:, j0], + -x1[:, j1], Ex_values[:, j1], Ex_ex_values[:, j1]) + fig_line.savefig(f'{visdir}/{name}_line_{tstr}_{rp_str}.png') + plt.close(fig_line) + name = 'Ey' + fig_line = plot_curve_along_s(name, 's', tstr, f'{theta0} and {theta1}', + x1[:, j0], Ey_values[:, j0], Ey_ex_values[:, j0], + -x1[:, j1], Ey_values[:, j1], Ey_ex_values[:, j1]) + fig_line.savefig(f'{visdir}/{name}_line_{tstr}_{rp_str}.png') + plt.close(fig_line) + name = 'Bz' + fig_line = plot_curve_along_s(name, 's', tstr, f'{theta0} and {theta1}', + x1[:, j0], Bz_values[:, j0], Bz_ex_values[:, j0], + -x1[:, j1], Bz_values[:, j1], Bz_ex_values[:, j1]) + fig_line.savefig(f'{visdir}/{name}_line_{tstr}_{rp_str}.png') + plt.close(fig_line) + + # Plot initial conditions + eval_data = build_eval_context() + if eval_data is not None: # when mpi_rank is 0, since this needs to be done in serial + V1_sx = eval_data['V1_sx'] + V1_sy = eval_data['V1_sy'] + V2_s = eval_data['V2_s'] + F_serial = eval_data['F_serial'] + x1 = eval_data['x1'] + x2 = eval_data['x2'] + x = eval_data['x'] + y = eval_data['y'] + gridlines = eval_data['gridlines'] + + Ex_serial, = V1_sx.import_fields('Ex.h5', 'Ex_field') + Ey_serial, = V1_sy.import_fields('Ey.h5', 'Ey_field') + B_serial, = V2_s.import_fields('B.h5', 'B_field') + + Ex_ex_values = np.empty_like(x1) + Ey_ex_values = np.empty_like(x1) + Bz_ex_values = np.empty_like(x1) + + Ex_values = np.empty_like(x1) + Ey_values = np.empty_like(x1) + Bz_values = np.empty_like(x1) + + for i, x1i in enumerate(x1[:, 0]): + for j, x2j in enumerate(x2[0, :]): + + Ex_values[i, j], Ey_values[i, j] = \ + push_2d_hcurl(Ex_serial, Ey_serial, x1i, x2j, F_serial) + + Bz_values[i, j] = push_2d_l2(B_serial, x1i, x2j, F_serial) + + xij, yij = F_serial(x1i, x2j) + Ex_ex_values[i, j], Ey_ex_values[i, j] = \ + Ex_ex_t(t, xij, yij), Ey_ex_t(t, xij, yij) + + Bz_ex_values[i, j] = Bz_ex_t(t, xij, yij) + + # fields along s for fixed theta + plot_fields_along_s(tstr='t0') + + # Electric field, x component + fig = plot_field_and_error(r'E^x', 0, x, y, Ex_values, Ex_ex_values, *gridlines) + fig.savefig(f'{visdir}/Ex_t0_{rp_str}.png') + plt.close(fig) + + # Electric field, y component + fig = plot_field_and_error(r'E^y', 0, x, y, Ey_values, Ey_ex_values, *gridlines) + fig.savefig(f'{visdir}/Ey_t0_{rp_str}.png') + plt.close(fig) + + # Magnetic field, z component + fig = plot_field_and_error(r'B^z', 0, x, y, Bz_values, Bz_ex_values, *gridlines) + fig.savefig(f'{visdir}/Bz_t0_{rp_str}.png') + plt.close(fig) + + if show_figs: + # Plot exact and approximate solutions at t = 0 + fig, axs = plt.subplots(3, 3, figsize=(12, 12)) + im0 = axs[0, 0].contourf(x, y, Ex_ex_values, 50) + im1 = axs[0, 1].contourf(x, y, Ey_ex_values, 50) + im2 = axs[0, 2].contourf(x, y, Bz_ex_values, 50) + im3 = axs[1, 0].contourf(x, y, Ex_values, 50) + im4 = axs[1, 1].contourf(x, y, Ey_values, 50) + im5 = axs[1, 2].contourf(x, y, Bz_values, 50) + im6 = axs[2, 0].contourf(x, y, Ex_values - Ex_ex_values, 50) + im7 = axs[2, 1].contourf(x, y, Ey_values - Ey_ex_values, 50) + im8 = axs[2, 2].contourf(x, y, Bz_values - Bz_ex_values, 50) + axs[0, 0].set_title(r'$E^x$ at t = 0') + axs[0, 1].set_title(r'$E^y$ at t = 0') + axs[0, 2].set_title(r'$B^z$ at t = 0') + axs[1, 0].set_title(r'$E_h^x$ at t = 0') + axs[1, 1].set_title(r'$E_h^y$ at t = 0') + axs[1, 2].set_title(r'$B_h^z$ at t = 0') + axs[2, 0].set_title(r'$E^x - E_h^x$ at t = 0') + axs[2, 1].set_title(r'$E^y - E_h^y$ at t = 0') + axs[2, 2].set_title(r'$B^z - B_h^z$ at t = 0') + for i in range(3): + for j in range(3): + axs[i, j].plot(*gridlines[0], color='k') + axs[i, j].plot(*gridlines[1], color='k') + axs[i, j].set_xlabel('x', fontsize=14) + axs[i, j].set_ylabel('y', fontsize=14, rotation='horizontal') + axs[i, j].set_aspect('equal') + add_colorbar(im0, axs[0, 0]) + add_colorbar(im1, axs[0, 1]) + add_colorbar(im2, axs[0, 2]) + add_colorbar(im3, axs[1, 0]) + add_colorbar(im4, axs[1, 1]) + add_colorbar(im5, axs[1, 2]) + add_colorbar(im6, axs[2, 0]) + add_colorbar(im7, axs[2, 1]) + add_colorbar(im8, axs[2, 2]) + fig.suptitle('Compare Exact Solution and Approximate solution at initial time') + fig.tight_layout() + + # Need a small pause to show the plot of the initial condition + plt.pause(.1) + + # L2 norms (of ref solution) + normx = lambda x1, x2: Ex_ex_t(t, *F(x1, x2)) + normy = lambda x1, x2: Ey_ex_t(t, *F(x1, x2)) + normz = lambda x1, x2: Bz_ex_t(t, *F(x1, x2)) + + norm_l2_Ex = l2_norm_of(normx) + norm_l2_Ey = l2_norm_of(normy) + norm_l2_Bz = l2_norm_of(normz) + + # L2 errors + errx = lambda x1, x2: push_2d_hcurl(E_log.fields[0], E_log.fields[1], x1, x2, F)[0] - Ex_ex_t(t, *F(x1, x2)) + erry = lambda x1, x2: push_2d_hcurl(E_log.fields[0], E_log.fields[1], x1, x2, F)[1] - Ey_ex_t(t, *F(x1, x2)) + errz = lambda x1, x2: push_2d_l2(B_log, x1, x2, F) - Bz_ex_t(t, *F(x1, x2)) + + error_l2_Ex = l2_norm_of(errx) / norm_l2_Ex + error_l2_Ey = l2_norm_of(erry) / norm_l2_Ey + error_l2_Bz = l2_norm_of(errz) / norm_l2_Bz + + print('L2 norm of rel. error on Ex(t,x,y) at initial time: {:.2e}'.format(error_l2_Ex)) + print('L2 norm of rel. error on Ey(t,x,y) at initial time: {:.2e}'.format(error_l2_Ey)) + print('L2 norm of rel. error on Bz(t,x,y) at initial time: {:.2e}'.format(error_l2_Bz)) + + # ============================================================================== + # SOLUTION + # ============================================================================== + + def Strang_update(dtau): + # Strang splitting, 2nd order + + # b := b - dt/2 * curl e + db = step_faraday_2d @ e + b.mul_iadd(- dtau / 2, db) + + # e := e + dt * curl b + de = step_ampere_2d @ b + e.mul_iadd(dtau, de) + + # b := b - dt/2 * curl e + db = step_faraday_2d @ e + b.mul_iadd(- dtau / 2, db) + + # weights for Suzuki-Yoshida composition (4th-order splitting) + + gamma_1 = 1 / (2 - 2 ** (1 / 3)) + gamma_2 = 1 - 2 * gamma_1 + + # Time loop + for ts in range(1, nsteps + 1): + + print(f'step = {ts}/{nsteps}') + + if splitting_order == 2: + + Strang_update(dt) + + elif splitting_order == 4: + + Strang_update(dt * gamma_1) + Strang_update(dt * gamma_2) + Strang_update(dt * gamma_1) + + else: + raise NotImplementedError('splitting_order must be 2 or 4') + + t += dt + + # diag + e.update_ghost_regions() + b.update_ghost_regions() + + e = P1 @ e + b = P2 @ b + + print('ts = {:4d}, t = {:8.4f}'.format(ts, t)) + + N = 10 + if show_figs: + V0_h.plot_2d_decomposition(F, refine=N) + + e = P1 @ e + b = P2 @ b + + Ex_field = FemField(V1_s, coeffs=e[0]) + Ey_field = FemField(V1_theta, coeffs=e[1]) + B_field = FemField(V2_h, coeffs=b) + V1_s.export_fields('Ex_final.h5', Ex_field=Ex_field) + V1_theta.export_fields('Ey_final.h5', Ey_field=Ey_field) + V2_h.export_fields('B_final.h5', B_field=B_field) + + if eval_data is not None: + Ex_serial, = V1_sx.import_fields('Ex_final.h5', 'Ex_field') + Ey_serial, = V1_sy.import_fields('Ey_final.h5', 'Ey_field') + B_serial, = V2_s.import_fields('B_final.h5', 'B_field') + + for i, x1i in enumerate(x1[:, 0]): + for j, x2j in enumerate(x2[0, :]): + + Ex_values[i, j], Ey_values[i, j] = \ + push_2d_hcurl(Ex_serial, Ey_serial, x1i, x2j, F_serial) + + Bz_values[i, j] = push_2d_l2(B_serial, x1i, x2j, F_serial) + + xij, yij = F_serial(x1i, x2j) + Ex_ex_values[i, j], Ey_ex_values[i, j] = \ + Ex_ex_t(t, xij, yij), Ey_ex_t(t, xij, yij) + + Bz_ex_values[i, j] = Bz_ex_t(t, xij, yij) + + # Error at final time + error_Ex = abs(Ex_ex_values - Ex_values).max() + error_Ey = abs(Ey_ex_values - Ey_values).max() + error_Bz = abs(Bz_ex_values - Bz_values).max() + print() + print('Max-norm of error on Ex(t,x) at final time: {:.2e}'.format(error_Ex)) + print('Max-norm of error on Ey(t,x) at final time: {:.2e}'.format(error_Ey)) + print('Max-norm of error on Bz(t,x) at final time: {:.2e}'.format(error_Bz)) + + # L2 norms (of ref solution) + normx = lambda x1, x2: Ex_ex_t(t, *F(x1, x2)) + normy = lambda x1, x2: Ey_ex_t(t, *F(x1, x2)) + normz = lambda x1, x2: Bz_ex_t(t, *F(x1, x2)) + + norm_l2_Ex = l2_norm_of(normx) + norm_l2_Ey = l2_norm_of(normy) + norm_l2_Bz = l2_norm_of(normz) + + E1_final = FemField(V1_s, coeffs=e[0]) + E2_final = FemField(V1_theta, coeffs=e[1]) + B_final = FemField(V2_h, coeffs=b) + + # L2 errors + errx = lambda x1, x2: push_2d_hcurl(E1_final, E2_final, x1, x2, F)[0] - Ex_ex_t(t, *F(x1, x2)) + erry = lambda x1, x2: push_2d_hcurl(E1_final, E2_final, x1, x2, F)[1] - Ey_ex_t(t, *F(x1, x2)) + errz = lambda x1, x2: push_2d_l2(B_final, x1, x2, F) - Bz_ex_t(t, *F(x1, x2)) + + error_l2_Ex = l2_norm_of(errx) / norm_l2_Ex + error_l2_Ey = l2_norm_of(erry) / norm_l2_Ey + error_l2_Bz = l2_norm_of(errz) / norm_l2_Bz + + print() + print('L2 norm of rel. error on Ex(t,x,y) at final time: {:.2e}'.format(error_l2_Ex)) + print('L2 norm of rel. error on Ey(t,x,y) at final time: {:.2e}'.format(error_l2_Ey)) + print('L2 norm of rel. error on Bz(t,x,y) at final time: {:.2e}'.format(error_l2_Bz)) + + if eval_data is not None and show_figs: + # Plot exact and approximate solution at final time + fig1, axs = plt.subplots(3, 3, figsize=(12, 12)) + im0 = axs[0, 0].contourf(x, y, Ex_ex_values, 50) + im1 = axs[0, 1].contourf(x, y, Ey_ex_values, 50) + im2 = axs[0, 2].contourf(x, y, Bz_ex_values, 50) + im3 = axs[1, 0].contourf(x, y, Ex_values, 50) + im4 = axs[1, 1].contourf(x, y, Ey_values, 50) + im5 = axs[1, 2].contourf(x, y, Bz_values, 50) + im6 = axs[2, 0].contourf(x, y, Ex_values - Ex_ex_values, 50) + im7 = axs[2, 1].contourf(x, y, Ey_values - Ey_ex_values, 50) + im8 = axs[2, 2].contourf(x, y, Bz_values - Bz_ex_values, 50) + axs[0, 0].set_title(r'$E^x$ at t = {:10.3e}'.format(t)) + axs[0, 1].set_title(r'$E^y$ at t = {:10.3e}'.format(t)) + axs[0, 2].set_title(r'$B$ at t = {:10.3e}'.format(t)) + axs[1, 0].set_title(r'$E_h^x$ at t = {:10.3e}'.format(t)) + axs[1, 1].set_title(r'$E_h^y$ at t = {:10.3e}'.format(t)) + axs[1, 2].set_title(r'$B_h$ at t = {:10.3e}'.format(t)) + axs[2, 0].set_title(r'$E^x - E^x_h$ at t = {:10.3e}'.format(t)) + axs[2, 1].set_title(r'$E^y - E^y_h$ at t = {:10.3e}'.format(t)) + axs[2, 2].set_title(r'$B - B_h$ at t = {:10.3e}'.format(t)) + for i in range(3): + for j in range(3): + axs[i, j].plot(*gridlines[0], color='k') + axs[i, j].plot(*gridlines[1], color='k') + axs[i, j].set_xlabel('x', fontsize=14) + axs[i, j].set_ylabel('y', fontsize=14, rotation='horizontal') + axs[i, j].set_aspect('equal') + add_colorbar(im0, axs[0, 0]) + add_colorbar(im1, axs[0, 1]) + add_colorbar(im2, axs[0, 2]) + add_colorbar(im3, axs[1, 0]) + add_colorbar(im4, axs[1, 1]) + add_colorbar(im5, axs[1, 2]) + add_colorbar(im6, axs[2, 0]) + add_colorbar(im7, axs[2, 1]) + add_colorbar(im8, axs[2, 2]) + fig1.suptitle('Compare Exact Solution and Approximate solution at final time') + fig1.tight_layout() + + # fields along s, final time + plot_fields_along_s(tstr='T') + + # Electric field, x component + fig = plot_field_and_error(r'E^x', tend, x, y, Ex_values, Ex_ex_values, *gridlines) + fig.savefig(f'{visdir}/Ex_T_{rp_str}.png') + plt.close(fig) # fig.clf() + + # Electric field, y component + fig = plot_field_and_error(r'E^y', tend, x, y, Ey_values, Ey_ex_values, *gridlines) + fig.savefig(f'{visdir}/Ey_T_{rp_str}.png') + plt.close(fig) # fig.clf() + + # Magnetic field, z component + fig = plot_field_and_error(r'B^z', tend, x, y, Bz_values, Bz_ex_values, *gridlines) + fig.savefig(f'{visdir}/Bz_T_{rp_str}.png') + plt.close(fig) + + return locals() + + +# ============================================================================== +# SCRIPT CAPABILITIES +# ============================================================================== +if __name__ == '__main__': + import argparse + + parser = argparse.ArgumentParser( + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + description="Solve Transverse Time Harmonic Maxwell system on analytical disk with CONGA polar spline method." + ) + + parser.add_argument('--study', + choices=('L2_proj', 'maxwell_bessel', 'maxwell_wave'), + default='maxwell_bessel', + dest='study', + help='Study to be performed' + ) + + parser.add_argument('-S', + action='store_true', + dest='use_spline_mapping', + help='Use spline mapping in finite element calculations' + ) + + parser.add_argument('-D', + type=float, + default=0, + dest='shift_D', + help='Shafranov shift for parametrization of Disk' + ) + + parser.add_argument('-n', '--ncells', + nargs=2, + type=int, + default=[10, 20], + dest='ncells', + help='Number of grid cells (elements) along each dimension' + ) + + parser.add_argument('-s', '--smoothness', + type=int, + default=0, + dest='smooth', + help='Smoothness at the pole. Only C0 and C1 possible. C0 as default.' + ) + + parser.add_argument('-d', '--degree', + nargs=2, + type=int, + default=[3, 3], + dest='degree', + help='Polynomial spline degrees' + ) + + parser.add_argument('-o', '--splitting_order', + type=int, + default=2, + choices=[2, 4], + help='Order of accuracy of operator splitting' + ) + + # ... + time_opts = parser.add_mutually_exclusive_group() + time_opts.add_argument('-t', + type=int, + default=1, + dest='nsteps', + metavar='NSTEPS', + help='Number of time-steps to be taken' + ) + time_opts.add_argument('-T', + type=float, + dest='tend', + metavar='END_TIME', + help='Run simulation until given final time' + ) + + parser.add_argument('--tol', + type=float, + default=1e-7, + help='Tolerance for iterative solver (L2-norm of residual)' + ) + + parser.add_argument('--scipy', + action='store_true', + dest='use_scipy', + help='use scipy matrices and direct inverses' + ) + + parser.add_argument('-v', '--verbose', + action='store_true', + help='Print convergence information of iterative solver' + ) + + # Read input arguments + args = parser.parse_args() + + print(f'running maxwell_2d_TE with args:') + print(f'{args}') + # Run simulation + namespace = run_maxwell_2d_TE(**vars(args)) + + # Keep matplotlib windows open + plt.show() diff --git a/psydac/feec/polar/examples/poisson_2d.py b/psydac/feec/polar/examples/poisson_2d.py new file mode 100644 index 000000000..29aba8734 --- /dev/null +++ b/psydac/feec/polar/examples/poisson_2d.py @@ -0,0 +1,792 @@ +""" +Solve manufactured 2D Poisson problems on polar mapped domains. + +The script builds analytical or spline-approximated polar domains (disk, target or Czarny), +assembles and solves the scalar Poisson system, applies optional treatments of the polar singularity +(C0/C1 CONGA or C1 polar projectors), solves the resulting linear system, +computes L2/H1 errors against the exact solution, and plots the result. + +Example of run: +mpirun -n 2 python poisson_2d.py -S -n 8 10 -d 2 2 -t disk -D 0.2 -m 'C0conga' +""" + +from dataclasses import dataclass + +import sympy +from mpi4py import MPI +from time import time, sleep + +import numpy as np +import matplotlib.pyplot as plt +from sympy import sin, cos, pi, Rational + +from sympde.topology.analytical_mapping import TargetMapping, CzarnyMapping +from sympde.topology.domain import Square, Domain +from sympde.topology import ScalarFunctionSpace, elements_of +from sympde.expr import BilinearForm, LinearForm, integral +from sympde.calculus import dot, grad +from sympde.topology.mapping import Mapping + +from psydac.api.discretization import discretize +from psydac.feec.polar.examples.utils_congapol import add_colorbar, create_tensor_spline_space +from psydac.linalg.stencil import StencilVector, StencilMatrix +from psydac.linalg.basic import LinearOperator +from psydac.linalg.solvers import inverse +from psydac.fem.tensor import TensorFemSpace +from psydac.fem.basic import FemField +from psydac.mapping.discrete import SplineMapping +from psydac.utilities.utils import refine_array_1d +from psydac.cad.geometry import Geometry +from psydac.ddm.cart import DomainDecomposition + +from psydac.polar.c1_projections import C1Projector + +from psydac.api.settings import PSYDAC_BACKENDS + +from psydac.feec.polar.conga_projections import C0PolarProjection_V0, C1PolarProjection_V0 + +backend = PSYDAC_BACKENDS['pyccel-gcc'] + + +# ============================================================================== +class Laplacian: + """ + Symbolic Laplace operator associated with a mapping F from logical to physical coordinates. + Builds Laplacian in logical coordinates using the metric induced by F. + """ + + def __init__(self, mapping): + assert isinstance(mapping, Mapping) + + self._eta = mapping.logical_coordinates + self._metric = mapping.metric_expr + self._metric_det = mapping.metric_det_expr + + # ... + def __call__(self, phi): + from sympy import sqrt, Matrix + + u = self._eta + G = self._metric + sqrt_g = sqrt(self._metric_det) + + # Store column vector of partial derivatives of phi w.r.t. uj + dphi_du = Matrix([phi.diff(uj) for uj in u]) + + # Compute gradient of phi in tangent basis: A = G^(-1) dphi_du + A = G.LUsolve(dphi_du) + + # Compute Laplacian of phi using formula for divergence of vector A + lapl = sum((sqrt_g * Ai).diff(ui) for ui, Ai in zip(u, A)) / sqrt_g + + return lapl + + +# ============================= EXACT SOLUTION ================================# +class Poisson2D: + r""" + Exact solution to the 2D Poisson equation with Dirichlet boundary + conditions, to be employed for the method of manufactured solutions. + + :code + $(\partial^2_{xx} + \partial^2_{yy}) \phi(x,y) = -\rho(x,y)$ + + """ + + def __init__(self, domain_log, mapping, phi_log, rho_log): + assert isinstance(mapping, Mapping) + + self._domain_log = domain_log + self._mapping = mapping + self._phi_log = phi_log + self._rho_log = rho_log + + s, t = mapping.logical_coordinates + self._phi_log_callable = sympy.lambdify([s, t], phi_log) + self._rho_log_callable = sympy.lambdify([s, t], rho_log) + + @staticmethod + def disk(R, shift_D): + r""" + Solve Poisson's equation on a disk of radius R centered at (x,y) = (0, 0), + with logical coordinates (s, theta): + + - The radial coordinate s belongs to the interval [0, R]; + - The angular coordinate theta belongs to the interval [0, 2 * pi). + + : code + $\phi(x,y) = (1 - ((x^2 + y^2) / R^2) ** 4) * sin(kx * x) * cos(ky * y)$. + """ + domain_log = Square('Omega', bounds1=(0, R), bounds2=(0, 2 * np.pi)) + params = dict(c1=shift_D * R * R, c2=0, k=0, D=shift_D) + mapping = TargetMapping('TM', **params) + + # physical field (cf use of physical ref solution in Maxwell case) + k = params['k'] + D = params['D'] + kx = 2 * pi / (R * (1 - k + D)) + ky = 2 * pi / (R * (1 + k)) + x, y = sympy.symbols('x, y') + phi_phys = (1 - ((x * x + y * y) / (R * R)) ** 4) * sin(kx * x) * cos(ky * y) + rho_phys = - phi_phys.diff(x, x) - phi_phys.diff(y, y) + + x_log, y_log = mapping.expressions + phi_log = phi_phys.subs({x: x_log, y: y_log}) + rho_log = rho_phys.subs({x: x_log, y: y_log}) + + return Poisson2D(domain_log, mapping, phi_log, rho_log) + + @staticmethod + def target(): + r""" + Solve Poisson's equation on a polar domain, with logical coordinates (s, theta): + + - The radial coordinate s belongs to the interval [0, 1]; + - The angular coordinate theta belongs to the interval [0, 2 * pi). + + The shape of the domain is set by parameter k in [0, 1): for k = 0 we have + a disk and for k --> 1 the disk is horizontally squeezed. + + The parameter D in (-(1 - k)/2, (1 - k)/2) moves horizontally the pole within + the shape. + + The parameter c1, c2 are just shifts on the plane of the domain. + + : code + $\phi(x,y) = (1 - s^8)\sin(k_x(x - 0.5))\cos(k_y y)$. + + """ + + domain_log = Square('Omega', bounds1=(0, 1), bounds2=(0, 2 * np.pi)) + params = dict(c1=0, c2=0, k=Rational(3, 10), D=Rational(2, 10)) + mapping = TargetMapping('F', **params) + + lapl = Laplacian(mapping) + s, t = mapping.logical_coordinates + x, y = mapping.expressions + + # Manufactured solution in logical coordinates + k = params['k'] + D = params['D'] + kx = 2 * pi / (1 - k + D) + ky = 2 * pi / (1 + k) + + phi_log = (1 - s ** 8) * sin(kx * (x - 0.5)) * cos(ky * y) + rho_log = - lapl(phi_log) + + return Poisson2D(domain_log, mapping, phi_log, rho_log) + + @staticmethod + def czarny(): + r""" + Solve Poisson's equation on a czarny domain, with logical coordinates (s, theta): + + - The radial coordinate s belongs to the interval [0, 1]; + - The angular coordinate theta belongs to the interval [0, 2 * pi). + + : code + $\phi(x,y) = (1 - s^8)\sin(\pi x)\cos(\pi y)$. + + """ + + domain_log = Square('Omega', bounds1=(0, 1), bounds2=(0, 2 * np.pi)) + params = dict(c1=0, c2=0, eps=Rational(1, 5), b=Rational(7, 5)) + mapping = CzarnyMapping('F', **params) + + lapl = Laplacian(mapping) + s, t = mapping.logical_coordinates + x, y = mapping.expressions + + # Manufactured solution in logical coordinates + phi_log = (1 - s ** 8) * sin(pi * x) * cos(pi * y) + rho_log = - lapl(phi_log) + + return Poisson2D(domain_log, mapping, phi_log, rho_log) + + @property + def domain_log(self): + return self._domain_log + + @property + def mapping(self): + return self._mapping + + @property + def phi_log(self): + return self._phi_log + + @property + def rho_log(self): + return self._rho_log + + @property + def phi_log_callable(self): + return self._phi_log_callable + + @property + def rho_log_callable(self): + return self._rho_log_callable + + +# ====================== CONGA (PENALIZED) POISSON ============================# + +class CongaLaplacian(LinearOperator): + + def __init__(self, S, M, P, alpha): + assert isinstance(S, StencilMatrix) + assert isinstance(M, StencilMatrix) + assert isinstance(P, (C0PolarProjection_V0, C1PolarProjection_V0)) + + W0 = P.W0.coeff_space + + assert S.domain is S.codomain is W0 + assert M.domain is M.codomain is W0 + + self.S = S + self.M = M + self.P = P + self.alpha = alpha + self.W0 = W0 + + def dot(self, x, out=None): + if out is None: + y = self.M._domain.zeros() + else: + assert isinstance(out, StencilVector) + assert out.space is self.M._domain + y = out + + y1 = self.P.T.dot(self.S.dot(self.P.dot(x))) + y2 = x - self.P.dot(x) + y2 = self.M.dot(y2) + y2 = self.alpha * (y2 - self.P.T.dot(y2)) + # y = y1 + y2 + y1.copy(out=y) + y += y2 + + y.update_ghost_regions() + return y + + def transpose(self, conjugate=False): + return CongaLaplacian(self.S.T, self.M.T, self.P.T, self.alpha) + + def tosparse(self): + S = self.S.tosparse() + M = self.M.tosparse() + P = self.P.tosparse() + alpha = self.alpha + n = self.W0.dimension + + from scipy.sparse import eye + I = eye(n) + + A = alpha * (I - P).T @ M @ (I - P) + P.T @ S @ P + + return A + + def toarray(self): + return self.tosparse().toarray() + + @property + def T(self): + return self.transpose() + + @property + def shape(self): + return (self.W0.dimension, self.W0.dimension) + + @property + def domain(self): + return self.W0 + + @property + def codomain(self): + return self.W0 + + @property + def dtype(self): + return float + +@dataclass +class ErrorDiagnostics: + ref_l2: float + ref_h1: float + rel_l2: float + rel_h1: float + +############################################################################### + +def run_poisson_2d(*, test_case, ncells, degree, + shift_D, R, use_spline_mapping, smooth_method, + cgtol, cgiter, alphaCONGA, verbose=False): + timing = {} + timing['assembly'] = 0.0 + timing['projection'] = 0.0 + timing['solution'] = 0.0 + timing['diagnostics'] = 0.0 + timing['export'] = 0.0 + + # Method of manufactured solution + if test_case == 'disk': + model = Poisson2D.disk(R=R, shift_D=shift_D) + elif test_case == 'target': + model = Poisson2D.target() + elif test_case == 'czarny': + model = Poisson2D.czarny() + else: + raise ValueError("Only available test-cases are 'disk', 'target' and 'czarny'") + + if smooth_method not in ('polar-std', 'polar-spec', 'C0conga', 'C1conga', 'None'): + raise ValueError( + "Only available options for pole smoothness are 'polar-spec', 'polar-std', 'C0conga', 'C1conga', 'None'") + + if smooth_method == 'polar-spec' and (not use_spline_mapping): + print('WARNING: C1 conforming discretization only available for spline mappings') + print('The domain will be approximated in the 0-forms spline space.') + print() + use_spline_mapping = True + + # Communicator, size, rank + mpi_comm = MPI.COMM_WORLD + mpi_size = mpi_comm.Get_size() + mpi_rank = mpi_comm.Get_rank() + + periodic = [False, True] + + if use_spline_mapping: + + # ==================== SPLINE SPACE FOR SPLINE MAPPINGS =======================# + + V = create_tensor_spline_space(ncells, degree, periodic, + (model.domain_log.bounds1, model.domain_log.bounds2), mpi_comm) + + #TODO: maybe define a parent class Model + + # ==================== MAPPING & PHYSICAL DOMAIN ==============================# + # Create spline mapping by interpolation of analytical mapping + map_analytic = model.mapping.get_callable_mapping() + map_discrete = SplineMapping.from_mapping(V, map_analytic) + # Create symbolic mapping with callable mapping as spline + mapping = Mapping('M', dim=2) + mapping.set_callable_mapping(map_discrete) + # In order to create a sympde.Domain object from this mapping we have + # to create first a HDF5 file and then load as sympde.Domain.fromfile + t0 = time() + geometry = Geometry.from_discrete_mapping(map_discrete, comm=mpi_comm) + geometry.export('geo.h5') + t1 = time() + timing['export'] += t1 - t0 + domain = Domain.from_file('geo.h5') + + #check_regular_ring_map(map_discrete) + + else: + # Only symbolic mapping is necessary + mapping = model.mapping + domain = mapping(model.domain_log) + + # ========================== SYMBOLIC DEFINITION ==============================# + + # Equations + V0 = ScalarFunctionSpace('V0', domain) + u0, v0 = elements_of(V0, names='u0, v0') + aM = BilinearForm((u0, v0), integral(domain, u0 * v0)) + aS = BilinearForm((u0, v0), integral(domain, dot(grad(u0), grad(v0)))) + + rhs = LinearForm(v0, integral(domain, model.rho_log * v0)) + + # ============================= DISCRETIZATION ================================# + if use_spline_mapping: + domain_h = discretize(domain, filename='geo.h5', comm=mpi_comm) + V0_h = discretize(V0, domain_h) + F = list(domain_h.mappings.values()).pop() + else: + domain_h = discretize(domain, ncells=ncells, periodic=periodic, comm=mpi_comm) + V0_h = discretize(V0, domain_h, degree=degree) + F = mapping.get_callable_mapping() + + aM_h = discretize(aM, domain_h, (V0_h, V0_h), backend=backend) + aS_h = discretize(aS, domain_h, (V0_h, V0_h), backend=backend) + rhs_h = discretize(rhs, domain_h, V0_h, backend=backend) + + M = aM_h.assemble() + S = aS_h.assemble() + b = rhs_h.assemble() + + S.update_ghost_regions() + b.update_ghost_regions() + M.update_ghost_regions() + + # =================== PROJECT THE EXACT SOLUTION =========================# + + from psydac.feec.global_geometric_projectors import GlobalGeometricProjectorH1 + + Pi0 = GlobalGeometricProjectorH1(V0_h) + phi_ref = Pi0(model.phi_log_callable) + phi_ref.coeffs.update_ghost_regions() + + + # ========================= HANDLING THE SINGULARITY ===========================# + + # If required by user, create C1 projector and then restrict + # stiffness/mass matrices and right-hand-side vector to C1 space + t0 = time() + bc = None + Sc = None + if smooth_method == 'polar-spec': + proj = C1Projector(F) + Sp = proj.change_matrix_basis(S) + bp = proj.change_rhs_basis(b) + alpha = 'None' + if smooth_method == 'polar-std': + # Build standard polar map from control points of standard polar map + n1, n2 = [W.nbasis for W in V0_h.spaces] + rho = np.array([i1 / (n1 - 1) for i1 in range(n1)]) + theta = np.array([i2 * 2 * np.pi / n2 for i2 in range(n2)]) + sin_theta = np.sin(theta) + cos_theta = np.cos(theta) + cp = np.zeros((n1, n2, 2)) + for i1 in range(n1): + for i2 in range(n2): + cp[i1, i2, 0] = rho[i1] * cos_theta[i2] + cp[i1, i2, 1] = rho[i1] * sin_theta[i2] + F_std = SplineMapping.from_control_points(V0_h, cp) + proj = C1Projector(F_std) + Sp = proj.change_matrix_basis(S) + bp = proj.change_rhs_basis(b) + alpha = 'None' + elif smooth_method == 'C1conga': + gamma = 1.0 # any value would be ok. + alpha = alphaCONGA + P0 = C1PolarProjection_V0(V0_h, gamma=gamma, hbc=True) # hbc imposes the boundary conditions + Sc = CongaLaplacian(S, M, P0, alpha) + bc = P0.T.dot(b) + elif smooth_method == 'C0conga': + alpha = alphaCONGA + P0 = C0PolarProjection_V0(V0_h, hbc=True) # hbc imposes the boundary conditions + Sc = CongaLaplacian(S, M, P0, alpha) + bc = P0.T.dot(b) + elif smooth_method == 'None': + alpha = 'None' + t1 = time() + timing['projection'] = t1 - t0 + + # Apply homogeneous Dirichlet boundary conditions for the conforming + # smooth_method case 'polar' and non-conforming case 'None' + # NOTE: this does not affect ghost regions + # For each angular index j on the last owned radial line, replace the assembled equation + # (S u)[e1, j] = b[e1, j] + # by the equation + # u[e1, j] = 0. + + e1 = V0_h.coeff_space.ends[0] + if e1 == V0_h.coeff_space.npts[0] - 1: + if smooth_method in ('polar-std', 'polar-spec'): + last = bp[1].space.npts[0] - 1 + Sp[1, 1][last, :, :, :] = 0. + Sp[1, 1][last, :, 0, 0] = 1. + bp[1][last, :] = 0. + elif smooth_method == 'None': + S[e1, :, :, :] = 0. + S[e1, :, 0, 0] = 1. #set diagonal entries to 1 + b[e1, :] = 0. #RHS is 0 at the outer radial boundary + + # ====================== SOLVE GALERKIN SYSTEM WITH CG ========================# + + # Solve linear system + t0 = time() + if smooth_method in ('polar-std', 'polar-spec'): + Sp_inv = inverse(Sp, 'cg', tol = cgtol, maxiter = cgiter, verbose=verbose) + xp = Sp_inv.dot(bp) + xsol = proj.convert_to_tensor_basis(xp) + info = Sp_inv.get_info() + elif smooth_method == 'C1conga': + Sc_inv = inverse(Sc, 'cg', tol=cgtol, maxiter=cgiter, verbose=verbose) + xsol = Sc_inv.dot(bc) + info = Sc_inv.get_info() + elif smooth_method == 'C0conga': + Sc_inv = inverse(Sc, 'cg', tol=cgtol, maxiter=cgiter, verbose=verbose) + xsol = Sc_inv.dot(bc) + info = Sc_inv.get_info() + elif smooth_method == 'None': + pc = S.diagonal(inverse=True) + S_inv = inverse(S, 'cg', pc=pc, tol=cgtol, maxiter=cgiter, verbose=verbose) + xsol = S_inv.dot(b) + info = S_inv.get_info() + t1 = time() + timing['solution'] = t1 - t0 + + # ========================= APPROXIMATION ERROR ===============================# + + # Create potential field for discrete solution + phi = FemField(V0_h, coeffs=xsol) + phi.coeffs.update_ghost_regions() + + t0 = time() + errors = compute_errors(phi, phi_ref, M, S) + t1 = time() + timing['diagnostics'] = t1 - t0 + + # Write solution to HDF5 file + t0 = time() + V0_h.export_fields('fields.h5', phi=phi) + t1 = time() + timing['export'] += t1 - t0 + # =============================== PRINTING INFO ===============================# + + # +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ + # Print some information to terminal + for i in range(mpi_size): + if i == mpi_rank: + print('--------------------------------------------------') + print(' RANK = {}'.format(mpi_rank)) + print('--------------------------------------------------') + print('> Grid :: [{ne1},{ne2}]'.format(ne1=ncells[0], ne2=ncells[1])) + print('> Degree :: [{p1},{p2}]'.format(p1=degree[0], p2=degree[1])) + print('> Penalization alpha :: {alpha} '.format(alpha=alpha)) + print( '> CG info :: ',info ) + print('> L2 norm solution :: {:.2e}'.format(errors.ref_l2)) + print('> H1 norm solution :: {:.2e}'.format(errors.ref_h1)) + print('> L2 error (relative) :: {:.2e}'.format(errors.rel_l2)) + print('> H1 error (relative) :: {:.2e}'.format(errors.rel_h1)) + print('') + print('> Assembly time :: {:.2e}'.format(timing['assembly'])) + if smooth_method: + print('> Project. time :: {:.2e}'.format(timing['projection'])) + print('> Solution time :: {:.2e}'.format(timing['solution'])) + print('> Evaluat. time :: {:.2e}'.format(timing['diagnostics'])) + print('> Export time :: {:.2e}'.format(timing['export'])) + print('', flush=True) + sleep(0.001) + mpi_comm.Barrier() + + # =============================== VISUALIZATION ===============================# + + N = 10 + V0_h.plot_2d_decomposition(mapping.get_callable_mapping(), refine=N) + + # Non-master processes stop here + if mpi_rank != 0: + return + + plot_solution(use_spline_mapping, model, ncells, periodic, V0_h, refine=N) + + return locals() + +def compute_errors(phi, phi_ref, M, S): + + # L2 and H1 norms + ref_l2_2 = phi_ref.coeffs.inner(M.dot(phi_ref.coeffs)) + ref_h1_semi2 = phi_ref.coeffs.inner(S.dot(phi_ref.coeffs)) + ref_l2 = np.sqrt(ref_l2_2) # L2 norm of ref solution + ref_h1 = np.sqrt(ref_h1_semi2 + ref_l2_2) # H1 norm of ref solution + + # L2 and H1 errors + phi_diff = phi_ref.coeffs - phi.coeffs + + err_l2_2 = phi_diff.inner(M.dot(phi_diff)) + err_h1_semi2 = phi_diff.inner(S.dot(phi_diff)) + err_l2 = np.sqrt(err_l2_2) + err_h1 = np.sqrt(err_h1_semi2 + err_l2_2) + + return ErrorDiagnostics( + ref_l2=ref_l2, + ref_h1=ref_h1, + rel_l2=err_l2 / ref_l2, + rel_h1=err_h1 / ref_h1, + ) + +# ============================================================================== +# Plotting +# ============================================================================== +def plot_solution(use_spline_mapping, model, ncells, periodic, V0_h, refine=10): + """ + Plot exact solution, numerical solution and error + """ + + if use_spline_mapping: + geometry = Geometry(filename='geo.h5', comm=MPI.COMM_SELF) + map_discrete = [*geometry.mappings.values()].pop() + Vnew = map_discrete.space + map_plot = map_discrete + + else: + dd = DomainDecomposition(ncells, periodic, comm=MPI.COMM_SELF) + Vnew = TensorFemSpace(dd, *V0_h.spaces) + map_plot = model.mapping.get_callable_mapping() + + # Import solution vector into new serial field + phi, = Vnew.import_fields('fields.h5', 'phi') + + # Callable exact solution in logical coordinates + phi_e = model.phi_log_callable + + # Compute numerical solution (and error) on refined logical grid + [sk1, sk2], [ek1, ek2] = Vnew.local_domain + V1_plot, V2_plot = Vnew.spaces + + eta1 = refine_array_1d(V1_plot.breaks[sk1:ek1 + 2], refine) + eta2 = refine_array_1d(V2_plot.breaks[sk2:ek2 + 2], refine) + num = np.array([[phi(e1, e2) for e2 in eta2] for e1 in eta1]) + ex = np.array([[phi_e(e1, e2) for e2 in eta2] for e1 in eta1]) + err = num - ex + + # Compute physical coordinates of logical grid + pcoords = np.array([[map_plot(e1, e2) for e2 in eta2] for e1 in eta1]) + xx = pcoords[:, :, 0] + yy = pcoords[:, :, 1] + + # Create figure with 3 subplots: + # 1. exact solution on exact domain + # 2. numerical solution on mapped domain (analytical or spline) + # 3. numerical error on mapped domain (analytical or spline) + fig, axes = plt.subplots(1, 3, figsize=(15, 4.8)) + + # Plot exact solution + ax = axes[0] + im = ax.contourf(xx, yy, ex, 40, cmap='jet') + add_colorbar(im, ax) + ax.set_xlabel(r'$x$', rotation='horizontal') + ax.set_ylabel(r'$y$', rotation='horizontal') + ax.set_title(r'$\phi_{ex}(x,y)$') + ax.plot(xx[:, ::refine], yy[:, ::refine], 'k') + ax.plot(xx[::refine, :].T, yy[::refine, :].T, 'k') + ax.set_aspect('equal') + + # Plot numerical solution + ax = axes[1] + im = ax.contourf(xx, yy, num, 40, cmap='jet') + add_colorbar(im, ax) + ax.set_xlabel(r'$x$', rotation='horizontal') + ax.set_ylabel(r'$y$', rotation='horizontal') + ax.set_title(r'$\phi(x,y)$') + ax.plot(xx[:, ::refine], yy[:, ::refine], 'k') + ax.plot(xx[::refine, :].T, yy[::refine, :].T, 'k') + ax.set_aspect('equal') + + # Plot numerical error + ax = axes[2] + im = ax.contourf(xx, yy, err, 40, cmap='jet') + add_colorbar(im, ax) + ax.set_xlabel(r'$x$', rotation='horizontal') + ax.set_ylabel(r'$y$', rotation='horizontal') + ax.set_title(r'$\phi(x,y) - \phi_{ex}(x,y)$') + ax.plot( xx[:,::refine] , yy[:,::refine] , 'k' ) + ax.plot( xx[::refine,:].T, yy[::refine,:].T, 'k' ) + ax.set_aspect('equal') + + # Show figure + fig.tight_layout() + fig.show() + +# ============================================================================== +# Parser +# ============================================================================== +def parse_input_arguments(): + import argparse + + parser = argparse.ArgumentParser( + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + description="Solve Poisson's equation on a 2D polar domain." + ) + + parser.add_argument('-t', + type=str, + choices=('disk', 'target', 'czarny'), + default='disk', + dest='test_case', + help='Test case' + ) + + parser.add_argument('-D', + type=float, + default=0, + dest='shift_D', + help='Shafranov shift for parametrization of Disk' + ) + + parser.add_argument('-R', + type=float, + default=1.0, + dest='R', + help='Radius of the disk' + ) + + parser.add_argument('-d', + type=int, + nargs=2, + default=[2, 2], + metavar=('P1', 'P2'), + dest='degree', + help='Spline degree along each dimension' + ) + + parser.add_argument('-n', '--ncells', + type=int, + nargs=2, + default=[10, 20], + metavar=('N1', 'N2'), + dest='ncells', + help='Number of grid cells (elements) along each dimension' + ) + + parser.add_argument('-S', + action='store_true', + dest='use_spline_mapping', + help='Use spline mapping in finite element calculations' + ) + + parser.add_argument('-m', + choices=('polar-spec', 'polar-std', 'C0conga', 'C1conga', 'None'), + default='C1conga', + dest='smooth_method', + help='Apply smoothing method at pole either C1-conforming geometry specific / C1-conforming standardized / C1-CONGA / C0-CONGA' + ) + + parser.add_argument('-v', + action='store_true', + dest='verbose', + help='See CG iterations and L2 norm of the residuals' + ) + + parser.add_argument('-a', + type=float, + default=1000, + dest='alphaCONGA', + help='Penalization constant for CONGA methods' + ) + + parser.add_argument('--cgtol', + type=float, + default=1e-10, + dest='cgtol', + help='absolute tol for the residual error to stop CG' + ) + + parser.add_argument('--maxiter', + type=int, + default=100000, + dest='cgiter', + help='max number of iterations for CG' + ) + + return parser.parse_args() + + +# ============================================================================== +# Script functionality +# ============================================================================== +if __name__ == '__main__': + + args = parse_input_arguments() + namespace = run_poisson_2d(**vars(args)) + + import __main__ + + if hasattr(__main__, '__file__'): + try: + __IPYTHON__ + except NameError: + import matplotlib.pyplot as plt + + plt.show() diff --git a/psydac/feec/polar/examples/utils_congapol.py b/psydac/feec/polar/examples/utils_congapol.py new file mode 100644 index 000000000..8a01c1821 --- /dev/null +++ b/psydac/feec/polar/examples/utils_congapol.py @@ -0,0 +1,152 @@ +import numpy as np + +from psydac.ddm.cart import DomainDecomposition +from psydac.fem.splines import SplineSpace +from psydac.fem.tensor import TensorFemSpace + + +def print_map_polar_coeffs(map_discrete): + """ + Used for debugging purposes. Prints information about discrete polar coefficients + """ + # Print spline mapping + print('Spline mapping:') + print(map_discrete) + print('vars(map_discrete):') + print(vars(map_discrete)) + print() + print('vars(map_discrete._fields[0]) :') + print(vars(map_discrete._fields[0])) + + print() + print('vars(map_discrete._fields[0]._space._spaces[0]) :') + print(vars(map_discrete._fields[0]._space._spaces[0])) + print('vars(map_discrete._fields[0]._space._spaces[1]) :') + print(vars(map_discrete._fields[0]._space._spaces[1])) + n_s = map_discrete._fields[0]._space._spaces[0]._nbasis + n_theta = map_discrete._fields[0]._space._spaces[1]._nbasis + print('n_s = ', n_s) + print('n_theta = ', n_theta) + print() + + map_0_c = map_discrete._fields[0]._coeffs.toarray() + map_1_c = map_discrete._fields[1]._coeffs.toarray() + + print('pole: ', map_0_c[0], map_1_c[0]) + map_0_c -= map_0_c[0] + map_1_c -= map_1_c[0] + + # print(map_0_c**2 + map_1_c**2) + radius_pole = np.sqrt(map_0_c[:n_theta] ** 2 + map_1_c[:n_theta] ** 2) + radius_first_ring = np.sqrt( + map_0_c[n_theta:2 * n_theta] ** 2 + map_1_c[n_theta:2 * n_theta] ** 2 + ) + rho_1 = radius_first_ring[0] + + cs = map_0_c[n_theta:2 * n_theta] / rho_1 + sn = map_1_c[n_theta:2 * n_theta] / rho_1 + theta = np.arctan2(sn, cs) + + print('radius_pole:', radius_pole) + print('radius first ring:', radius_first_ring) + print('D theta_j:', np.mod(theta[1:] - theta[:-1], 2 * np.pi)) + + +def check_regular_ring_map(map_discrete, verbose=False): + """ + Check the first two radial rings of a 2D spline mapping + + Performs 3 checks: + 1. that the pole ring (i=0) collapses to a single point + 2. the 1st ring (i=1) has constant radius + 3. the 1st ring is uniformly spaced in angle + and prints diagnostic information. Used for debugging purposes. + + """ + n_s = map_discrete._fields[0]._space._spaces[0]._nbasis + n_theta = map_discrete._fields[0]._space._spaces[1]._nbasis + assert n_s == map_discrete._fields[1]._space._spaces[0]._nbasis + assert n_theta == map_discrete._fields[1]._space._spaces[1]._nbasis + + if verbose: + print('n_s = ', n_s) + print('n_theta = ', n_theta) + print() + + # read mapping coefficients + map_0_c = map_discrete._fields[0]._coeffs.toarray() + map_1_c = map_discrete._fields[1]._coeffs.toarray() + + # shift + map_0_c -= map_0_c[0] + map_1_c -= map_1_c[0] + + # distance of every coefficient in the pole ring from the first pole coefficient (should be 0) + radius_pole = np.sqrt(map_0_c[:n_theta] ** 2 + map_1_c[:n_theta] ** 2) + # distance from the pole of the first radial ring + radius_first_ring = np.sqrt( + map_0_c[n_theta:2 * n_theta] ** 2 + map_1_c[n_theta:2 * n_theta] ** 2 + ) + rho_1 = radius_first_ring[0] + + # compute angles of the 1st ring + cs = map_0_c[n_theta:2 * n_theta] / rho_1 + sn = map_1_c[n_theta:2 * n_theta] / rho_1 + theta = np.arctan2(sn, cs) + + delta_theta = np.mod(theta[1:] - theta[:-1], 2 * np.pi) + if verbose: + print('radius_pole:', radius_pole) + print('radius first ring:', radius_first_ring) + print('D theta_j:', delta_theta) + print() + + circle_error = np.linalg.norm(radius_pole[1:] - radius_pole[:-1]) + #check that 1st ring points are uniformly spaced in angle + regularity_error = np.linalg.norm(delta_theta[1:] - delta_theta[:-1]) + first_ring_radius_error = np.linalg.norm(radius_first_ring[1:] - radius_first_ring[:-1]) + regularity_check = circle_error + regularity_error + first_ring_radius_error < 1e-12 + + print('CHECK: spline mapping is regular: ', regularity_check) + if not regularity_check: + print(' - circle error on pole ring:', circle_error) + print(' - regularity error on 1st ring:', regularity_error) + print(' - first ring radius error:', first_ring_radius_error) + print() + +def add_colorbar(im, ax, **kwargs): + from mpl_toolkits.axes_grid1 import make_axes_locatable + divider = make_axes_locatable(ax) + cax = divider.append_axes("right", size=0.2, pad=0.3) + cbar = ax.get_figure().colorbar(im, cax=cax, **kwargs) + return cbar + +def create_tensor_spline_space(ncells, spline_degrees, periodic, bounds, mpi_comm=None): + """ + Create a 2D tensor-product spline finite element space on a rectangular + logical domain (e.g. with bounds ``[[0, R], [0, 2*pi]]``). + + Returns + ------- + TensorFemSpace + The 2D tensor-product spline finite element space. + """ + + # Number of elements and spline degree + ne1, ne2 = ncells + p1, p2 = spline_degrees + + # Create uniform grid + grid_1 = np.linspace(*bounds[0], num=ne1 + 1) + grid_2 = np.linspace(*bounds[1], num=ne2 + 1) + + # Create 1D finite element spaces + S1 = SplineSpace(p1, grid=grid_1, periodic=periodic[0]) + S2 = SplineSpace(p2, grid=grid_2, periodic=periodic[1]) + + # Create 2D tensor product finite element space + domain_decomposition = DomainDecomposition(ncells, periodic, comm=mpi_comm) + V = TensorFemSpace(domain_decomposition, S1, S2) + return V + + diff --git a/psydac/feec/polar/examples/waveTE.py b/psydac/feec/polar/examples/waveTE.py new file mode 100644 index 000000000..9f151396a --- /dev/null +++ b/psydac/feec/polar/examples/waveTE.py @@ -0,0 +1,206 @@ +from sympde.topology import Square +from sympde.topology.analytical_mapping import PolarMapping + +from psydac.feec.polar.examples.utils_congapol import add_colorbar +from psydac.feec.pull_push import push_2d_hcurl, push_2d_l2 +from numpy import pi +import numpy as np +import matplotlib.pyplot as plt + +class GaussianSolution: + """ + Initial Gaussian circular wave for the TE Maxwell test. + + This class defines the initial condition used for the Gaussian wave + propagation experiment. It is not an exact time-dependent Maxwell solution. + The electric field is initialized as a localized rotational Gaussian pulse, + + E0(x, y) = scale * (y - y0, -(x - x0)) + * exp(-((x - x0)^2 + (y - y0)^2) / (2 sigma^2)), + + and the magnetic field is initialized as + + B0 = curl E0 = d_x Ey - d_y Ex. + + Parameters + ---------- + sigma : float + Width of the Gaussian pulse. + + x0, y0 : float + Center of the Gaussian pulse in physical coordinates. + + scale : float, optional + Amplitude scaling factor for the initial fields. + """ + + def __init__(self, sigma, x0, y0, scale=1): + + self.x0 = x0 + self.y0 = y0 + self.sigma = sigma + self.scale = scale + + def _gaussian(self, x, y): + from numpy import exp + X = x - self.x0 + Y = y - self.y0 + sig2 = self.sigma**2 + return self.scale * exp(-(X*X + Y*Y)/(2*sig2)) + + def Ex_ex(self, t, x, y): + Y = y - self.y0 + return Y * self._gaussian(x, y) + + def Ey_ex(self, t, x, y): + X = x - self.x0 + return -X * self._gaussian(x, y) + + def Bz_ex(self, t, x, y): + """ + Bz = curl E = d_x Ey - d_y Ex + """ + X = x - self.x0 + Y = y - self.y0 + sig2 = self.sigma**2 + r2 = X*X + Y*Y + return (r2/sig2 - 2.0) * self._gaussian(x, y) + +def main(): + """ + This function is not currently used. It is kept for possible future development. + """ + + # Logical domain is rectangle [0, R] x [0, 2pi] + R = 2.0 + + # Speed of light equal c and scaling of the fields by a scale factor + sigma = 0.1 + scale = 1 + + # Exact solution + exact_solution = GaussianSolution(sigma=sigma, x0=0, y0=0, scale=scale) + + # Exact fields, as callable functions of (t, s, theta) + Es_ex = exact_solution.Es_ex + Et_ex = exact_solution.Et_ex + B_ex = exact_solution.B_ex + Bt_ex = exact_solution.Bt_ex + + # Logical domain: [0, R] x [0, 2pi] + logical_domain = Square('Omega', bounds1=[0, R], bounds2=[0, 2 * pi]) + + # Physical domain: disk of radius R obtained as image of the logical_domain + # with the analytical mapping of a circle + mapping = PolarMapping('F', c1=0, c2=0, rmin=0, rmax=1) + domain = mapping(logical_domain) + F = mapping.get_callable_mapping() + + # Set time + t = 0 + + Es = lambda x, y: Es_ex(t, x, y) + Et = lambda x, y: Et_ex(t, x, y) + B = lambda x, y: B_ex(t, x, y) + Bt = lambda x, y: Bt_ex(t, x, y) + + # Plot of fields + N = 100 + + rho = np.linspace(1e-20, R, N) + theta = np.linspace(0, 2 * pi, N) + rho, theta = np.meshgrid(rho, theta, indexing='ij') + x, y = F(rho, theta) + + Ex_values = np.empty_like(rho) + Ey_values = np.empty_like(rho) + B_values = np.empty_like(rho) + + for i, x1i in enumerate(rho[:, 0]): + for j, x2j in enumerate(theta[0, :]): + Ex_values[i, j], Ey_values[i, j] = \ + push_2d_hcurl(Es, Et, x1i, x2j, F) + + B_values[i, j] = push_2d_l2(B, x1i, x2j, F) + + fig, axs = plt.subplots(2, 2, figsize=(12, 12)) + im0 = axs[0, 0].contourf(x, y, Ex_values, 50) + im1 = axs[0, 1].contourf(x, y, Ey_values, 50) + im2 = axs[1, 0].contourf(x, y, np.sqrt(Ex_values ** 2 + Ey_values ** 2), 50) + im3 = axs[1, 1].contourf(x, y, B_values, 50) + axs[0, 0].set_title(r'$E_x$') + axs[0, 1].set_title(r'$E_y$') + axs[1, 0].set_title(r'$||\mathbf{E}||$') + axs[1, 1].set_title('$B$') + add_colorbar(im0, axs[0, 0]) + add_colorbar(im1, axs[0, 1]) + add_colorbar(im2, axs[1, 0]) + add_colorbar(im3, axs[1, 1]) + + # Test: curl E = - d_t B_z with finite Differences + # on a tensor grid in the square inscribed the disk + N = 160 + l = R / np.sqrt(2) # edge length + + x, dx = np.linspace(-l, l, N, retstep=True) + y, dy = np.linspace(-l, l, N, retstep=True) + x, y = np.meshgrid(x, y, indexing='ij') + rho = np.sqrt(x ** 2 + y ** 2) + theta = np.arctan2(y, x) % (2 * pi) + + Ex_values = np.empty_like(rho) + Ey_values = np.empty_like(rho) + Bt_values = np.empty_like(rho) + + ni, nj = rho.shape + for i in range(ni): + for j in range(nj): + x1_ij = rho[i, j] + x2_ij = theta[i, j] + Ex_values[i, j], Ey_values[i, j] = \ + push_2d_hcurl(Es, Et, x1_ij, x2_ij, F) + + Bt_values[i, j] = push_2d_l2(Bt, x1_ij, x2_ij, F) + + fig, axs = plt.subplots(2, 3, figsize=(15, 15)) + im1 = axs[0, 0].contourf(x, y, np.sqrt(Ex_values ** 2 + Ey_values ** 2)) + im2 = axs[0, 1].contourf(x, y, -Bt_values) + axs[0, 0].set_title(r'$||\mathbf{E}||$') + axs[0, 1].set_title(r'$-\partial_t B$') + for axi in axs.flat: + axi.set_aspect('equal') + add_colorbar(im1, axs[0, 0]) + add_colorbar(im2, axs[0, 1]) + + curlE_values = np.zeros_like(rho) + valerr = 0 + for i in range(1, ni - 1): + for j in range(1, nj - 1): + curlE_values[i, j] = (Ey_values[i + 1, j] - Ey_values[i - 1, j]) / (2 * dx) \ + - (Ex_values[i, j + 1] - Ex_values[i, j - 1]) / (2 * dy) + + d = np.abs(Bt_values[i, j] + curlE_values[i, j]) + + if valerr < d: + valerr = d + + skip = (slice(None, None, int(N / 20)), slice(None, None, int(N / 20))) + im3 = axs[1, 0].contourf(x, y, curlE_values) + axs[1, 0].quiver(x[skip], y[skip], Ex_values[skip], Ey_values[skip]) + add_colorbar(im3, axs[1, 0]) + axs[1, 0].set_title(r'curl $\mathbf{E}$') + im4 = axs[1, 1].contourf(x, y, -Bt_values) + axs[1, 1].quiver(x[skip], y[skip], Ex_values[skip], Ey_values[skip]) + add_colorbar(im4, axs[1, 1]) + im5 = axs[0, 2].contourf(x, y, Ex_values) + im6 = axs[1, 2].contourf(x, y, Ey_values) + add_colorbar(im5, axs[0, 2]) + add_colorbar(im6, axs[1, 2]) + print('|curlE + d_t B| <= ', valerr) + + +if __name__ == "__main__": + main() + plt.show() + + diff --git a/psydac/feec/polar/tests/__init__.py b/psydac/feec/polar/tests/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/psydac/feec/polar/tests/test_conga_projections.py b/psydac/feec/polar/tests/test_conga_projections.py new file mode 100644 index 000000000..0b41d64a6 --- /dev/null +++ b/psydac/feec/polar/tests/test_conga_projections.py @@ -0,0 +1,312 @@ +import numpy as np +from numpy import pi + +import pytest +from sympde import ScalarFunctionSpace, VectorFunctionSpace + +from mpi4py import MPI + +from sympde.topology import Square +from sympde.topology.analytical_mapping import PolarMapping + +from psydac.api.discretization import discretize +from psydac.feec.polar.conga_projections import C0PolarProjection_V0, C0PolarProjection_V2, C0PolarProjection_V1, \ + C1PolarProjection_V0, C1PolarProjection_V1 +from psydac.utilities.gather_variable_len_arrays import gather_vlen_array, gather_vlen_arrays +from psydac.fem.basic import FemField +from psydac.linalg.block import BlockVector + +ATOL = 1e-12 +RTOL = 1e-12 + + +def get_domain(R): + # Build physical domain - disk of radius R + logical_domain = Square('Omega', bounds1=[0, R], bounds2=[0, 2 * pi]) + mapping = PolarMapping('F', c1=0, c2=0, rmin=0, rmax=1) + domain = mapping(logical_domain) + + return domain + +def get_random_vector(space): + [s1, s2] = space.coeff_space.starts + [e1, e2] = space.coeff_space.ends + + phi = FemField(space) + x = phi.coeffs + x[s1:e1 + 1, s2:e2 + 1] = np.random.random([e1 - s1 + 1, e2 - s2 + 1]) + x.update_ghost_regions() + return x + +def get_random_block_vector(space): + x = BlockVector(space.coeff_space) + for i in (0, 1): + [s1, s2] = space.coeff_space[i].starts + [e1, e2] = space.coeff_space[i].ends + x[i][s1:e1 + 1, s2:e2 + 1] = np.random.random([e1 - s1 + 1, e2 - s2 + 1]) + return x + + +@pytest.mark.parametrize('transposed', [False, True]) +@pytest.mark.parametrize('hbc', [True, False]) +@pytest.mark.parametrize('degree', [[2, 2], [2, 3]]) +@pytest.mark.parametrize('ncells', [[8, 10], [15, 12]]) +@pytest.mark.parametrize('R', [1]) +@pytest.mark.parametrize('Projector', [C0PolarProjection_V0, C1PolarProjection_V0]) +@pytest.mark.mpi + +def test_PolarProjection_V0(Projector, R, ncells, degree, hbc, transposed, verbose=False): + mpi_comm = MPI.COMM_WORLD + domain = get_domain(R) + + # Discrete physical domain and discrete space + domain_h = discretize(domain, ncells=ncells, periodic=[False, True], comm=mpi_comm) + V0 = ScalarFunctionSpace('V0', domain) + V0_h = discretize(V0, domain_h, degree=degree) + + P0 = Projector(V0_h, hbc=hbc, transposed=transposed) + + x = get_random_vector(V0_h) + phiC = FemField(V0_h) + y = phiC.coeffs + P0.dot(x, out=y) + + # Checking projection property P0(P0(phi)) = P0(phi) + assert np.allclose(P0.dot(y)[:, :], y[:, :], atol=ATOL, rtol=RTOL) + + sp_P0 = P0.tosparse().tocoo() + if verbose: + print("Projection P0 in matrix form on rank", mpi_comm.rank) + print(sp_P0.toarray()) + + + [n1, n2] = V0_h.coeff_space.npts + assert sp_P0.shape == (n1 * n2, n1 * n2) + + # gather sparse matrix entries (rows, columns, data) on root process + data = gather_vlen_array(sp_P0.data, mpi_comm) + cols_rows = gather_vlen_arrays((sp_P0.col, sp_P0.row), mpi_comm) + + if mpi_comm.rank == 0: + + cols = cols_rows[0] + rows = cols_rows[1] + + # Check the number of non-zero entries in the sparse matrix + if Projector == C0PolarProjection_V0: + assert len(data) == n2 * n2 + n2 * (n1 - (2 if hbc else 1)) + else: + assert len(data) == 3 * n2 * n2 + n2 * (n1 - (3 if hbc else 2)) + + if transposed: + rows, cols = cols, rows + + # Check all non-zero entries + for i, j, v in zip(rows, cols, data): + if i < n2 and j < n2: + assert np.isclose(v, 1 / n2, atol=ATOL, rtol=RTOL) + elif Projector == C1PolarProjection_V0 and j < n2 <= i < 2 * n2: + assert np.isclose(v, 1 / n2, atol=ATOL, rtol=RTOL) + elif Projector == C1PolarProjection_V0 and n2 <= i < 2 * n2 and n2 <= j < 2 * n2: + assert np.isclose(v, 2 / n2 * np.cos( (i - j) * 2 * np.pi / n2), atol=ATOL, rtol=RTOL) + else: + # identity + assert i == j + assert np.isclose(v, 1.0, atol=ATOL, rtol=RTOL) + + + # Comparing results of dot and tosparse + x_global = mpi_comm.allreduce(x.toarray(), op=MPI.SUM) + y_sp = sp_P0 @ x_global + y = mpi_comm.allreduce(y.toarray(), op=MPI.SUM) + y_sp = mpi_comm.allreduce(y_sp, op=MPI.SUM) + + assert np.allclose(y_sp, y, atol=ATOL, rtol=RTOL) + + + +@pytest.mark.parametrize('transposed', [False, True]) +@pytest.mark.parametrize('hbc', [True, False]) +@pytest.mark.parametrize('degree', [[2, 2], [2, 3]]) +@pytest.mark.parametrize('ncells', [[8, 10], [15, 12]]) +@pytest.mark.parametrize('R', [1]) +@pytest.mark.parametrize('Projector', [C0PolarProjection_V1, C1PolarProjection_V1]) +@pytest.mark.mpi + +def test_PolarProjection_V1(Projector, R, ncells, degree, hbc, transposed, verbose=False): + mpi_comm = MPI.COMM_WORLD + domain = get_domain(R) + + # Discrete physical domain and discrete space + domain_h = discretize(domain, ncells=ncells, periodic=[False, True], comm=mpi_comm) + V1 = VectorFunctionSpace('V1', domain, kind='hcurl') + V1_h = discretize(V1, domain_h, degree=degree) + + P1 = Projector(V1_h, hbc=hbc) + if transposed: + P1 = P1.T + + x = get_random_block_vector(V1_h) + y = BlockVector(V1_h.coeff_space) + P1.dot(x, out=y) + z = P1.dot(y) + # Checking projection property P1(P1(phi)) = P1(phi) + assert np.allclose(z[0][:, :], y[0][:, :], atol=ATOL, rtol=RTOL) + assert np.allclose(z[1][:, :], y[1][:, :], atol=ATOL, rtol=RTOL) + + sp_P1 = P1.tosparse().tocoo() + if verbose: + print("Projection P1 in matrix form on rank", mpi_comm.rank) + print(sp_P1.toarray()) + + [n01, n02] = V1_h.coeff_space[0].npts + [n11, n12] = V1_h.coeff_space[1].npts + + assert sp_P1.shape == (n01 * n02 + n11 * n12, n01 * n02 + n11 * n12) + + # gather sparse matrix entries (rows, columns, data) on root process + data = gather_vlen_array(sp_P1.data, mpi_comm) + cols_rows = gather_vlen_arrays((sp_P1.col, sp_P1.row), mpi_comm) + + if mpi_comm.rank == 0: + + cols = cols_rows[0] + rows = cols_rows[1] + + # Check the number of non-zero entries in the sparse matrix + if Projector == C0PolarProjection_V1: + assert len(data) == n01 * n02 + 2 * n02 + (n11 - (3 if hbc else 2)) * n12 + else: + assert len(data) == 3 * n02 * n02 + n02 * (n01 - 1) + (n11 - (3 if hbc else 2)) * n12 + + if transposed: + rows, cols = cols, rows + + # Check all non-zero entries + for i, j, v in zip(rows, cols, data): + if Projector == C0PolarProjection_V1: + # Block P1_00 + if i < n01 * n02: + # identity + assert i == j + assert np.isclose(v, 1.0, atol=ATOL, rtol=RTOL) + # Block P1_10 + elif n01 * n02 + n02 <= i < n01 * n02 + 2 * n02: + # d matrix + if j == i - n02 * (n01 + 1): + assert np.isclose(v, -1.0, atol=ATOL, rtol=RTOL) + else: + assert np.isclose(v, 1.0, atol=ATOL, rtol=RTOL) + assert j == (i - n02 * (n01 + 1) + 1) % n02 + # Block P1_11 + else: + assert i >= n01 * n02 + 2 * n02 + assert j == i + assert np.isclose(v, 1.0, atol=ATOL, rtol=RTOL) + + if Projector == C1PolarProjection_V1: + p_ij = 2 / n02 * np.cos((i - j) * 2 * np.pi / n02) + # Block P1_00 + if i < n02: + # matrix p + assert j < n02 + assert np.isclose(v, p_ij, atol=ATOL, rtol=RTOL) + elif 2 * n02 > i >= n02 > j: + # matrix I - p + if i == j + n02: + assert np.isclose(v, 1 - p_ij, atol=ATOL, rtol=RTOL) + else: + assert np.isclose(v, - p_ij, atol=ATOL, rtol=RTOL) + elif n02 <= i < n01 * n02: + assert j == i + assert np.isclose(v, 1.0, atol=ATOL, rtol=RTOL) + + # Block P1_10 + elif j < n02: + # matrix q + q_ij = (2 / n02 * np.cos((i + 1 - j) * 2 * np.pi / n02) - p_ij) + assert np.isclose(v, q_ij, atol=ATOL, rtol=RTOL) + + # Block P1_11 + else: + assert i == j + assert i >= n01 * n02 + 2 * n02 + assert np.isclose(v, 1.0, atol=ATOL, rtol=RTOL) + + + # Comparing results of dot and tosparse + x_global = mpi_comm.allreduce(x.toarray(), op=MPI.SUM) + y_sp = sp_P1 @ x_global + y = mpi_comm.allreduce(y.toarray(), op=MPI.SUM) + y_sp = mpi_comm.allreduce(y_sp, op=MPI.SUM) + + assert np.allclose(y_sp, y, atol=ATOL, rtol=RTOL) + + +@pytest.mark.parametrize( 'transposed', [True, False]) +@pytest.mark.parametrize('degree', [[2, 2], [2, 3]]) +@pytest.mark.parametrize('ncells', [[8, 10], [15, 12]]) +@pytest.mark.parametrize( 'R', [1]) +@pytest.mark.mpi + +def test_PolarProjection_V2(R, ncells, degree, transposed, verbose=False): + mpi_comm = MPI.COMM_WORLD + domain = get_domain(R) + + # Discrete physical domain and discrete space + domain_h = discretize(domain, ncells=ncells, periodic=[False, True], comm=mpi_comm) + V2 = ScalarFunctionSpace('V2', domain) + V2_h = discretize(V2, domain_h, degree=degree) + + P2 = C0PolarProjection_V2(V2_h, transposed=transposed) + + x = get_random_vector(V2_h) + phiC = FemField(V2_h) + y = phiC.coeffs + P2.dot(x, out=y) + + # Checking projection property P0(P0(phi)) = P0(phi) + assert np.allclose(P2.dot(y)[:, :], y[:, :], atol=ATOL, rtol=RTOL) + + # Comparing the global sparse matrix to reference file + sp_P2 = P2.tosparse().tocoo() + if verbose: + print("Projection P2 in matrix form on rank", mpi_comm.rank) + print(sp_P2.toarray()) + + [n1, n2] = V2_h.coeff_space.npts + assert sp_P2.shape == (n1 * n2, n1 * n2) + + data = gather_vlen_array(sp_P2.data, mpi_comm) + cols_rows = gather_vlen_arrays((sp_P2.col, sp_P2.row), mpi_comm) + + if mpi_comm.rank == 0: + + cols = cols_rows[0] + rows = cols_rows[1] + + # Check the number of non-zero entries in the sparse matrix + assert len(data) == n1 * n2 + + if transposed: + rows, cols = cols, rows + + # Check all non-zero entries + for i, j, v in zip(rows, cols, data): + assert np.isclose(v, 1.0, atol=ATOL, rtol=RTOL) + if j < n2: + assert i == j + n2 + else: + assert i == j + + + # Comparing results of dot and tosparse + x_global = mpi_comm.allreduce(x.toarray(), op=MPI.SUM) + y_sp = sp_P2 @ x_global + + assert np.allclose(y_sp, y.toarray(), atol=ATOL, rtol=RTOL) + +if __name__ == '__main__': + test_PolarProjection_V0(C0PolarProjection_V0, 1, [3, 3], [1, 1], False, False, verbose=True) + test_PolarProjection_V1(C0PolarProjection_V1, 1, [2, 3], [1, 1], False, False, verbose=True) + test_PolarProjection_V2(1, [8, 10], [2, 2], False, verbose=True) diff --git a/psydac/utilities/gather_variable_len_arrays.py b/psydac/utilities/gather_variable_len_arrays.py new file mode 100644 index 000000000..1f2e2635d --- /dev/null +++ b/psydac/utilities/gather_variable_len_arrays.py @@ -0,0 +1,78 @@ +import numpy as np +from mpi4py.util.dtlib import from_numpy_dtype + +def gather_vlen_arrays(arrays, mpi_comm, mpi_root=0): + + """ + Gather several 1D NumPy arrays of the same local length onto root process. Length may vary between processes. + + Parameters + ---------- + arrays : iterable of 1D arrays + All local arrays must have the same length and dtype + All ranks must pass the same number of arrays + mpi_comm : MPI communicator + mpi_root : int, default=0 + + Returns + ------- + On root: + tuple of gathered 1D arrays + On other ranks: + None + """ + + rank = mpi_comm.rank + arrays = tuple(np.asarray(a) for a in arrays) + dtype = arrays[0].dtype + arr_len = arrays[0].size + + for a in arrays: + if a.dtype != dtype: + raise TypeError("all arrays must have same dtype") + if a.size != arr_len: + raise ValueError("all arrays must have same size") + + # reshape as (arr_len, num of arrays) and flatten + arrays_flat = np.ascontiguousarray(np.column_stack(arrays)).ravel(order="C") + + local_len = np.array(arrays_flat.size, dtype=np.int64) + + # Gather local flattened array sizes on root process + recvcounts = np.empty(mpi_comm.size, dtype=np.int64) if rank == mpi_root else None + mpi_comm.Gather(local_len, recvcounts, root=mpi_root) + + if rank == mpi_root: + displs = np.empty(mpi_comm.size, dtype=np.int64) + displs[0] = 0 + displs[1:] = np.cumsum(recvcounts[:-1]) + + total_len = int(np.sum(recvcounts)) + global_array = np.empty(total_len, dtype=dtype) + + else: + displs = None + global_array = None + mpi_type = from_numpy_dtype(dtype) + + mpi_comm.Gatherv(arrays_flat, [global_array, recvcounts, displs, mpi_type], root=mpi_root) + + if rank != mpi_root: + return None + + num_arr = len(arrays) + gathered = global_array.reshape(-1, num_arr, order="C") + return tuple(gathered[:, j].copy() for j in range(num_arr)) + +#=============================================================================== +def gather_vlen_array(v, mpi_comm, mpi_root=0): + + """ Gather 1D arrays of possibly different lengths onto root process + Other processes return None + Local arrays must be of the same type + """ + + result = gather_vlen_arrays((v,), mpi_comm, mpi_root=mpi_root) + if mpi_comm.rank == mpi_root: + return result[0] + return None diff --git a/pyproject.toml b/pyproject.toml index 2a47c1260..1240e8807 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -31,7 +31,7 @@ dependencies = [ 'termcolor', # Our packages from PyPi - 'sympde == 0.19.2', + 'sympde @ https://github.com/pyccel/sympde/archive/refs/heads/fix_check_linearity.zip', 'pyccel >= 2.2.3', 'gelato == 0.12',