diff --git a/.gitignore b/.gitignore index 6d8aa6d8e..e30d36a73 100644 --- a/.gitignore +++ b/.gitignore @@ -115,3 +115,7 @@ pyvenv.cfg *kernels.txt struphy.log struphy.log.* +polarfeec.jl/ + +# output of the polar-splines tutorial (spectrum explorer) +tutorials/spectrum_explorer/ diff --git a/CUDA_STRATEGY.md b/CUDA_STRATEGY.md index c2c09d408..ab91cfed7 100644 --- a/CUDA_STRATEGY.md +++ b/CUDA_STRATEGY.md @@ -714,8 +714,8 @@ Every mapping runs on the GPU: the spline mappings (`kind_map` 0–2) join the a NumPy backend, and the control points are copied to the active backend once. Found and tested without a GPU with cunumpy's fake CuPy (`CUNUMPY_FAKE_CUPY=1`), which rejects host/device mixing like CuPy. - **Still host-only.** `EQDSKequilibrium` cannot be created on CuPy (a `Tokamak` on CuPy takes an equilibrium created - on the NumPy backend, or builds its default one there; fixed in #696). Polar splines: `PolarExtractionBlocksC1` builds SciPy sparse - matrices from the control points, and the polar extraction operators apply them to the stencil data, which lives on + on the NumPy backend, or builds its default one there; fixed in #696). Polar splines: the polar extraction blocks (then `PolarExtractionBlocksC1`) are SciPy sparse + matrices, and the polar extraction operators apply them to the stencil data, which lives on the device on CuPy; `Derham` raised `NotImplementedError` for `polar_splines=True` on CuPy (resolved in #695, see [Polar splines on CuPy](#polar-splines-on-cupy-695)). - **Device helpers.** `b_splines_slim`, `b_der_splines_slim` (`bsplines/bsplines_kernels.cuh`), `evaluation_kernel_2d` @@ -828,8 +828,8 @@ the [porting order](#porting-order); part of #688). ## Polar splines on CuPy (#695) -- **Blocks on the host.** `PolarExtractionBlocksC1` builds every block (extraction, DOF extraction, polar - derivatives) with NumPy/SciPy from a host copy of the control points (`xp.to_numpy(domain.cx)`), on every backend: +- **Blocks on the host.** `PolarExtractionBlocksHomotopy` builds every block (extraction, duals, polar + derivatives) with NumPy/SciPy from host evaluations of the domain, on every backend: they are setup data, and `.toarray()`, `.T` and the restart left inverse (`SplineFunction._restart_extraction_op`) keep working on them. - **Products on the device.** `PolarExtractionOperator` and `PolarLinearOperator` keep their `blocks_*` as host diff --git a/doc/sections/dev_guide.rst b/doc/sections/dev_guide.rst index cbe2941e2..024f4d819 100644 --- a/doc/sections/dev_guide.rst +++ b/doc/sections/dev_guide.rst @@ -54,6 +54,7 @@ This command is only available when Struphy is installed in editable mode, and i ../_collections/tutorials/dev_tutorial_feec_basics ../_collections/tutorials/dev_tutorial_feec_bcs ../_collections/tutorials/dev_tutorial_data_structs + ../_collections/tutorials/dev_tutorial_polar_homotopy .. toctree:: diff --git a/doc/sections/subsections/models-all.rst b/doc/sections/subsections/models-all.rst index ce68d4394..eb7ad2697 100644 --- a/doc/sections/subsections/models-all.rst +++ b/doc/sections/subsections/models-all.rst @@ -204,6 +204,11 @@ Pure fields :undoc-members: :show-inheritance: +.. autoclass:: struphy.models.ScalarWave + :members: doc_pde, doc_scalar_quantities, doc_normalization, doc_discretization, doc_long_description, doc_examples, doc_use_cases, doc_cannot_be_used_for + :undoc-members: + :show-inheritance: + Particle tracing ^^^^^^^^^^^^^^^^ diff --git a/doc/sections/subsections/numerics-geomFE.rst b/doc/sections/subsections/numerics-geomFE.rst index cece56401..9f9491eb6 100644 --- a/doc/sections/subsections/numerics-geomFE.rst +++ b/doc/sections/subsections/numerics-geomFE.rst @@ -266,15 +266,21 @@ See Fig. 5 for an example with :math:`p=4`. Polar splines ^^^^^^^^^^^^^ -For domains with a polar singularity at :math:`\eta_1 = 0` (e.g. ``IGAPolarCylinder``, ``IGAPolarTorus``), -``DerhamOptions(polar_splines=True)`` builds the C1 polar spline spaces of -`Holderied and Possanner (JCP 2022) `_. -The first two rings of tensor-product splines around the pole are replaced by a few polar basis functions, -via extraction operators :math:`\mathbb E` (basis) and :math:`\mathbb P` (degrees of freedom) that act in the -poloidal plane :math:`(\eta_1, \eta_2)` only. - -Hence, the inter-/histopolation matrices of the commuting projectors, :math:`\mathbb P \mathcal I \mathbb E^T`, -and the mass matrices with weights independent of :math:`\eta_3` have a 2D x 1D Kronecker structure +For domains with a polar singularity at :math:`\eta_1 = 0` (e.g. ``HollowCylinder(a1=0)``, ``IGAPolarCylinder``, +``IGAPolarTorus``), ``DerhamOptions(polar_splines=True, polar_regularity=r)`` builds polar spline spaces that are +:math:`C^{\mathfrak r}` at the pole (see :ref:`polar_splines_homotopy` for the construction). Regularity +:math:`\mathfrak r = 1` (the default) spans exactly the C1 polar spline spaces of +`Toshniwal et al. (CMAME 2017) `_ and +`Holderied and Possanner (JCP 2022) `_, which Struphy used before +(``src/struphy/polar/tests/test_polar_regularity_one.py`` compares all degrees with the former C1 extraction +matrices); unlike the former construction, it needs neither a spline mapping nor a multiple of 3 elements in +:math:`\eta_2`. The innermost rings of tensor-product splines around the pole are replaced by polar basis +functions, via extraction operators :math:`\mathbb E` (basis) and dual coefficients :math:`\tilde{\mathbb F}` that act +in the poloidal plane :math:`(\eta_1, \eta_2)` only. The commuting projectors are the tensor-product projections +followed by the polar duals, :math:`\bar\pi = \tilde{\mathbb F}^T \pi`, so that their polar inter-/histopolation +matrix is the identity. + +The mass matrices with weights independent of :math:`\eta_3` have a 2D x 1D Kronecker structure :math:`\mathbb A_\textnormal{pol} \otimes \mathbb A_3` (:mod:`struphy.polar.kron_2dx1d`). They are solved with a 1d direct solver in :math:`\eta_3` and one of the following solvers for the sparse poloidal matrix :math:`\mathbb A_\textnormal{pol}`: @@ -284,7 +290,6 @@ sparse poloidal matrix :math:`\mathbb A_\textnormal{pol}`: * ``"lu"``: sparse LU factorization, * ``"ilu"``: incomplete LU factorization, used as preconditioner of an iterative solver. -The solver of the projectors is set with ``DerhamOptions(polar_poloidal_solver=...)``. For weighted mass matrices, use :class:`~struphy.feec.preconditioner.PolarMassMatrixPreconditioner` (``precond="PolarMassMatrixPreconditioner"`` in the propagator options); it is exact for axisymmetric weights and mappings. @@ -305,11 +310,134 @@ and are included in the polar stiffness preconditioner (the average over :math:` On the CuPy backend, extraction, derivatives, Kronecker operators and these solves keep their coefficients on the device. Factorization and mass-preconditioner assembly run once on the host; cached device factors apply subsequent solves. Identical factors are shared across polar spaces -and projector/preconditioner instances of the same Derham. Matrix contents and solver settings +and preconditioner instances of the same Derham. Matrix contents and solver settings are matched exactly; updated weights use new factors, and unused entries are released. Polar MPI collectives support CUDA-aware device buffers or host staging through ``cunumpy.mpi.mpi_buffer``. The feectools halo exchanges require CUDA-aware MPI on real GPUs. ``SplineFunction.__call__`` dispatches CUDA kernels for marker and dense/sparse grid evaluation. The fixed local-grid helper ``eval_tp_fixed_loc`` remains host-only; distributed field-evaluation MPI handling and coverage are separate follow-ups. -.. image:: ../../pics/polar_derham.png \ No newline at end of file +.. image:: ../../pics/polar_derham.png + + +.. _polar_splines_homotopy: + +Construction of the polar splines (homotopy operators) +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +``DerhamOptions(polar_splines=True, polar_regularity=r)`` builds the polar subspaces of R. Remmerswaal (*High-order FEEC discretization on 2D and 3D domains with a disk-like coordinate +singularity*, draft, and the Julia package `PolarFEEC.jl `_). +For a regularity :math:`\mathfrak r` the 2D polar spaces of the cross-section are + +.. math:: + + \bar V^{0,\mathfrak r} = \pi^0 T^{\mathfrak r} \oplus V^{0,\mathfrak r}_\textnormal{ext}\,,\qquad + \bar V^{1,\mathfrak r} = \nabla \pi^0 T^{\mathfrak r} \oplus \textnormal{curl}^\dagger \tilde\pi^{2,1} T^{\mathfrak r-2} + \oplus V^{1,\mathfrak r}_\textnormal{ext}\,,\qquad + \bar V^{2,\mathfrak r} = \tilde\pi^{2,1} T^{\mathfrak r-2} \oplus V^{2,\mathfrak r}_\textnormal{ext}\,, + +where :math:`T^{\mathfrak r}` are the Cartesian polynomials of total degree :math:`\mathfrak r`, :math:`\pi^k` the +commuting tensor-product projections, :math:`\tilde\pi^{2,1}` a conforming 2-form projection and +:math:`\textnormal{curl}^\dagger` the Poincaré homotopy operator (a radial integration in logical coordinates). +The external spaces consist of the tensor-product splines of the rings :math:`i > \mathfrak r` (resp. the +corresponding rings of the 1- and 2-forms). The spaces are conforming for :math:`\mathfrak r \geq 1`, +:math:`C^{\mathfrak r}`, :math:`C^{\mathfrak r-1}` and :math:`C^{\mathfrak r-2}` at the pole (exactly when the mapping +is :math:`\mathfrak r`-compatible with the logical spaces), they form a cochain and a chain complex with the +cohomology of the disk, and the commuting projections :math:`\bar\pi^k = \bar\Pi^k \circ \pi^k` commute with the +derivatives and with the homotopies. In 3D the polar complex is the total complex +:math:`\bar W^m = \bigoplus_{p+q=m} \bar V^{p,\mathfrak r} \otimes V^q_3` with the 1D complex in :math:`\eta_3`; +it has the cohomology of the solid torus. + +The polar spaces are built by :class:`~struphy.polar.homotopy.complex2d.PolarComplex2D` (extraction matrices +:math:`E` and duals :math:`F` with :math:`F^T E = I`, homotopies and cohomology projections) from the 1D spline +spaces of the ``Derham`` complex itself (:class:`~struphy.polar.homotopy.derham_spaces.DerhamSpace1D`: feectools +basis values, the derivative ``grad_1d_matrix`` and the projection grids of ``get_pts_and_wts``), in Struphy's +coefficient ordering (:mod:`struphy.polar.homotopy.layout`) and on the cross-section of the Struphy domain +(:class:`~struphy.polar.homotopy.derham_blocks.DomainCrossSection`). The central polynomials are scaled to the size +of the polar rings, which keeps the polar matrices as well conditioned as the tensor-product ones. +:class:`~struphy.polar.homotopy.derham_blocks.PolarExtractionBlocksHomotopy` provides the blocks of the polar +operators of :mod:`struphy.polar.linear_operators`: + +* per tensor-product component the first ``n_rings`` rings hold the polar coefficients: + :math:`\mathfrak r + 2` rings for 0-forms, :math:`(\mathfrak r + 1, \mathfrak r + 2)` for the two components of + 2D 1-forms and :math:`\mathfrak r + 1` for 2D 2-forms, so that the extraction and the duals are the identity on + the outer rings and the derivatives map no tensor-product coefficient into the polar coefficients; +* the duals are those of the paper with the non-conforming ring-0 :math:`\eta_2`-component removed; on the + coefficients of physical fields both coincide; +* the commuting projectors apply the tensor-product projection followed by the polar duals, so that the polar + inter-/histopolation matrix is the identity (no shift of the projection points off the pole is needed); +* the polynomials are projected on the reference cross-section ``polar_cross_section``: the (planar) domain at + :math:`\eta_3 = 0` (default) or the unit disk; for regularity :math:`\mathfrak r \geq 2` a warning is given if the + cross-sections change shape along :math:`\eta_3` (the splines are then only approximately + :math:`C^{\mathfrak r}` away from :math:`\eta_3 = 0`). + +The options ``polar_enrichment`` (``"cochain_chain"``, ``"cochain"``, ``"none"``) and +``polar_polynomial_projection`` (``"commuting"``, ``"l2"``) select the variants of the paper. The requirements are +:math:`\mathfrak r \leq p_1` and ``num_elements[1]`` :math:`\geq 2\mathfrak r + 1`; any domain with a pole can be +used, e.g. ``HollowCylinder(a1=0)``, ``HollowTorus(a1=0)``, ``IGAPolarTorus`` and the analytic +``PerturbedDiskCylinder`` and ``RoundedSquareCylinder``. See the tutorials ``tutorial_polar_splines`` and +``dev_tutorial_polar_homotopy``. + +**Trigonometric functions in** :math:`\eta_2`. With ``DerhamOptions(angular_basis="trig")`` the :math:`\eta_2`-direction +uses the ``n = num_elements[1]`` (odd) nodal trigonometric functions of :mod:`struphy.feec.trig_space` instead of +periodic B-splines: the 0-forms are the cardinal functions of the nodes :math:`\eta_2 = i/n`, the 1-forms their cell +histopolants, so that the derivative is the forward difference and the inter- and histopolation matrices are the +identity. These are the compatible angular spaces of the paper: the polar splines are then exactly +:math:`C^{\mathfrak r}` at the pole for every :math:`\mathfrak r \leq (n - 1)/2`. Mass matrices, global projectors, +derivatives, field evaluation and the polar splines support them; :math:`\eta_2` must not be MPI-decomposed, and +particles, local projectors and ``BasisProjectionOperator`` need B-splines in every direction (a clear error is +raised). + +======================================================================== ========================================================== +Paper / PolarFEEC.jl Struphy +======================================================================== ========================================================== +1D spaces, scaled spaces, derivative and histopolation matrices :class:`~struphy.polar.homotopy.derham_spaces.DerhamSpace1D` (the 1D spaces of ``Derham``) +trigonometric spaces :mod:`struphy.feec.trig_space` (``angular_basis="trig"``) +tensor-product and physical form spaces, Piola maps, GDOFs :mod:`struphy.polar.homotopy.forms2d` +disk-like mappings Struphy domains (:class:`~struphy.polar.homotopy.derham_blocks.DomainCrossSection`) +Zernike basis, centre-shifted polynomial basis :mod:`struphy.polar.homotopy.polynomials` +``PolarComplex``: extraction, duals, :math:`\tilde\Pi^{2,\mathfrak r}`, :class:`~struphy.polar.homotopy.complex2d.PolarComplex2D` +homotopies, boundary homotopies, mean projection :math:`\mathfrak B` +``stiffness_solver``, ``range_projection`` (2D curl-curl "free lunch") :mod:`struphy.polar.homotopy.solvers` +``TensorProductComplex``, ``PeriodicComplex``, ``PolarComplex3D`` :mod:`struphy.polar.homotopy.complex3d` +3D polar complex of a domain ``DerhamOptions(polar_splines=True)`` +======================================================================== ========================================================== + +The studies of PolarFEEC.jl and of the NMPP seminar talk (2026) are reproduced by the demos +:mod:`struphy.polar.homotopy.demos` (used by the tutorials), the tests ``src/struphy/polar/tests/test_homotopy_*.py`` +and the examples: + +================================================================== ================================================================== +Study Struphy +================================================================== ================================================================== +1D spaces, forms on the square and on a mapped domain (jl 01-03) ``demos.spaces`` +polar bases, regularity at the centre (talk; jl 04) ``demos.basis``, ``demos.regularity`` +convergence and commutativity of the polar projections (talk) ``demos.projections`` +Poisson and curl-curl on a rounded square (jl 05) ``demos.pde``; ``examples/Poisson/polar_rounded_square``, + ``examples/CurlCurl/polar_rounded_square`` +Maxwell spectrum on the disk (jl 06) ``demos.maxwell`` +field lines in a twisted tube, :math:`\iota(r)` (jl 07) ``demos.fieldlines`` with the domain ``HelicalTube("islands")`` +assembly cost of the polar masses (jl 08) ``demos.performance`` +0- and 1-form wave equations, regularity and convergence (talk) ``demos.waves``; model ``ScalarWave`` and ``Maxwell`` with + ``algo="leapfrog"``; ``examples/ScalarWave/gaussian_perturbed_disk``, + ``examples/Maxwell/gaussian_perturbed_disk`` +3D 1-form wave in a helical tube (talk) ``demos.waves3d``; ``examples/Maxwell/gaussian_helical_tube`` + (``HelicalTube("stellarator")``) +guiding-centre orbits near the magnetic axis (talk) ``demos.orbits``; ``examples/GuidingCenter/vacuum_near_axis`` with + the equilibrium ``HelicalVacuumField`` +spectrum explorer of the curl-curl operator (talk) ``demos.spectrum`` (``write_explorer(folder)``: an HTML page) +de Rham diagrams, coordinate transformation, boundary homotopy ``demos.schematics`` (diagrams with ``tikzfigure``) +all figures and videos of the talk ``python -m struphy.polar.homotopy.demos.talk_figures OUT`` + (``--full``: the talk's resolution, hours; ``--videos``) +================================================================== ================================================================== + +The wave and guiding-centre scripts of the talk were not available; these studies follow the methods stated in the +talk and in the wave script of the PolarFEEC.jl history (staggered leapfrog at :math:`0.9` times the stability limit +:math:`2/\sqrt{\lambda_\textnormal{max}}`, Gaussian initial data next to the axis, maximum relative +:math:`L^\infty` error during the final 10%, RK4 orbits with the relative energy error +:math:`\max_t |E(t) - E(0)| / (\mu |B|)`). The reference solutions (the finest mesh of the highest regularity), the +sampling of the :math:`L^\infty` norm, the final time :math:`T = 1`, the helical-tube geometries and the vacuum field +(a single helical harmonic :math:`B = B_0 \nabla(z + \delta\, I_m(kr)\sin(m\theta - kz)/k)`) are choices of Struphy. +With these choices the orbit study does not show the strong dependence of the energy error on :math:`\mathfrak r` +reported in the talk; the spectrum explorer, the regularity study and the projection studies reproduce the talk +(the eigenvalues and dimensions of the explorer agree with the talk's data to :math:`10^{-13}`). \ No newline at end of file diff --git a/doc/sections/subsections/propagators-avail.rst b/doc/sections/subsections/propagators-avail.rst index 829e82423..bc6def472 100644 --- a/doc/sections/subsections/propagators-avail.rst +++ b/doc/sections/subsections/propagators-avail.rst @@ -26,6 +26,11 @@ Field solvers :exclude-members: options, allocate :show-inheritance: +.. automodule:: struphy.propagators.scalar_wave_leapfrog + :members: + :exclude-members: options, allocate + :show-inheritance: + Linear MHD solvers ^^^^^^^^^^^^^^^^^^ diff --git a/doc/sections/tutorials.rst b/doc/sections/tutorials.rst index 09a21e2b2..407a22ca9 100644 --- a/doc/sections/tutorials.rst +++ b/doc/sections/tutorials.rst @@ -52,6 +52,8 @@ It is recommended to use the same Python environment as for Struphy, e.g., by in ../_collections/tutorials/tutorial_domains ../_collections/tutorials/tutorial_mhd_equilibria + ../_collections/tutorials/tutorial_polar_derham + ../_collections/tutorials/tutorial_polar_splines .. toctree:: diff --git a/examples/CurlCurl/gaussian_on_axis_cylinder/params_gaussian_on_axis_cylinder.py b/examples/CurlCurl/gaussian_on_axis_cylinder/params_gaussian_on_axis_cylinder.py index e53020185..6c890bd49 100644 --- a/examples/CurlCurl/gaussian_on_axis_cylinder/params_gaussian_on_axis_cylinder.py +++ b/examples/CurlCurl/gaussian_on_axis_cylinder/params_gaussian_on_axis_cylinder.py @@ -49,7 +49,7 @@ "--poloidal-solver", default="schur", choices=["schur", "lu", "ilu"], - help="Poloidal solver of the polar preconditioners and projectors.", + help="Poloidal solver of the polar preconditioners.", ) parser.add_argument("--num-elements", type=int, nargs=3, default=[64, 66, 64], help="Number of elements per direction.") parser.add_argument( @@ -173,7 +173,7 @@ # Time stepping time_opts = Time() -# Grid (polar splines need a multiple of 3 elements in eta2) +# Grid num_elements = tuple(args.num_elements) degree = (3, 3, 3) grid = grids.TensorProductGrid(num_elements=num_elements, mpi_dims_mask=(True, True, True)) @@ -188,7 +188,6 @@ degree=degree, bcs=(("free", "dirichlet"), None, None), polar_splines=True, - polar_poloidal_solver=args.poloidal_solver, ) else: a1 = 1e-2 diff --git a/examples/CurlCurl/gaussian_on_axis_torus/params_gaussian_on_axis_torus.py b/examples/CurlCurl/gaussian_on_axis_torus/params_gaussian_on_axis_torus.py index 490e93f9c..f7bc7efcf 100644 --- a/examples/CurlCurl/gaussian_on_axis_torus/params_gaussian_on_axis_torus.py +++ b/examples/CurlCurl/gaussian_on_axis_torus/params_gaussian_on_axis_torus.py @@ -49,7 +49,7 @@ "--poloidal-solver", default="schur", choices=["schur", "lu", "ilu"], - help="Poloidal solver of the polar preconditioners and projectors.", + help="Poloidal solver of the polar preconditioners.", ) parser.add_argument("--num-elements", type=int, nargs=3, default=[64, 66, 64], help="Number of elements per direction.") parser.add_argument( @@ -173,7 +173,7 @@ # Time stepping time_opts = Time() -# Grid (polar splines need a multiple of 3 elements in eta2) +# Grid num_elements = tuple(args.num_elements) degree = (3, 3, 3) grid = grids.TensorProductGrid(num_elements=num_elements, mpi_dims_mask=(True, True, True)) @@ -190,7 +190,6 @@ degree=degree, bcs=(("free", "dirichlet"), None, None), polar_splines=True, - polar_poloidal_solver=args.poloidal_solver, ) else: a1 = 1e-2 diff --git a/examples/CurlCurl/polar_rounded_square/params_polar_rounded_square.py b/examples/CurlCurl/polar_rounded_square/params_polar_rounded_square.py new file mode 100644 index 000000000..acc3ee585 --- /dev/null +++ b/examples/CurlCurl/polar_rounded_square/params_polar_rounded_square.py @@ -0,0 +1,121 @@ +"""Curl-curl problem curl curl A = J on a rounded square with high-order (homotopy) polar splines of regularity r, the +Dirichlet case (A x n = 0) of PolarFEEC.jl example 05; the pproc reports the L2 error of curl A (rate p).""" + +import argparse +from pathlib import Path + +parser = argparse.ArgumentParser(add_help=__name__ == "__main__") +parser.add_argument( + "--out-folder", + default=str(Path(__file__).resolve().parent), + help="Folder in which the output folder is created (default: beside this file).", +) +parser.add_argument("--sim-folder", default="sim_1", help="Name of the output folder.") +parser.add_argument("--degree", type=int, default=3, help="Spline degree p in eta1 and eta2.") +parser.add_argument("--regularity", type=int, default=None, help="Regularity r of the polar splines (default: p).") +parser.add_argument("--cells", type=int, default=6, help="Number n of radial cells (6n angular cells).") +parser.add_argument("--sigma", type=float, default=1e-10, help="Coefficient of the stabilizing mass term.") +parser.add_argument("--test-run", action="store_true", help="Small resolution (n = 4, p = 2), for the tests.") +args, _ = parser.parse_known_args() +if args.test_run: + args.cells, args.degree = 4, 2 +if args.regularity is None: + args.regularity = args.degree + +name = f"Curl-curl on a rounded square, homotopy polar splines (p = {args.degree}, r = {args.regularity})" +description = f""" +Curl-curl problem (curl Phi, curl A) + sigma (Phi, A) = (Phi, J) for all Phi in H0(curl) on the rounded square +g > 0, g = 1 - x^2 - y^2 - eps Re((x + i y)^4), eps = 0.15, with A x n = 0 on the boundary and high-order polar +splines of regularity r = {args.regularity} at the centre. + +Physics. The curl-curl problem determines the vector potential A of a current J (magnetostatics); A is unique only +up to gradients, so the gauge-independent field B = curl A is compared. J = curl* B = (d_y B, -d_x B) with the +manufactured B = sin(2x) cos(y) (odd in x, zero mean: the compatibility condition of the Dirichlet problem). The +domain is parametrized by a disk-like map with its pole at the centre of the cross-section; the polar splines of +R. Remmerswaal (homotopy construction) are conforming and C^(r-1) at the pole. In 2D the paper solves this problem +with one 2-form mass solve and the homotopy operator ("a free two-dimensional lunch", struphy.polar.homotopy); here +the stabilized stiffness system of the CurlCurl model is solved (sigma = {args.sigma}), which gives the same curl A. + +Setup. RoundedSquareCylinder(eps=0.15), {args.cells} x {6 * args.cells} x 1 cells, degree {args.degree}, +DerhamOptions(polar_splines=True, polar_regularity={args.regularity}), Dirichlet conditions at eta1 = 1. + +Expected result. The L2 error of curl A converges like h^p (PolarFEEC.jl example 05). +""" + +import logging + +import numpy as np + +from struphy import set_logging_level + +set_logging_level(logging.INFO) + +from struphy import BaseUnits, DerhamOptions, EnvironmentOptions, Simulation, Time, domains, grids +from struphy.models import CurlCurl + +model = CurlCurl(base_units=BaseUnits()) +model.em_fields.e_field.save_data = True + +env = EnvironmentOptions(out_folders=args.out_folder, sim_folder=args.sim_folder, restart=False) +grid = grids.TensorProductGrid(num_elements=(args.cells, 6 * args.cells, 1)) +EPS = 0.15 +domain = domains.RoundedSquareCylinder(a=1.0, eps=EPS, Lz=1.0) +derham_opts = DerhamOptions( + degree=(args.degree, args.degree, 1), + bcs=(("free", "dirichlet"), None, None), + polar_splines=True, + polar_regularity=args.regularity, + nquads=(args.degree + 3, args.degree + 3, 2), +) + +sim = Simulation( + model=model, + name=name, + description=description, + params_path=__file__, + env=env, + time_opts=Time(), + domain=domain, + equil=None, + grid=grid, + derham_opts=derham_opts, +) + +from struphy.linear_algebra.solver import SolverParameters + +model.propagators.curl_curl.options = model.propagators.curl_curl.Options( + sigma=args.sigma, + solver="pcg", + precond="StiffnessPreconditioner", + solver_params=SolverParameters(tol=1e-12, maxiter=5000, info=True, recycle=False), +) + + +def exact_b(x, y, z): + """B_z = curl A of the Dirichlet problem.""" + return np.sin(2 * x) * np.cos(y) + 0.0 * z + + +def current(x, y): + """J = (d_y B, -d_x B).""" + return -np.sin(2 * x) * np.sin(y), -2 * np.cos(2 * x) * np.cos(y) + + +def _current_logical(i): + """Logical 1-form component i of J: (DF^T J)_i.""" + + def component(e1, e2, e3): + x, y, _ = domain(e1, e2, e3) + DF = np.asarray(domain.jacobian(e1, e2, e3)).reshape((3, 3) + np.shape(x)) + jx, jy = current(x, y) + return DF[0, i] * jx + DF[1, i] * jy + + return component + + +# the load is integrated against the Dirichlet test functions (a source in V1 would lose its boundary coefficients) +model.propagators.curl_curl.j = tuple(_current_logical(i) for i in range(3)) + +if __name__ == "__main__": + # post-processing: pproc_polar_rounded_square.py + sim.run(one_time_step=True) diff --git a/examples/CurlCurl/polar_rounded_square/pproc_polar_rounded_square.py b/examples/CurlCurl/polar_rounded_square/pproc_polar_rounded_square.py new file mode 100644 index 000000000..5c51c8bd6 --- /dev/null +++ b/examples/CurlCurl/polar_rounded_square/pproc_polar_rounded_square.py @@ -0,0 +1,86 @@ +"""Post-process the CurlCurl polar_rounded_square example: L2 error of curl A (computed with the polar curl of the +saved A) and plots of curl A and its error. Run as ``python pproc_polar_rounded_square.py [sim_folder] [flags]``.""" + +import argparse +import json +import os +import sys +from pathlib import Path + +import numpy as np +from matplotlib import pyplot as plt + +from struphy import Output + +SCRIPT_DIR = Path(__file__).resolve().parent +DEFAULT_OUTPUT = SCRIPT_DIR / "sim_1" + +sys.path.insert(0, str(SCRIPT_DIR)) +from params_polar_rounded_square import args, exact_b, parser, sim # noqa: E402 + + +def gauss_points(cells, n): + x, w = np.polynomial.legendre.leggauss(n) + a = np.arange(cells)[:, None] / cells + return (a + (x[None, :] + 1) / (2 * cells)).ravel(), np.tile(w / (2 * cells), cells) + + +def curl_a(out): + """Spline function of curl A (the polar curl of the saved 1-form A).""" + a = out.spline_fields(t=-1)["em_fields"]["e_field"] + derham = out._spline_derham + b = derham.create_spline_function("curl_a", "Hdiv") + b.vector = derham.curl.dot(a.vector) + return b + + +def bz(b, e1, e2, e3): + """Physical B_z = b_3 L_z / det(DF) on the tensor grid e1 x e2 x e3 (single e3).""" + values = b(e1, e2, e3)[2].reshape(e1.size, e2.size) + detj = sim.domain.jacobian_det(e1, e2, e3).reshape(e1.size, e2.size) + return values * sim.domain.params["Lz"] / detj + + +def main(path_out=DEFAULT_OUTPUT, save_figs=True): + out = Output(path_out) + b = curl_a(out) + cells = sim.grid.num_elements + e1, w1 = gauss_points(cells[0], args.degree + 3) + e2, w2 = gauss_points(cells[1], args.degree + 3) + e3 = np.array([0.5]) + x, y, z = (c.reshape(e1.size, e2.size) for c in sim.domain(e1, e2, e3)) + detj = np.abs(sim.domain.jacobian_det(e1, e2, e3)).reshape(e1.size, e2.size) / sim.domain.params["Lz"] + weights = np.outer(w1, w2) * detj + error = np.sqrt(np.sum(weights * (bz(b, e1, e2, e3) - exact_b(x, y, z)) ** 2)) + print(f"L2 error of curl A: {error:.3e}") + + results_dir = os.path.join(out.path_out, "results") + os.makedirs(results_dir, exist_ok=True) + with open(os.path.join(results_dir, "errors.json"), "w") as f: + json.dump(dict(l2_error=float(error), cells=args.cells, degree=args.degree, regularity=args.regularity), f) + + p1, p2 = np.linspace(1e-3, 1, 81), np.linspace(0, 1, 241) + x, y, z = (c.reshape(p1.size, p2.size) for c in sim.domain(p1, p2, e3)) + values = bz(b, p1, p2, e3) + fig, axes = plt.subplots(1, 2, figsize=(10, 4.5)) + for ax, field, title in ((axes[0], values, r"$\mathrm{curl}\,A_h$"), (axes[1], values - exact_b(x, y, z), "error")): + mesh = ax.pcolormesh(x, y, field, shading="gouraud", cmap="RdBu_r") + plt.colorbar(mesh, ax=ax, shrink=0.8) + ax.set_aspect("equal") + ax.set_title(title) + fig.tight_layout() + if save_figs: + fig.savefig(os.path.join(results_dir, "curl_a.png")) + plt.close(fig) + else: + plt.show() + return error + + +if __name__ == "__main__": + _, other_flags = parser.parse_known_args() + pproc_parser = argparse.ArgumentParser(description="Post-process a saved simulation run.") + pproc_parser.add_argument("path_out", nargs="?", default=DEFAULT_OUTPUT, help="Simulation output folder.") + pproc_parser.add_argument("--show", action="store_true", help="Show the figures instead of saving them.") + pproc_args = pproc_parser.parse_args(other_flags) + main(pproc_args.path_out, save_figs=not pproc_args.show) diff --git a/examples/GuidingCenter/screw_pinch_equilibrium/params_screw_pinch_equilibrium.py b/examples/GuidingCenter/screw_pinch_equilibrium/params_screw_pinch_equilibrium.py index d71b4d832..f650e074d 100644 --- a/examples/GuidingCenter/screw_pinch_equilibrium/params_screw_pinch_equilibrium.py +++ b/examples/GuidingCenter/screw_pinch_equilibrium/params_screw_pinch_equilibrium.py @@ -153,7 +153,7 @@ time_opts = Time(dt=args.dt, Tend=args.tend, split_algo="LieTrotter") -# Grid (polar splines need a multiple of 3 elements in eta2); eta1 is not decomposed, such that the process +# Grid; eta1 is not decomposed, such that the process # with the pole has all rings of the polar splines and each process has at least degree elements per direction num_elements = tuple(args.num_elements) degree = (3, 3, 3) diff --git a/examples/GuidingCenter/tokamak_equilibrium/params_tokamak_equilibrium.py b/examples/GuidingCenter/tokamak_equilibrium/params_tokamak_equilibrium.py index a803b72f6..0531b6327 100644 --- a/examples/GuidingCenter/tokamak_equilibrium/params_tokamak_equilibrium.py +++ b/examples/GuidingCenter/tokamak_equilibrium/params_tokamak_equilibrium.py @@ -167,7 +167,7 @@ time_opts = Time(dt=args.dt, Tend=args.tend, split_algo="LieTrotter") -# Grid (polar splines need a multiple of 3 elements in eta2); eta1 is not decomposed, such that the process +# Grid; eta1 is not decomposed, such that the process # with the pole has all rings of the polar splines and each process has at least degree elements per direction num_elements = tuple(args.num_elements) degree = (3, 3, 3) diff --git a/examples/GuidingCenter/vacuum_near_axis/convergence_vacuum_near_axis.py b/examples/GuidingCenter/vacuum_near_axis/convergence_vacuum_near_axis.py new file mode 100644 index 000000000..f413a96e6 --- /dev/null +++ b/examples/GuidingCenter/vacuum_near_axis/convergence_vacuum_near_axis.py @@ -0,0 +1,63 @@ +"""Energy conservation of RK4 guiding-centre orbits near the axis against the time step, for polar splines of +regularity r = 0, ..., 4 (the orbit study of the polar FEEC talk): runs params_vacuum_near_axis.py for every (r, dt), +post-processes the runs and plots the maximum (over the eight orbits) of max_t |E(t) - E(0)| / (mu |B|) against dt, +with a dotted O(dt^4) reference. + +Run as ``python convergence_vacuum_near_axis.py [--regularities 0 1 2 3 4] [--dts 0.04 0.02 0.01 0.005 0.0025] +[--periods 20]``. The default (20 transits, dt down to 0.0025) is a long study. +""" + +import argparse +import json +import subprocess +import sys +from pathlib import Path + +import numpy as np +from matplotlib import pyplot as plt + +SCRIPT_DIR = Path(__file__).resolve().parent + + +def run(r, dt, periods, mpi): + folder = f"conv_r{r}_dt{dt:g}" + flags = ["--regularity", str(r), "--dt", str(dt), "--periods", str(periods), "--sim-folder", folder] + flags += ["--save-step", str(max(1, int(round(0.04 / dt))))] + command = [sys.executable, str(SCRIPT_DIR / "params_vacuum_near_axis.py"), *flags] + if mpi > 1: + command = ["mpirun", "-n", str(mpi), *command] + subprocess.run(command, check=True) + subprocess.run( + [sys.executable, str(SCRIPT_DIR / "pproc_vacuum_near_axis.py"), str(SCRIPT_DIR / folder), *flags], check=True + ) + with open(SCRIPT_DIR / folder / "results" / "energy_errors.json") as f: + return max(json.load(f)["errors"]) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--regularities", type=int, nargs="+", default=[0, 1, 2, 3, 4]) + parser.add_argument("--dts", type=float, nargs="+", default=[0.04, 0.02, 0.01, 0.005, 0.0025]) + parser.add_argument("--periods", type=float, default=20) + parser.add_argument("--mpi", type=int, default=1) + args = parser.parse_args() + + dts = np.array(args.dts) + table = {r: np.array([run(r, dt, args.periods, args.mpi) for dt in dts]) for r in args.regularities} + fig, ax = plt.subplots(figsize=(6, 5)) + for i, (r, errors) in enumerate(table.items()): + print(f"r = {r}: {errors}") + ax.loglog(dts, errors, "o-", color=f"C{i}", label=rf"$\mathfrak{{r}} = {r}$") + reference = table[max(table)] + ax.loglog(dts, reference[0] * (dts / dts[0]) ** 4, "k:", label=r"$O(\Delta t^4)$") + ax.set_xlabel(r"$\Delta t$") + ax.set_ylabel(r"$\max_t |E(t) - E(0)| / (\mu |B|)$") + ax.set_title("RK4 guiding-centre orbits near the axis") + ax.legend() + fig.tight_layout() + fig.savefig(SCRIPT_DIR / "convergence_vacuum_near_axis.png", dpi=150) + print("wrote", SCRIPT_DIR / "convergence_vacuum_near_axis.png") + + +if __name__ == "__main__": + main() diff --git a/examples/GuidingCenter/vacuum_near_axis/params_vacuum_near_axis.py b/examples/GuidingCenter/vacuum_near_axis/params_vacuum_near_axis.py new file mode 100644 index 000000000..9ee48e1d2 --- /dev/null +++ b/examples/GuidingCenter/vacuum_near_axis/params_vacuum_near_axis.py @@ -0,0 +1,140 @@ +"""Guiding-centre orbits near the magnetic axis of a helical vacuum field with high-order (homotopy) polar splines of +regularity r: energy conservation of the explicit RK4 pusher with the projected background (the orbit study of the +polar FEEC talk of R. Remmerswaal).""" + +import argparse +from pathlib import Path + +parser = argparse.ArgumentParser(add_help=__name__ == "__main__") +parser.add_argument( + "--out-folder", + default=str(Path(__file__).resolve().parent), + help="Folder in which the output folder is created (default: beside this file).", +) +parser.add_argument("--sim-folder", default="sim_1", help="Name of the output folder.") +parser.add_argument("--regularity", type=int, default=2, help="Regularity r of the polar splines at the axis.") +parser.add_argument("--degree", type=int, default=4, help="Spline degree in eta1 and eta2.") +parser.add_argument("--num-elements", type=int, nargs=3, default=[8, 12, 16], help="Number of elements.") +parser.add_argument("--dt", type=float, default=0.04, help="Time step of RK4.") +parser.add_argument("--periods", type=float, default=20, help="Number of toroidal periods (transits of the cylinder).") +parser.add_argument("--epsilon", type=float, default=0.01, help="Normalized cyclotron period (size of the drifts).") +parser.add_argument("--save-step", type=int, default=5, help="Save the orbits every save_step time steps.") +parser.add_argument("--test-run", action="store_true", help="Small resolution and a few time steps, for the tests.") +args, _ = parser.parse_known_args() +if args.test_run: + args.num_elements, args.degree, args.regularity, args.periods, args.save_step = [4, 9, 4], 3, 2, 0.02, 1 + +import numpy as np + +LZ = 2 * np.pi +A = 0.3 +tend = args.periods * LZ + +name = f"Guiding-centre orbits near the axis of a helical vacuum field, polar splines r = {args.regularity}" +description = f""" +Guiding-centre orbits near the magnetic axis of a straight helical vacuum field B = B0 grad(z + delta/k I_2(k r) +sin(2 theta - k z)) (current-free, elliptical rotating flux surfaces, three periods along the cylinder), traced with +the explicit RK4 pusher of the GuidingCenter model in the background projected into the polar spline spaces. + +Physics. The guiding centres move along the field lines and drift with the grad-B and curvature drifts; the energy +E = v_par^2 / 2 + mu |B| and the magnetic moment mu are invariants. The model evaluates the background (|B|, the unit +vector b, grad |B| and (curl b) . b) from its commuting projections into the spline spaces. Near the axis (the pole +of the polar coordinates of the cross-section), the regularity of these splines at the axis determines how smooth +the projected field is, and thereby how well RK4 conserves the energy of orbits that pass close to or through the +axis. The talk shows that the energy error of RK4 decreases like dt^4 only for sufficiently regular (and +r-compatible) splines. + +Setup. Cylinder of radius {A} and length 2 pi (HollowCylinder(a1=0)), {args.num_elements} elements, degree +{args.degree}, DerhamOptions(polar_splines=True, polar_regularity={args.regularity}), HelicalVacuumField(B0=1, +delta=0.5, m=2, k=3). Eight markers at distances 0.01 ... 0.04 from the axis, co- and counter-passing (v_par = +-1) +with mu = 1e-3, epsilon = {args.epsilon}. RK4 with dt = {args.dt} for {args.periods} transits (t = {tend:.4g}). + +Expected result. The pproc plots the orbits projected onto the cross-section (overview and close-up) and the +relative energy error max_t |E(t) - E(0)| / (mu |B|) of each orbit. For the dt-convergence of the talk run +convergence_vacuum_near_axis.py (dt = 0.04 ... 0.0025, r = 0 ... 4; a long study). + +Not stated in the talk and chosen here: the vacuum field, the geometry and the markers. + +References: +- R. Remmerswaal, High-order FEEC discretization on 2D and 3D domains with a disk-like coordinate singularity + (draft); talk at the NMPP seminar 2026; PolarFEEC.jl (https://gitlab.mpcdf.mpg.de/rwr/polarfeec.jl). +- R. G. Littlejohn, "Variational principles of guiding centre motion", J. Plasma Phys. 29, 111 (1983). +""" + +import logging + +from struphy import set_logging_level + +set_logging_level(logging.INFO) + +from struphy import ( + BaseUnits, + BoundaryParameters, + DerhamOptions, + EnvironmentOptions, + LoadingParameters, + SavingParameters, + Simulation, + Time, + WeightsParameters, + domains, + equils, + grids, + maxwellians, +) +from struphy.models import GuidingCenter +from struphy.ode.utils import ButcherTableau + +model = GuidingCenter(base_units=BaseUnits(), epsilon=args.epsilon) +model.kinetic_ions.var.save_data = True + +env = EnvironmentOptions(out_folders=args.out_folder, sim_folder=args.sim_folder, save_step=args.save_step) +time_opts = Time(dt=args.dt, Tend=tend, split_algo="LieTrotter") +grid = grids.TensorProductGrid(num_elements=tuple(args.num_elements), mpi_dims_mask=(False, True, True)) +domain = domains.HollowCylinder(a1=0.0, a2=A, Lz=LZ) +derham_opts = DerhamOptions( + degree=(args.degree, args.degree, 3), + bcs=(("free", "free"), None, None), + polar_splines=True, + polar_regularity=args.regularity, +) +equil = equils.HelicalVacuumField(B0=1.0, delta=0.5, m=2, k=3.0) + +sim = Simulation( + model=model, + name=name, + description=description, + params_path=__file__, + env=env, + time_opts=time_opts, + domain=domain, + equil=equil, + grid=grid, + derham_opts=derham_opts, +) + +# eight markers (eta1, eta2, eta3, v_par, mu) at distances 0.01 ... 0.04 from the axis +MU = 1e-3 +initial = tuple( + (r / A, angle, 0.0, sign, MU) + for r, angle in ((0.01, 0.0), (0.02, 0.125), (0.03, 0.25), (0.04, 0.375)) + for sign in (1.0, -1.0) +) + +model.kinetic_ions.set_markers( + loading_params=LoadingParameters(Np=len(initial), seed=1234, specific_markers=initial), + weights_params=WeightsParameters(), + boundary_params=BoundaryParameters(bc=("remove", "periodic", "periodic")), + saving_params=SavingParameters(n_markers=len(initial)), + bufsize=2.0, +) + +butcher = ButcherTableau(algo="rk4") +model.propagators.push_bxe.options = model.propagators.push_bxe.Options(algo="explicit", butcher=butcher) +model.propagators.push_parallel.options = model.propagators.push_parallel.Options(algo="explicit", butcher=butcher) + +model.kinetic_ions.var.add_background(maxwellians.GyroMaxwellian2D(n=(1.0, None), B0=1.0)) + +if __name__ == "__main__": + # post-processing: pproc_vacuum_near_axis.py + sim.run() diff --git a/examples/GuidingCenter/vacuum_near_axis/pproc_vacuum_near_axis.py b/examples/GuidingCenter/vacuum_near_axis/pproc_vacuum_near_axis.py new file mode 100644 index 000000000..7e7f01243 --- /dev/null +++ b/examples/GuidingCenter/vacuum_near_axis/pproc_vacuum_near_axis.py @@ -0,0 +1,91 @@ +"""Post-process the GuidingCenter vacuum_near_axis example: orbits projected onto the cross-section (overview and +close-up) and the relative energy error max_t |E(t) - E(0)| / (mu |B|) of each orbit, with E = v_par^2 / 2 + mu |B|_h +and |B|_h the projected field strength used by the pusher. Run as +``python pproc_vacuum_near_axis.py [sim_folder] [variant flags]``.""" + +import argparse +import json +import os +import sys +from pathlib import Path + +import numpy as np +from matplotlib import pyplot as plt + +from struphy import Output + +SCRIPT_DIR = Path(__file__).resolve().parent +DEFAULT_OUTPUT = SCRIPT_DIR / "sim_1" + +sys.path.insert(0, str(SCRIPT_DIR)) +from params_vacuum_near_axis import LZ, A, args, parser, sim # noqa: E402 + + +def projected_field_strength(): + """|B|_h: the commuting projection of |B| into the polar 0-forms (as in the GuidingCenter model).""" + from struphy.feec.psydac_derham import Derham + + derham = Derham(sim.grid, sim.derham_opts, domain=sim.domain) + sim.equil.domain = sim.domain + absB = derham.create_spline_function("absB", "H1") + absB.vector = derham.P0(sim.equil.absB0) + return absB + + +def logical(x, y, z): + """Logical coordinates of points in the cylinder HollowCylinder(a1=0, a2=A, Lz=LZ).""" + return np.hypot(x, y) / A, (np.arctan2(y, x) / (2 * np.pi)) % 1.0, (z / LZ) % 1.0 + + +def energy_errors(orbits, absB): + x, y, z = (np.asarray(orbits[c]) for c in ("x", "y", "z")) + v_par, mu = np.asarray(orbits["v_par"]), np.asarray(orbits["mu"]) + e1, e2, e3 = logical(x, y, z) + markers = np.column_stack([e1.ravel(), e2.ravel(), e3.ravel()]) + B = np.asarray(absB(markers)).reshape(x.shape) + energy = v_par**2 / 2 + mu * B + return np.max(np.abs(energy - energy[0]), axis=0) / (mu[0] * B[0]) + + +def main(path_out=DEFAULT_OUTPUT, save_figs=True): + out = Output(path_out) + orbits = out.evaluate("kinetic_ions/orbits") + results_dir = os.path.join(out.path_out, "results") + os.makedirs(results_dir, exist_ok=True) + + errors = energy_errors(orbits, projected_field_strength()) + print("relative energy errors max_t |E(t) - E(0)| / (mu |B|):", np.array2string(errors, precision=2)) + with open(os.path.join(results_dir, "energy_errors.json"), "w") as f: + json.dump(dict(errors=errors.tolist(), dt=args.dt, regularity=args.regularity), f) + + fig, axes = plt.subplots(1, 2, figsize=(10, 4.8)) + x, y = np.asarray(orbits["x"]), np.asarray(orbits["y"]) + for ax, extent in ((axes[0], A), (axes[1], 0.05)): + for i in range(x.shape[1]): + ax.plot(x[:, i], y[:, i], ".", ms=1, color=f"C{i % 10}") + ax.add_patch(plt.Circle((0, 0), A, color="k", fill=False)) + ax.set_xlim(-extent, extent) + ax.set_ylim(-extent, extent) + ax.set_aspect("equal") + ax.set_xlabel("$x$") + ax.set_ylabel("$y$") + axes[0].set_title("orbits projected onto a cross-section") + axes[1].set_title("close-up") + fig.suptitle(rf"guiding centres near the axis, $\mathfrak{{r}} = {args.regularity}$, $\Delta t = {args.dt}$") + fig.tight_layout() + + if save_figs: + fig.savefig(os.path.join(results_dir, "orbits.png")) + plt.close("all") + else: + plt.show() + return errors + + +if __name__ == "__main__": + _, other_flags = parser.parse_known_args() + pproc_parser = argparse.ArgumentParser(description="Post-process a saved simulation run.") + pproc_parser.add_argument("path_out", nargs="?", default=DEFAULT_OUTPUT, help="Simulation output folder.") + pproc_parser.add_argument("--show", action="store_true", help="Show the figures instead of saving them.") + pproc_args = pproc_parser.parse_args(other_flags) + main(pproc_args.path_out, save_figs=not pproc_args.show) diff --git a/examples/GyrokineticPoisson/gaussian_on_axis_cylinder/params_gaussian_on_axis_cylinder.py b/examples/GyrokineticPoisson/gaussian_on_axis_cylinder/params_gaussian_on_axis_cylinder.py index 1102d60b9..59d5d0991 100644 --- a/examples/GyrokineticPoisson/gaussian_on_axis_cylinder/params_gaussian_on_axis_cylinder.py +++ b/examples/GyrokineticPoisson/gaussian_on_axis_cylinder/params_gaussian_on_axis_cylinder.py @@ -48,7 +48,7 @@ "--poloidal-solver", default="schur", choices=["schur", "lu", "ilu"], - help="Poloidal solver of the polar preconditioners and projectors.", + help="Poloidal solver of the polar preconditioners.", ) parser.add_argument("--num-elements", type=int, nargs=3, default=[64, 66, 64], help="Number of elements per direction.") parser.add_argument( @@ -176,7 +176,7 @@ # Time stepping time_opts = Time() -# Grid (polar splines need a multiple of 3 elements in eta2) +# Grid num_elements = tuple(args.num_elements) degree = (3, 3, 3) grid = grids.TensorProductGrid(num_elements=num_elements, mpi_dims_mask=(True, True, True)) @@ -192,7 +192,6 @@ degree=degree, bcs=(("free", "dirichlet"), None, None), polar_splines=True, - polar_poloidal_solver=args.poloidal_solver, ) else: a1 = 1e-2 diff --git a/examples/GyrokineticPoisson/gaussian_on_axis_torus/params_gaussian_on_axis_torus.py b/examples/GyrokineticPoisson/gaussian_on_axis_torus/params_gaussian_on_axis_torus.py index cc6103ba3..6a940ce06 100644 --- a/examples/GyrokineticPoisson/gaussian_on_axis_torus/params_gaussian_on_axis_torus.py +++ b/examples/GyrokineticPoisson/gaussian_on_axis_torus/params_gaussian_on_axis_torus.py @@ -48,7 +48,7 @@ "--poloidal-solver", default="schur", choices=["schur", "lu", "ilu"], - help="Poloidal solver of the polar preconditioners and projectors.", + help="Poloidal solver of the polar preconditioners.", ) parser.add_argument("--num-elements", type=int, nargs=3, default=[64, 66, 64], help="Number of elements per direction.") parser.add_argument( @@ -177,7 +177,7 @@ # Time stepping time_opts = Time() -# Grid (polar splines need a multiple of 3 elements in eta2) +# Grid num_elements = tuple(args.num_elements) degree = (3, 3, 3) grid = grids.TensorProductGrid(num_elements=num_elements, mpi_dims_mask=(True, True, True)) @@ -194,7 +194,6 @@ degree=degree, bcs=(("free", "dirichlet"), None, None), polar_splines=True, - polar_poloidal_solver=args.poloidal_solver, ) else: a1 = 1e-2 diff --git a/examples/Maxwell/gaussian_helical_tube/params_gaussian_helical_tube.py b/examples/Maxwell/gaussian_helical_tube/params_gaussian_helical_tube.py new file mode 100644 index 000000000..fee3ef259 --- /dev/null +++ b/examples/Maxwell/gaussian_helical_tube/params_gaussian_helical_tube.py @@ -0,0 +1,150 @@ +"""3D 1-form wave equation (Maxwell) in a stellarator-like tube with high-order (homotopy) polar splines: a localized +magnetic pulse next to the magnetic axis (the 3D wave study of the polar FEEC talk of R. Remmerswaal).""" + +import argparse +from pathlib import Path + +parser = argparse.ArgumentParser(add_help=__name__ == "__main__") +parser.add_argument( + "--out-folder", + default=str(Path(__file__).resolve().parent), + help="Folder in which the output folder is created (default: beside this file).", +) +parser.add_argument("--sim-folder", default="sim_1", help="Name of the output folder.") +parser.add_argument("--regularity", type=int, default=2, help="Regularity r of the polar splines at the axis.") +parser.add_argument("--degree", type=int, default=4, help="Spline degree p in all directions.") +parser.add_argument("--n", type=int, default=5, help="Mesh parameter: n x 2n x 8n cells.") +parser.add_argument("--final-time", type=float, default=0.2, help="Final time.") +parser.add_argument("--dt", type=float, default=2e-3, help="Time step (leapfrog: below 2 / sqrt(lambda_max)).") +parser.add_argument("--test-run", action="store_true", help="Small resolution and a few time steps, for the tests.") +args, _ = parser.parse_known_args() +if args.test_run: + args.n, args.degree, args.regularity, args.final_time, args.dt = 3, 2, 1, 0.004, 0.002 + +name = f"3D 1-form wave in a stellarator-like tube, polar splines p = {args.degree}, r = {args.regularity}" +description = f""" +Maxwell's equations dE/dt = curl B, dB/dt = -curl E (the 1-form wave equation for E = -dA/dt, B = curl A) in the +stellarator-like tube of the 3D wave and curl-curl studies of the talk: a helical magnetic axis with three periods, +a three-petal cross-section of radius 0.3 that turns once around the torus. E x n = 0 on the boundary. + +Physics. The initial field B = curl A_0 is a magnetic pulse next to the magnetic axis (the line eta1 = 0 of the polar +coordinates), A_0 = a(x) grad(zeta)-like with the Gaussian a of width 0.04 * 0.3 around the cross-section point of +logical coordinates (0.05, 0) and of width 0.1 along the axis (as the pulse of the talk). The wave crosses the axis, +where the tensor-product splines are not conforming; the high-order polar splines (DerhamOptions(polar_splines= +True, polar_regularity=r)) are the total complex of the 2D polar splines of the cross-section and the splines along the axis. +Since B_0 is the exact curl of a 1-form given as logical components, its commuting projection is the discrete curl +of the projection of A_0: div B = 0 holds exactly in polar coefficients. + +Setup. HelicalTube("stellarator"), {args.n} x {2 * args.n} x {8 * args.n} cells, splines of degree {args.degree}, +polar regularity r = {args.regularity}, Maxwell with the symplectic leapfrog (algo="leapfrog"), dt = {args.dt}. + +Expected result. The pproc plots |B| on the poloidal cut through the pulse at the first, middle and last snapshot, +and the energies. The talk compares r = 1, ..., 4 for degree-4 splines (convergence in h of the maximum error over +the final 10%; the resolutions n = 8 ... 32 are long runs). + +Not stated in the talk and chosen here: the initial field is curl A_0 of the Gaussian pulse (the talk uses the +L2-orthogonal projection of a 2-form pulse onto the range of curl). + +References: +- R. Remmerswaal, High-order FEEC discretization on 2D and 3D domains with a disk-like coordinate singularity + (draft); talk at the NMPP seminar 2026; PolarFEEC.jl (https://gitlab.mpcdf.mpg.de/rwr/polarfeec.jl). +""" + +import logging + +import numpy as np + +from struphy import set_logging_level + +set_logging_level(logging.INFO) + +from struphy import BaseUnits, DerhamOptions, EnvironmentOptions, Simulation, Time, domains, grids +from struphy.initial.base import GenericPerturbation +from struphy.models import Maxwell + +base_units = BaseUnits() +model = Maxwell(base_units=base_units) +model.em_fields.e_field.save_data = True +model.em_fields.b_field.save_data = True + +num_elements = (args.n, 2 * args.n, 8 * args.n) +degree = (args.degree,) * 3 +env = EnvironmentOptions(out_folders=args.out_folder, sim_folder=args.sim_folder, restart=False) +time_opts = Time(dt=args.dt, Tend=args.final_time) +grid = grids.TensorProductGrid(num_elements=num_elements) +domain = domains.HelicalTube("stellarator", num_elements=(8, 24, 64), degree=(4, 4, 4)) +derham_opts = DerhamOptions( + degree=degree, + bcs=(("free", "dirichlet"), None, None), + polar_splines=True, + polar_regularity=args.regularity, +) + +sim = Simulation( + model=model, + name=name, + description=description, + params_path=__file__, + env=env, + time_opts=time_opts, + domain=domain, + equil=None, + grid=grid, + derham_opts=derham_opts, +) + +from struphy.linear_algebra.solver import SolverParameters + +model.propagators.maxwell.options = model.propagators.maxwell.Options( + algo="leapfrog", + precond="PolarMassMatrixPreconditioner", + solver_params=SolverParameters(tol=1e-10, maxiter=500), +) + +# ------------------ +# Initial condition: B_0 = curl A_0, A_0 = (0, 0, a) in logical 1-form components +# ------------------ +g = domain.geometry +WIDTH, AXIAL_WIDTH = 0.04 * 0.3, 0.1 + + +def section(e1, e2): + radius = e1 * (1 + g["eps"] * e1 * np.cos(2 * np.pi * g["k"] * e2)) + return radius * g["a"] * np.cos(2 * np.pi * e2), radius * g["b"] * np.sin(2 * np.pi * e2) + + +CENTER = section(0.05, 0.0) +_speed = 2 * np.pi * np.sqrt((g["R0"] + g["delta"]) ** 2 + (g["periods"] * g["delta_z"]) ** 2) + + +def potential(e1, e2, e3): + """Logical component a_3 of A_0: the Gaussian pulse (cross-section and arc length along the axis).""" + x2, y2 = section(e1, e2) + arclength = _speed * (np.mod(e3 + 0.5, 1.0) - 0.5) + return np.exp( + -((x2 - CENTER[0]) ** 2 + (y2 - CENTER[1]) ** 2) / (2 * WIDTH**2) - arclength**2 / (2 * AXIAL_WIDTH**2) + ) + + +def _derivative(i, h=1e-7): + def d(e1, e2, e3): + plus, minus = [e1, e2, e3], [e1, e2, e3] + plus[i] = plus[i] + h + minus[i] = minus[i] - h + return (potential(*plus) - potential(*minus)) / (2 * h) + + return d + + +# curl of (0, 0, a) in logical components: (d_2 a, -d_1 a, 0) +_d1, _d2 = _derivative(0), _derivative(1) +model.em_fields.b_field.add_perturbation( + GenericPerturbation(lambda e1, e2, e3: _d2(e1, e2, e3), given_in_basis="2", comp=0) +) +model.em_fields.b_field.add_perturbation( + GenericPerturbation(lambda e1, e2, e3: -_d1(e1, e2, e3), given_in_basis="2", comp=1) +) + +if __name__ == "__main__": + # post-processing: pproc_gaussian_helical_tube.py + sim.run() diff --git a/examples/Maxwell/gaussian_helical_tube/pproc_gaussian_helical_tube.py b/examples/Maxwell/gaussian_helical_tube/pproc_gaussian_helical_tube.py new file mode 100644 index 000000000..b59c21fc0 --- /dev/null +++ b/examples/Maxwell/gaussian_helical_tube/pproc_gaussian_helical_tube.py @@ -0,0 +1,81 @@ +"""Post-process the Maxwell gaussian_helical_tube example: |B| on the poloidal cut eta3 = 0 (through the pulse) at the +first, middle and last snapshot, and the energies. Run as ``python pproc_gaussian_helical_tube.py [sim_folder] [flags]``. +""" + +import argparse +import os +import sys +from pathlib import Path + +import numpy as np +from matplotlib import pyplot as plt + +from struphy import Output + +SCRIPT_DIR = Path(__file__).resolve().parent +DEFAULT_OUTPUT = SCRIPT_DIR / "sim_1" + +sys.path.insert(0, str(SCRIPT_DIR)) +from params_gaussian_helical_tube import args, parser, sim # noqa: E402 + + +def cut(out, step, eta3=0.0, n1=61, n2=121): + """|B| on the poloidal cut eta3 and the coordinates of the cut in its own plane (pole at the origin).""" + e1, e2, e3 = np.linspace(1e-3, 1, n1), np.linspace(0, 1, n2), np.array([eta3]) + b = out.evaluate("em_fields/b_field", t=step, eta1=e1, eta2=e2, eta3=e3, representation="2") + b = np.stack([b.sel(component=c).values.reshape(n1, n2) for c in range(3)]) + # physical B = DF b / det(DF) of the logical 2-form b + DF = np.asarray(sim.domain.jacobian(e1, e2, e3)).reshape(3, 3, n1, n2) + detj = np.asarray(sim.domain.jacobian_det(e1, e2, e3)).reshape(n1, n2) + magnitude = np.linalg.norm(np.einsum("ijkl,jkl->ikl", DF, b), axis=0) / np.abs(detj) + X = np.array(sim.domain(e1, e2, e3)).reshape(3, n1, n2) + pole = np.array(sim.domain(np.array([0.0]), np.array([0.0]), e3)).reshape(3, 1, 1) + d = X - pole + q1 = d[:, -1, 0] / np.linalg.norm(d[:, -1, 0]) + q2 = d[:, -1, n2 // 4] - (d[:, -1, n2 // 4] @ q1) * q1 + q2 /= np.linalg.norm(q2) + return np.einsum("i,ijk->jk", q1, d), np.einsum("i,ijk->jk", q2, d), magnitude + + +def main(path_out=DEFAULT_OUTPUT, save_figs=True): + out = Output(path_out) + snapshots = out.standard_snapshots() + times = [float(out.time[n]) for n in snapshots] + results_dir = os.path.join(out.path_out, "results") + os.makedirs(results_dir, exist_ok=True) + + fig, axes = plt.subplots(1, len(snapshots), figsize=(4 * len(snapshots), 4), squeeze=False) + for ax, n, t in zip(axes[0], snapshots, times): + u, v, magnitude = cut(out, n) + ax.pcolormesh(u, v, magnitude, cmap="magma", shading="gouraud") + ax.set_aspect("equal") + ax.set_axis_off() + ax.set_title(f"$t = {t:.3f}$") + fig.suptitle( + rf"$|B|$ on the poloidal cut $\zeta = 0$, polar splines $p = {args.degree}$, $\mathfrak{{r}} = {args.regularity}$" + ) + fig.tight_layout() + + scalars = out.evaluate("scalars") + fig_e, ax = plt.subplots(figsize=(5, 3.5)) + for key in ("electric_energy", "magnetic_energy", "total_energy"): + ax.plot(scalars["t"], scalars[key], label=key.replace("_", " ")) + ax.set_xlabel("$t$") + ax.legend() + fig_e.tight_layout() + + if save_figs: + fig.savefig(os.path.join(results_dir, "b_snapshots.png")) + fig_e.savefig(os.path.join(results_dir, "energies.png")) + plt.close("all") + else: + plt.show() + + +if __name__ == "__main__": + _, other_flags = parser.parse_known_args() + pproc_parser = argparse.ArgumentParser(description="Post-process a saved simulation run.") + pproc_parser.add_argument("path_out", nargs="?", default=DEFAULT_OUTPUT, help="Simulation output folder.") + pproc_parser.add_argument("--show", action="store_true", help="Show the figures instead of saving them.") + pproc_args = pproc_parser.parse_args(other_flags) + main(pproc_args.path_out, save_figs=not pproc_args.show) diff --git a/examples/Maxwell/gaussian_perturbed_disk/params_gaussian_perturbed_disk.py b/examples/Maxwell/gaussian_perturbed_disk/params_gaussian_perturbed_disk.py new file mode 100644 index 000000000..c9b09d2ff --- /dev/null +++ b/examples/Maxwell/gaussian_perturbed_disk/params_gaussian_perturbed_disk.py @@ -0,0 +1,151 @@ +"""1-form wave equation (Maxwell in 2D) on the perturbed disk with high-order (homotopy) polar splines of regularity r: +a Gaussian pulse of B_z = curl A next to the polar axis crosses the axis (the 1-form wave study of the polar FEEC talk +of R. Remmerswaal).""" + +import argparse +from pathlib import Path + +parser = argparse.ArgumentParser(add_help=__name__ == "__main__") +parser.add_argument( + "--out-folder", + default=str(Path(__file__).resolve().parent), + help="Folder in which the output folder is created (default: beside this file).", +) +parser.add_argument("--sim-folder", default="sim_1", help="Name of the output folder.") +parser.add_argument("--regularity", type=int, default=3, help="Regularity r of the polar splines at the axis.") +parser.add_argument("--degree", type=int, default=5, help="Spline degree p in eta1 and eta2.") +parser.add_argument("--cells", type=int, default=8, help="Number of radial cells N^x.") +parser.add_argument("--angular-cells", type=int, default=16, help="Number of angular cells N^y.") +parser.add_argument( + "--enrichment", + default="cochain_chain", + choices=["cochain_chain", "cochain", "none"], + help="Enrichment of the polar spaces (the talk compares cochain_chain with cochain).", +) +parser.add_argument("--final-time", type=float, default=1.0, help="Final time.") +parser.add_argument( + "--cfl", type=float, default=0.9, help="Fraction of the leapfrog stability limit 2/sqrt(lambda_max)." +) +parser.add_argument("--test-run", action="store_true", help="Small resolution and a few time steps, for the tests.") +args, _ = parser.parse_known_args() +if args.test_run: + args.cells, args.angular_cells, args.degree, args.regularity, args.final_time = 4, 12, 3, 2, 0.05 + +name = f"1-form wave on the perturbed disk, polar splines p = {args.degree}, r = {args.regularity}" +description = f""" +Maxwell's equations dE/dt = curl B, dB/dt = -curl E in the cross-section of a cylinder, i.e. the 1-form wave +equation d^2 A / dt^2 + curl curl A = 0 for E = -dA/dt, B = curl A, on the perturbed disk +R = eta1 (1 + 0.2 eta1 cos(6 pi eta2)), with E x n = 0 on the boundary, E = 0 and +B_z = exp(-((x - 0.05)^2 + y^2) / (2 * 0.04^2)) - mean at t = 0. + +Physics. The pulse of B_z = curl A is centred next to the axis of the cylinder (the pole of the polar coordinates), +so the wave crosses the coordinate singularity. The tensor-product 1-form space is not conforming there; the +high-order polar splines (R. Remmerswaal, homotopy construction, DerhamOptions(polar_splines=True, polar_regularity=r)) are +conforming for r >= 1 and C^(r-1) at the axis. The talk compares r = 0 (non-conforming) with higher regularities and +the enrichments "cochain" (gradients of polynomials only, too poor) and "cochain_chain" (with the homotopy images of +polynomial 2-forms). The mean of B_z is conserved (E x n = 0) and is subtracted, since it does not propagate. + +Setup. PerturbedDiskCylinder(eps=0.2, k=3), {args.cells} x {args.angular_cells} x 1 cells, splines of degree +{args.degree}, polar regularity r = {args.regularity}, enrichment {args.enrichment}. Model Maxwell with the +symplectic leapfrog (algo="leapfrog") of the talk; the time step is {args.cfl} times the stability limit +2 / sqrt(lambda_max) of these spaces (computed with struphy.polar.homotopy). + +Expected result. The pproc plots B_z at the first, middle and last snapshot and the energies; the discrete +electromagnetic energy is conserved up to the leapfrog oscillation of order dt^2. + +References: +- R. Remmerswaal, High-order FEEC discretization on 2D and 3D domains with a disk-like coordinate singularity + (draft); talk at the NMPP seminar 2026; PolarFEEC.jl (https://gitlab.mpcdf.mpg.de/rwr/polarfeec.jl). +""" + +import logging + +import numpy as np + +from struphy import set_logging_level + +set_logging_level(logging.INFO) + +from struphy import BaseUnits, DerhamOptions, EnvironmentOptions, Simulation, Time, domains, grids +from struphy.initial.base import GenericPerturbation +from struphy.models import Maxwell +from struphy.polar.homotopy.demos.waves import WaveProblem, gaussian + +base_units = BaseUnits() +model = Maxwell(base_units=base_units) +model.em_fields.e_field.save_data = True +model.em_fields.b_field.save_data = True + +# time step: a fraction of the leapfrog stability limit of these polar spaces +problem = WaveProblem( + 1, + args.regularity, + degree=args.degree, + radial_cells=args.cells, + angular="bspline", + angular_cells=args.angular_cells, + enrichment=args.enrichment, +) +dt = args.cfl / 0.9 * problem.dt +dt = args.final_time / np.ceil(args.final_time / dt) + +env = EnvironmentOptions(out_folders=args.out_folder, sim_folder=args.sim_folder, restart=False) +time_opts = Time(dt=dt, Tend=args.final_time) +grid = grids.TensorProductGrid(num_elements=(args.cells, args.angular_cells, 1)) +domain = domains.PerturbedDiskCylinder(R1=1.0, eps=0.2, k=3, Lz=1.0) +derham_opts = DerhamOptions( + degree=(args.degree, args.degree, 1), + bcs=(("free", "dirichlet"), None, None), + polar_splines=True, + polar_regularity=args.regularity, + polar_enrichment=args.enrichment, +) + +sim = Simulation( + model=model, + name=name, + description=description, + params_path=__file__, + env=env, + time_opts=time_opts, + domain=domain, + equil=None, + grid=grid, + derham_opts=derham_opts, +) + +from struphy.linear_algebra.solver import SolverParameters + +model.propagators.maxwell.options = model.propagators.maxwell.Options( + algo="leapfrog", + precond="PolarMassMatrixPreconditioner", + solver_params=SolverParameters(tol=1e-12, maxiter=500), +) + +# mean of the Gaussian over the cross-section (Gauss quadrature in the logical coordinates) +_x, _w = np.polynomial.legendre.leggauss(8) +_e1 = ((np.arange(64)[:, None] + (_x[None, :] + 1) / 2) / 64).ravel() +_w1 = np.tile(_w / 128, 64) +_e3 = np.array([0.5]) +_X, _Y, _ = (c.reshape(_e1.size, _e1.size) for c in domain(_e1, _e1, _e3)) +_detj = np.abs(domain.jacobian_det(_e1, _e1, _e3)).reshape(_e1.size, _e1.size) +_weights = np.outer(_w1, _w1) * _detj +MEAN = float(np.sum(_weights * gaussian(_X, _Y)) / np.sum(_weights)) + + +def exact_initial_bz(x, y, z): + """Initial B_z (physical) with zero mean.""" + return gaussian(x, y) - MEAN + 0.0 * z + + +def initial_b_logical_z(e1, e2, e3): + """Logical 2-form component b_3 = (det(DF) DF^{-1} B)_3 = det(DF) B_z / L_z (finite at the axis).""" + x, y, z = domain(e1, e2, e3) + return domain.jacobian_det(e1, e2, e3) / domain.params["Lz"] * exact_initial_bz(x, y, z) + + +model.em_fields.b_field.add_perturbation(GenericPerturbation(initial_b_logical_z, given_in_basis="2", comp=2)) + +if __name__ == "__main__": + # post-processing: pproc_gaussian_perturbed_disk.py + sim.run() diff --git a/examples/Maxwell/gaussian_perturbed_disk/pproc_gaussian_perturbed_disk.py b/examples/Maxwell/gaussian_perturbed_disk/pproc_gaussian_perturbed_disk.py new file mode 100644 index 000000000..2366346b9 --- /dev/null +++ b/examples/Maxwell/gaussian_perturbed_disk/pproc_gaussian_perturbed_disk.py @@ -0,0 +1,76 @@ +"""Post-process the Maxwell gaussian_perturbed_disk example: B_z = curl A on the cross-section at the first, middle +and last snapshot, and the energies. Run as ``python pproc_gaussian_perturbed_disk.py [sim_folder] [variant flags]``.""" + +import argparse +import os +import sys +from pathlib import Path + +import numpy as np +from matplotlib import pyplot as plt + +from struphy import Output + +SCRIPT_DIR = Path(__file__).resolve().parent +DEFAULT_OUTPUT = SCRIPT_DIR / "sim_1" + +sys.path.insert(0, str(SCRIPT_DIR)) +from params_gaussian_perturbed_disk import args, parser, sim # noqa: E402 + + +def bz_cross_section(out, step, n1=81, n2=181): + """Physical B_z = b_3 L_z / det(DF) on the cross-section eta3 = 0.5 (eta1 > 0).""" + e1, e2, e3 = np.linspace(1e-3, 1, n1), np.linspace(0, 1, n2), np.array([0.5]) + b = out.evaluate("em_fields/b_field", t=step, eta1=e1, eta2=e2, eta3=e3, representation="2") + b3 = b.sel(component=2).values.reshape(n1, n2) + detj = sim.domain.jacobian_det(e1, e2, e3).reshape(n1, n2) + x, y, _ = (c.reshape(n1, n2) for c in sim.domain(e1, e2, e3)) + return x, y, b3 * sim.domain.params["Lz"] / detj + + +def main(path_out=DEFAULT_OUTPUT, save_figs=True): + out = Output(path_out) + snapshots = out.standard_snapshots() + times = [float(out.time[n]) for n in snapshots] + results_dir = os.path.join(out.path_out, "results") + os.makedirs(results_dir, exist_ok=True) + + fig, axes = plt.subplots(1, len(snapshots), figsize=(4 * len(snapshots), 4), squeeze=False) + for ax, n, t in zip(axes[0], snapshots, times): + x, y, bz = bz_cross_section(out, n) + limit = np.max(np.abs(bz)) or 1.0 + ax.pcolormesh(x, y, bz, cmap="RdBu_r", vmin=-limit, vmax=limit, shading="gouraud") + ax.set_aspect("equal") + ax.set_axis_off() + ax.set_title(f"$t = {t:.3f}$") + fig.suptitle( + rf"$B_z = \mathrm{{curl}}\,A$, polar splines $p = {args.degree}$, $\mathfrak{{r}} = {args.regularity}$, " + f"{args.enrichment}" + ) + fig.tight_layout() + + scalars = out.evaluate("scalars") + fig_e, ax = plt.subplots(figsize=(5, 3.5)) + for key in ("electric_energy", "magnetic_energy", "total_energy"): + ax.plot(scalars["t"], scalars[key], label=key.replace("_", " ")) + ax.set_xlabel("$t$") + ax.legend() + fig_e.tight_layout() + total = np.asarray(scalars["total_energy"]) + print(f"max relative variation of the total energy: {np.max(np.abs(total - total[0])) / total[0]:.2e}") + + if save_figs: + fig.savefig(os.path.join(results_dir, "bz_snapshots.png")) + fig_e.savefig(os.path.join(results_dir, "energies.png")) + plt.close("all") + else: + plt.show() + + +if __name__ == "__main__": + _, other_flags = parser.parse_known_args() + pproc_parser = argparse.ArgumentParser(description="Post-process a saved simulation run.") + pproc_parser.add_argument("path_out", nargs="?", default=DEFAULT_OUTPUT, help="Simulation output folder.") + pproc_parser.add_argument("--show", action="store_true", help="Show the figures instead of saving them.") + pproc_args = pproc_parser.parse_args(other_flags) + main(pproc_args.path_out, save_figs=not pproc_args.show) diff --git a/examples/Poisson/gaussian_on_axis_cylinder/params_gaussian_on_axis_cylinder.py b/examples/Poisson/gaussian_on_axis_cylinder/params_gaussian_on_axis_cylinder.py index ee91d083c..68b082c76 100644 --- a/examples/Poisson/gaussian_on_axis_cylinder/params_gaussian_on_axis_cylinder.py +++ b/examples/Poisson/gaussian_on_axis_cylinder/params_gaussian_on_axis_cylinder.py @@ -47,7 +47,7 @@ "--poloidal-solver", default="schur", choices=["schur", "lu", "ilu"], - help="Poloidal solver of the polar preconditioners and projectors.", + help="Poloidal solver of the polar preconditioners.", ) parser.add_argument("--num-elements", type=int, nargs=3, default=[64, 66, 64], help="Number of elements per direction.") parser.add_argument( @@ -159,7 +159,7 @@ # Time stepping time_opts = Time() -# Grid (polar splines need a multiple of 3 elements in eta2) +# Grid num_elements = tuple(args.num_elements) degree = (3, 3, 3) grid = grids.TensorProductGrid(num_elements=num_elements, mpi_dims_mask=(True, True, True)) @@ -174,7 +174,6 @@ degree=degree, bcs=(("free", "dirichlet"), None, None), polar_splines=True, - polar_poloidal_solver=args.poloidal_solver, ) else: a1 = 1e-2 diff --git a/examples/Poisson/gaussian_on_axis_torus/params_gaussian_on_axis_torus.py b/examples/Poisson/gaussian_on_axis_torus/params_gaussian_on_axis_torus.py index 38a22bd00..58abde107 100644 --- a/examples/Poisson/gaussian_on_axis_torus/params_gaussian_on_axis_torus.py +++ b/examples/Poisson/gaussian_on_axis_torus/params_gaussian_on_axis_torus.py @@ -48,7 +48,7 @@ "--poloidal-solver", default="schur", choices=["schur", "lu", "ilu"], - help="Poloidal solver of the polar preconditioners and projectors.", + help="Poloidal solver of the polar preconditioners.", ) parser.add_argument("--num-elements", type=int, nargs=3, default=[64, 66, 64], help="Number of elements per direction.") parser.add_argument( @@ -160,7 +160,7 @@ # Time stepping time_opts = Time() -# Grid (polar splines need a multiple of 3 elements in eta2) +# Grid num_elements = tuple(args.num_elements) degree = (3, 3, 3) grid = grids.TensorProductGrid(num_elements=num_elements, mpi_dims_mask=(True, True, True)) @@ -177,7 +177,6 @@ degree=degree, bcs=(("free", "dirichlet"), None, None), polar_splines=True, - polar_poloidal_solver=args.poloidal_solver, ) else: a1 = 1e-2 diff --git a/examples/Poisson/polar_rounded_square/convergence_polar_rounded_square.py b/examples/Poisson/polar_rounded_square/convergence_polar_rounded_square.py new file mode 100644 index 000000000..cafc93963 --- /dev/null +++ b/examples/Poisson/polar_rounded_square/convergence_polar_rounded_square.py @@ -0,0 +1,59 @@ +"""h-convergence of the Poisson polar_rounded_square example (struphy counterpart of the Poisson panel of +PolarFEEC.jl example 05): runs params_polar_rounded_square.py for n radial cells and degrees p (regularity r = p), +post-processes every run and plots the L2 error against h = 1/n with dotted h^(p + 1) reference slopes. + +Run as ``python convergence_polar_rounded_square.py [--degrees 2 3 4] [--cells 4 6 8 12] [--mpi N]``. +""" + +import argparse +import json +import subprocess +import sys +from pathlib import Path + +import numpy as np +from matplotlib import pyplot as plt + +SCRIPT_DIR = Path(__file__).resolve().parent + + +def run(p, n, mpi): + folder = f"conv_p{p}_n{n}" + flags = ["--degree", str(p), "--cells", str(n), "--sim-folder", folder] + command = [sys.executable, str(SCRIPT_DIR / "params_polar_rounded_square.py"), *flags] + if mpi > 1: + command = ["mpirun", "-n", str(mpi), *command] + subprocess.run(command, check=True) + subprocess.run([sys.executable, str(SCRIPT_DIR / "pproc_polar_rounded_square.py"), str(SCRIPT_DIR / folder), + *flags], check=True) + with open(SCRIPT_DIR / folder / "results" / "errors.json") as f: + return json.load(f)["l2_error"] + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--degrees", type=int, nargs="+", default=[2, 3, 4]) + parser.add_argument("--cells", type=int, nargs="+", default=[4, 6, 8, 12]) + parser.add_argument("--mpi", type=int, default=1, help="Number of MPI processes per run.") + args = parser.parse_args() + + hs = 1 / np.array(args.cells) + table = {p: np.array([run(p, n, args.mpi) for n in args.cells]) for p in args.degrees} + fig, ax = plt.subplots(figsize=(6, 5)) + for i, (p, errors) in enumerate(table.items()): + rates = np.log(errors[:-1] / errors[1:]) / np.log(hs[:-1] / hs[1:]) + print(f"p = {p}: errors {errors}, observed rates {rates}") + ax.loglog(hs, errors, "o-", color=f"C{i}", label=f"$p = r = {p}$") + ax.loglog(hs[-3:], 0.3 * errors[-1] * (hs[-3:] / hs[-1]) ** (p + 1), ":", color=f"C{i}") + ax.set_xticks(hs, [f"1/{n}" for n in args.cells]) + ax.set_xlabel("$h = 1/n$") + ax.set_ylabel(r"$\Vert \phi - \phi_h \Vert_{L^2}$") + ax.set_title("Poisson on a rounded square, homotopy polar splines\n(dotted: $h^{p+1}$)") + ax.legend() + fig.tight_layout() + fig.savefig(SCRIPT_DIR / "convergence_polar_rounded_square.png", dpi=150) + print("wrote", SCRIPT_DIR / "convergence_polar_rounded_square.png") + + +if __name__ == "__main__": + main() diff --git a/examples/Poisson/polar_rounded_square/params_polar_rounded_square.py b/examples/Poisson/polar_rounded_square/params_polar_rounded_square.py new file mode 100644 index 000000000..e329eb9a8 --- /dev/null +++ b/examples/Poisson/polar_rounded_square/params_polar_rounded_square.py @@ -0,0 +1,166 @@ +"""Poisson equation -Laplace phi = rho on a rounded square with high-order (homotopy) polar splines of regularity r. +Manufactured solution of PolarFEEC.jl example 05; the pproc reports the L2 error, convergence_polar_rounded_square.py +runs the h-convergence study (rate p + 1).""" + +# ----------------------------- +# Description of the simulation +# ----------------------------- + +import argparse +from pathlib import Path + +# no -h here when imported by the pproc script, which has its own help +parser = argparse.ArgumentParser(add_help=__name__ == "__main__") +parser.add_argument( + "--out-folder", + default=str(Path(__file__).resolve().parent), + help="Folder in which the output folder is created (default: beside this file).", +) +parser.add_argument("--sim-folder", default="sim_1", help="Name of the output folder.") +parser.add_argument("--degree", type=int, default=3, help="Spline degree p in eta1 and eta2.") +parser.add_argument( + "--regularity", type=int, default=None, help="Regularity r of the polar splines at the pole (default: p)." +) +parser.add_argument("--cells", type=int, default=6, help="Number n of radial cells (6n angular cells).") +parser.add_argument( + "--precond", + default="StiffnessPreconditioner", + choices=["None", "MassMatrixPreconditioner", "StiffnessPreconditioner"], + help="Preconditioner of the CG solver.", +) +parser.add_argument("--test-run", action="store_true", help="Small resolution (n = 4, p = 2), for the tests.") +args, _ = parser.parse_known_args() +if args.test_run: + args.cells, args.degree = 4, 2 +if args.regularity is None: + args.regularity = args.degree + +precond = None if args.precond == "None" else args.precond + +name = f"Poisson on a rounded square, homotopy polar splines (p = {args.degree}, r = {args.regularity})" +description = f""" +Poisson equation -Laplace phi = rho on the rounded square g > 0, g = 1 - x^2 - y^2 - eps Re((x + i y)^4), +eps = 0.15, with high-order polar splines of regularity r = {args.regularity} at the centre. + +Physics. Poisson's equation determines the electrostatic potential phi of a charge density rho. The domain is +parametrized by a disk-like map whose logical edge eta1 = 0 collapses to the centre of the cross-section (the axis +of the cylinder); tensor-product splines are not smooth there. The polar splines of R. Remmerswaal (homotopy +construction, see struphy.polar.homotopy) replace the r + 1 innermost rings of coefficients by projected polynomials +of degree r: the 0-forms are C^r at the centre, conforming, and form a de Rham complex with commuting projections. + +Setup. RoundedSquareCylinder(eps=0.15), {args.cells} x {6 * args.cells} x 1 cells, splines of degree {args.degree} +(one periodic cell in eta3: a 2D problem), DerhamOptions(polar_splines=True, +polar_regularity={args.regularity}), homogeneous Dirichlet conditions at eta1 = 1. The manufactured solution of +PolarFEEC.jl example 05 is phi = g (1 + sin(2x + y)), which vanishes on the boundary; rho = -Laplace phi. +CG with preconditioner {precond}. + +Expected result. The pproc script prints the L2 error of phi; the error converges like h^(p + 1) +(convergence_polar_rounded_square.py), as with the homotopy solver of PolarFEEC.jl. + +References: +- R. Remmerswaal, High-order FEEC discretization on 2D and 3D domains with a disk-like coordinate singularity + (draft), and the Julia package PolarFEEC.jl (https://gitlab.mpcdf.mpg.de/rwr/polarfeec.jl). +- P. Jiang et al., Polar spline discretizations with high-order regularity at the origin, J. Comput. Phys. 564, + 115129 (2026). +""" + +import logging + +from struphy import set_logging_level + +set_logging_level(logging.INFO) + +# ------------------ +# Import Struphy API +# ------------------ +from struphy import BaseUnits, DerhamOptions, EnvironmentOptions, Simulation, Time, domains, grids +from struphy.models import Poisson + +base_units = BaseUnits() +model = Poisson(base_units=base_units) +model.em_fields.phi.save_data = True +model.em_fields.source.save_data = True + +env = EnvironmentOptions(out_folders=args.out_folder, sim_folder=args.sim_folder, restart=False) +time_opts = Time() + +num_elements = (args.cells, 6 * args.cells, 1) +degree = (args.degree, args.degree, 1) +grid = grids.TensorProductGrid(num_elements=num_elements) + +EPS = 0.15 +domain = domains.RoundedSquareCylinder(a=1.0, eps=EPS, Lz=1.0) +derham_opts = DerhamOptions( + degree=degree, + bcs=(("free", "dirichlet"), None, None), + polar_splines=True, + polar_regularity=args.regularity, + nquads=(args.degree + 3, args.degree + 3, 2), +) + +sim = Simulation( + model=model, + name=name, + description=description, + params_path=__file__, + env=env, + time_opts=time_opts, + domain=domain, + equil=None, + grid=grid, + derham_opts=derham_opts, +) + +# ------------------ +# Propagator options +# ------------------ +from struphy.linear_algebra.solver import SolverParameters + +model.propagators.poisson.options = model.propagators.poisson.Options( + stab_eps=0.0, + solver="pcg", + precond=precond, + solver_params=SolverParameters(tol=1e-12, maxiter=5000, info=True, recycle=False), +) + +# ------------------ +# Initial conditions +# ------------------ +import numpy as np + +from struphy.initial.base import GenericPerturbation + + +def level(x, y): + """g, its gradient and its Laplacian (-4, since Re((x + i y)^4) is harmonic).""" + g = 1 - x**2 - y**2 - EPS * (x**4 - 6 * x**2 * y**2 + y**4) + return g, (-2 * x - 4 * EPS * (x**3 - 3 * x * y**2), -2 * y - 4 * EPS * (y**3 - 3 * x**2 * y)), -4.0 + + +def exact_solution(x, y, z): + g, _, _ = level(x, y) + return g * (1 + np.sin(2 * x + y)) + 0.0 * z + + +def rhs_fun(x, y, z): + g, (gx, gy), lap_g = level(x, y) + s, sx, sy, lap_s = 1 + np.sin(2 * x + y), 2 * np.cos(2 * x + y), np.cos(2 * x + y), -5 * np.sin(2 * x + y) + return -(g * lap_s + s * lap_g + 2 * (gx * sx + gy * sy)) + 0.0 * z + + +# The source variable is saved for plotting. The solver gets the source as a callable (a logical 0-form proxy): its +# L2 load is integrated against the test functions with homogeneous Dirichlet conditions, whereas a source in the +# space V0 would have its boundary coefficients removed (rho does not vanish on the boundary). +model.em_fields.source.add_perturbation(GenericPerturbation(rhs_fun, given_in_basis="physical")) + + +def rhs_logical(e1, e2, e3): + return rhs_fun(*domain(e1, e2, e3)) + + +model.propagators.poisson.rho = rhs_logical + + +if __name__ == "__main__": + # post-processing: pproc_polar_rounded_square.py + sim.run(one_time_step=True) diff --git a/examples/Poisson/polar_rounded_square/pproc_polar_rounded_square.py b/examples/Poisson/polar_rounded_square/pproc_polar_rounded_square.py new file mode 100644 index 000000000..970310d05 --- /dev/null +++ b/examples/Poisson/polar_rounded_square/pproc_polar_rounded_square.py @@ -0,0 +1,93 @@ +"""Post-process the Poisson polar_rounded_square example: L2 error of phi and plots of phi and of the error. + +Run as ``python pproc_polar_rounded_square.py [sim_folder] [variant flags of params_polar_rounded_square.py]``. +Errors and figures are written into ``/results``. +""" + +import argparse +import json +import os +import sys +from pathlib import Path + +import numpy as np +from matplotlib import pyplot as plt + +from struphy import Output + +SCRIPT_DIR = Path(__file__).resolve().parent +DEFAULT_OUTPUT = SCRIPT_DIR / "sim_1" + +sys.path.insert(0, str(SCRIPT_DIR)) +from params_polar_rounded_square import args, exact_solution, parser, sim # noqa: E402 + + +def gauss_points(cells, n): + """Gauss-Legendre points and weights of n points in each of `cells` cells of [0, 1].""" + x, w = np.polynomial.legendre.leggauss(n) + a = np.arange(cells)[:, None] / cells + return (a + (x[None, :] + 1) / (2 * cells)).ravel(), np.tile(w / (2 * cells), cells) + + +def l2_error(out, domain, n_quad): + """Physical L2 error of phi on the cross-section (the problem does not depend on eta3).""" + cells = sim.grid.num_elements + e1, w1 = gauss_points(cells[0], n_quad) + e2, w2 = gauss_points(cells[1], n_quad) + e3 = np.array([0.5]) + phi = out.evaluate("em_fields/phi", t=-1, eta1=e1, eta2=e2, eta3=e3).values.reshape(e1.size, e2.size) + x, y, z = (c.reshape(e1.size, e2.size) for c in domain(e1, e2, e3)) + detj = np.abs(domain.jacobian_det(e1, e2, e3)).reshape(e1.size, e2.size) / domain.params["Lz"] + weights = np.outer(w1, w2) * detj + error = phi - exact_solution(x, y, z) + norm = np.sqrt(np.sum(weights * exact_solution(x, y, z) ** 2)) + return np.sqrt(np.sum(weights * error**2)), norm + + +def main(path_out=DEFAULT_OUTPUT, save_figs=True): + out = Output(path_out) + domain = sim.domain + error, norm = l2_error(out, domain, sim.derham_opts.degree[0] + 3) + print(f"L2 error of phi: {error:.3e} (relative {error / norm:.3e})") + + results_dir = os.path.join(out.path_out, "results") + os.makedirs(results_dir, exist_ok=True) + with open(os.path.join(results_dir, "errors.json"), "w") as f: + json.dump( + dict( + l2_error=float(error), + relative_l2_error=float(error / norm), + cells=args.cells, + degree=args.degree, + regularity=args.regularity, + ), + f, + ) + + # phi and the error on the cross-section + e1, e2, e3 = np.linspace(0, 1, 81), np.linspace(0, 1, 241), np.array([0.5]) + phi = out.evaluate("em_fields/phi", t=-1, eta1=e1, eta2=e2, eta3=e3).values.reshape(e1.size, e2.size) + x, y, z = (c.reshape(e1.size, e2.size) for c in domain(e1, e2, e3)) + fig, axes = plt.subplots(1, 2, figsize=(10, 4.5)) + for ax, values, title in ((axes[0], phi, r"$\phi_h$"), (axes[1], phi - exact_solution(x, y, z), "error")): + mesh = ax.pcolormesh(x, y, values, shading="gouraud", cmap="RdBu_r" if title == "error" else "viridis") + plt.colorbar(mesh, ax=ax, shrink=0.8) + ax.set_aspect("equal") + ax.set_title(title) + fig.suptitle(f"Poisson, homotopy polar splines p = {args.degree}, r = {args.regularity}, n = {args.cells}") + fig.tight_layout() + if save_figs: + fig.savefig(os.path.join(results_dir, "phi.png")) + plt.close(fig) + else: + plt.show() + return error + + +if __name__ == "__main__": + _, other_flags = parser.parse_known_args() + pproc_parser = argparse.ArgumentParser(description="Post-process a saved simulation run.") + pproc_parser.add_argument("path_out", nargs="?", default=DEFAULT_OUTPUT, help="Simulation output folder.") + pproc_parser.add_argument("--show", action="store_true", help="Show the figures instead of saving them.") + pproc_args = pproc_parser.parse_args(other_flags) + main(pproc_args.path_out, save_figs=not pproc_args.show) diff --git a/examples/ScalarWave/gaussian_perturbed_disk/params_gaussian_perturbed_disk.py b/examples/ScalarWave/gaussian_perturbed_disk/params_gaussian_perturbed_disk.py new file mode 100644 index 000000000..318171029 --- /dev/null +++ b/examples/ScalarWave/gaussian_perturbed_disk/params_gaussian_perturbed_disk.py @@ -0,0 +1,123 @@ +"""0-form wave equation on the perturbed disk with high-order (homotopy) polar splines of regularity r: a Gaussian +pulse next to the polar axis crosses the axis (the 0-form wave study of the polar FEEC talk of R. Remmerswaal).""" + +import argparse +from pathlib import Path + +parser = argparse.ArgumentParser(add_help=__name__ == "__main__") +parser.add_argument( + "--out-folder", + default=str(Path(__file__).resolve().parent), + help="Folder in which the output folder is created (default: beside this file).", +) +parser.add_argument("--sim-folder", default="sim_1", help="Name of the output folder.") +parser.add_argument("--regularity", type=int, default=3, help="Regularity r of the polar splines at the axis.") +parser.add_argument("--degree", type=int, default=5, help="Spline degree p in eta1 and eta2.") +parser.add_argument("--cells", type=int, default=8, help="Number of radial cells N^x.") +parser.add_argument("--angular-cells", type=int, default=16, help="Number of angular cells N^y.") +parser.add_argument("--final-time", type=float, default=1.0, help="Final time.") +parser.add_argument( + "--cfl", type=float, default=0.9, help="Fraction of the leapfrog stability limit 2/sqrt(lambda_max)." +) +parser.add_argument("--test-run", action="store_true", help="Small resolution and a few time steps, for the tests.") +args, _ = parser.parse_known_args() +if args.test_run: + args.cells, args.angular_cells, args.degree, args.regularity, args.final_time = 4, 12, 3, 2, 0.05 + +name = f"0-form wave on the perturbed disk, polar splines p = {args.degree}, r = {args.regularity}" +description = f""" +0-form (scalar) wave equation d^2 phi / dt^2 - Laplace phi = 0 on the perturbed disk +R = eta1 (1 + 0.2 eta1 cos(6 pi eta2)), with phi = 0 on the boundary and a Gaussian pulse +exp(-((x - 0.05)^2 + y^2) / (2 * 0.04^2)) at rest next to the polar axis. + +Physics. The pulse is centred next to the axis of the cylinder (the pole of the polar coordinates of the +cross-section), so the expanding wave crosses the coordinate singularity. Tensor-product splines are discontinuous +there; the high-order polar splines (R. Remmerswaal, homotopy construction, DerhamOptions(polar_splines=True, polar_regularity=r)) +are C^r at the axis. The talk compares r = 0 (continuity) with higher regularities: low regularity leaves spurious +waves and kinks at the axis, and limits the convergence of nearly resolved waves. + +Setup. PerturbedDiskCylinder(eps=0.2, k=3), {args.cells} x {args.angular_cells} x 1 cells, splines of degree +{args.degree}, polar regularity r = {args.regularity}, Dirichlet condition at eta1 = 1. Model ScalarWave with the +symplectic leapfrog (velocity Verlet) of the talk; the time step is {args.cfl} times the stability limit +2 / sqrt(lambda_max) of these spaces (computed with struphy.polar.homotopy). + +Expected result. The pproc plots phi at the first, middle and last snapshot and the energy; the staggered discrete +energy is conserved. For the convergence in h of the talk (maximum error during the final 10%), run several +resolutions and regularities (see struphy.polar.homotopy.demos.waves). + +References: +- R. Remmerswaal, High-order FEEC discretization on 2D and 3D domains with a disk-like coordinate singularity + (draft); talk at the NMPP seminar 2026; PolarFEEC.jl (https://gitlab.mpcdf.mpg.de/rwr/polarfeec.jl). +""" + +import logging + +import numpy as np + +from struphy import set_logging_level + +set_logging_level(logging.INFO) + +from struphy import BaseUnits, DerhamOptions, EnvironmentOptions, Simulation, Time, domains, grids +from struphy.initial.base import GenericPerturbation +from struphy.models import ScalarWave +from struphy.polar.homotopy.demos.waves import WaveProblem, gaussian + +base_units = BaseUnits() +model = ScalarWave(base_units=base_units) +model.em_fields.phi.save_data = True +model.em_fields.psi.save_data = True + +# time step: a fraction of the leapfrog stability limit of these polar spaces +problem = WaveProblem( + 0, + args.regularity, + degree=args.degree, + radial_cells=args.cells, + angular="bspline", + angular_cells=args.angular_cells, +) +dt = args.cfl / 0.9 * problem.dt +dt = args.final_time / np.ceil(args.final_time / dt) + +env = EnvironmentOptions(out_folders=args.out_folder, sim_folder=args.sim_folder, restart=False) +time_opts = Time(dt=dt, Tend=args.final_time) +grid = grids.TensorProductGrid(num_elements=(args.cells, args.angular_cells, 1)) +domain = domains.PerturbedDiskCylinder(R1=1.0, eps=0.2, k=3, Lz=1.0) +derham_opts = DerhamOptions( + degree=(args.degree, args.degree, 1), + bcs=(("free", "dirichlet"), None, None), + polar_splines=True, + polar_regularity=args.regularity, +) + +sim = Simulation( + model=model, + name=name, + description=description, + params_path=__file__, + env=env, + time_opts=time_opts, + domain=domain, + equil=None, + grid=grid, + derham_opts=derham_opts, +) + +from struphy.linear_algebra.solver import SolverParameters + +model.propagators.wave.options = model.propagators.wave.Options( + precond="PolarMassMatrixPreconditioner", + solver_params=SolverParameters(tol=1e-12, maxiter=500), +) + + +def initial_phi(x, y, z): + return gaussian(x, y) + 0.0 * z + + +model.em_fields.phi.add_perturbation(GenericPerturbation(initial_phi, given_in_basis="physical")) + +if __name__ == "__main__": + # post-processing: pproc_gaussian_perturbed_disk.py + sim.run() diff --git a/examples/ScalarWave/gaussian_perturbed_disk/pproc_gaussian_perturbed_disk.py b/examples/ScalarWave/gaussian_perturbed_disk/pproc_gaussian_perturbed_disk.py new file mode 100644 index 000000000..a4e045bf8 --- /dev/null +++ b/examples/ScalarWave/gaussian_perturbed_disk/pproc_gaussian_perturbed_disk.py @@ -0,0 +1,68 @@ +"""Post-process the ScalarWave gaussian_perturbed_disk example: phi on the cross-section at the first, middle and last +snapshot, and the kinetic energy. Run as ``python pproc_gaussian_perturbed_disk.py [sim_folder] [variant flags]``.""" + +import argparse +import os +import sys +from pathlib import Path + +import numpy as np +from matplotlib import pyplot as plt + +from struphy import Output + +SCRIPT_DIR = Path(__file__).resolve().parent +DEFAULT_OUTPUT = SCRIPT_DIR / "sim_1" + +sys.path.insert(0, str(SCRIPT_DIR)) +from params_gaussian_perturbed_disk import args, parser, sim # noqa: E402 + + +def cross_section(out, name, step, n1=81, n2=181): + """Values of a scalar field on the cross-section eta3 = 0.5 and the physical coordinates.""" + e1, e2, e3 = np.linspace(0, 1, n1), np.linspace(0, 1, n2), np.array([0.5]) + values = out.evaluate(name, t=step, eta1=e1, eta2=e2, eta3=e3).values.reshape(n1, n2) + x, y, _ = (c.reshape(n1, n2) for c in sim.domain(e1, e2, e3)) + return x, y, values + + +def main(path_out=DEFAULT_OUTPUT, save_figs=True): + out = Output(path_out) + snapshots = out.standard_snapshots() + times = [float(out.time[n]) for n in snapshots] + results_dir = os.path.join(out.path_out, "results") + os.makedirs(results_dir, exist_ok=True) + + fig, axes = plt.subplots(1, len(snapshots), figsize=(4 * len(snapshots), 4), squeeze=False) + for ax, n, t in zip(axes[0], snapshots, times): + x, y, phi = cross_section(out, "em_fields/phi", n) + limit = np.max(np.abs(phi)) or 1.0 + ax.pcolormesh(x, y, phi, cmap="RdBu_r", vmin=-limit, vmax=limit, shading="gouraud") + ax.set_aspect("equal") + ax.set_axis_off() + ax.set_title(f"$t = {t:.3f}$") + fig.suptitle(rf"$\phi$, polar splines $p = {args.degree}$, $\mathfrak{{r}} = {args.regularity}$") + fig.tight_layout() + + scalars = out.evaluate("scalars") + fig_e, ax = plt.subplots(figsize=(5, 3.5)) + ax.plot(scalars["t"], scalars["kinetic_energy"]) + ax.set_xlabel("$t$") + ax.set_ylabel("kinetic energy") + fig_e.tight_layout() + + if save_figs: + fig.savefig(os.path.join(results_dir, "phi_snapshots.png")) + fig_e.savefig(os.path.join(results_dir, "kinetic_energy.png")) + plt.close("all") + else: + plt.show() + + +if __name__ == "__main__": + _, other_flags = parser.parse_known_args() + pproc_parser = argparse.ArgumentParser(description="Post-process a saved simulation run.") + pproc_parser.add_argument("path_out", nargs="?", default=DEFAULT_OUTPUT, help="Simulation output folder.") + pproc_parser.add_argument("--show", action="store_true", help="Show the figures instead of saving them.") + pproc_args = pproc_parser.parse_args(other_flags) + main(pproc_args.path_out, save_figs=not pproc_args.show) diff --git a/pyproject.toml b/pyproject.toml index 62c46b7b2..0f72d4a28 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -69,6 +69,7 @@ phys = [ ] pproc = [ "plasma-plots>=0.1.0", + "tikzfigure>=0.5.0", ] dev = [ "struphy[mpi]", @@ -138,6 +139,9 @@ struphy = "struphy.console.main:struphy" "run_03/*.ini", "run_03/*.nc", ] +"struphy.polar.tests" = [ + "data/*.npz", +] struphy = [ "compile_struphy.mk", "**/*.cu", diff --git a/src/struphy/api/models/__init__.py b/src/struphy/api/models/__init__.py index 17bf65569..0e4ccbb05 100644 --- a/src/struphy/api/models/__init__.py +++ b/src/struphy/api/models/__init__.py @@ -22,6 +22,7 @@ from struphy.models.poisson import Poisson from struphy.models.pressure_less_sph import PressureLessSPH from struphy.models.random_particle_diffusion import RandomParticleDiffusion +from struphy.models.scalar_wave import ScalarWave from struphy.models.shear_alfven import ShearAlfven from struphy.models.two_fluid_quasi_neutral_toy import TwoFluidQuasiNeutralToy from struphy.models.variational_barotropic_fluid import VariationalBarotropicFluid @@ -63,6 +64,7 @@ "Poisson", "PressureLessSPH", "RandomParticleDiffusion", + "ScalarWave", "ShearAlfven", "TwoFluidQuasiNeutralToy", "VariationalBarotropicFluid", diff --git a/src/struphy/eigenvalue_solvers/derivatives.py b/src/struphy/eigenvalue_solvers/derivatives.py index 0e34cceb1..00573464b 100644 --- a/src/struphy/eigenvalue_solvers/derivatives.py +++ b/src/struphy/eigenvalue_solvers/derivatives.py @@ -33,9 +33,10 @@ def grad_1d_matrix(spl_kind, NbaseN): grad = xp.zeros((NbaseD, NbaseN), dtype=float) + # accumulate: for a single periodic basis function both entries coincide and the derivative is zero for i in range(NbaseD): - grad[i, i] = -1.0 - grad[i, (i + 1) % NbaseN] = 1.0 + grad[i, i] += -1.0 + grad[i, (i + 1) % NbaseN] += 1.0 return grad @@ -126,35 +127,10 @@ def discrete_derivatives_3d(space): D = [[D1, D2, D3]] D = spa.bmat(D, format="csr") - # C^k polar derivatives else: - # discrete grad (full space) - G1 = spa.kron(space.polar_splines.G1.copy(), spa.identity(grad_1d_3.shape[1])) - G2 = spa.kron(space.polar_splines.G2.copy(), spa.identity(grad_1d_3.shape[1])) - G3 = spa.kron(spa.identity(space.polar_splines.Nbase0), grad_1d_3) - - G = [[G1], [G2], [G3]] - G = spa.bmat(G, format="csr") - - # discrete curl (full space) - C12 = spa.kron(space.polar_splines.VC.copy(), spa.identity(grad_1d_3.shape[0])) - C21 = spa.kron(space.polar_splines.SC.copy(), spa.identity(grad_1d_3.shape[1])) - - C11_12 = -spa.kron(spa.identity(space.polar_splines.Nbase1C_2), grad_1d_3) - C11_21 = spa.kron(spa.identity(space.polar_splines.Nbase1C_1), grad_1d_3) - - C11 = spa.bmat([[None, C11_12], [C11_21, None]]) - - C = [[C11, C12], [C21, None]] - C = spa.bmat(C, format="csr") - - # discrete div (full space) - D1 = spa.kron(space.polar_splines.D1.copy(), spa.identity(grad_1d_3.shape[0])) - D2 = spa.kron(space.polar_splines.D2.copy(), spa.identity(grad_1d_3.shape[0])) - D3 = spa.kron(spa.identity(space.polar_splines.Nbase2), grad_1d_3) - - D = [[D1, D2, D3]] - D = spa.bmat(D, format="csr") + raise NotImplementedError( + "polar splines are provided by struphy.feec.psydac_derham.Derham (polar_splines=True)" + ) # apply boundary operators G0 = space.B1.dot(G.dot(space.B0.T)) diff --git a/src/struphy/feec/basis_projection_ops.py b/src/struphy/feec/basis_projection_ops.py index 56ec67fbf..1d6da6c65 100644 --- a/src/struphy/feec/basis_projection_ops.py +++ b/src/struphy/feec/basis_projection_ops.py @@ -1674,6 +1674,14 @@ def __init__( polar_shift: bool = False, use_cache: bool = False, ): + if getattr(P.space, "spaces", None) is not None and any( + getattr(sp, "is_trigonometric", False) + for comp in getattr(P.space, "spaces", []) + for sp in getattr(comp, "spaces", [comp]) + ): + raise NotImplementedError( + "BasisProjectionOperator needs B-splines in every direction (not angular_basis='trig')" + ) # only for M1 Mac users PSYDAC_BACKEND_GPYCCEL["flags"] = "-O3 -march=native -mtune=native -ffast-math -ffree-line-length-none" diff --git a/src/struphy/feec/projectors.py b/src/struphy/feec/projectors.py index 1f87674ac..64cd6b57e 100644 --- a/src/struphy/feec/projectors.py +++ b/src/struphy/feec/projectors.py @@ -41,14 +41,6 @@ ) from struphy.kernel_arguments.local_projectors_args_cuda import CudaLocalProjectorsArguments from struphy.kernel_arguments.local_projectors_args_kernels import LocalProjectorsArguments -from struphy.polar.kron_2dx1d import ( - PolarKron2Dx1DOperator, - PolarKron2Dx1DSolver, - bc_masks, - projector_groups, - spline_bases, -) -from struphy.polar.linear_operators import PolarExtractionOperator logger = logging.getLogger("struphy") @@ -354,20 +346,10 @@ class CommutingProjector: boundary_op : BoundaryOperator The boundary operator applying essential boundary conditions to a vector. If not given, is set to identity. - polar_rtol : float - Relative tolerance (w.r.t. the norm of the right-hand side) of the iterative solver for polar projections - (only used if the poloidal solver is not exact, i.e. for ``"ilu"``). - - polar_maxiter : int - Maximum number of iterations of the iterative solver for polar projections. - - poloidal_solver : str - Solver for the poloidal (eta1-eta2) matrices of the 2D x 1D Kronecker structure of polar projections: - ``"schur"`` (exact, block elimination of the pole block), ``"lu"`` (exact, sparse LU) or ``"ilu"`` - (incomplete LU, used as preconditioner of an iterative solver), see :mod:`struphy.polar.kron_2dx1d`. - - poloidal_solver_kwargs : dict | None - Further options of the poloidal solver (``drop_tol``, ``fill_factor`` of the incomplete LU factorization). + polar : bool + Whether the projector maps to polar splines. The polar projection is the tensor-product projection followed + by the polar dual coefficients, :math:`\bar\pi = \tilde F^\top \pi`: the ``dofs_extraction_op`` is then + :math:`\tilde F^\top \mathcal I^{-1}` and the polar inter-/histopolation matrix is the identity. """ def __init__( @@ -376,10 +358,7 @@ def __init__( dofs_extraction_op=None, base_extraction_op=None, boundary_op=None, - polar_rtol: float = 1e-13, - polar_maxiter: int = 1000, - poloidal_solver: str = "schur", - poloidal_solver_kwargs: dict | None = None, + polar: bool = False, ): self._projector_tensor = projector_tensor @@ -402,7 +381,7 @@ def __init__( else: self._boundary_op = IdentityOperator(self.space.coeff_space) - self._is_polar = isinstance(self.dofs_extraction_op, PolarExtractionOperator) + self._is_polar = polar self._imat = projector_tensor.imat_kronecker self._imatT = self._imat.T @@ -411,43 +390,15 @@ def __init__( E = self._base_extraction_op B = self._boundary_op - self._polar_rtol = polar_rtol - self._polar_maxiter = polar_maxiter - self._poloidal_solver = poloidal_solver - self._poloidal_solver_kwargs = poloidal_solver_kwargs or {} - - if self._is_polar: - # 2D x 1D Kronecker structure (P_pol (I_1 x I_2) E_pol^T) x I_3 of the polar inter-/histopolation matrix - polar_space = P.codomain - groups = projector_groups(polar_space, self.space, P, E) - bc = getattr(B, "bc", None) - bases = spline_bases(self.space) - masks = [bc_masks(polar_space, g.comps, bases, bc) for g in groups] - - # factors of the operators to invert, key (transposed, apply_bc); boundary conditions as D A D + I - D - groups_bc = [g.masked(*m, True, polar_space) for g, m in zip(groups, masks)] - self._polar_groups = { - (False, False): groups, - (False, True): groups_bc, - (True, False): [g.transpose() for g in groups], - (True, True): [g.transpose() for g in groups_bc], - } - - # inter-/histopolation matrices I = P * I * E^T and I0 = B * P * I * E^T * B^T - self._I = PolarKron2Dx1DOperator(polar_space, groups) - self._I0 = PolarKron2Dx1DOperator( - polar_space, [g.masked(*m, False, polar_space) for g, m in zip(groups, masks)] - ) - self._IT = self._I.transpose() - self._I0T = self._I0.transpose() - - # solvers are built on first use - self._polar_solvers = {} - self._polar_krylov = {} - - # transposed boundary operator and buffer for the right-hand side with boundary conditions - self._boundary_op_T = self._boundary_op.T - self._polar_rhs_bc = self._boundary_op_T.codomain.zeros() + if polar: + # polar splines: the dofs extraction operator is F~^T I^{-1}, so that the polar "DOFs" are the polar + # coefficients of the commuting projection pi_bar = F~^T pi and P I E^T is the identity + polar_space = self._dofs_extraction_op.codomain + self._I = IdentityOperator(polar_space) + self._I0 = B @ B.T + self._IT = self._I + self._I0T = self._I0 + self._pc = self._pc0 = self._pcT = self._pc0T = IdentityOperator(polar_space) else: # build inter-/histopolation matrix I = ID * P * I * E^T * ID^T and I0 = B * P * I * E^T * B^T as ComposedLinearOperator self._I = P @ self._imat @ E.T @@ -463,8 +414,6 @@ def __init__( self._pcT = ProjectorPreconditioner(self, transposed=True, apply_bc=False) self._pc0T = ProjectorPreconditioner(self, transposed=True, apply_bc=True) - self._polar_info = None - @property def projector_tensor(self): """Tensor product projector.""" @@ -495,11 +444,6 @@ def is_polar(self): """Whether the projector maps to polar splines (True) or pure tensor product splines.""" return self._is_polar - @property - def polar_info(self): - """Info dict of the iterative solver of the last polar projection (None before the first one).""" - return self._polar_info - @property def I(self): """Inter-/histopolation matrix ID * P * I * E^T * ID^T as ComposedLinearOperator (ID = IdentityOperator).""" @@ -522,40 +466,23 @@ def I0T(self): @property def pc(self): - """(Approximate) inverse of I: Kronecker solver (tensor product), resp. 2D x 1D Kronecker solver (polar).""" - return self._polar_solver(False, False) if self.is_polar else self._pc + """Inverse of I: Kronecker solver (tensor product), resp. the identity (polar).""" + return self._pc @property def pc0(self): """(Approximate) inverse of I0 (with boundary conditions).""" - return self._polar_solver(False, True) if self.is_polar else self._pc0 + return self._pc0 @property def pcT(self): """(Approximate) inverse of the transposed I.""" - return self._polar_solver(True, False) if self.is_polar else self._pcT + return self._pcT @property def pc0T(self): """(Approximate) inverse of the transposed I0 (with boundary conditions).""" - return self._polar_solver(True, True) if self.is_polar else self._pc0T - - @property - def poloidal_solver(self) -> str: - """Solver for the poloidal matrices of polar projections ("schur", "lu" or "ilu").""" - return self._poloidal_solver - - def _polar_solver(self, transposed: bool, apply_bc: bool) -> PolarKron2Dx1DSolver: - """2D x 1D Kronecker solver of the polar inter-/histopolation matrix (built on first use).""" - key = (transposed, apply_bc) - if key not in self._polar_solvers: - self._polar_solvers[key] = PolarKron2Dx1DSolver( - self._I.domain, - self._polar_groups[key], - self._poloidal_solver, - **self._poloidal_solver_kwargs, - ) - return self._polar_solvers[key] + return self._pc0T def solve(self, rhs, transposed=False, apply_bc=False, out=None, x0=None): """ @@ -584,50 +511,22 @@ def solve(self, rhs, transposed=False, apply_bc=False, out=None, x0=None): assert isinstance(rhs, Vector) assert rhs.space == self._I.domain - # standard (tensor product) case (Kronecker solver) - if not self.is_polar: - if transposed: - pc = self.pc0T if apply_bc else self.pcT - else: - pc = self.pc0 if apply_bc else self.pc - return pc.solve(rhs, out=out) - - # polar case: 2D x 1D Kronecker solver - solver = self._polar_solver(transposed, apply_bc) - - # boundary conditions are only imposed if requested (as in the tensor product case) - if apply_bc: - rhs = self._boundary_op_T.dot(rhs, out=self._polar_rhs_bc) - - if solver.exact: - return solver.dot(rhs, out=out) - - # approximate poloidal solver (ILU): preconditioner of an iterative solve - key = (transposed, apply_bc) - if key not in self._polar_krylov: - if transposed: - op = self._I0T if apply_bc else self._IT + # polar splines: the polar inter-/histopolation matrix is the identity + if self.is_polar: + if out is None: + out = rhs.copy() else: - op = self._I0 if apply_bc else self._I - self._polar_krylov[key] = inverse( - op, "pbicgstab", pc=solver, tol=self._polar_rtol, maxiter=self._polar_maxiter, verbose=False - ) - krylov = self._polar_krylov[key] - - # relative tolerance w.r.t. the norm of the right-hand side - rhs_norm = float(xp.sqrt(rhs.inner(rhs))) - krylov.set_options(x0=x0, tol=self._polar_rtol * rhs_norm if rhs_norm > 0.0 else self._polar_rtol) - x = krylov.solve(rhs, out=out) - - info = krylov.get_info() - self._polar_info = info - if not info["success"]: - logger.warning( - f"Polar projection did not converge: {info['niter']} iterations, " - f"residual norm {info['res_norm']:.2e} (rhs norm {rhs_norm:.2e}, rtol {self._polar_rtol:.1e}).", - ) + rhs.copy(out=out) + if apply_bc: + out = self._boundary_op.dot(out) + return out - return x + # standard (tensor product) case (Kronecker solver) + if transposed: + pc = self.pc0T if apply_bc else self.pcT + else: + pc = self.pc0 if apply_bc else self.pc + return pc.solve(rhs, out=out) def get_dofs(self, fun, dofs=None, apply_bc=False): """ diff --git a/src/struphy/feec/psydac_derham.py b/src/struphy/feec/psydac_derham.py index 31fca09c3..2773f0fb0 100644 --- a/src/struphy/feec/psydac_derham.py +++ b/src/struphy/feec/psydac_derham.py @@ -37,6 +37,7 @@ from struphy.feec.linear_operators import BoundaryOperator from struphy.feec.local_projectors_kernels import get_local_problem_size, select_quasi_points from struphy.feec.projectors import CommutingProjector, CommutingProjectorLocal +from struphy.feec.trig_space import TrigonometricSpace, tensor_fem_space from struphy.feec.utilities import get_quad_grids from struphy.fields_background.base import FluidEquilibrium, MHDequilibrium from struphy.geometry.base import Domain @@ -50,7 +51,6 @@ from struphy.kernel_arguments.spline_args_cuda import CudaSplineArguments from struphy.kernel_arguments.spline_args_kernels import SplineArguments from struphy.polar.basic import PolarDerhamSpace, PolarVector -from struphy.polar.extraction_operators import PolarExtractionBlocksC1 from struphy.polar.linear_operators import PolarExtractionOperator, PolarLinearOperator from struphy.topology.grids import TensorProductGrid from struphy.utils.utils import check_option @@ -147,8 +147,8 @@ def __init__(self, V0: TensorFemSpace, V1: VectorFemSpace, V2: VectorFemSpace, V [L2_space.ncells], [L2_space.periodic], ) - H1_space_tensor = TensorFemSpace(H1_domain_decomp, H1_space) - L2_space_tensor = TensorFemSpace(L2_domain_decomp, L2_space) + H1_space_tensor = tensor_fem_space(H1_domain_decomp, H1_space) + L2_space_tensor = tensor_fem_space(L2_domain_decomp, L2_space) H1_space_tensor.symbolic_space = f"H1_1d_eta{d + 1}" L2_space_tensor.symbolic_space = f"L2_1d_eta{d + 1}" @@ -284,11 +284,6 @@ class SplineAttributes1D: Number of quadrature points in each logical direction used to build projection point/weight sets. - polar_splines : bool, default=False - Polar spline regularity flag. ``False`` selects standard tensor-product - splines, while ``True`` enables C1 polar splines (including the optional - shift of points near ``eta_1 = 0`` where required). - local_projectors : bool, default=False If ``True``, also computes quasi-interpolation point/weight grids used by local commuting projectors. @@ -313,7 +308,6 @@ def __init__( femspace: TensorFemSpace | VectorFemSpace, nquads: tuple[int, int, int], nquads_proj: tuple[int, int, int], - polar_splines: bool = False, local_projectors: bool = False, ): # inputs @@ -322,7 +316,6 @@ def __init__( self._nquads = nquads self._nq_pr = nquads_proj self._local_projectors = local_projectors - self._polar_splines = polar_splines # grid attributes self._nbasis = [] @@ -388,10 +381,7 @@ def __init__( ] if local_projectors: - ptsloc, wtsloc = get_pts_and_wts_quasi( - space, - polar_shift=d == 0 and self.polar_splines, - ) + ptsloc, wtsloc = get_pts_and_wts_quasi(space) self._proj_loc_grid_pts[-1] += [ptsloc] self._proj_loc_grid_wts[-1] += [wtsloc] @@ -400,7 +390,6 @@ def __init__( s, e, n_quad=self.nquads_proj[d], - polar_shift=d == 0 and self.polar_splines, ) self._proj_grid_subs[-1] += [subs] @@ -459,11 +448,6 @@ def nquads_proj(self) -> tuple[int, int, int]: """The number of quadrature points in each direction for the projection.""" return self._nq_pr - @property - def polar_splines(self) -> bool: - """The polar splines flag.""" - return self._polar_splines - @property def local_projectors(self) -> bool: """Whether to use local projectors.""" @@ -550,7 +534,7 @@ class Derham: """High-level discrete de Rham complex on the 3D logical unit cube. The class builds a tensor-product base complex through ``DiscreteDerham`` - and augments it with optional boundary enforcement, optional C1 polar + and augments it with optional boundary enforcement, optional polar extraction operators, and commuting projectors. It also stores convenient metadata (1D spline attributes, decomposition helpers, index arrays) used by initialization, evaluation, and kernel interfaces. @@ -668,11 +652,25 @@ def __init__( assert len(nquads_proj) == 3 self._nq_pr = tuple(nquads_proj) - # smoothness at the polar singularity: default False for standard tensor product splines, or True for C1 polar splines + # smoothness at the polar singularity: False for standard tensor product splines, True for polar splines of + # regularity polar_regularity self._polar_splines = polar_splines - # solver for the poloidal matrices in the 2D x 1D Kronecker solves of polar projections - self._polar_poloidal_solver = options.polar_poloidal_solver + # basis in eta2: periodic B-splines or trigonometric functions (num_elements[1] of them, odd) + self._angular_basis = options.angular_basis + if self._angular_basis == "trig": + assert self.spl_kind[1], "trigonometric functions need a periodic eta2 (bcs[1] = None)" + assert num_elements[1] % 2 == 1 and num_elements[1] >= 3, ( + "with angular_basis='trig', num_elements[1] is the (odd) number of trigonometric functions" + ) + assert mpi_dims_mask is None or not mpi_dims_mask[1], ( + "trigonometric functions in eta2 need an undecomposed eta2: use mpi_dims_mask = (., False, .)" + ) + assert not local_projectors, "local projectors do not support trigonometric functions" + self._polar_regularity = options.polar_regularity + self._polar_enrichment = options.polar_enrichment + self._polar_polynomial_projection = options.polar_polynomial_projection + self._polar_cross_section = options.polar_cross_section # Other derived properties self._mpi_dims_mask = mpi_dims_mask @@ -715,35 +713,30 @@ def __init__( derham.V0, self.nquads, self.nquads_proj, - polar_splines=self.polar_splines, local_projectors=local_projectors, ) self._V1splines = SplineAttributes1D( derham.V1, self.nquads, self.nquads_proj, - polar_splines=self.polar_splines, local_projectors=local_projectors, ) self._V2splines = SplineAttributes1D( derham.V2, self.nquads, self.nquads_proj, - polar_splines=self.polar_splines, local_projectors=local_projectors, ) self._V3splines = SplineAttributes1D( derham.V3, self.nquads, self.nquads_proj, - polar_splines=self.polar_splines, local_projectors=local_projectors, ) self._Vvsplines = SplineAttributes1D( derham.Vv, self.nquads, self.nquads_proj, - polar_splines=self.polar_splines, local_projectors=local_projectors, ) @@ -755,7 +748,6 @@ def __init__( H1_space, (self.nquads[n],), (self.nquads_proj[n],), - polar_splines=False, # the 1d spaces are not polar, even if the 3d space is local_projectors=False, ) ] @@ -764,7 +756,6 @@ def __init__( L2_space, (self.nquads[n],), (self.nquads_proj[n],), - polar_splines=False, # the 1d spaces are not polar, even if the 3d space is local_projectors=False, ) ] @@ -807,8 +798,16 @@ def __init__( if not self.polar_splines: self._ck_blocks = None else: - assert self.domain is not None - self._ck_blocks = PolarExtractionBlocksC1(self.domain, self) + from struphy.polar.homotopy.derham_blocks import PolarExtractionBlocksHomotopy + + self._ck_blocks = PolarExtractionBlocksHomotopy( + self.domain, + self, + regularity=self.polar_regularity, + enrichment=self.polar_enrichment, + polynomial_projection=self.polar_polynomial_projection, + cross_section=self.polar_cross_section, + ) # Polar spaces share live factors; obsolete mass-weight factors are not retained. self._polar_factorization_cache = WeakValueDictionary() @@ -1010,13 +1009,34 @@ def with_projectors(self) -> bool: @property def polar_splines(self) -> bool: - """C^k smoothness at eta_1=0. Is False for standard tensor product splines and True for C^1 polar splines.""" + """Smoothness at eta_1=0: False for standard tensor product splines, True for polar splines of regularity + ``polar_regularity``.""" return self._polar_splines @property - def polar_poloidal_solver(self) -> str: - """Solver for the poloidal matrices in the 2D x 1D Kronecker solves of polar projections ("schur", "lu" or "ilu").""" - return self._polar_poloidal_solver + def angular_basis(self) -> str: + """Basis in eta2: ``"bspline"`` (periodic B-splines) or ``"trig"`` (trigonometric functions).""" + return getattr(self, "_angular_basis", "bspline") + + @property + def polar_regularity(self) -> int: + """Regularity at the pole of the "homotopy" polar splines.""" + return self._polar_regularity + + @property + def polar_enrichment(self) -> str: + """Enrichment of the "homotopy" polar spaces.""" + return self._polar_enrichment + + @property + def polar_polynomial_projection(self) -> str: + """Projection of the polynomials in the "homotopy" polar construction.""" + return self._polar_polynomial_projection + + @property + def polar_cross_section(self) -> str: + """Reference cross-section of the "homotopy" polar complex.""" + return self._polar_cross_section @property def with_local_projectors(self) -> bool: @@ -1219,8 +1239,9 @@ def polar_coeff_spaces(self) -> dict[str, StencilVectorSpace | BlockVectorSpace } @property - def ck_blocks(self) -> PolarExtractionBlocksC1 | None: - """Polar extraction blocks for C1 polar splines. Is None if polar_splines is False (standard tensor product splines).""" + def ck_blocks(self): + """Polar extraction blocks (:class:`~struphy.polar.homotopy.derham_blocks.PolarExtractionBlocksHomotopy`). Is + None if polar_splines is False (standard tensor product splines).""" return self._ck_blocks @property @@ -1520,8 +1541,17 @@ def div_tp(self): @property def args_derham(self) -> DerhamArguments | CudaDerhamArguments: """Mandatory pusher kernel arguments for the backend used at initialization.""" + self._require_splines("particles (pusher and accumulation kernels)") return self._args_derham + def _require_splines(self, feature): + """Raise for features whose kernels evaluate B-splines in every direction.""" + if self.angular_basis == "trig": + raise NotImplementedError( + f"{feature} need B-splines in every direction; trigonometric functions in eta2 " + "(DerhamOptions(angular_basis='trig')) support FEEC operations only" + ) + def kernels(self) -> tuple: """The kernels called on the data of this complex. @@ -1735,6 +1765,7 @@ def prepare_eval_tp_fixed(self, grids_1d): bds : 3-tuple of 3d float arrays Values of degree non-zero D-Splines at quadrature points in format (n, nq, basis). """ + self._require_splines("prepare_eval_tp_fixed (fixed-grid spline evaluation)") # spline degree and knot vectors must come from N-spline spaces (V0 space) spans, bns, bds = [], [], [] @@ -1868,10 +1899,13 @@ def _discretize_space( SplineSpace(degree, multiplicity=m, grid=grid, periodic=P) for degree, m, grid, P in zip(degree_i, multiplicity_i, grids, periodic) ] + # trigonometric functions in eta2 (num_elements[1] functions, see struphy.feec.trig_space) + if getattr(self, "_angular_basis", "bspline") == "trig": + spaces_1d[1] = TrigonometricSpace(ncells[1], "B") carts = create_cart([ddm], [spaces_1d]) - Vh = TensorFemSpace(ddm, *spaces_1d, cart=carts[0]) + Vh = tensor_fem_space(ddm, *spaces_1d, cart=carts[0]) if V == "H1": Wh = Vh @@ -1929,7 +1963,7 @@ def _assemble_polar_extraction_operators( pol_space, ) - # C^1 polar spline case + # polar spline case else: pol_space = PolarDerhamSpace(self, sp_id) @@ -1956,13 +1990,18 @@ def _assemble_polar_extraction_operators( def _assemble_projectors(self, *projectors: GlobalGeometricProjector) -> tuple[CommutingProjector, ...]: tmp = [] for (sp_id, sp_form), projector in zip(self.space_to_form.items(), projectors): + if self.polar_splines: + # the polar commuting projection is pi_bar = F~^T pi (paper): the polar "DOFs" are the polar + # coefficients F~^T I^{-1} dofs, and the polar inter-/histopolation matrix is the identity + self._dofs_extraction_ops[sp_form] = self._dofs_extraction_ops[sp_form] @ projector.solver + self._dofs_extraction_ops[sp_id] = self._dofs_extraction_ops[sp_form] tmp.append( CommutingProjector( projector, dofs_extraction_op=self._dofs_extraction_ops[sp_form], base_extraction_op=self._extraction_ops[sp_form], boundary_op=self.boundary_ops[sp_form], - poloidal_solver=self.polar_poloidal_solver, + polar=self.polar_splines, ) ) return tuple(tmp) @@ -2889,6 +2928,43 @@ def eval_tp_fixed_loc(self, spans, bases, out=None): return out + def _eval_with_1d_bases(self, *etas, squeeze_out=False): + """Evaluation by contracting the global coefficients with the 1D basis values of each direction (B-/M-splines + or trigonometric functions); markers ``(N, 3)`` or tensor grids, as :meth:`__call__`.""" + import numpy + + stencils = [self._vector_stencil] if isinstance(self._vector_stencil, StencilVector) else self._vector_stencil + fem = self.fem_space + components = [fem] if isinstance(fem, TensorFemSpace) else list(fem.spaces) + + def basis_1d(space, x): + from struphy.polar.homotopy.derham_spaces import basis_values_1d + + return basis_values_1d(space, numpy.asarray(xp.to_numpy(x), dtype=float)) + + if len(etas) == 1: + markers = numpy.asarray(xp.to_numpy(etas[0]))[:, :3] + axes = None + else: + E1, E2, E3, _ = Domain.prepare_eval_pts(*etas) + axes = [numpy.asarray(xp.to_numpy(E))[:, 0, 0] for E in (E1,)] + axes += [numpy.asarray(xp.to_numpy(E2))[0, :, 0], numpy.asarray(xp.to_numpy(E3))[0, 0, :]] + values = [] + for space, stencil in zip(components, stencils): + coefficients = numpy.asarray(xp.to_numpy(stencil.toarray())).reshape(stencil.space.npts) + if self.derham.comm is not None and self.derham.comm.Get_size() > 1: + coefficients = self.derham.comm.allreduce(coefficients) + if axes is None: + B = [basis_1d(V, markers[:, d]) for d, V in enumerate(space.spaces)] + values.append(numpy.einsum("ijk,mi,mj,mk->m", coefficients, *B)) + else: + B = [basis_1d(V, axes[d]) for d, V in enumerate(space.spaces)] + values.append(numpy.einsum("ijk,ai,bj,ck->abc", coefficients, *B)) + if squeeze_out: + values = [numpy.squeeze(v) for v in values] + values = [xp.asarray(v) for v in values] + return values[0] if len(values) == 1 else values + def __call__(self, *etas, out=None, tmp=None, squeeze_out=False, local=False): """ Evaluates the spline function on the global domain, unless local=True, @@ -2924,6 +3000,10 @@ def __call__(self, *etas, out=None, tmp=None, squeeze_out=False, local=False): # extract coefficients and update ghost regions self.extract_coeffs(update_ghost_regions=True) + # trigonometric functions in eta2: evaluation with the 1D bases (the spline kernels assume B-splines) + if self.derham.angular_basis == "trig": + return self._eval_with_1d_bases(*etas, squeeze_out=squeeze_out) + # marker evaluation if len(etas) == 1: marker_evaluation = True diff --git a/src/struphy/feec/tests/test_trig_space.py b/src/struphy/feec/tests/test_trig_space.py new file mode 100644 index 000000000..9af16c8ee --- /dev/null +++ b/src/struphy/feec/tests/test_trig_space.py @@ -0,0 +1,134 @@ +"""Trigonometric functions in eta2 of the Derham complex (``DerhamOptions(angular_basis="trig")``).""" + +import numpy as np +import pytest + + +def _derham(n=9, polar=False, regularity=2, domain=None, num_elements=(6, None, 1), degree=(3, 3, 1)): + from struphy.feec.psydac_derham import Derham + from struphy.geometry import domains + from struphy.io.options import DerhamOptions + from struphy.topology.grids import TensorProductGrid + + domain = domains.HollowCylinder(a1=0.0, a2=1.0, Lz=1.0) if domain is None else domain + options = DerhamOptions( + degree=degree, + bcs=(("free", "free"), None, None), + angular_basis="trig", + polar_splines=polar, + polar_regularity=regularity, + nquads=(degree[0] + 1, n + 2, 2), + ) + grid = TensorProductGrid(num_elements=(num_elements[0], n, num_elements[2]), mpi_dims_mask=(True, False, False)) + return Derham(grid, options, domain=domain), domain + + +def test_nodal_bases(): + from struphy.feec.trig_space import cardinal_values, histopolant_values + + n = 9 + nodes = np.arange(n) / n + assert np.allclose(cardinal_values(nodes, n)[0], np.eye(n)) + u, w = np.polynomial.legendre.leggauss(20) + integrals = np.array([(w / (2 * n)) @ histopolant_values((j + 0.5 + 0.5 * u) / n, n)[0] for j in range(n)]) + assert np.allclose(integrals, np.eye(n), atol=1e-13) + x = np.linspace(0, 1, 37) + dL, H = cardinal_values(x, n, 1)[1], histopolant_values(x, n)[0] + assert np.allclose(dL, np.roll(H, 1, axis=1) - H) + assert np.allclose(cardinal_values(x, n)[0].sum(axis=1), 1.0) + + +def test_masses_equal_a_dense_assembly(): + """Mass matrices assembled by Struphy's kernels (all functions active on every cell, duplicate stencil offsets) + equal a dense assembly with the 1D bases.""" + from struphy.feec.mass import WeightedMassOperators + from struphy.polar.homotopy.derham_spaces import basis_values_1d + + derham, domain = _derham(n=7, num_elements=(4, None, 1)) + mass = WeightedMassOperators(derham, domain) + M0 = mass.M0.matrix.tosparse().toarray() + # dense: int B_i B_j det(DF) on Gauss points of the cells + V = derham.V0fem.spaces + x1, w1 = _gauss_on(V[0].breaks, 6) + x2, w2 = _gauss_on(V[1].breaks, 12) + B1, B2 = basis_values_1d(V[0], x1), basis_values_1d(V[1], x2) + det = np.asarray(domain.jacobian_det(x1, x2, np.array([0.5])))[:, :, 0] + weights = np.outer(w1, w2) * det + dense = np.einsum("ab,ai,bj,ak,bl->ijkl", weights, B1, B2, B1, B2).reshape(B1.shape[1] * B2.shape[1], -1) + assert np.allclose(M0, dense, atol=1e-12 * np.abs(dense).max()) + assert np.allclose(mass.M0.matrix.transpose().tosparse().toarray(), dense.T, atol=1e-12 * np.abs(dense).max()) + + +def _gauss_on(breaks, q): + u, w = np.polynomial.legendre.leggauss(q) + a, b = breaks[:-1, None], breaks[1:, None] + return ((a + b) / 2 + (b - a) / 2 * u).ravel(), ((b - a) / 2 * w).ravel() + + +def test_projections_derivatives_and_evaluation(): + derham, domain = _derham() + + def f(a, b, c): + return (1 + a**2) * np.cos(2 * np.pi * b) + 0.3 * a * np.sin(4 * np.pi * b) + 0 * c + + def grad_f(a, b, c): + return ( + 2 * a * np.cos(2 * np.pi * b) + 0.3 * np.sin(4 * np.pi * b) + 0 * c, + -2 * np.pi * (1 + a**2) * np.sin(2 * np.pi * b) + 1.2 * np.pi * a * np.cos(4 * np.pi * b) + 0 * c, + 0 * a, + ) + + u = derham.P0(f) + g = derham.P1([lambda a, b, c, i=i: grad_f(a, b, c)[i] for i in range(3)]) + r = derham.grad.dot(u) - g + assert np.sqrt(r.inner(r) / g.inner(g)) < 1e-7 + # trigonometric polynomials of degree <= 4 in eta2 are interpolated exactly along eta2 + F = derham.create_spline_function("u", "H1") + F.vector = u + e1, e2 = np.array(derham.V0fem.spaces[0].greville), np.linspace(0, 1, 13) + values = np.asarray(F(e1, e2, np.array([0.5])))[:, :, 0] + assert np.allclose(values, f(*np.meshgrid(e1, e2, indexing="ij"), 0), atol=1e-12) + markers = np.column_stack([np.full(13, e1[2]), e2, np.full(13, 0.5)]) + assert np.allclose(np.asarray(F(markers)), values[2]) + + +def test_polar_splines_with_trigonometric_functions_equal_the_reference(): + """The polar complex on the Derham's trigonometric spaces spans the spaces of the reference construction + (PolarFEEC.jl's modal trigonometric spaces) in every degree. (The polar duals are pseudo-inverses in coefficient + space, so the commuting projections depend on the angular coordinates for r >= 2.)""" + from struphy.polar.homotopy import PolarComplex2D + from struphy.polar.homotopy.demos import common + from struphy.polar.tests.reference_core import ClampedBSplineSpace, TrigonometricSpace + + e1, e2 = np.linspace(0.02, 1, 23), np.linspace(0, 1, 41) + for r in (1, 2, 3): + struphy = common.polar_complex(r, 5, (6, 13), domain=common.star_disk(), angular="trig") + reference = PolarComplex2D( + ClampedBSplineSpace(5, 6), TrigonometricSpace(6), common.cross_section(common.star_disk()), r + ) + assert [struphy.dimension(k) for k in range(3)] == [reference.dimension(k) for k in range(3)] + for k in range(3): + values = [] + for c in (struphy, reference): + E = c.extraction(k).toarray() + columns = [c.spaces[k].evaluate(E[:, j], e1, e2) for j in range(E.shape[1])] + values.append(np.array([np.ravel(v) for v in columns])) + A, B = values + ranks = [np.linalg.matrix_rank(M, tol=1e-8 * np.abs(M).max()) for M in (A, B, np.vstack([A, B]))] + assert ranks == [A.shape[0]] * 3, (r, k, ranks) + assert struphy.cohomology_dimensions() == [1, 0, 0] + + +def test_trigonometric_eta2_limits(): + derham, _ = _derham() + with pytest.raises(NotImplementedError, match="particles"): + derham.args_derham + from struphy.feec.psydac_derham import Derham + from struphy.io.options import DerhamOptions + from struphy.topology.grids import TensorProductGrid + + with pytest.raises(AssertionError, match="odd"): + Derham( + TensorProductGrid(num_elements=(4, 8, 1), mpi_dims_mask=(True, False, False)), + DerhamOptions(degree=(2, 2, 1), bcs=(("free", "free"), None, None), angular_basis="trig"), + ) diff --git a/src/struphy/feec/trig_space.py b/src/struphy/feec/trig_space.py new file mode 100644 index 000000000..e2e766ce2 --- /dev/null +++ b/src/struphy/feec/trig_space.py @@ -0,0 +1,224 @@ +r"""Trigonometric 1D spaces for the Derham complex (periodic directions). + +The space of trigonometric polynomials of degree ``M`` on ``[0, 1)`` has dimension ``n = 2 M + 1``. Struphy uses the +*nodal* bases on the points :math:`x_i = i / n`: + +* 0-forms (``basis="B"``): the cardinal functions (Dirichlet kernels) + :math:`L_i(x) = \sin(n \pi (x - x_i)) / (n \sin(\pi (x - x_i)))`, :math:`L_i(x_j) = \delta_{ij}`; +* 1-forms (``basis="M"``): the cell histopolants + :math:`H_i(x) = 1 + \sum_{k=1}^{M} \frac{2 \pi k}{n \sin(\pi k / n)} \cos(2 \pi k (x - x_i - \tfrac{1}{2n}))`, + :math:`\int_{x_j}^{x_{j+1}} H_i = \delta_{ij}`. + +They satisfy :math:`L_i' = H_{i-1} - H_i`, so the coefficient derivative is the forward difference of the B-splines +and the interpolation and histopolation matrices are the identity: the trigonometric spaces form the same 1D +complex as periodic splines, with the geometric degrees of freedom at the nodes and on the cells. + +For the FEEC data structures the spaces present themselves as periodic splines on ``n`` cells (the cells +``[x_i, x_{i+1}]``) whose basis functions are all active on every cell: the 0-form space has "degree" ``n`` (``n + 1`` +local slots per cell, one of them a zero duplicate), the 1-form space degree ``n - 1``, both with ``n`` pads. The +assembly grids (:class:`TrigonometricAssemblyGrid`) provide the basis values, so the mass-matrix kernels of Struphy +work unchanged. Particles, local projectors and spline-kernel evaluation do not support trigonometric spaces. +""" + +import numpy as np +from feectools.ddm.cart import CartDecomposition, DomainDecomposition +from feectools.fem.grid import FemAssemblyGrid +from feectools.fem.partitioning import partition_coefficients +from feectools.fem.splines import SplineSpace +from feectools.fem.tensor import TensorFemSpace +from feectools.linalg.direct_solvers import SparseSolver +from feectools.linalg.stencil import StencilVectorSpace +from scipy.sparse import csc_matrix, identity + + +# ---------------------------------------------------------------------------------------------------------- bases +def cardinal_values(x, n, nderiv=0): + """Values (and derivatives) of the cardinal trigonometric functions ``L_i``, shape ``(nderiv + 1, len(x), n)``.""" + x = np.atleast_1d(np.asarray(x, dtype=float)) + modes = (n - 1) // 2 + k = np.arange(1, modes + 1) + t = x[:, None, None] - np.arange(n)[None, :, None] / n # (points, functions, 1) + phase = 2 * np.pi * k[None, None, :] * t + out = np.empty((nderiv + 1, x.size, n)) + for d in range(nderiv + 1): + factor = (2 * np.pi * k) ** d + # d-th derivative of cos is cos(. + d pi / 2) + values = np.sum(factor * np.cos(phase + d * np.pi / 2), axis=2) + out[d] = (2 * values + (1.0 if d == 0 else 0.0)) / n + return out + + +def histopolant_values(x, n, nderiv=0): + """Values (and derivatives) of the cell histopolants ``H_i``, shape ``(nderiv + 1, len(x), n)``.""" + x = np.atleast_1d(np.asarray(x, dtype=float)) + modes = (n - 1) // 2 + k = np.arange(1, modes + 1) + weight = 2 * np.pi * k / (n * np.sin(np.pi * k / n)) + t = x[:, None, None] - (np.arange(n)[None, :, None] + 0.5) / n + phase = 2 * np.pi * k[None, None, :] * t + out = np.empty((nderiv + 1, x.size, n)) + for d in range(nderiv + 1): + factor = weight * (2 * np.pi * k) ** d + out[d] = np.sum(factor * np.cos(phase + d * np.pi / 2), axis=2) + (1.0 if d == 0 else 0.0) + return out + + +# ---------------------------------------------------------------------------------------------------------- spaces +class TrigonometricSpace(SplineSpace): + """Nodal trigonometric space with ``n`` (odd) functions on the periodic interval, a 0-form (``basis="B"``) or + 1-form (``basis="M"``) space of the Derham complex (see the module documentation).""" + + is_trigonometric = True + + def __init__(self, n, basis="B"): + if n % 2 == 0 or n < 3: + raise ValueError(f"the number of trigonometric functions must be odd and at least 3 (got {n})") + if basis not in ("B", "M"): + raise ValueError("basis must be 'B' (0-forms) or 'M' (1-forms)") + n = int(n) + self.n_functions = n + self.modes = (n - 1) // 2 + nodes = np.arange(n + 1) / n + degree = n if basis == "B" else n - 1 + self._degree = degree + self._pads = n + # knots of a uniform periodic spline of the same "degree" (only for code that reads them) + self._knots = np.arange(-degree, n + degree + 1) / n + self._periodic = True + self._multiplicity = 1 + self._parent_multiplicity = 1 + self._dirichlet = (False, False) + self._basis = basis + self._nbasis = n + self._breaks = nodes + self._ncells = n + self._greville = nodes[:-1].copy() + self._ext_greville = nodes.copy() + self._scaling_array = None + self._histopolation_grid = nodes.copy() + domain_decomposition = DomainDecomposition([n], [True]) + cart = CartDecomposition(domain_decomposition, [n], [np.array([0])], [np.array([n - 1])], [self._pads], [1]) + self._coeff_space = StencilVectorSpace(cart) + self._interpolation_ready = False + self._histopolation_ready = False + self._symbolic_space = None + + def __repr__(self): + return f"TrigonometricSpace(n={self.n_functions}, basis={self._basis!r})" + + def basis_values(self, x, nderiv=0): + """Basis values and derivatives at the points ``x``, shape ``(nderiv + 1, len(x), n)``.""" + values = cardinal_values if self._basis == "B" else histopolant_values + return values(np.asarray(x, dtype=float) % 1.0, self.n_functions, nderiv) + + def partner(self): + """The other space of the 1D complex (1-forms of a 0-form space and vice versa).""" + return TrigonometricSpace(self.n_functions, "M" if self._basis == "B" else "B") + + def init_interpolation(self, dtype=float): + # the cardinal functions interpolate at the nodes + self.imat = np.eye(self.n_functions) + self._interpolator = SparseSolver(csc_matrix(identity(self.n_functions))) + self._interpolation_ready = True + + def init_histopolation(self, dtype=float): + # the histopolants histopolate on the cells + self.hmat = np.eye(self.n_functions) + self._histopolator = SparseSolver(csc_matrix(identity(self.n_functions))) + self._histopolation_ready = True + + +def is_trigonometric(space): + return getattr(space, "is_trigonometric", False) + + +class TrigonometricAssemblyGrid(FemAssemblyGrid): + """Quadrature data of a trigonometric space on the cells ``[x_i, x_{i+1}]`` (the interface of feectools' + ``FemAssemblyGrid``): all functions are active on every cell; the 0-form space has one zero duplicate slot.""" + + def __init__(self, space, start, end, *, nquads, nderiv=1): + n = space.n_functions + u, w = np.polynomial.legendre.leggauss(nquads) + breaks = space.breaks + a, b = breaks[:-1, None], breaks[1:, None] + points = (a + b) / 2 + (b - a) / 2 * u[None, :] + weights = (b - a) / 2 * w[None, :] + degree = space.degree + slots = degree + 1 + values = space.basis_values(points.ravel(), nderiv).reshape(nderiv + 1, n, nquads, n) + basis = np.zeros((n, slots, nderiv + 1, nquads)) + for cell in range(n): + for slot in range(min(slots, n)): + function = (cell + slot) % n + basis[cell, slot] = values[:, cell, :, function] + spans = np.arange(n) + degree + cells = slice(start, end + 1) + self._num_elements = end - start + 1 + self._num_quad_pts = nquads + self._spans = spans[cells].copy() + self._basis = basis[cells].copy() + self._points = points[cells].copy() + self._weights = weights[cells].copy() + self._indices = tuple(range(start, end + 1)) + self._quad_rule_x, self._quad_rule_w = u, w + self._local_element_start = 0 + self._local_element_end = self._num_elements - 1 + + num_elements = property(lambda self: self._num_elements) + num_quad_pts = property(lambda self: self._num_quad_pts) + spans = property(lambda self: self._spans) + basis = property(lambda self: self._basis) + points = property(lambda self: self._points) + weights = property(lambda self: self._weights) + indices = property(lambda self: self._indices) + quad_rule_x = property(lambda self: self._quad_rule_x) + quad_rule_w = property(lambda self: self._quad_rule_w) + local_element_start = property(lambda self: self._local_element_start) + local_element_end = property(lambda self: self._local_element_end) + + +class TrigonometricTensorFemSpace(TensorFemSpace): + """A ``TensorFemSpace`` with trigonometric directions: assembly grids and degree reduction (0-forms to 1-forms) + of the trigonometric spaces.""" + + def get_assembly_grids(self, *nquads): + assert len(nquads) == self.ldim + grids = [] + for i, (V, nq) in enumerate(zip(self.spaces, nquads)): + cache = self._assembly_grids[i] + if nq not in cache: + s, e = int(self._element_starts[i]), int(self._element_ends[i]) + if is_trigonometric(V): + cache[nq] = TrigonometricAssemblyGrid(V, s, e, nderiv=1, nquads=nq) + else: + cache[nq] = FemAssemblyGrid(V, s, e, nderiv=V.degree, nquads=nq) + grids.append(cache[nq]) + return tuple(grids) + + def reduce_degree(self, axes, multiplicity=None, basis="B"): + if isinstance(axes, int): + axes = [axes] + trig_axes = [a for a in axes if is_trigonometric(self.spaces[a])] + other_axes = [a for a in axes if a not in trig_axes] + if multiplicity is None: + multiplicity = [self.multiplicity[i] for i in axes] + elif isinstance(multiplicity, int): + multiplicity = [multiplicity] + other_multiplicity = [m for m, a in zip(multiplicity, axes) if a in other_axes] + reduced = super().reduce_degree(other_axes, other_multiplicity, basis) if other_axes else self + spaces = list(reduced.spaces) + for axis in trig_axes: + spaces[axis] = TrigonometricSpace(self.spaces[axis].n_functions, basis="M") + v = self.coeff_space + npts = [s.nbasis for s in spaces] + global_starts, global_ends = partition_coefficients(v.cart.domain_decomposition, spaces) + cart = v.cart.reduce_npts(npts, global_starts, global_ends, shifts=[s.multiplicity for s in spaces]) + space = TrigonometricTensorFemSpace(self._domain_decomposition, *spaces, cart=cart, dtype=v.dtype) + space._interpolation_ready = False + return space + + +def tensor_fem_space(domain_decomposition, *spaces, **kwargs): + """``TensorFemSpace`` of the 1D spaces (with trigonometric support if one of them is trigonometric).""" + cls = TrigonometricTensorFemSpace if any(is_trigonometric(s) for s in spaces) else TensorFemSpace + return cls(domain_decomposition, *spaces, **kwargs) diff --git a/src/struphy/fields_background/base.py b/src/struphy/fields_background/base.py index 426fd8f28..5ab9194d5 100644 --- a/src/struphy/fields_background/base.py +++ b/src/struphy/fields_background/base.py @@ -18,6 +18,15 @@ logger = logging.getLogger("struphy") +def _cartesian_vector(function, domain, etas, squeeze_out): + """Cartesian vector field ``function(x, y, z)`` at the logical points, without pull-back and push-forward (which + need the inverse Jacobian and are singular on a polar axis). Returns also (x, y, z).""" + xyz = domain(*etas, squeeze_out=squeeze_out) + values = function(xyz[0], xyz[1], xyz[2]) + out = xp.array([xp.broadcast_to(xp.asarray(c, dtype=float), xp.shape(xyz[0])) for c in values], dtype=float) + return out, xyz + + class FluidEquilibrium(metaclass=ABCMeta): """ Abstract base class for callable fluid equilibria on arbitrary domains. @@ -818,6 +827,24 @@ def gradB_xyz(self, x, y, z): def domain(self, new_domain): super(CartesianFluidEquilibriumWithB, type(self)).domain.fset(self, new_domain) + def b_cart(self, *etas, squeeze_out=False): + """Cartesian components of magnetic field evaluated on logical cube [0, 1]^3 (directly from ``b_xyz``, also on + a polar axis). Returns also (x,y,z).""" + return _cartesian_vector(self.b_xyz, self.domain, etas, squeeze_out) + + def gradB_cart(self, *etas, squeeze_out=False): + """Cartesian components of gradient of magnetic field strength (directly from ``gradB_xyz``). Returns also + (x,y,z).""" + return _cartesian_vector(self.gradB_xyz, self.domain, etas, squeeze_out) + + def b1(self, *etas, squeeze_out=False): + """1-form components of magnetic field (pulled back from ``b_xyz``, also on a polar axis).""" + return self.domain.pull(self.b_cart(*etas, squeeze_out=False)[0], *etas, kind="1", squeeze_out=squeeze_out) + + def b2(self, *etas, squeeze_out=False): + """2-form components of magnetic field (pulled back from ``b_xyz``, also on a polar axis).""" + return self.domain.pull(self.b_cart(*etas, squeeze_out=False)[0], *etas, kind="2", squeeze_out=squeeze_out) + class LogicalFluidEquilibriumWithB(LogicalFluidEquilibrium, FluidEquilibriumWithB): r""" @@ -1456,6 +1483,36 @@ def gradB_xyz(self, x, y, z): def domain(self, new_domain): super(CartesianMHDequilibrium, type(self)).domain.fset(self, new_domain) + def b_cart(self, *etas, squeeze_out=False): + """Cartesian components of magnetic field evaluated on logical cube [0, 1]^3 (directly from ``b_xyz``, also on + a polar axis). Returns also (x,y,z).""" + return _cartesian_vector(self.b_xyz, self.domain, etas, squeeze_out) + + def j_cart(self, *etas, squeeze_out=False): + """Cartesian components of current (directly from ``j_xyz``). Returns also (x,y,z).""" + return _cartesian_vector(self.j_xyz, self.domain, etas, squeeze_out) + + def gradB_cart(self, *etas, squeeze_out=False): + """Cartesian components of gradient of magnetic field strength (directly from ``gradB_xyz``). Returns also + (x,y,z).""" + return _cartesian_vector(self.gradB_xyz, self.domain, etas, squeeze_out) + + def b1(self, *etas, squeeze_out=False): + """1-form components of magnetic field (pulled back from ``b_xyz``, also on a polar axis).""" + return self.domain.pull(self.b_cart(*etas, squeeze_out=False)[0], *etas, kind="1", squeeze_out=squeeze_out) + + def b2(self, *etas, squeeze_out=False): + """2-form components of magnetic field (pulled back from ``b_xyz``, also on a polar axis).""" + return self.domain.pull(self.b_cart(*etas, squeeze_out=False)[0], *etas, kind="2", squeeze_out=squeeze_out) + + def j1(self, *etas, squeeze_out=False): + """1-form components of current (pulled back from ``j_xyz``, also on a polar axis).""" + return self.domain.pull(self.j_cart(*etas, squeeze_out=False)[0], *etas, kind="1", squeeze_out=squeeze_out) + + def j2(self, *etas, squeeze_out=False): + """2-form components of current (pulled back from ``j_xyz``, also on a polar axis).""" + return self.domain.pull(self.j_cart(*etas, squeeze_out=False)[0], *etas, kind="2", squeeze_out=squeeze_out) + class AxisymmMHDequilibrium(CartesianMHDequilibrium): r""" diff --git a/src/struphy/fields_background/equils.py b/src/struphy/fields_background/equils.py index cde0e9ec7..033110aba 100644 --- a/src/struphy/fields_background/equils.py +++ b/src/struphy/fields_background/equils.py @@ -10,6 +10,7 @@ from typing import TYPE_CHECKING import cunumpy as xp +import numpy as np from line_profiler import profile import struphy @@ -824,6 +825,94 @@ def gradB_xyz(self, x, y, z): return gradBx, gradBy, gradBz +class HelicalVacuumField(CartesianMHDequilibrium): + r"""Straight helical vacuum field (stellarator-like near the axis), :math:`\mathbf B = B_0 \nabla \Phi` with the + harmonic potential + + .. math:: + + \Phi = z + \frac{\delta}{k} I_m(k r) \sin(m \theta - k z)\,, + + where :math:`I_m` is the modified Bessel function. The field is current-free (:math:`\nabla \times \mathbf B = 0`), + with constant density and pressure. Near the axis the :math:`m = 2` field has elliptical, rotating flux surfaces + and a rotational transform of order :math:`\delta^2`: the vacuum field of the guiding-centre orbit study of the + high-order polar splines (orbits near the magnetic axis, energy conservation). + + Parameters + ---------- + B0 : float + Field strength on the axis. + delta : float + Amplitude of the helical harmonic. + m : int + Poloidal mode number. + k : float + Axial wave number (``k = 2 pi N / L_z`` for ``N`` periods on a cylinder of length ``L_z``). + n0, p0 : float + Constant density and pressure. + """ + + @classmethod + def doc_formula(cls): + r""" + .. math:: + + \mathbf B = B_0 \left( \delta I_m'(kr) \sin\varphi\, \mathbf e_r + + \delta \frac{m}{kr} I_m(kr) \cos\varphi\, \mathbf e_\theta + + (1 - \delta I_m(kr) \cos\varphi)\, \mathbf e_z \right)\,,\qquad \varphi = m\theta - kz\,. + """ + + def __init__( + self, B0: float = 1.0, delta: float = 0.5, m: int = 2, k: float = 3.0, n0: float = 1.0, p0: float = 1.0 + ): + self.params = copy.deepcopy(locals()) + + def _cylindrical(self, x, y, z): + from scipy.special import iv, ivp + + x, y, z = (np.asarray(xp.to_numpy(xp.asarray(c)), dtype=float) for c in (x, y, z)) + B0, delta, m, k = (self.params[key] for key in ("B0", "delta", "m", "k")) + r = np.sqrt(x**2 + y**2) + theta = np.arctan2(y, x) + phase = m * theta - k * z + b_r = B0 * delta * ivp(m, k * r) * np.sin(phase) + # (m / (k r)) I_m(k r) = (I_{m-1} - I_{m+1}) / 2, regular at the axis + b_theta = B0 * delta * (iv(m - 1, k * r) - iv(m + 1, k * r)) / 2 * np.cos(phase) + b_z = B0 * (1 - delta * iv(m, k * r) * np.cos(phase)) + return theta, b_r, b_theta, b_z + + def b_xyz(self, x, y, z): + """Magnetic field.""" + theta, b_r, b_theta, b_z = self._cylindrical(x, y, z) + c, s = np.cos(theta), np.sin(theta) + return xp.asarray(b_r * c - b_theta * s), xp.asarray(b_r * s + b_theta * c), xp.asarray(b_z) + + def j_xyz(self, x, y, z): + """Current density (zero: vacuum field).""" + return 0 * x, 0 * x, 0 * x + + def p_xyz(self, x, y, z): + """Pressure.""" + return self.params["p0"] + 0 * x + + def n_xyz(self, x, y, z): + """Number density.""" + return self.params["n0"] + 0 * x + + def gradB_xyz(self, x, y, z, h=1e-6): + """Gradient of the field strength (central differences).""" + + def absB(x, y, z): + bx, by, bz = self.b_xyz(x, y, z) + return (bx**2 + by**2 + bz**2) ** 0.5 + + return ( + (absB(x + h, y, z) - absB(x - h, y, z)) / (2 * h), + (absB(x, y + h, z) - absB(x, y - h, z)) / (2 * h), + (absB(x, y, z + h) - absB(x, y, z - h)) / (2 * h), + ) + + class AdhocTorus(AxisymmMHDequilibrium): r""" Ad hoc tokamak MHD equilibrium with circular concentric flux surfaces. diff --git a/src/struphy/fields_background/tests/test_equils_cupy.py b/src/struphy/fields_background/tests/test_equils_cupy.py index 8c51674bc..0662a5b4d 100644 --- a/src/struphy/fields_background/tests/test_equils_cupy.py +++ b/src/struphy/fields_background/tests/test_equils_cupy.py @@ -51,14 +51,15 @@ "CurrentSheet": ("CurrentSheet", {}, "Cuboid", {}), "GenericCartesianFluidEquilibrium": ("GenericCartesianFluidEquilibrium", {}, "Cuboid", {}), "GenericCartesianFluidEquilibriumWithB": ("GenericCartesianFluidEquilibriumWithB", {}, "Cuboid", {}), + "HelicalVacuumField": ("HelicalVacuumField", {}, "HollowCylinder", {"a1": 0.0, "a2": 0.3, "Lz": 2 * np.pi}), } # equilibria that need an optional package (skipped without it) NEEDS = {"GVECequilibrium": "gvec", "DESCequilibrium": "desc"} -# equilibria evaluated through their native code on the host (one round trip per call); all others are evaluated on -# the device only, without host/device transfers -HOST_EVALUATION = {"GVECequilibrium", "DESCequilibrium"} +# equilibria evaluated through their native code on the host (one round trip per call; HelicalVacuumField: the Bessel +# functions of SciPy); all others are evaluated on the device only, without host/device transfers +HOST_EVALUATION = {"GVECequilibrium", "DESCequilibrium", "HelicalVacuumField"} # methods that models and projections call on the logical domain; those an equilibrium does not provide are skipped METHODS = ( diff --git a/src/struphy/geometry/domains/__init__.py b/src/struphy/geometry/domains/__init__.py index 270fa55dd..663c23e23 100644 --- a/src/struphy/geometry/domains/__init__.py +++ b/src/struphy/geometry/domains/__init__.py @@ -17,12 +17,15 @@ "Cuboid": "struphy.geometry.domains.cuboid", "DESCunit": "struphy.geometry.domains.desc_unit", "GVECunit": "struphy.geometry.domains.gvec_unit", + "HelicalTube": "struphy.geometry.domains.helical_tube", "HollowCylinder": "struphy.geometry.domains.hollow_cylinder", "HollowTorus": "struphy.geometry.domains.hollow_torus", "IGAPolarCylinder": "struphy.geometry.domains.iga_polar_cylinder", "IGAPolarTorus": "struphy.geometry.domains.iga_polar_torus", "Orthogonal": "struphy.geometry.domains.orthogonal", + "PerturbedDiskCylinder": "struphy.geometry.domains.perturbed_disk_cylinder", "PoweredEllipticCylinder": "struphy.geometry.domains.powered_elliptic_cylinder", + "RoundedSquareCylinder": "struphy.geometry.domains.rounded_square_cylinder", "ShafranovDshapedCylinder": "struphy.geometry.domains.shafranov_dshaped_cylinder", "ShafranovShiftCylinder": "struphy.geometry.domains.shafranov_shift_cylinder", "ShafranovSqrtCylinder": "struphy.geometry.domains.shafranov_sqrt_cylinder", @@ -34,12 +37,15 @@ from struphy.geometry.domains.cuboid import Cuboid from struphy.geometry.domains.desc_unit import DESCunit from struphy.geometry.domains.gvec_unit import GVECunit + from struphy.geometry.domains.helical_tube import HelicalTube from struphy.geometry.domains.hollow_cylinder import HollowCylinder from struphy.geometry.domains.hollow_torus import HollowTorus from struphy.geometry.domains.iga_polar_cylinder import IGAPolarCylinder from struphy.geometry.domains.iga_polar_torus import IGAPolarTorus from struphy.geometry.domains.orthogonal import Orthogonal + from struphy.geometry.domains.perturbed_disk_cylinder import PerturbedDiskCylinder from struphy.geometry.domains.powered_elliptic_cylinder import PoweredEllipticCylinder + from struphy.geometry.domains.rounded_square_cylinder import RoundedSquareCylinder from struphy.geometry.domains.shafranov_dshaped_cylinder import ShafranovDshapedCylinder from struphy.geometry.domains.shafranov_shift_cylinder import ShafranovShiftCylinder from struphy.geometry.domains.shafranov_sqrt_cylinder import ShafranovSqrtCylinder @@ -50,12 +56,15 @@ "Cuboid", "DESCunit", "GVECunit", + "HelicalTube", "HollowCylinder", "HollowTorus", "IGAPolarCylinder", "IGAPolarTorus", "Orthogonal", + "PerturbedDiskCylinder", "PoweredEllipticCylinder", + "RoundedSquareCylinder", "ShafranovDshapedCylinder", "ShafranovShiftCylinder", "ShafranovSqrtCylinder", diff --git a/src/struphy/geometry/domains/helical_tube/__init__.py b/src/struphy/geometry/domains/helical_tube/__init__.py new file mode 100644 index 000000000..a751e7fe7 --- /dev/null +++ b/src/struphy/geometry/domains/helical_tube/__init__.py @@ -0,0 +1,14 @@ +"""Domain package for :class:`struphy.geometry.domains.helical_tube.helical_tube.HelicalTube`. + +The class is imported lazily on first attribute access (see :mod:`struphy.geometry.domains.hollow_cylinder`). +""" + +__all__ = ["HelicalTube"] + + +def __getattr__(name): + if name == "HelicalTube": + from struphy.geometry.domains.helical_tube.helical_tube import HelicalTube + + return HelicalTube + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/src/struphy/geometry/domains/helical_tube/helical_tube.py b/src/struphy/geometry/domains/helical_tube/helical_tube.py new file mode 100644 index 000000000..01386b34c --- /dev/null +++ b/src/struphy/geometry/domains/helical_tube/helical_tube.py @@ -0,0 +1,116 @@ +import copy + +import cunumpy as xp +import numpy as np + +from struphy.geometry.base import Spline, interp_mapping + + +def helical_axis(zeta, R0, delta, delta_z, periods): + """Centre curve ``((R0 + delta cos 2 pi m zeta) cos 2 pi zeta, ... sin 2 pi zeta, delta_z sin 2 pi m zeta)``.""" + rho = R0 + delta * np.cos(2 * np.pi * periods * zeta) + return np.stack( + [rho * np.cos(2 * np.pi * zeta), rho * np.sin(2 * np.pi * zeta), delta_z * np.sin(2 * np.pi * periods * zeta)] + ) + + +def axis_derivative(zeta, R0, delta, delta_z, periods): + m = periods + rho = R0 + delta * np.cos(2 * np.pi * m * zeta) + drho = -2 * np.pi * m * delta * np.sin(2 * np.pi * m * zeta) + c, s = np.cos(2 * np.pi * zeta), np.sin(2 * np.pi * zeta) + return np.stack( + [ + drho * c - 2 * np.pi * rho * s, + drho * s + 2 * np.pi * rho * c, + 2 * np.pi * m * delta_z * np.cos(2 * np.pi * m * zeta), + ] + ) + + +def rotating_frame(zeta, R0, delta, delta_z, periods, rotation): + """The vertical normal frame ``N1 = T x e_z / |T x e_z|``, ``N2 = T x N1``, rotated by ``pi * rotation * zeta``.""" + T = axis_derivative(zeta, R0, delta, delta_z, periods) + T = T / np.linalg.norm(T, axis=0) + N1 = np.stack([T[1], -T[0], np.zeros_like(T[0])]) + N1 = N1 / np.linalg.norm(N1, axis=0) + N2 = np.cross(T, N1, axis=0) + c, s = np.cos(np.pi * rotation * zeta), np.sin(np.pi * rotation * zeta) + return c * N1 + s * N2, c * N2 - s * N1 + + +class HelicalTube(Spline): + r"""Stellarator-like tube :math:`F(\eta_1, \eta_2, \eta_3) = \gamma(\eta_3) + N(\eta_3)\,\Phi_\textnormal{2D}(\eta_1, \eta_2)` + around a helical axis, the 3D domain of the high-order polar-spline studies of R. Remmerswaal (PolarFEEC.jl + ``TubeMapping``). + + * axis :math:`\gamma(\zeta) = ((R_0 + \delta\cos 2\pi m\zeta)\cos 2\pi\zeta, (R_0 + \delta\cos 2\pi m\zeta)\sin 2\pi\zeta, + \delta_z \sin 2\pi m\zeta)` with ``m = periods``, + * frame: the vertical normal frame :math:`N_1 = T \times e_z / |T \times e_z|`, :math:`N_2 = T \times N_1`, rotated by + the angle :math:`\pi\,\textnormal{rotation}\,\zeta` (``rotation`` must be even for a periodic frame), + * cross-section :math:`\Phi_\textnormal{2D} = R\,(a\cos 2\pi\eta_2, b\sin 2\pi\eta_2)`, + :math:`R = \eta_1 (1 + \varepsilon\,\eta_1 \cos(2\pi k \eta_2))`. + + The presets of the talk: ``"stellarator"`` (the 3D wave and curl-curl studies: three periods, three-petal section of + radius 0.3 turning once) and ``"islands"`` (PolarFEEC.jl example 07: four periods, ellipse 0.28 x 0.14 turning by + half a turn per period). The mapping is interpolated by splines of the given degree on the given grid. + + Parameters + ---------- + preset : str or None + ``"stellarator"``, ``"islands"`` or ``None`` (use the remaining parameters). + num_elements, degree : tuple[int, int, int] + Spline grid of the interpolated mapping. + R0, delta, delta_z, periods, rotation, a, b, eps, k : float + Geometry parameters (ignored for a preset). + """ + + PRESETS = { + "stellarator": dict(R0=1.0, delta=0.1, delta_z=0.1, periods=3, rotation=2, a=0.3, b=0.3, eps=0.2, k=3), + "islands": dict(R0=1.0, delta=0.12, delta_z=0.12, periods=4, rotation=4, a=0.28, b=0.14, eps=0.0, k=0), + } + + def __init__( + self, + preset: str | None = "stellarator", + num_elements: tuple[int, int, int] = (16, 48, 96), + degree: tuple[int, int, int] = (4, 4, 4), + R0: float = 1.0, + delta: float = 0.1, + delta_z: float = 0.1, + periods: int = 3, + rotation: int = 2, + a: float = 0.3, + b: float = 0.3, + eps: float = 0.2, + k: int = 3, + ): + geometry = dict(R0=R0, delta=delta, delta_z=delta_z, periods=periods, rotation=rotation, a=a, b=b, eps=eps, k=k) + if preset is not None: + geometry = dict(self.PRESETS[preset]) + self.geometry = geometry + g = geometry + axis = (g["R0"], g["delta"], g["delta_z"], g["periods"]) + + def section(e1, e2): + radius = e1 * (1 + g["eps"] * e1 * np.cos(2 * np.pi * g["k"] * e2)) + return radius * g["a"] * np.cos(2 * np.pi * e2), radius * g["b"] * np.sin(2 * np.pi * e2) + + def component(i): + def X(e1, e2, e3): + if np.ndim(e1) == np.ndim(e2) == np.ndim(e3) == 1: + e1, e2, e3 = np.meshgrid(e1, e2, e3, indexing="ij") + x2, y2 = section(e1, e2) + N1, N2 = rotating_frame(e3, *axis, g["rotation"]) + return helical_axis(e3, *axis)[i] + N1[i] * x2 + N2[i] * y2 + + return X + + spl_kind = (False, True, True) + cx, cy, cz = interp_mapping(num_elements, degree, spl_kind, component(0), component(1), component(2)) + # the first ring of control points is the axis: make it exactly the same point for all eta2 + for c in (cx, cy, cz): + c[0] = xp.mean(c[0], axis=0)[None, :] + + self._preset_params = copy.deepcopy(dict(preset=preset, **geometry)) + super().__init__(num_elements=num_elements, degree=degree, spl_kind=spl_kind, cx=cx, cy=cy, cz=cz) diff --git a/src/struphy/geometry/domains/perturbed_disk_cylinder/__init__.py b/src/struphy/geometry/domains/perturbed_disk_cylinder/__init__.py new file mode 100644 index 000000000..e487eef50 --- /dev/null +++ b/src/struphy/geometry/domains/perturbed_disk_cylinder/__init__.py @@ -0,0 +1,14 @@ +"""Domain package for :class:`struphy.geometry.domains.perturbed_disk_cylinder.perturbed_disk_cylinder.PerturbedDiskCylinder`. + +The class is imported lazily on first attribute access (see :mod:`struphy.geometry.domains.hollow_cylinder`). +""" + +__all__ = ["PerturbedDiskCylinder"] + + +def __getattr__(name): + if name == "PerturbedDiskCylinder": + from struphy.geometry.domains.perturbed_disk_cylinder.perturbed_disk_cylinder import PerturbedDiskCylinder + + return PerturbedDiskCylinder + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/src/struphy/geometry/domains/perturbed_disk_cylinder/perturbed_disk_cylinder.py b/src/struphy/geometry/domains/perturbed_disk_cylinder/perturbed_disk_cylinder.py new file mode 100644 index 000000000..5e4a5e029 --- /dev/null +++ b/src/struphy/geometry/domains/perturbed_disk_cylinder/perturbed_disk_cylinder.py @@ -0,0 +1,58 @@ +import copy + +from struphy.geometry.base import Domain + + +class PerturbedDiskCylinder(Domain): + r"""Cylinder whose cross-section is the perturbed (``k``-petal) disk of the polar FEEC paper. + + The cross-section is the disk-like map :math:`R(\eta_1, \eta_2) = R_1 \eta_1 (1 + \varepsilon \eta_1^q \cos(2\pi k + \eta_2))` (``q = 1``: the perturbed disk, ``q = 0``: a star-shaped disk) of R. Remmerswaal, *High-order FEEC discretization on 2D and 3D domains with a disk-like coordinate + singularity*, which is exactly representable by the polar splines (``r``-compatible) only with trigonometric + functions in :math:`\eta_2`. With :math:`\varepsilon = 0` it is the circular cylinder of radius :math:`R_1`. + + Parameters + ---------- + R1 : float + Radius of the unperturbed disk (default: 1.0). + eps : float + Amplitude of the perturbation (default: 0.2). + k : int + Number of petals (default: 3). + Lz : float + Length of the cylinder (default: 1.0). + power : float + Power ``q`` of ``eta1`` in the perturbation (default: 1.0; 0.0 gives the star-shaped disk + ``R = R_1 eta1 (1 + eps cos(2 pi k eta2))`` of the regularity study of the talk). + """ + + @classmethod + def doc_mapping(cls): + r""" + .. math:: + + F: \begin{bmatrix}\eta_1\\ \eta_2\\ \eta_3\end{bmatrix}\mapsto \begin{bmatrix} + \,\,x= &R_1\,\eta_1\,\left(1 + \varepsilon\,\eta_1^q\cos(2\pi k\,\eta_2)\right)\cos(2\pi\,\eta_2)\,\,\\ + \,\,y= &R_1\,\eta_1\,\left(1 + \varepsilon\,\eta_1^q\cos(2\pi k\,\eta_2)\right)\sin(2\pi\,\eta_2)\,\,\\ + \,\,z= &L_z\,\eta_3\,\,\end{bmatrix} + """ + + def __init__( + self, + R1: float = 1.0, + eps: float = 0.2, + k: int = 3, + Lz: float = 1.0, + power: float = 1.0, + ): + self.kind_map = 23 + + # use params setter + self.params = copy.deepcopy(locals()) + self.params_numpy = self.get_params_numpy() + + # periodicity in eta3-direction and pole at eta1=0 + self.periodic_eta3 = False + self.pole = True + + super().__init__() diff --git a/src/struphy/geometry/domains/perturbed_disk_cylinder/perturbed_disk_cylinder_cuda.cuh b/src/struphy/geometry/domains/perturbed_disk_cylinder/perturbed_disk_cylinder_cuda.cuh new file mode 100644 index 000000000..ade2900f5 --- /dev/null +++ b/src/struphy/geometry/domains/perturbed_disk_cylinder/perturbed_disk_cylinder_cuda.cuh @@ -0,0 +1,58 @@ +#pragma once +#include "struphy/geometry/domains/constants_cuda.cuh" +namespace struphy_cuda::perturbed_disk_cylinder_kernels { +/** + * Evaluate the perturbed disk cylinder mapping, as in perturbed_disk_cylinder_kernels.perturbed_disk. + * + * R = r1 eta1 (1 + eps eta1^q cos(2 pi k eta2)), F_x = R cos(2 pi eta2), F_y = R sin(2 pi eta2), F_z = lz eta3. + * + * @param eta1 Logical coordinate along the first axis. + * @param eta2 Logical coordinate along the second axis. + * @param eta3 Logical coordinate along the third axis. + * @param r1 Radius of the unperturbed disk. + * @param eps Amplitude of the perturbation. + * @param k Number of petals. + * @param lz Length in z-direction. + * @param q Power of eta1 in the perturbation. + * @param f_out Output buffer for the three physical coordinates. + */ +__device__ inline void perturbed_disk(double eta1, double eta2, double eta3, double r1, double eps, double k, + double lz, double q, double* f_out) { + double radius = r1 * eta1 * (1.0 + eps * pow(eta1, q) * cos(2 * pi * k * eta2)); + f_out[0] = radius * cos(2 * pi * eta2); + f_out[1] = radius * sin(2 * pi * eta2); + f_out[2] = lz * eta3; +} + +/** + * Evaluate the Jacobian of the perturbed disk cylinder mapping, as in perturbed_disk_cylinder_kernels.perturbed_disk_df. + * + * @param eta1 Logical coordinate along the first axis. + * @param eta2 Logical coordinate along the second axis. + * @param r1 Radius of the unperturbed disk. + * @param eps Amplitude of the perturbation. + * @param k Number of petals. + * @param lz Length in z-direction. + * @param q Power of eta1 in the perturbation. + * @param df_out Output 3x3 matrix stored as nine row-major entries. + */ +__device__ inline void perturbed_disk_df(double eta1, double eta2, double r1, double eps, double k, double lz, + double q, double* df_out) { + double ck = cos(2 * pi * k * eta2); + double sk = sin(2 * pi * k * eta2); + double radius = r1 * eta1 * (1.0 + eps * pow(eta1, q) * ck); + double dr1 = r1 * (1.0 + (q + 1.0) * eps * pow(eta1, q) * ck); + double dr2 = -2 * pi * k * r1 * eps * pow(eta1, q + 1.0) * sk; + double c = cos(2 * pi * eta2); + double s = sin(2 * pi * eta2); + df_out[0] = dr1 * c; + df_out[1] = dr2 * c - 2 * pi * radius * s; + df_out[2] = 0.0; + df_out[3] = dr1 * s; + df_out[4] = dr2 * s + 2 * pi * radius * c; + df_out[5] = 0.0; + df_out[6] = 0.0; + df_out[7] = 0.0; + df_out[8] = lz; +} +} // namespace struphy_cuda::perturbed_disk_cylinder_kernels diff --git a/src/struphy/geometry/domains/perturbed_disk_cylinder/perturbed_disk_cylinder_kernels.py b/src/struphy/geometry/domains/perturbed_disk_cylinder/perturbed_disk_cylinder_kernels.py new file mode 100644 index 000000000..6e6c9af5d --- /dev/null +++ b/src/struphy/geometry/domains/perturbed_disk_cylinder/perturbed_disk_cylinder_kernels.py @@ -0,0 +1,73 @@ +"""Pyccel kernels for the mapping and Jacobian matrix of :class:`struphy.geometry.domains.PerturbedDiskCylinder`.""" + +from numpy import cos, pi, sin +from pyccel.decorators import pure + + +@pure +def perturbed_disk( + eta1: float, eta2: float, eta3: float, r1: float, eps: float, k: float, lz: float, q: float, f_out: "float[:]" +): + r"""Point-wise evaluation of + + .. math:: + + F_x &= R(\eta_1, \eta_2) \cos(2\pi\,\eta_2)\,, \quad F_y = R(\eta_1, \eta_2) \sin(2\pi\,\eta_2)\,, + \quad F_z = L_z\,\eta_3\,, + + R &= R_1\,\eta_1\,\left(1 + \varepsilon\,\eta_1^q \cos(2\pi k\,\eta_2)\right)\,. + + Parameters + ---------- + eta1, eta2, eta3 : float + Logical coordinate in [0, 1]. + + r1 : float + Radius of the unperturbed disk. + + eps : float + Amplitude of the perturbation. + + k : float + Number of petals. + + lz : float + Length in third direction. + + q : float + Power of eta1 in the perturbation (1: the perturbed disk of the paper, 0: a star-shaped disk). + + f_out : array[float] + Output: (x, y, z) = F(eta1, eta2, eta3). + """ + + radius = r1 * eta1 * (1.0 + eps * eta1**q * cos(2 * pi * k * eta2)) + + f_out[0] = radius * cos(2 * pi * eta2) + f_out[1] = radius * sin(2 * pi * eta2) + f_out[2] = lz * eta3 + + +@pure +def perturbed_disk_df( + eta1: float, eta2: float, r1: float, eps: float, k: float, lz: float, q: float, df_out: "float[:,:]" +): + """Jacobian matrix for :meth:`perturbed_disk_cylinder_kernels.perturbed_disk`.""" + + ck = cos(2 * pi * k * eta2) + sk = sin(2 * pi * k * eta2) + radius = r1 * eta1 * (1.0 + eps * eta1**q * ck) + dr1 = r1 * (1.0 + (q + 1.0) * eps * eta1**q * ck) + dr2 = -2 * pi * k * r1 * eps * eta1 ** (q + 1.0) * sk + c = cos(2 * pi * eta2) + s = sin(2 * pi * eta2) + + df_out[0, 0] = dr1 * c + df_out[0, 1] = dr2 * c - 2 * pi * radius * s + df_out[0, 2] = 0.0 + df_out[1, 0] = dr1 * s + df_out[1, 1] = dr2 * s + 2 * pi * radius * c + df_out[1, 2] = 0.0 + df_out[2, 0] = 0.0 + df_out[2, 1] = 0.0 + df_out[2, 2] = lz diff --git a/src/struphy/geometry/domains/rounded_square_cylinder/__init__.py b/src/struphy/geometry/domains/rounded_square_cylinder/__init__.py new file mode 100644 index 000000000..93903d26d --- /dev/null +++ b/src/struphy/geometry/domains/rounded_square_cylinder/__init__.py @@ -0,0 +1,14 @@ +"""Domain package for :class:`struphy.geometry.domains.rounded_square_cylinder.rounded_square_cylinder.RoundedSquareCylinder`. + +The class is imported lazily on first attribute access (see :mod:`struphy.geometry.domains.hollow_cylinder`). +""" + +__all__ = ["RoundedSquareCylinder"] + + +def __getattr__(name): + if name == "RoundedSquareCylinder": + from struphy.geometry.domains.rounded_square_cylinder.rounded_square_cylinder import RoundedSquareCylinder + + return RoundedSquareCylinder + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/src/struphy/geometry/domains/rounded_square_cylinder/rounded_square_cylinder.py b/src/struphy/geometry/domains/rounded_square_cylinder/rounded_square_cylinder.py new file mode 100644 index 000000000..f19f7b11d --- /dev/null +++ b/src/struphy/geometry/domains/rounded_square_cylinder/rounded_square_cylinder.py @@ -0,0 +1,50 @@ +import copy + +from struphy.geometry.base import Domain + + +class RoundedSquareCylinder(Domain): + r"""Cylinder whose cross-section is a rounded square, the zero set of + :math:`g = 1 - x^2 - y^2 - \varepsilon (x^4 - 6x^2y^2 + y^4)` (for :math:`a = 1`). + + This is the domain of the Poisson and curl-curl examples of PolarFEEC.jl (``examples/05_polar_pde.jl``). + + Parameters + ---------- + a : float + Scale of the cross-section (default: 1.0). + eps : float + Squareness of the cross-section (default: 0.15). + Lz : float + Length of the cylinder (default: 1.0). + """ + + @classmethod + def doc_mapping(cls): + r""" + .. math:: + + F: \begin{bmatrix}\eta_1\\ \eta_2\\ \eta_3\end{bmatrix}\mapsto \begin{bmatrix} + \,\,x= &a\,\eta_1\,\rho(\eta_2)\cos(2\pi\,\eta_2)\,\,\\ + \,\,y= &a\,\eta_1\,\rho(\eta_2)\sin(2\pi\,\eta_2)\,\,\\ + \,\,z= &L_z\,\eta_3\,\,\end{bmatrix}\,, \qquad + \rho(\eta_2) = \sqrt{\frac{2}{1 + \sqrt{1 + 4\varepsilon\cos(8\pi\,\eta_2)}}} + """ + + def __init__( + self, + a: float = 1.0, + eps: float = 0.15, + Lz: float = 1.0, + ): + self.kind_map = 24 + + # use params setter + self.params = copy.deepcopy(locals()) + self.params_numpy = self.get_params_numpy() + + # periodicity in eta3-direction and pole at eta1=0 + self.periodic_eta3 = False + self.pole = True + + super().__init__() diff --git a/src/struphy/geometry/domains/rounded_square_cylinder/rounded_square_cylinder_cuda.cuh b/src/struphy/geometry/domains/rounded_square_cylinder/rounded_square_cylinder_cuda.cuh new file mode 100644 index 000000000..c6f7cd928 --- /dev/null +++ b/src/struphy/geometry/domains/rounded_square_cylinder/rounded_square_cylinder_cuda.cuh @@ -0,0 +1,56 @@ +#pragma once +#include "struphy/geometry/domains/constants_cuda.cuh" +namespace struphy_cuda::rounded_square_cylinder_kernels { +/** + * Evaluate the rounded square cylinder mapping, as in rounded_square_cylinder_kernels.rounded_square. + * + * rho = sqrt(2 / (1 + sqrt(1 + 4 eps cos(8 pi eta2)))), F_x = a eta1 rho cos(2 pi eta2), + * F_y = a eta1 rho sin(2 pi eta2), F_z = lz eta3. + * + * @param eta1 Logical coordinate along the first axis. + * @param eta2 Logical coordinate along the second axis. + * @param eta3 Logical coordinate along the third axis. + * @param a Scale of the cross-section. + * @param eps Squareness of the cross-section. + * @param lz Length in z-direction. + * @param f_out Output buffer for the three physical coordinates. + */ +__device__ inline void rounded_square(double eta1, double eta2, double eta3, double a, double eps, double lz, + double* f_out) { + double q = sqrt(1.0 + 4.0 * eps * cos(8 * pi * eta2)); + double rho = sqrt(2.0 / (1.0 + q)); + f_out[0] = a * eta1 * rho * cos(2 * pi * eta2); + f_out[1] = a * eta1 * rho * sin(2 * pi * eta2); + f_out[2] = lz * eta3; +} + +/** + * Evaluate the Jacobian of the rounded square cylinder mapping, as in rounded_square_cylinder_kernels.rounded_square_df. + * + * @param eta1 Logical coordinate along the first axis. + * @param eta2 Logical coordinate along the second axis. + * @param a Scale of the cross-section. + * @param eps Squareness of the cross-section. + * @param lz Length in z-direction. + * @param df_out Output 3x3 matrix stored as nine row-major entries. + */ +__device__ inline void rounded_square_df(double eta1, double eta2, double a, double eps, double lz, double* df_out) { + double q = sqrt(1.0 + 4.0 * eps * cos(8 * pi * eta2)); + double rho = sqrt(2.0 / (1.0 + q)); + double drho = 16 * pi * eps * sin(8 * pi * eta2) / (rho * (1.0 + q) * (1.0 + q) * q); + double radius = a * eta1 * rho; + double dr1 = a * rho; + double dr2 = a * eta1 * drho; + double c = cos(2 * pi * eta2); + double s = sin(2 * pi * eta2); + df_out[0] = dr1 * c; + df_out[1] = dr2 * c - 2 * pi * radius * s; + df_out[2] = 0.0; + df_out[3] = dr1 * s; + df_out[4] = dr2 * s + 2 * pi * radius * c; + df_out[5] = 0.0; + df_out[6] = 0.0; + df_out[7] = 0.0; + df_out[8] = lz; +} +} // namespace struphy_cuda::rounded_square_cylinder_kernels diff --git a/src/struphy/geometry/domains/rounded_square_cylinder/rounded_square_cylinder_kernels.py b/src/struphy/geometry/domains/rounded_square_cylinder/rounded_square_cylinder_kernels.py new file mode 100644 index 000000000..b3d7fd704 --- /dev/null +++ b/src/struphy/geometry/domains/rounded_square_cylinder/rounded_square_cylinder_kernels.py @@ -0,0 +1,67 @@ +"""Pyccel kernels for the mapping and Jacobian matrix of :class:`struphy.geometry.domains.RoundedSquareCylinder`.""" + +from numpy import cos, pi, sin, sqrt +from pyccel.decorators import pure + + +@pure +def rounded_square(eta1: float, eta2: float, eta3: float, a: float, eps: float, lz: float, f_out: "float[:]"): + r"""Point-wise evaluation of + + .. math:: + + F_x &= a\,\eta_1\,\rho(\eta_2) \cos(2\pi\,\eta_2)\,, \quad F_y = a\,\eta_1\,\rho(\eta_2) \sin(2\pi\,\eta_2)\,, + \quad F_z = L_z\,\eta_3\,, + + \rho &= \sqrt{2 / (1 + \sqrt{1 + 4\varepsilon \cos(8\pi\,\eta_2)})}\,, + + the rounded square :math:`1 - x^2 - y^2 - \varepsilon (x^4 - 6 x^2 y^2 + y^4) = 0` (for :math:`a = 1`). + + Parameters + ---------- + eta1, eta2, eta3 : float + Logical coordinate in [0, 1]. + + a : float + Scale of the cross-section. + + eps : float + Squareness of the cross-section. + + lz : float + Length in third direction. + + f_out : array[float] + Output: (x, y, z) = F(eta1, eta2, eta3). + """ + + q = sqrt(1.0 + 4.0 * eps * cos(8 * pi * eta2)) + rho = sqrt(2.0 / (1.0 + q)) + + f_out[0] = a * eta1 * rho * cos(2 * pi * eta2) + f_out[1] = a * eta1 * rho * sin(2 * pi * eta2) + f_out[2] = lz * eta3 + + +@pure +def rounded_square_df(eta1: float, eta2: float, a: float, eps: float, lz: float, df_out: "float[:,:]"): + """Jacobian matrix for :meth:`rounded_square_cylinder_kernels.rounded_square`.""" + + q = sqrt(1.0 + 4.0 * eps * cos(8 * pi * eta2)) + rho = sqrt(2.0 / (1.0 + q)) + drho = 16 * pi * eps * sin(8 * pi * eta2) / (rho * (1.0 + q) ** 2 * q) + radius = a * eta1 * rho + dr1 = a * rho + dr2 = a * eta1 * drho + c = cos(2 * pi * eta2) + s = sin(2 * pi * eta2) + + df_out[0, 0] = dr1 * c + df_out[0, 1] = dr2 * c - 2 * pi * radius * s + df_out[0, 2] = 0.0 + df_out[1, 0] = dr1 * s + df_out[1, 1] = dr2 * s + 2 * pi * radius * c + df_out[1, 2] = 0.0 + df_out[2, 0] = 0.0 + df_out[2, 1] = 0.0 + df_out[2, 2] = lz diff --git a/src/struphy/geometry/evaluation_kernels.cuh b/src/struphy/geometry/evaluation_kernels.cuh index eafebf6b9..37db362a1 100644 --- a/src/struphy/geometry/evaluation_kernels.cuh +++ b/src/struphy/geometry/evaluation_kernels.cuh @@ -11,6 +11,8 @@ #include "struphy/geometry/domains/orthogonal/orthogonal_cuda.cuh" #include "struphy/geometry/domains/colella/colella_cuda.cuh" #include "struphy/geometry/domains/hollow_cylinder/hollow_cylinder_cuda.cuh" +#include "struphy/geometry/domains/perturbed_disk_cylinder/perturbed_disk_cylinder_cuda.cuh" +#include "struphy/geometry/domains/rounded_square_cylinder/rounded_square_cylinder_cuda.cuh" #include "struphy/geometry/domains/powered_elliptic_cylinder/powered_elliptic_cylinder_cuda.cuh" #include "struphy/geometry/domains/hollow_torus/hollow_torus_cuda.cuh" #include "struphy/geometry/domains/shafranov_shift_cylinder/shafranov_shift_cylinder_cuda.cuh" @@ -67,6 +69,14 @@ __device__ inline void f(double eta1, double eta2, double eta3, const DomainArgs hollow_torus_kernels::hollow_torus(eta1, eta2, eta3, args.params[0], args.params[1], args.params[2], args.params[3], args.params[4], args.params[5], f_out); return; + case 23: + perturbed_disk_cylinder_kernels::perturbed_disk(eta1, eta2, eta3, args.params[0], args.params[1], + args.params[2], args.params[3], args.params[4], f_out); + return; + case 24: + rounded_square_cylinder_kernels::rounded_square(eta1, eta2, eta3, args.params[0], args.params[1], + args.params[2], f_out); + return; case 30: shafranov_shift_cylinder_kernels::shafranov_shift(eta1, eta2, eta3, args.params[0], args.params[1], args.params[2], args.params[3], f_out); @@ -135,6 +145,14 @@ __device__ inline void df(double eta1, double eta2, double eta3, const DomainArg hollow_torus_kernels::hollow_torus_df(eta1, eta2, eta3, args.params[0], args.params[1], args.params[2], args.params[3], args.params[4], args.params[5], df_out); return; + case 23: + perturbed_disk_cylinder_kernels::perturbed_disk_df(eta1, eta2, args.params[0], args.params[1], + args.params[2], args.params[3], args.params[4], df_out); + return; + case 24: + rounded_square_cylinder_kernels::rounded_square_df(eta1, eta2, args.params[0], args.params[1], + args.params[2], df_out); + return; case 30: shafranov_shift_cylinder_kernels::shafranov_shift_df(eta1, eta2, eta3, args.params[0], args.params[1], args.params[2], args.params[3], df_out); @@ -231,6 +249,18 @@ __device__ inline void df_inv(double eta1, double eta2, double eta3, const Domai dfinv_out[6] = 0.0; // [2, 0] dfinv_out[7] = 0.0; // [2, 1] break; + case 23: + dfinv_out[2] = 0.0; // [0, 2] + dfinv_out[5] = 0.0; // [1, 2] + dfinv_out[6] = 0.0; // [2, 0] + dfinv_out[7] = 0.0; // [2, 1] + break; + case 24: + dfinv_out[2] = 0.0; // [0, 2] + dfinv_out[5] = 0.0; // [1, 2] + dfinv_out[6] = 0.0; // [2, 0] + dfinv_out[7] = 0.0; // [2, 1] + break; case 22: dfinv_out[8] = 0.0; // [2, 2] break; @@ -325,6 +355,18 @@ __device__ inline void g(double eta1, double eta2, double eta3, const DomainArgs g_out[6] = 0.0; // [2, 0] g_out[7] = 0.0; // [2, 1] break; + case 23: + g_out[2] = 0.0; // [0, 2] + g_out[5] = 0.0; // [1, 2] + g_out[6] = 0.0; // [2, 0] + g_out[7] = 0.0; // [2, 1] + break; + case 24: + g_out[2] = 0.0; // [0, 2] + g_out[5] = 0.0; // [1, 2] + g_out[6] = 0.0; // [2, 0] + g_out[7] = 0.0; // [2, 1] + break; case 22: // straight field line coordinates if (args.params[3] == 1.0) { @@ -434,6 +476,18 @@ __device__ inline void g_inv(double eta1, double eta2, double eta3, const Domain ginv_out[6] = 0.0; // [2, 0] ginv_out[7] = 0.0; // [2, 1] break; + case 23: + ginv_out[2] = 0.0; // [0, 2] + ginv_out[5] = 0.0; // [1, 2] + ginv_out[6] = 0.0; // [2, 0] + ginv_out[7] = 0.0; // [2, 1] + break; + case 24: + ginv_out[2] = 0.0; // [0, 2] + ginv_out[5] = 0.0; // [1, 2] + ginv_out[6] = 0.0; // [2, 0] + ginv_out[7] = 0.0; // [2, 1] + break; case 22: // straight field line coordinates if (args.params[3] == 1.0) { diff --git a/src/struphy/geometry/evaluation_kernels.py b/src/struphy/geometry/evaluation_kernels.py index 9089ea182..5dd4c9df3 100644 --- a/src/struphy/geometry/evaluation_kernels.py +++ b/src/struphy/geometry/evaluation_kernels.py @@ -11,7 +11,9 @@ import struphy.geometry.domains.hollow_cylinder.hollow_cylinder_kernels as hollow_cylinder_kernels import struphy.geometry.domains.hollow_torus.hollow_torus_kernels as hollow_torus_kernels import struphy.geometry.domains.orthogonal.orthogonal_kernels as orthogonal_kernels +import struphy.geometry.domains.perturbed_disk_cylinder.perturbed_disk_cylinder_kernels as perturbed_disk_cylinder_kernels import struphy.geometry.domains.powered_elliptic_cylinder.powered_elliptic_cylinder_kernels as powered_elliptic_cylinder_kernels +import struphy.geometry.domains.rounded_square_cylinder.rounded_square_cylinder_kernels as rounded_square_cylinder_kernels import struphy.geometry.domains.shafranov_dshaped_cylinder.shafranov_dshaped_cylinder_kernels as shafranov_dshaped_cylinder_kernels import struphy.geometry.domains.shafranov_shift_cylinder.shafranov_shift_cylinder_kernels as shafranov_shift_cylinder_kernels import struphy.geometry.domains.shafranov_sqrt_cylinder.shafranov_sqrt_cylinder_kernels as shafranov_sqrt_cylinder_kernels @@ -150,6 +152,28 @@ def f( args.params[5], f_out, ) + elif args.kind_map == 23: + perturbed_disk_cylinder_kernels.perturbed_disk( + eta1, + eta2, + eta3, + args.params[0], + args.params[1], + args.params[2], + args.params[3], + args.params[4], + f_out, + ) + elif args.kind_map == 24: + rounded_square_cylinder_kernels.rounded_square( + eta1, + eta2, + eta3, + args.params[0], + args.params[1], + args.params[2], + f_out, + ) elif args.kind_map == 30: shafranov_shift_cylinder_kernels.shafranov_shift( eta1, @@ -308,6 +332,26 @@ def df( args.params[5], df_out, ) + elif args.kind_map == 23: + perturbed_disk_cylinder_kernels.perturbed_disk_df( + eta1, + eta2, + args.params[0], + args.params[1], + args.params[2], + args.params[3], + args.params[4], + df_out, + ) + elif args.kind_map == 24: + rounded_square_cylinder_kernels.rounded_square_df( + eta1, + eta2, + args.params[0], + args.params[1], + args.params[2], + df_out, + ) elif args.kind_map == 30: shafranov_shift_cylinder_kernels.shafranov_shift_df( eta1, @@ -454,6 +498,16 @@ def df_inv( dfinv_out[1, 2] = 0.0 dfinv_out[2, 0] = 0.0 dfinv_out[2, 1] = 0.0 + elif args.kind_map == 23: + dfinv_out[0, 2] = 0.0 + dfinv_out[1, 2] = 0.0 + dfinv_out[2, 0] = 0.0 + dfinv_out[2, 1] = 0.0 + elif args.kind_map == 24: + dfinv_out[0, 2] = 0.0 + dfinv_out[1, 2] = 0.0 + dfinv_out[2, 0] = 0.0 + dfinv_out[2, 1] = 0.0 elif args.kind_map == 22: dfinv_out[2, 2] = 0.0 elif args.kind_map == 30: @@ -555,6 +609,16 @@ def g( g_out[1, 2] = 0.0 g_out[2, 0] = 0.0 g_out[2, 1] = 0.0 + elif args.kind_map == 23: + g_out[0, 2] = 0.0 + g_out[1, 2] = 0.0 + g_out[2, 0] = 0.0 + g_out[2, 1] = 0.0 + elif args.kind_map == 24: + g_out[0, 2] = 0.0 + g_out[1, 2] = 0.0 + g_out[2, 0] = 0.0 + g_out[2, 1] = 0.0 elif args.kind_map == 22: # straight field line coordinates if args.params[3] == 1.0: @@ -674,6 +738,16 @@ def g_inv( ginv_out[1, 2] = 0.0 ginv_out[2, 0] = 0.0 ginv_out[2, 1] = 0.0 + elif args.kind_map == 23: + ginv_out[0, 2] = 0.0 + ginv_out[1, 2] = 0.0 + ginv_out[2, 0] = 0.0 + ginv_out[2, 1] = 0.0 + elif args.kind_map == 24: + ginv_out[0, 2] = 0.0 + ginv_out[1, 2] = 0.0 + ginv_out[2, 0] = 0.0 + ginv_out[2, 1] = 0.0 elif args.kind_map == 22: # straight field line coordinates if args.params[3] == 1.0: diff --git a/src/struphy/initial/eigenfunctions.py b/src/struphy/initial/eigenfunctions.py index 8cf791a89..93d95901a 100644 --- a/src/struphy/initial/eigenfunctions.py +++ b/src/struphy/initial/eigenfunctions.py @@ -66,6 +66,14 @@ def __init__(self, derham, **params): nnz_pol = derham.boundary_ops["2"].dim_nz_pol nnz_tor = derham.boundary_ops["2"].dim_nz_tor + # the eigenvectors must have the layout of the 2-forms of this Derham complex (spectra of the former C1 polar + # splines, which had fewer polar rings, cannot be read with the polar splines of regularity >= 1) + assert U2_eig.shape[0] == sum(nnz_pol), ( + f"the spectrum {spec_path} has {U2_eig.shape[0]} coefficients per eigenvector, but the 2-forms of this " + f"Derham complex have {sum(nnz_pol)} poloidal coefficients; recompute the spectrum with the same " + "polar splines, resolution and boundary conditions." + ) + eig_vec_1 = U2_eig[ 0 * nnz_pol[0] + 0 * nnz_pol[1] + 0 * nnz_pol[2] : 1 * nnz_pol[0] + 0 * nnz_pol[1] + 0 * nnz_pol[2], mode, diff --git a/src/struphy/io/options.py b/src/struphy/io/options.py index 4f8919ed6..9d9596572 100644 --- a/src/struphy/io/options.py +++ b/src/struphy/io/options.py @@ -1,5 +1,6 @@ import logging import os +import warnings from dataclasses import dataclass, fields from typing import Any, Callable, Literal @@ -79,6 +80,10 @@ class LiteralOptions: OptsMassPrecond = Literal["MassMatrixPreconditioner", "PolarMassMatrixPreconditioner", None] OptsShearAlfvenPrecond = Literal["MassMatrixPreconditioner", "PolarMassMatrixPreconditioner", "MultiGrid", None] OptsPoloidalSolver = Literal["schur", "lu", "ilu"] + OptsPolarMethod = Literal["c1", "homotopy"] + OptsAngularBasis = Literal["bspline", "trig"] + OptsPolarEnrichment = Literal["cochain_chain", "cochain", "none"] + OptsPolarPolynomialProjection = Literal["commuting", "l2"] OptsDiffusionPrecond = Literal[ "MultiGrid", "StiffnessPreconditioner", @@ -243,20 +248,50 @@ class DerhamOptions(OptionsBase): projector/histopolation integrals. If ``None``, backend defaults are used. + angular_basis : str + Basis in the periodic direction ``eta_2``: ``"bspline"`` (periodic B-splines of degree ``degree[1]`` on + ``num_elements[1]`` cells) or ``"trig"`` (``num_elements[1]`` trigonometric functions, odd, ``degree[1]`` is + ignored; see :mod:`struphy.feec.trig_space`). With trigonometric functions the polar splines of a disk-like + mapping that is exactly representable (``r``-compatible) are exactly ``C^r`` at the pole. FEEC only: + particles, local projectors and spline-kernel field evaluation are not supported, and ``eta_2`` must not be + MPI-decomposed. + polar_splines : bool Smoothness at a possible polar singularity at ``eta_1 = 0``. - ``False`` gives standard tensor-product splines, ``True`` gives C1 polar - splines. + ``False`` gives standard tensor-product splines, ``True`` gives polar + splines of regularity ``polar_regularity`` (built with homotopy + operators, see :mod:`struphy.polar.homotopy`; regularity 1 spans the C1 + polar splines of Toshniwal et al. 2017 and Holderied & Possanner 2022). local_projectors : bool Whether to build local commuting projectors based on quasi-inter-/histopolation. - polar_poloidal_solver : str - Solver for the poloidal (eta_1-eta_2) matrices in the 2D x 1D Kronecker - solves of the polar commuting projectors (only used with ``polar_splines=True``): - ``"schur"`` (exact, block elimination of the pole block), ``"lu"`` - (exact, sparse LU) or ``"ilu"`` (incomplete LU as preconditioner). + polar_poloidal_solver : str | None + Deprecated and ignored: the polar commuting projectors need no poloidal solver (their polar + inter-/histopolation matrix is the identity). + + polar_method : str | None + Deprecated: there is one polar construction. ``"c1"`` gives ``polar_regularity=1`` (the same spaces as the + former C1 polar splines), ``"homotopy"`` keeps ``polar_regularity``. + + polar_regularity : int + Regularity ``r`` of the polar splines at the pole: 0-forms are ``C^r``, 1-forms ``C^{r-1}`` + and 2-forms ``C^{r-2}`` at the pole (conforming for ``r >= 1``). Requires ``r <= degree[0]`` and + ``num_elements[1] >= 2r + 1``. + + polar_enrichment : str + Enrichment of the polar spaces: ``"cochain_chain"`` (the paper's construction), + ``"cochain"`` (no compact 2-forms) or ``"none"`` (external tensor-product functions only). + + polar_polynomial_projection : str + Projection of the polynomials onto the tensor-product spaces in the polar construction: + ``"commuting"`` (geometric degrees of freedom, the paper's choice) or ``"l2"`` (a research option, well-defined + only when the polynomials are exactly representable, i.e. not with B-splines in ``eta_2``). + + polar_cross_section : str + Reference cross-section of the polar complex: ``"domain"`` (the domain at ``eta_3 = 0``) + or ``"unit_disk"``. """ degree: tuple[int, int, int] = (1, 1, 1) @@ -269,10 +304,40 @@ class DerhamOptions(OptionsBase): nquads_proj: tuple[int, int, int] | None = None polar_splines: bool = False local_projectors: bool = False - polar_poloidal_solver: LiteralOptions.OptsPoloidalSolver = "schur" + angular_basis: LiteralOptions.OptsAngularBasis = "bspline" + polar_poloidal_solver: LiteralOptions.OptsPoloidalSolver | None = None + polar_method: LiteralOptions.OptsPolarMethod | None = None + polar_regularity: int = 1 + polar_enrichment: LiteralOptions.OptsPolarEnrichment = "cochain_chain" + polar_polynomial_projection: LiteralOptions.OptsPolarPolynomialProjection = "commuting" + polar_cross_section: str = "domain" def __post_init__(self): - check_option(self.polar_poloidal_solver, LiteralOptions.OptsPoloidalSolver) + check_option(self.angular_basis, LiteralOptions.OptsAngularBasis) + if self.polar_poloidal_solver is not None: + check_option(self.polar_poloidal_solver, LiteralOptions.OptsPoloidalSolver) + warnings.warn( + "DerhamOptions.polar_poloidal_solver is deprecated and ignored: the polar commuting projectors need " + "no poloidal solver.", + DeprecationWarning, + stacklevel=3, + ) + if self.polar_method is not None: + check_option(self.polar_method, LiteralOptions.OptsPolarMethod) + warnings.warn( + "DerhamOptions.polar_method is deprecated: there is one polar construction, of regularity " + "polar_regularity ('c1' is polar_regularity=1, which spans the former C1 polar splines).", + DeprecationWarning, + stacklevel=3, + ) + if self.polar_method == "c1": + self.polar_regularity = 1 + check_option(self.polar_enrichment, LiteralOptions.OptsPolarEnrichment) + check_option(self.polar_polynomial_projection, LiteralOptions.OptsPolarPolynomialProjection) + assert self.polar_regularity >= 0, "polar_regularity must be nonnegative" + assert self.polar_cross_section in ("domain", "unit_disk"), ( + "polar_cross_section must be 'domain' or 'unit_disk'" + ) for bc in self.bcs: if bc is not None: assert isinstance(bc, tuple) and len(bc) == 2, ( diff --git a/src/struphy/models/__init__.py b/src/struphy/models/__init__.py index fdb67f440..77370e369 100644 --- a/src/struphy/models/__init__.py +++ b/src/struphy/models/__init__.py @@ -29,6 +29,7 @@ "Poisson": "struphy.models.poisson", "PressureLessSPH": "struphy.models.pressure_less_sph", "RandomParticleDiffusion": "struphy.models.random_particle_diffusion", + "ScalarWave": "struphy.models.scalar_wave", "ShearAlfven": "struphy.models.shear_alfven", "ToyDrift": "struphy.models.toy_drift", "TwoFluidQuasiNeutralToy": "struphy.models.two_fluid_quasi_neutral_toy", @@ -70,6 +71,7 @@ from struphy.models.poisson import Poisson from struphy.models.pressure_less_sph import PressureLessSPH from struphy.models.random_particle_diffusion import RandomParticleDiffusion + from struphy.models.scalar_wave import ScalarWave from struphy.models.shear_alfven import ShearAlfven from struphy.models.toy_drift import ToyDrift from struphy.models.two_fluid_quasi_neutral_toy import TwoFluidQuasiNeutralToy @@ -110,6 +112,7 @@ "Poisson", "PressureLessSPH", "RandomParticleDiffusion", + "ScalarWave", "ShearAlfven", "ToyDrift", "TwoFluidQuasiNeutralToy", diff --git a/src/struphy/models/scalar_wave.py b/src/struphy/models/scalar_wave.py new file mode 100644 index 000000000..5436772be --- /dev/null +++ b/src/struphy/models/scalar_wave.py @@ -0,0 +1,129 @@ +import copy + +from struphy import BaseUnits +from struphy.io.options import LiteralOptions +from struphy.models.base import StruphyModel +from struphy.models.scalars import BilinearEnergyFEEC, Scalars +from struphy.models.species import FieldSpecies +from struphy.models.variables import FEECVariable +from struphy.propagators.scalar_wave_leapfrog import ScalarWaveLeapfrog + + +class ScalarWave(StruphyModel): + """Scalar (0-form) wave equation, the model problem of the polar-spline wave studies. + + Parameters + ---------- + base_units: BaseUnits + Base units for normalization (default: BaseUnits()) + """ + + ## species + + class EMFields(FieldSpecies): + def __init__(self): + self.phi = FEECVariable(space="H1") + self.psi = FEECVariable(space="H1") + self.init_variables() + + ## propagators + + class Propagators: + def __init__(self): + self.wave = ScalarWaveLeapfrog() + + ## abstract methods + + def __init__(self, base_units: BaseUnits = BaseUnits()): + # 0. store input parameters + self.params = copy.deepcopy(locals()) + + # 1. instantiate all species + self.em_fields = self.EMFields() + + # 2. derive units + self.setup_equation_params(base_units=base_units) + + # 3. instantiate all propagators + self.propagators = self.Propagators() + + # 4. assign variables to propagators + self.propagators.wave.variables.phi = self.em_fields.phi + self.propagators.wave.variables.psi = self.em_fields.psi + + # 5. scalars: kinetic energy 1/2 |psi|^2 (the potential energy 1/2 |grad phi|^2 is computed in the pproc) + self.scalars = Scalars(kinetic_energy=BilinearEnergyFEEC(self.em_fields.psi)) + + @property + def bulk_species(self): + return None + + @property + def velocity_scale(self): + return None + + def post_allocate(self): + pass + + ## abstract methods for documentation + + @classmethod + def model_type(cls) -> LiteralOptions.ModelTypes: + return "Toy" + + @classmethod + def doc_pde(cls): + r"""**PDEs solved by model:** + + .. math:: + + \frac{\partial^2 \phi}{\partial t^2} - c^2 \Delta \phi = 0 + """ + + @classmethod + def doc_normalization(cls): + r"""All quantities are dimensionless; the wave speed :math:`c` is an option of the propagator.""" + + @classmethod + def doc_scalar_quantities(cls): + r"""**The following scalars are tracked during simulation:** + + - Kinetic energy: :math:`\frac{1}{2} \int |\partial_t \phi|^2 \, dV`""" + + @classmethod + def doc_discretization(cls): + """Time integration is performed by the following propagators (in sequence): + + 1. :class:`~struphy.propagators.scalar_wave_leapfrog.ScalarWaveLeapfrog` + """ + return rf"""**1. propagators.wave.ScalarWaveLeapfrog:** + +{ScalarWaveLeapfrog.__doc__} +""" + + @classmethod + def doc_long_description(cls): + r"""The 0-form wave equation in H1 with the field :math:`\phi` and its time derivative :math:`\psi`, integrated + by the symplectic leapfrog. It is the model problem of the high-order polar-spline studies (smoothness of the + waves crossing the polar axis, see ``DerhamOptions(polar_splines=True, polar_regularity=r)``).""" + + @classmethod + def doc_examples(cls): + r"""Create and initialize a ScalarWave model: + + .. code-block:: python + + from struphy.models import ScalarWave + + model = ScalarWave() + model.em_fields.phi # wave field + model.em_fields.psi # its time derivative + """ + + @classmethod + def doc_use_cases(cls): + """Verification of FEEC spaces (e.g. polar splines) with a simple hyperbolic problem.""" + + @classmethod + def doc_cannot_be_used_for(cls): + """Plasma dynamics; the model has no particles and no coupling to other fields.""" diff --git a/src/struphy/polar/basic.py b/src/struphy/polar/basic.py index 7704193dc..1335cbbbd 100644 --- a/src/struphy/polar/basic.py +++ b/src/struphy/polar/basic.py @@ -27,9 +27,10 @@ class PolarDerhamSpace(VectorSpace): def __init__(self, derham: "Derham", space_id: str): assert not derham.spl_kind[0], "Spline basis in eta1 must be clamped" assert derham.spl_kind[1], "Spline basis in eta2 must be periodic" - assert (derham.num_elements[1] / 3) % 1 == 0.0, "Number of elements in eta2 must be a multiple of 3" - assert derham.degree[0] > 1 and derham.degree[1] > 1, "Spline degrees in (eta1, eta2) must be at least two" + # the polar extraction blocks define the layout of the polar coefficients and rings + assert derham.ck_blocks is not None, "the Derham complex has no polar splines (polar_splines=False)" + layout = derham.ck_blocks.space_layout # other properties self._dtype = float @@ -48,75 +49,41 @@ def __init__(self, derham: "Derham", space_id: str): self._ends = self._parent_space.ends # polar properties - if space_id == "H1": - self._n_polar = (3,) - self._n_rings = (2,) - self._dimension = ((self.n[0] - self.n_rings[0]) * self.n[1] + self.n_polar[0]) * self.n[2] - self._n2 = (self.n[1],) - self._n3 = (self.n[2],) - self._type_of_basis_3 = (derham.V0splines.spline_types[0][2],) - elif space_id == "Hcurl": - self._n_polar = (0, 2, 3) - self._n_rings = (1, 2, 2) - dim1 = ((self.d[0] - self.n_rings[0]) * self.n[1] + self.n_polar[0]) * self.n[2] - dim2 = ((self.n[0] - self.n_rings[1]) * self.d[1] + self.n_polar[1]) * self.n[2] - dim3 = ((self.n[0] - self.n_rings[2]) * self.n[1] + self.n_polar[2]) * self.d[2] - self._dimension = dim1 + dim2 + dim3 - self._n2 = (self.n[1], self.d[1], self.n[1]) - self._n3 = (self.n[2], self.n[2], self.d[2]) - self._type_of_basis_3 = ( - derham.V1splines.spline_types[0][2], - derham.V1splines.spline_types[1][2], - derham.V1splines.spline_types[2][2], - ) - elif space_id == "Hdiv": - self._n_polar = (2, 0, 0) - self._n_rings = (2, 1, 1) - dim1 = ((self.n[0] - self.n_rings[0]) * self.d[1] + self.n_polar[0]) * self.d[2] - dim2 = ((self.d[0] - self.n_rings[1]) * self.n[1] + self.n_polar[1]) * self.d[2] - dim3 = ((self.d[0] - self.n_rings[2]) * self.d[1] + self.n_polar[2]) * self.n[2] - self._dimension = dim1 + dim2 + dim3 - self._n2 = (self.d[1], self.n[1], self.d[1]) - self._n3 = (self.d[2], self.d[2], self.n[2]) - self._type_of_basis_3 = ( - derham.V2splines.spline_types[0][2], - derham.V2splines.spline_types[1][2], - derham.V2splines.spline_types[2][2], - ) - elif space_id == "L2": - self._n_polar = (0,) - self._n_rings = (1,) - self._dimension = ((self.d[0] - self.n_rings[0]) * self.d[1] + self.n_polar[0]) * self.d[2] - self._n2 = (self.d[1],) - self._n3 = (self.d[2],) - self._type_of_basis_3 = (derham.V3splines.spline_types[0][2],) - elif space_id == "H1vec": - self._n_polar = (3, 3, 3) - self._n_rings = (2, 2, 2) - self._dimension = (((self.n[0] - self.n_rings[0]) * self.n[1] + self.n_polar[0]) * self.n[2]) * 3 - self._n2 = (self.n[1], self.n[1], self.n[1]) - self._n3 = (self.n[2], self.n[2], self.n[2]) - self._type_of_basis_3 = ( - derham.Vvsplines.spline_types[0][2], - derham.Vvsplines.spline_types[1][2], - derham.Vvsplines.spline_types[2][2], - ) - else: - raise ValueError("Space not supported.") + assert space_id in ("H1", "Hcurl", "Hdiv", "L2", "H1vec"), "Space not supported." + form = derham.space_to_form[space_id] + self._n_polar, self._n_rings = layout(form) + spaces = { + "H1": derham.V0splines, + "Hcurl": derham.V1splines, + "Hdiv": derham.V2splines, + "L2": derham.V3splines, + "H1vec": derham.Vvsplines, + }[space_id] + nbasis = spaces.nbasis # per component, per direction + self._n2 = tuple(nb[1] for nb in nbasis) + self._n3 = tuple(nb[2] for nb in nbasis) + self._type_of_basis_3 = tuple(st[2] for st in spaces.spline_types) + self._dimension = sum( + (nb[0] - nr) * nb[1] * nb[2] + npol * nb[2] for nb, nr, npol in zip(nbasis, self._n_rings, self._n_polar) + ) self._n_comps = len(self.n_polar) + # the process(es) at the pole must hold all polar rings and the first tensor-product ring; the check is + # collective, such that all processes raise (otherwise the processes without the pole would wait forever) if self.n_comps == 1: self._is_pole_rank = self.starts[0] == 0 - if self.is_pole_rank: - assert self.ends[0] > self.n_rings[0], "MPI coeff decomposition in eta_1 too small for polar splines!" + decomposition_ok = not self.is_pole_rank or self.ends[0] > self.n_rings[0] else: self._is_pole_rank = self.starts[0][0] == 0 - for n in range(3): - if self.starts[n][0] == 0: - assert self.ends[n][0] > self.n_rings[n], ( - "MPI coeff decomposition in eta_1 too small for polar splines!" - ) + decomposition_ok = all(self.starts[n][0] != 0 or self.ends[n][0] > self.n_rings[n] for n in range(3)) + if self.comm is not None and self.comm.Get_size() > 1: + decomposition_ok = self.comm.allreduce(bool(decomposition_ok), op=MPI.LAND) + assert decomposition_ok, ( + f"MPI coeff decomposition in eta_1 too small for polar splines: the process at the pole needs at least " + f"{max(self.n_rings) + 2} coefficients in eta_1, the polar rings and two tensor-product rings (use more " + f"elements in eta_1 or mpi_dims_mask=(False, True, True))." + ) # sub-communicator of the processes at the pole (the only ones holding the inner tp rings) if self.comm is not None and self.comm.Get_size() > 1: diff --git a/src/struphy/polar/extraction_operators.py b/src/struphy/polar/extraction_operators.py deleted file mode 100644 index a6caa9e14..000000000 --- a/src/struphy/polar/extraction_operators.py +++ /dev/null @@ -1,943 +0,0 @@ -import cunumpy as xp -import numpy as np - -from struphy.eigenvalue_solvers.derivatives import grad_1d_matrix - - -# ============================= 2D polar splines (C1) =================================== -class PolarExtractionBlocksC1: - """ - Class for operator blocks defining C1 basis extraction operators and DOF extraction operators. - - / / / / / - k = 1 / / / / / - k = 0 / / / / / - ------------------------------- ----------------- - j = 0 | | | | | - j = 1 | | | | | - | | | | | - | i = 0 | ... | i = n_rings | i > n_rings | - | | | | | / - | | | | | / - j = n2-1 | | | | | / - ------------------------------------------------- - | polar rings | first tp ring | outer tp zone | - - Fig. : Indexing of 3d spline tensor-product (tp) basis functions/DOFs. - A linear combination of n_rings rings is used to construct n_polar polar basis functions/DOFs - and "first tp ring" basis functions/DOFs. - - The blocks are setup data: they are built on the host (NumPy, SciPy sparse) on every array backend, - from a host copy of the control points. On the CuPy backend the polar operators - (:mod:`struphy.polar.linear_operators`) copy them to the device once and apply them there. - - Parameters - ---------- - domain : struphy.geometry.base.Domain - Spline mapping (IGAPolarCylinder, IGAPolarTorus, ...) whose control points define the polar splines. - - derham : struphy.feec.psydac_derham.Derham - Discrete Derham complex. - """ - - def __init__(self, domain, derham): - from scipy.sparse import csr_matrix as csr - - # get control points (host copies; they live on the device on the CuPy backend) - cx = xp.to_numpy(domain.cx[:, :, 0]) - cy = xp.to_numpy(domain.cy[:, :, 0]) - - self._cx = cx - self._cy = cy - - self._pole = (cx[0, 0], cy[0, 0]) - - assert np.all(cx[0] == self.pole[0]) - assert np.all(cy[0] == self.pole[1]) - - self._n0 = cx.shape[0] - self._n1 = cx.shape[1] - self._n2 = derham.V0splines.nbasis[0][2] - - assert derham.spl_kind[1], "Use of poalr splines requires periodic splines in eta2." - assert self.n1 == derham.num_elements[1], ( - f"Polar splines: number of control points {self.n1} in eta2 direction is not consistent with the grid (with {derham.num_elements[1]})." - ) - - self._d0 = self.n0 - 1 - self._d1 = self.n1 - 0 - self._d2 = derham.V3splines.nbasis[0][2] - - self._n_rings = [(2,), (1, 2), (2, 1), (1,)] - self._n_polar = [(3,), (0, 2), (2, 0), (0,)] - - # size of control triangle - self._tau = max( - [ - ((self.cx[1] - self.pole[0]) * (-2)).max(), - ((self.cx[1] - self.pole[0]) - np.sqrt(3) * (self.cy[1] - self.pole[1])).max(), - ((self.cx[1] - self.pole[0]) + np.sqrt(3) * (self.cy[1] - self.pole[1])).max(), - ], - ) - - # barycentric coordinates - self._xi_0 = np.zeros((3, self.n1), dtype=float) - self._xi_1 = np.zeros((3, self.n1), dtype=float) - - self._xi_0[:, :] = 1 / 3 - - self._xi_1[0, :] = 1 / 3 + 2 / (3 * self.tau) * (self.cx[1] - self.pole[0]) - self._xi_1[1, :] = ( - 1 / 3 - - 1 / (3 * self.tau) * (self.cx[1] - self.pole[0]) - + np.sqrt(3) / (3 * self.tau) * (self.cy[1] - self.pole[1]) - ) - self._xi_1[2, :] = ( - 1 / 3 - - 1 / (3 * self.tau) * (self.cx[1] - self.pole[0]) - - np.sqrt(3) / (3 * self.tau) * (self.cy[1] - self.pole[1]) - ) - - # remove small values - self._xi_1[abs(self._xi_1) < 1e-14] = 0.0 - - # dictionary for basis extraction blocks ten_to_pol - self._e_ten_to_pol = {} - - # ============= basis extraction operator for discrete 0-forms ================ - - # first n_rings tp rings --> "polar coeffs" - e0_blocks_ten_to_pol = np.block([self.xi_0, self.xi_1]) - self._e_ten_to_pol["0"] = [[csr(e0_blocks_ten_to_pol)]] - - # ============ basis extraction operator for discrete 1-forms (Hcurl) ========= - - # first n_rings tp rings --> "polar coeffs" - e1_11_blocks_ten_to_pol = np.zeros((self.n_polar[1][0], self.n_rings[1][0] * self.n1), dtype=float) - e1_12_blocks_ten_to_pol = np.zeros((self.n_polar[1][0], self.n_rings[1][1] * self.d1), dtype=float) - - e1_21_blocks_ten_to_pol = np.zeros((self.n_polar[1][1], self.n_rings[1][0] * self.n1), dtype=float) - e1_22_blocks_ten_to_pol = np.zeros((self.n_polar[1][1], self.n_rings[1][1] * self.d1), dtype=float) - - # 1st component - for l in range(2): - for j in range(self.n1): - e1_21_blocks_ten_to_pol[l, j] = self.xi_1[l + 1, j] - self.xi_0[l + 1, j] - - # 2nd component - for l in range(2): - for j in range(1 * self.d1, 2 * self.d1): - e1_22_blocks_ten_to_pol[l, j] = ( - self.xi_1[l + 1, (j - self.d1 + 1) % self.d1] - self.xi_1[l + 1, j - self.d1] - ) - - self._e_ten_to_pol["1"] = [ - [csr(e1_11_blocks_ten_to_pol), csr(e1_12_blocks_ten_to_pol), None], - [csr(e1_21_blocks_ten_to_pol), csr(e1_22_blocks_ten_to_pol), None], - [None, None, csr(e0_blocks_ten_to_pol)], - ] - - # =============== basis extraction operator for discrete 1-forms (Hdiv) ========= - - # first n_rings tp rings --> "polar coeffs" - e3_blocks_ten_to_pol = np.zeros((self.n_polar[3][0], self.n_rings[3][0] * self.d1), dtype=float) - - self._e_ten_to_pol["2"] = [ - [csr(e1_22_blocks_ten_to_pol), csr(-e1_21_blocks_ten_to_pol), None], - [csr(e1_12_blocks_ten_to_pol), csr(e1_11_blocks_ten_to_pol), None], - [None, None, csr(e3_blocks_ten_to_pol)], - ] - - # ================== basis extraction operator for discrete 2-forms ============= - - # first n_rings tp rings --> "polar coeffs" - self._e_ten_to_pol["3"] = [[csr(e3_blocks_ten_to_pol)]] - - # ================ basis extraction operator for discrete vector fields ========= - - self._e_ten_to_pol["v"] = [ - [csr(e0_blocks_ten_to_pol), None, None], - [None, csr(e0_blocks_ten_to_pol), None], - [None, None, csr(e0_blocks_ten_to_pol)], - ] - - # ======= projection extraction operator for discrete 0-forms ==================== - self._p_ten_to_pol = {} - self._p_ten_to_ten = {} - - # first n_rings tp rings --> "polar coeffs" - p0_blocks_ten_to_pol = np.zeros((self.n_polar[0][0], self.n_rings[0][0] * self.n1), dtype=float) - - # !! NOTE: for odd spline degrees and periodic splines the first Greville point sometimes does NOT start at zero!! - if domain.degree[1] % 2 != 0 and not (abs(derham.V0fem.spaces[1].interpolation_grid[0]) < 1e-14): - p0_blocks_ten_to_pol[0, self.n1 + 3 * self.n1 // 3 - 1] = 1.0 - p0_blocks_ten_to_pol[1, self.n1 + 1 * self.n1 // 3 - 1] = 1.0 - p0_blocks_ten_to_pol[2, self.n1 + 2 * self.n1 // 3 - 1] = 1.0 - else: - p0_blocks_ten_to_pol[0, self.n1 + 0 * self.n1 // 3] = 1.0 - p0_blocks_ten_to_pol[1, self.n1 + 1 * self.n1 // 3] = 1.0 - p0_blocks_ten_to_pol[2, self.n1 + 2 * self.n1 // 3] = 1.0 - - self._p_ten_to_pol["0"] = [[csr(p0_blocks_ten_to_pol)]] - - # first n_rings + 1 tp rings --> "first tp ring" - p0_blocks_ten_to_ten = np.block([0 * np.identity(self.n1)] * self.n_rings[0][0] + [np.identity(self.n1)]) - - self._p_ten_to_ten["0"] = [[csr(p0_blocks_ten_to_ten)]] - - # =========== projection extraction operator for discrete 1-forms (Hcurl) ======== - - # first n_rings tp rings --> "polar coeffs" - p1_11_blocks_ten_to_pol = np.zeros((self.n_polar[1][0], self.n_rings[1][0] * self.n1), dtype=float) - p1_22_blocks_ten_to_pol = np.zeros((self.n_polar[1][1], self.n_rings[1][1] * self.d1), dtype=float) - - # !! NOTE: PSYDAC's first integration interval sometimes start at < 0 !! - if derham.V3fem.spaces[1].histopolation_grid[0] < -1e-14: - p1_22_blocks_ten_to_pol[0, (self.d1 + 0 * self.d1 // 3 + 1) : (self.d1 + 1 * self.d1 // 3 + 1)] = 1.0 - p1_22_blocks_ten_to_pol[1, (self.d1 + 0 * self.d1 // 3 + 1) : (self.d1 + 1 * self.d1 // 3 + 1)] = 1.0 - p1_22_blocks_ten_to_pol[1, (self.d1 + 1 * self.d1 // 3 + 1) : (self.d1 + 2 * self.d1 // 3 + 1)] = 1.0 - else: - p1_22_blocks_ten_to_pol[0, (self.d1 + 0 * self.d1 // 3) : (self.d1 + 1 * self.d1 // 3)] = 1.0 - p1_22_blocks_ten_to_pol[1, (self.d1 + 0 * self.d1 // 3) : (self.d1 + 1 * self.d1 // 3)] = 1.0 - p1_22_blocks_ten_to_pol[1, (self.d1 + 1 * self.d1 // 3) : (self.d1 + 2 * self.d1 // 3)] = 1.0 - - p1_12_blocks_ten_to_pol = np.zeros((self.n_polar[1][0], self.n_rings[1][1] * self.d1), dtype=float) - p1_21_blocks_ten_to_pol = np.zeros((self.n_polar[1][1], self.n_rings[1][0] * self.d1), dtype=float) - - self._p_ten_to_pol["1"] = [ - [csr(p1_11_blocks_ten_to_pol), csr(p1_12_blocks_ten_to_pol), None], - [csr(p1_21_blocks_ten_to_pol), csr(p1_22_blocks_ten_to_pol), None], - [None, None, csr(p0_blocks_ten_to_pol)], - ] - - # first n_rings + 1 tp rings --> "first tp ring" - p1_11_blocks_ten_to_ten = np.zeros((self.n1, self.n1), dtype=float) - - # !! NOTE: for odd spline degrees and periodic splines the first Greville point sometimes does NOT start at zero!! - if domain.degree[1] % 2 != 0 and not (abs(derham.V0fem.spaces[1].interpolation_grid[0]) < 1e-14): - p1_11_blocks_ten_to_ten[:, 3 * self.n1 // 3 - 1] = -np.roll(self.xi_1[0], -1) - p1_11_blocks_ten_to_ten[:, 1 * self.n1 // 3 - 1] = -np.roll(self.xi_1[1], -1) - p1_11_blocks_ten_to_ten[:, 2 * self.n1 // 3 - 1] = -np.roll(self.xi_1[2], -1) - else: - p1_11_blocks_ten_to_ten[:, 0 * self.n1 // 3] = -self.xi_1[0] - p1_11_blocks_ten_to_ten[:, 1 * self.n1 // 3] = -self.xi_1[1] - p1_11_blocks_ten_to_ten[:, 2 * self.n1 // 3] = -self.xi_1[2] - - p1_11_blocks_ten_to_ten += np.identity(self.n1) - - p1_11_blocks_ten_to_ten = np.block([p1_11_blocks_ten_to_ten, np.identity(self.n1)]) - - p1_22_blocks_ten_to_ten = np.block([0 * np.identity(self.d1)] * self.n_rings[1][1] + [np.identity(self.d1)]) - - p1_12_blocks_ten_to_ten = np.zeros((self.d1, (self.n_rings[1][1] + 1) * self.d1), dtype=float) - p1_21_blocks_ten_to_ten = np.zeros((self.n1, (self.n_rings[1][0] + 1) * self.n1), dtype=float) - - self._p_ten_to_ten["1"] = [ - [csr(p1_11_blocks_ten_to_ten), csr(p1_12_blocks_ten_to_ten), None], - [csr(p1_21_blocks_ten_to_ten), csr(p1_22_blocks_ten_to_ten), None], - [None, None, csr(p0_blocks_ten_to_ten)], - ] - - # ========== projection extraction operator for discrete 1-forms (Hdiv) ========== - - # first n_rings tp rings --> "polar coeffs" - p3_blocks_ten_to_pol = np.zeros((self.n_polar[3][0], self.n_rings[3][0] * self.d1), dtype=float) - - self._p_ten_to_pol["2"] = [ - [csr(p1_22_blocks_ten_to_pol), csr(p1_21_blocks_ten_to_pol), None], - [csr(p1_12_blocks_ten_to_pol), csr(p1_11_blocks_ten_to_pol), None], - [None, None, csr(p3_blocks_ten_to_pol)], - ] - - # first n_rings + 1 tp rings --> "first tp ring" - p3_blocks_ten_to_ten = np.zeros((self.d1, self.d1), dtype=float) - - a0 = np.diff(self.xi_1[1], append=self.xi_1[1, 0]) - a1 = np.diff(self.xi_1[2], append=self.xi_1[2, 0]) - - # !! NOTE: PSYDAC's first integration interval sometimes start at < 0 !! - if derham.V3fem.spaces[1].histopolation_grid[0] < -1e-14: - p3_blocks_ten_to_ten[:, (0 * self.n1 // 3 + 1) : (1 * self.n1 // 3 + 1)] = ( - -np.roll(a0, +1)[:, None] - np.roll(a1, +1)[:, None] - ) - p3_blocks_ten_to_ten[:, (1 * self.n1 // 3 + 1) : (2 * self.n1 // 3 + 1)] = -np.roll(a1, +1)[:, None] - else: - p3_blocks_ten_to_ten[:, 0 * self.n1 // 3 : 1 * self.n1 // 3] = -a0[:, None] - a1[:, None] - p3_blocks_ten_to_ten[:, 1 * self.n1 // 3 : 2 * self.n1 // 3] = -a1[:, None] - - p3_blocks_ten_to_ten += np.identity(self.d1) - - p3_blocks_ten_to_ten = np.block([p3_blocks_ten_to_ten, np.identity(self.d1)]) - - self._p_ten_to_ten["2"] = [ - [csr(p1_22_blocks_ten_to_ten), csr(p1_21_blocks_ten_to_ten), None], - [csr(p1_12_blocks_ten_to_ten), csr(p1_11_blocks_ten_to_ten), None], - [None, None, csr(p3_blocks_ten_to_ten)], - ] - - # =============== projection extraction operator for discrete 2-forms ============ - - # first n_rings tp rings --> "polar coeffs" - self._p_ten_to_pol["3"] = [[csr(p3_blocks_ten_to_pol)]] - - # first n_rings + 1 tp rings --> "first tp ring" - self._p_ten_to_ten["3"] = [[csr(p3_blocks_ten_to_ten)]] - - # ============ projection extraction operator for discrete vector fields ========= - - self._p_ten_to_pol["v"] = [ - [csr(p0_blocks_ten_to_pol), None, None], - [None, csr(p0_blocks_ten_to_pol), None], - [None, None, csr(p0_blocks_ten_to_pol)], - ] - - self._p_ten_to_ten["v"] = [ - [csr(p0_blocks_ten_to_ten), None, None], - [None, csr(p0_blocks_ten_to_ten), None], - [None, None, csr(p0_blocks_ten_to_ten)], - ] - - # ======================= discrete gradient ====================================== - - # "polar coeffs" to "polar coeffs" - grad_pol_to_pol_1 = np.zeros((self.n_polar[1][0], self.n_polar[0][0]), dtype=float) - grad_pol_to_pol_2 = np.array([[-1.0, 1.0, 0.0], [-1.0, 0.0, 1.0]]) - grad_pol_to_pol_3 = np.identity(self.n_polar[0][0], dtype=float) - - self._grad_pol_to_pol = [[csr(grad_pol_to_pol_1)], [csr(grad_pol_to_pol_2)], [csr(grad_pol_to_pol_3)]] - - # "polar coeffs" to "first tp ring" - grad_pol_to_ten_1 = np.zeros(((self.n_rings[1][0] + 1) * self.n1, self.n_polar[0][0])) - grad_pol_to_ten_2 = np.zeros(((self.n_rings[1][1] + 1) * self.d1, self.n_polar[0][0])) - grad_pol_to_ten_3 = np.zeros(((self.n_rings[0][0] + 1) * self.n1, self.n_polar[0][0])) - - grad_pol_to_ten_1[-self.n1 :, :] = -self.xi_1.T - - self._grad_pol_to_ten = [[csr(grad_pol_to_ten_1)], [csr(grad_pol_to_ten_2)], [csr(grad_pol_to_ten_3)]] - - # eta_3 direction - grad_e3_1 = np.identity(self.n2, dtype=float) - grad_e3_2 = np.identity(self.n2, dtype=float) - grad_e3_3 = xp.to_numpy(grad_1d_matrix(derham.spl_kind[2], self.n2)) - - self._grad_e3 = [[csr(grad_e3_1)], [csr(grad_e3_2)], [csr(grad_e3_3)]] - - # =========================== discrete curl ====================================== - - # "polar coeffs" to "polar coeffs" - curl_pol_to_pol_12 = np.identity(self.n_polar[1][1], dtype=float) - curl_pol_to_pol_13 = np.array([[-1.0, 1.0, 0.0], [-1.0, 0.0, 1.0]]) - - curl_pol_to_pol_21 = np.identity(self.n_polar[1][0], dtype=float) - curl_pol_to_pol_23 = np.zeros((self.n_polar[2][1], self.n_polar[0][0]), dtype=float) - - curl_pol_to_pol_31 = np.zeros((self.n_polar[3][0], self.n_polar[1][0]), dtype=float) - curl_pol_to_pol_32 = np.zeros((self.n_polar[3][0], self.n_polar[1][1]), dtype=float) - - self._curl_pol_to_pol = [ - [None, csr(-curl_pol_to_pol_12), csr(curl_pol_to_pol_13)], - [csr(curl_pol_to_pol_21), None, csr(-curl_pol_to_pol_23)], - [csr(-curl_pol_to_pol_31), csr(curl_pol_to_pol_32), None], - ] - - # "polar coeffs" to "first tp ring" - curl_pol_to_ten_12 = np.zeros(((self.n_rings[2][0] + 1) * self.d1, self.n_polar[1][1])) - curl_pol_to_ten_13 = np.zeros(((self.n_rings[2][0] + 1) * self.d1, self.n_polar[0][0])) - - curl_pol_to_ten_21 = np.zeros(((self.n_rings[2][1] + 1) * self.n1, self.n_polar[1][0])) - curl_pol_to_ten_23 = np.zeros(((self.n_rings[2][1] + 1) * self.n1, self.n_polar[0][0])) - - curl_pol_to_ten_31 = np.zeros(((self.n_rings[3][0] + 1) * self.n1, self.n_polar[1][0])) - curl_pol_to_ten_32 = np.zeros(((self.n_rings[3][0] + 1) * self.d1, self.n_polar[1][1])) - - curl_pol_to_ten_23[-self.n1 :, :] = -self.xi_1.T - - for l in range(2): - for j in range(self.d1, 2 * self.d1): - curl_pol_to_ten_32[j, l] = -( - self.xi_1[l + 1, (j - self.d1 + 1) % self.d1] - self.xi_1[l + 1, j - self.d1] - ) - - self._curl_pol_to_ten = [ - [None, csr(-curl_pol_to_ten_12), csr(curl_pol_to_ten_13)], - [csr(curl_pol_to_ten_21), None, csr(-curl_pol_to_ten_23)], - [csr(-curl_pol_to_ten_31), csr(curl_pol_to_ten_32), None], - ] - - # eta_3 direction - curl_e3_12 = xp.to_numpy(grad_1d_matrix(derham.spl_kind[2], self.n2)) - curl_e3_13 = np.identity(self.d2) - - curl_e3_21 = xp.to_numpy(grad_1d_matrix(derham.spl_kind[2], self.n2)) - curl_e3_23 = np.identity(self.d2) - - curl_e3_31 = np.identity(self.n2) - curl_e3_32 = np.identity(self.n2) - - self._curl_e3 = [ - [None, csr(curl_e3_12), csr(curl_e3_13)], - [csr(curl_e3_21), None, csr(curl_e3_23)], - [csr(curl_e3_31), csr(curl_e3_32), None], - ] - - # =========================== discrete div ====================================== - - # "polar coeffs" to "polar coeffs" - div_pol_to_pol_1 = np.zeros((self.n_polar[3][0], self.n_polar[2][0]), dtype=float) - div_pol_to_pol_2 = np.zeros((self.n_polar[3][0], self.n_polar[2][1]), dtype=float) - div_pol_to_pol_3 = np.identity(self.n_polar[3][0], dtype=float) - - self._div_pol_to_pol = [[csr(div_pol_to_pol_1), csr(div_pol_to_pol_2), csr(div_pol_to_pol_3)]] - - # "polar coeffs" to "first tp ring" - div_pol_to_ten_1 = np.zeros(((self.n_rings[3][0] + 1) * self.d1, self.n_polar[2][0])) - div_pol_to_ten_2 = np.zeros(((self.n_rings[3][0] + 1) * self.d1, self.n_polar[2][1])) - div_pol_to_ten_3 = np.zeros(((self.n_rings[3][0] + 1) * self.d1, self.n_polar[3][0])) - - for l in range(2): - for j in range(self.d1, 2 * self.d1): - div_pol_to_ten_1[j, l] = -( - self.xi_1[l + 1, (j - self.d1 + 1) % self.d1] - self.xi_1[l + 1, j - self.d1] - ) - - self._div_pol_to_ten = [[csr(div_pol_to_ten_1), csr(div_pol_to_ten_2), csr(div_pol_to_ten_3)]] - - # eta_3 direction - div_e3_1 = np.identity(self.d2, dtype=float) - div_e3_2 = np.identity(self.d2, dtype=float) - div_e3_3 = xp.to_numpy(grad_1d_matrix(derham.spl_kind[2], self.n2)) - - self._div_e3 = [[csr(div_e3_1), csr(div_e3_2), csr(div_e3_3)]] - - @property - def cx(self): - return self._cx - - @property - def cy(self): - return self._cy - - @property - def pole(self): - return self._pole - - @property - def n0(self): - return self._n0 - - @property - def n1(self): - return self._n1 - - @property - def n2(self): - return self._n2 - - @property - def d0(self): - return self._d0 - - @property - def d1(self): - return self._d1 - - @property - def d2(self): - return self._d2 - - @property - def n_rings(self): - return self._n_rings - - @property - def n_polar(self): - return self._n_polar - - @property - def tau(self): - return self._tau - - @property - def xi_0(self): - return self._xi_0 - - @property - def xi_1(self): - return self._xi_1 - - @property - def e_ten_to_pol(self): - return self._e_ten_to_pol - - @property - def p_ten_to_pol(self): - return self._p_ten_to_pol - - @property - def p_ten_to_ten(self): - return self._p_ten_to_ten - - @property - def grad_pol_to_pol(self): - return self._grad_pol_to_pol - - @property - def grad_pol_to_ten(self): - return self._grad_pol_to_ten - - @property - def grad_e3(self): - return self._grad_e3 - - @property - def curl_pol_to_pol(self): - return self._curl_pol_to_pol - - @property - def curl_pol_to_ten(self): - return self._curl_pol_to_ten - - @property - def curl_e3(self): - return self._curl_e3 - - @property - def div_pol_to_pol(self): - return self._div_pol_to_pol - - @property - def div_pol_to_ten(self): - return self._div_pol_to_ten - - @property - def div_e3(self): - return self._div_e3 - - -# ============================= 2D polar splines (C0) =================================== -class PolarSplines_C0_2D: - def __init__(self, n0, n1): - import scipy.sparse as spa - - d0 = n0 - 1 - d1 = n1 - 0 - - # number of polar basis functions in V0 (NN) - self.Nbase0 = (n0 - 1) * n1 + 1 - - # number of polar basis functions in V1_C (DN ND) (1st and 2nd component) - self.Nbase1C_1 = d0 * n1 - self.Nbase1C_2 = (n0 - 1) * d1 - - self.Nbase1C = self.Nbase1C_1 + self.Nbase1C_2 - - # number of polar basis functions in V1_D (ND DN) (1st and 2nd component) - self.Nbase1D_1 = (n0 - 1) * d1 - self.Nbase1D_2 = d0 * n1 - - self.Nbase1D = self.Nbase1D_1 + self.Nbase1D_2 - - # number of polar basis functions in V2 (DD) - self.Nbase2 = d0 * d1 - - # =========== extraction operators for discrete 0-forms ================== - # extraction operator for basis functions - self.E0_11 = spa.csr_matrix(xp.ones((1, n1), dtype=float)) - self.E0_22 = spa.identity((n0 - 1) * n1, format="csr") - - self.E0 = spa.bmat([[self.E0_11, None], [None, self.E0_22]], format="csr") - - # global projection extraction operator for interpolation points - self.P0_11 = xp.zeros((1, n1), dtype=float) - - self.P0_11[0, 0] = 1.0 - - self.P0_11 = spa.csr_matrix(self.P0_11) - - self.P0_22 = spa.identity((n0 - 1) * n1, format="csr") - - self.P0 = spa.bmat([[self.P0_11, None], [None, self.P0_22]], format="csr") - # ======================================================================= - - # =========== extraction operators for discrete 1-forms (H_curl) ======== - self.E1C_1 = spa.identity(d0 * n1, format="csr") - self.E1C_2 = spa.identity(n0 * d1, format="lil")[d1:, :].tocsr() - - # combined first and second component - self.E1C = spa.bmat([[self.E1C_1, None], [None, self.E1C_2]], format="csr") - - # extraction operator for interpolation/histopolation in global projector - self.P1C_1 = spa.identity(d0 * n1, format="csr") - self.P1C_2 = spa.identity(n0 * d1, format="lil")[d1:, :].tocsr() - - # combined first and second component - self.P1C = spa.bmat([[self.P1C_1, None], [None, self.P1C_2]], format="csr") - # ======================================================================== - - # =========== extraction operators for discrete 1-forms (H_div) ========== - self.E1D_1 = spa.identity(n0 * d1, format="lil")[d1:, :].tocsr() - self.E1D_2 = spa.identity(d0 * n1, format="csr") - - # combined first and second component - self.E1D = spa.bmat([[self.E1D_1, None], [None, self.E1D_2]], format="csr") - - # extraction operator for interpolation/histopolation in global projector - self.P1D_1 = spa.identity(n0 * d1, format="lil")[d1:, :].tocsr() - self.P1D_2 = spa.identity(d0 * n1, format="csr") - - # combined first and second component - self.P1D = spa.bmat([[self.P1D_1, None], [None, self.P1D_2]], format="csr") - # ======================================================================== - - # =========== extraction operators for discrete 2-forms ================== - self.E2 = spa.identity(d0 * d1, format="csr") - self.P2 = spa.identity(d0 * d1, format="csr") - # ======================================================================== - - # ========================= 1D discrete derivatives ====================== - grad_1d_1 = spa.csc_matrix(grad_1d_matrix(False, n0)) - grad_1d_2 = spa.csc_matrix(grad_1d_matrix(True, n1)) - # ======================================================================== - - # ========= discrete polar gradient matrix =============================== - # radial dofs (DN) - G11 = xp.zeros(((d0 - 0) * n1, 1), dtype=float) - G11[:n1, 0] = -1.0 - - G12 = spa.kron(grad_1d_1[:, 1:], spa.identity(n1)) - - self.G1 = spa.bmat([[G11, G12]], format="csr") - - # angular dofs (ND) - G21 = xp.zeros(((n0 - 1) * d1, 1), dtype=float) - G22 = spa.kron(spa.identity(n0 - 1), grad_1d_2, format="csr") - - self.G2 = spa.bmat([[G21, G22]], format="csr") - - # combined 1st and 2nd component - self.G = spa.bmat([[self.G1], [self.G2]], format="csr") - # ======================================================================== - - # ========= discrete polar curl matrix =================================== - # 2D vector curl (NN --> ND DN) - - # angular dofs (ND) - VC11 = xp.zeros(((n0 - 1) * d1, 1), dtype=float) - VC12 = spa.kron(spa.identity(n0 - 1), grad_1d_2, format="csr") - - self.VC1 = spa.bmat([[VC11, VC12]], format="csr") - - # radial dofs (DN) - VC21 = xp.zeros(((d0 - 0) * n1, 1), dtype=float) - VC21[:n1, 0] = 1.0 - - VC22 = -spa.kron(grad_1d_1[:, 1:], spa.identity(n1)) - - self.VC2 = spa.bmat([[VC21, VC22]], format="csr") - - # combined 1st and 2nd component - self.VC = spa.bmat([[self.VC1], [self.VC2]], format="csr") - - # 2D scalar curl (DN ND --> DD) - self.SC1 = -spa.kron(spa.identity(d0), grad_1d_2, format="csr") - self.SC2 = spa.kron(grad_1d_1[:, 1:], spa.identity(d1), format="csr") - - # combined 1st and 2nd component - self.SC = spa.bmat([[self.SC1, self.SC2]], format="csr") - # ======================================================================== - - # =============== discrete polar div matrix ============================== - self.D1 = spa.kron(grad_1d_1[:, 1:], spa.identity(d1), format="csr") - self.D2 = spa.kron(spa.identity(d0), grad_1d_2, format="csr") - - # combined 1st and 2nd component - self.D = spa.bmat([[self.D1, self.D2]], format="csr") - # ======================================================================== - - -# ============================= 2D polar splines (C1) =================================== -class PolarSplines_C1_2D: - def __init__(self, cx, cy): - import scipy.sparse as spa - - n0, n1 = cx.shape - - d0 = n0 - 1 - d1 = n1 - 0 - - # location of pole - self.x0 = cx[0, 0] - self.y0 = cy[0, 0] - - # number of polar basis functions in V0 (NN) - self.Nbase0 = (n0 - 2) * n1 + 3 - - # number of polar basis functions in V1_C (DN ND) (1st and 2nd component) - self.Nbase1C_1 = (d0 - 1) * n1 - self.Nbase1C_2 = (n0 - 2) * d1 + 2 - - self.Nbase1C = self.Nbase1C_1 + self.Nbase1C_2 - - # number of polar basis functions in V1_D (ND DN) (1st and 2nd component) - self.Nbase1D_1 = (n0 - 2) * d1 + 2 - self.Nbase1D_2 = (d0 - 1) * n1 - - self.Nbase1D = self.Nbase1D_1 + self.Nbase1D_2 - - # number of polar basis functions in V2 (DD) - self.Nbase2 = (d0 - 1) * d1 - - # size of control triangle - self.tau = xp.array( - [ - (-2 * (cx[1] - self.x0)).max(), - ((cx[1] - self.x0) - xp.sqrt(3) * (cy[1] - self.y0)).max(), - ((cx[1] - self.x0) + xp.sqrt(3) * (cy[1] - self.y0)).max(), - ], - ).max() - - self.Xi_0 = xp.zeros((3, n1), dtype=float) - self.Xi_1 = xp.zeros((3, n1), dtype=float) - - # barycentric coordinates - self.Xi_0[:, :] = 1 / 3 - - self.Xi_1[0, :] = 1 / 3 + 2 / (3 * self.tau) * (cx[1] - self.x0) - self.Xi_1[1, :] = ( - 1 / 3 - 1 / (3 * self.tau) * (cx[1] - self.x0) + xp.sqrt(3) / (3 * self.tau) * (cy[1] - self.y0) - ) - self.Xi_1[2, :] = ( - 1 / 3 - 1 / (3 * self.tau) * (cx[1] - self.x0) - xp.sqrt(3) / (3 * self.tau) * (cy[1] - self.y0) - ) - - # remove small values - self.Xi_1[abs(self.Xi_1) < 1e-15] = 0.0 - - # =========== extraction operators for discrete 0-forms ================== - # extraction operator for basis functions - self.E0_11 = spa.csr_matrix(xp.hstack((self.Xi_0, self.Xi_1))) - self.E0_22 = spa.identity((n0 - 2) * n1, format="csr") - - self.E0 = spa.bmat([[self.E0_11, None], [None, self.E0_22]], format="csr") - - # global projection extraction operator for interpolation points - self.P0_11 = xp.zeros((3, 2 * n1), dtype=float) - - self.P0_11[0, n1 + 0 * n1 // 3] = 1.0 - self.P0_11[1, n1 + 1 * n1 // 3] = 1.0 - self.P0_11[2, n1 + 2 * n1 // 3] = 1.0 - - self.P0_11 = spa.csr_matrix(self.P0_11) - - self.P0_22 = spa.identity((n0 - 2) * n1) - - self.P0 = spa.bmat([[self.P0_11, None], [None, self.P0_22]], format="csr") - # ======================================================================= - - # =========== extraction operators for discrete 1-forms (H_curl) ======== - self.E1C_12 = spa.identity((d0 - 1) * n1) - self.E1C_34 = spa.identity((n0 - 2) * d1) - - self.E1C_21 = xp.zeros((2, 1 * n1), dtype=float) - self.E1C_23 = xp.zeros((2, 2 * d1), dtype=float) - - # 1st component - for s in range(2): - for j in range(n1): - self.E1C_21[s, j] = self.Xi_1[s + 1, j] - self.Xi_0[s + 1, j] - - # 2nd component - for s in range(2): - for j in range(d1): - self.E1C_23[s, j] = 0.0 - self.E1C_23[s, n1 + j] = self.Xi_1[s + 1, (j + 1) % n1] - self.Xi_1[s + 1, j] - - # combined first and second component - self.E1C = spa.bmat( - [[None, self.E1C_12, None, None], [self.E1C_21, None, self.E1C_23, None], [None, None, None, self.E1C_34]], - format="csr", - ) - - # extraction operator for interpolation/histopolation in global projector - - # 1st component - self.P1C_11 = xp.zeros((n1, n1), dtype=float) - self.P1C_12 = spa.identity(n1) - self.P1C_23 = spa.identity((d0 - 2) * n1) - - self.P1C_11[:, 0 * n1 // 3] = -self.Xi_1[0] - self.P1C_11[:, 1 * n1 // 3] = -self.Xi_1[1] - self.P1C_11[:, 2 * n1 // 3] = -self.Xi_1[2] - self.P1C_11 += xp.identity(n1) - - # 2nd component - self.P1C_34 = xp.zeros((2, 2 * d1), dtype=float) - self.P1C_45 = spa.identity((n0 - 2) * d1) - - self.P1C_34[0, (d1 + 0 * d1 // 3) : (d1 + 1 * d1 // 3)] = xp.ones(d1 // 3, dtype=float) - self.P1C_34[1, (d1 + 0 * d1 // 3) : (d1 + 1 * d1 // 3)] = xp.ones(d1 // 3, dtype=float) - self.P1C_34[1, (d1 + 1 * d1 // 3) : (d1 + 2 * d1 // 3)] = xp.ones(d1 // 3, dtype=float) - - # combined first and second component - self.P1C = spa.bmat( - [ - [self.P1C_11, self.P1C_12, None, None, None], - [None, None, self.P1C_23, None, None], - [None, None, None, self.P1C_34, None], - [None, None, None, None, self.P1C_45], - ], - format="csr", - ) - # ========================================================================= - - # ========= extraction operators for discrete 1-forms (H_div) ============= - self.E1D_11 = xp.zeros((2, 2 * d1), dtype=float) - self.E1D_13 = xp.zeros((2, 1 * n1), dtype=float) - - self.E1D_22 = spa.identity((n0 - 2) * d1) - self.E1D_34 = spa.identity((d0 - 1) * n1) - - # 1st component - for s in range(2): - for j in range(d1): - self.E1D_11[s, j] = 0.0 - self.E1D_11[s, n1 + j] = self.Xi_1[s + 1, (j + 1) % n1] - self.Xi_1[s + 1, j] - - # 2nd component - for s in range(2): - for j in range(n1): - self.E1D_13[s, j] = -(self.Xi_1[s + 1, j] - self.Xi_0[s + 1, j]) - - # combined first and second component - self.E1D = spa.bmat( - [[self.E1D_11, None, self.E1D_13, None], [None, self.E1D_22, None, None], [None, None, None, self.E1C_34]], - format="csr", - ) - - # extraction operator for interpolation/histopolation in global projector - self.P1D_11 = self.P1C_34.copy() - self.P1D_22 = self.P1C_45.copy() - - self.P1D_33 = self.P1C_11.copy() - self.P1D_34 = self.P1C_12.copy() - self.P1D_45 = self.P1C_23.copy() - - # combined first and second component - self.P1D = spa.bmat( - [ - [self.P1D_11, None, None, None, None], - [None, self.P1D_22, None, None, None], - [None, None, self.P1D_33, self.P1D_34, None], - [None, None, None, None, self.P1D_45], - ], - format="csr", - ) - # ========================================================================= - - # =========== extraction operators for discrete 2-forms =================== - self.E2_1 = xp.zeros(((d0 - 1) * d1, d1), dtype=float) - self.E2_2 = spa.identity((d0 - 1) * d1) - - self.E2 = spa.bmat([[self.E2_1, self.E2_2]], format="csr") - - # extraction operator for histopolation in global projector - self.P2_11 = xp.zeros((d1, d1), dtype=float) - self.P2_12 = spa.identity(d1) - self.P2_23 = spa.identity((d0 - 2) * d1) - - for i in range(d1): - # block A - self.P2_11[i, 0 * n1 // 3 : 1 * n1 // 3] = -(self.Xi_1[1, (i + 1) % n1] - self.Xi_1[1, i]) - ( - self.Xi_1[2, (i + 1) % n1] - self.Xi_1[2, i] - ) - - # block B - self.P2_11[i, 1 * n1 // 3 : 2 * n1 // 3] = -(self.Xi_1[2, (i + 1) % n1] - self.Xi_1[2, i]) - - self.P2_11 += xp.identity(d1) - - self.P2 = spa.bmat([[self.P2_11, self.P2_12, None], [None, None, self.P2_23]], format="csr") - # ========================================================================= - - # ========================= 1D discrete derivatives ======================= - grad_1d_1 = spa.csc_matrix(grad_1d_matrix(False, n0)) - grad_1d_2 = spa.csc_matrix(grad_1d_matrix(True, n1)) - # ========================================================================= - - # ========= discrete polar gradient matrix ================================ - self.G1_1 = xp.zeros(((d0 - 1) * n1, 3), dtype=float) - self.G1_1[:n1, :] = -self.Xi_1.T - - self.G1_2 = spa.kron(grad_1d_1[1:, 2:], spa.identity(n1)) - - self.G1 = spa.bmat([[self.G1_1, self.G1_2]], format="csr") - - self.G2_11 = xp.zeros((2, 3), dtype=float) - - self.G2_11[0, 0] = -1.0 - self.G2_11[0, 1] = 1.0 - - self.G2_11[1, 0] = -1.0 - self.G2_11[1, 2] = 1.0 - - self.G2_22 = spa.kron(spa.identity(n0 - 2), grad_1d_2) - - self.G2 = spa.bmat([[self.G2_11, None], [None, self.G2_22]], format="csr") - - self.G = spa.bmat([[self.G1], [self.G2]], format="csr") - # ======================================================================= - - # ========= discrete polar curl matrix =================================== - # 2D vector curl - self.VC1_11 = xp.zeros((2, 3), dtype=float) - - self.VC1_11[0, 0] = -1.0 - self.VC1_11[0, 1] = 1.0 - - self.VC1_11[1, 0] = -1.0 - self.VC1_11[1, 2] = 1.0 - - self.VC1_22 = spa.kron(spa.identity(n0 - 2), grad_1d_2) - - self.VC1 = spa.bmat([[self.VC1_11, None], [None, self.VC1_22]], format="csr") - - self.VC2_11 = xp.zeros(((d0 - 1) * n1, 3), dtype=float) - self.VC2_11[:n1, :] = -self.Xi_1.T - - self.VC2_22 = spa.kron(grad_1d_1[1:, 2:], spa.identity(n1)) - - self.VC2 = -spa.bmat([[self.VC2_11, self.VC2_22]], format="csr") - - self.VC = spa.bmat([[self.VC1], [self.VC2]], format="csr") - - # 2D scalar curl - self.SC1 = -spa.kron(spa.identity(d0 - 1), grad_1d_2) - - self.SC2_1 = xp.zeros(((d0 - 1) * d1, 2), dtype=float) - - for s in range(2): - for j in range(d1): - self.SC2_1[j, s] = -(self.Xi_1[s + 1, (j + 1) % n1] - self.Xi_1[s + 1, j]) - - self.SC2_2 = spa.kron(grad_1d_1[1:, 2:], spa.identity(d1)) - - self.SC2 = spa.bmat([[self.SC2_1, self.SC2_2]], format="csr") - - self.SC = spa.bmat([[self.SC1, self.SC2]], format="csr") - # ========================================================================= - - # ========= discrete polar div matrix ===================================== - self.D1_1 = xp.zeros(((d0 - 1) * d1, 2), dtype=float) - - for s in range(2): - for j in range(d1): - self.D1_1[j, s] = -(self.Xi_1[s + 1, (j + 1) % n1] - self.Xi_1[s + 1, j]) - - self.D1_2 = spa.kron(grad_1d_1[1:, 2:], spa.identity(d1)) - - self.D1 = spa.bmat([[self.D1_1, self.D1_2]], format="csr") - - self.D2 = spa.kron(spa.identity(d0 - 1), grad_1d_2) - - self.D = spa.bmat([[self.D1, self.D2]], format="csr") - # ========================================================================= diff --git a/src/struphy/polar/homotopy/__init__.py b/src/struphy/polar/homotopy/__init__.py new file mode 100644 index 000000000..89eecaadf --- /dev/null +++ b/src/struphy/polar/homotopy/__init__.py @@ -0,0 +1,36 @@ +"""High-order polar FEEC by homotopy operators. + +A port of PolarFEEC.jl (R. Remmerswaal) to Struphy: the high-order polar subcomplex of regularity ``r`` with commuting +projections and (boundary-compatible) homotopy operators, built on the 1D spaces of a Struphy Derham complex +(B-splines, or trigonometric functions in ``eta2``) and on Struphy domains, and the homotopy-based solvers. +""" + +from struphy.polar.homotopy.complex2d import PolarComplex2D, PolarFormSpace, cumulative_sum_matrix +from struphy.polar.homotopy.derham_blocks import DomainCrossSection, derham_polar_complex, poloidal_derham +from struphy.polar.homotopy.derham_spaces import DerhamSpace1D, derham_factors +from struphy.polar.homotopy.forms2d import FormSpace2D, PhysicalFormSpace2D +from struphy.polar.homotopy.polynomials import ( + PolynomialBasis, + center_shifted_basis, + monomial_basis, + polynomial_dimension, + zernike_basis, +) + +__all__ = [ + "DerhamSpace1D", + "DomainCrossSection", + "FormSpace2D", + "PhysicalFormSpace2D", + "PolarComplex2D", + "PolarFormSpace", + "PolynomialBasis", + "center_shifted_basis", + "cumulative_sum_matrix", + "derham_factors", + "derham_polar_complex", + "monomial_basis", + "poloidal_derham", + "polynomial_dimension", + "zernike_basis", +] diff --git a/src/struphy/polar/homotopy/complex2d.py b/src/struphy/polar/homotopy/complex2d.py new file mode 100644 index 000000000..dad7b4023 --- /dev/null +++ b/src/struphy/polar/homotopy/complex2d.py @@ -0,0 +1,610 @@ +"""The high-order polar subcomplex of the 2D de Rham sequence, built with homotopy operators. + +Port of ``src/polar.jl`` of PolarFEEC.jl (R. Remmerswaal, *High-order FEEC discretization on 2D and 3D domains with a +disk-like coordinate singularity*). With regularity ``r`` the polar spaces are + +* ``V0 = pi0 T^r (+) external 0-forms`` (rings ``> r + 1``), +* ``V1 = grad pi0 T^r (+) curl^dagger Pi~^{2,1} pi2 T^{r-2} (+) external 1-forms``, +* ``V2 = Pi~^{2,1} pi2 T^{r-2} (+) external 2-forms`` (rings ``> r``), + +where ``T^r`` are the Cartesian polynomials of total degree ``r``, ``pi^k`` the tensor-product commuting projections +and ``curl^dagger`` the radial Poincare homotopy. Each degree is represented by an extraction matrix ``E`` (columns are +the full tensor-product coefficients of the polar basis functions, central columns first) and dual matrix ``F`` with +``F^T E = I``; ``E F^T`` is the coefficient projector of the commuting polar projections. + +All indices are zero-based; ring ``i`` of the paper is index ``i - 1``. +""" + +from functools import cached_property + +import numpy as np +import scipy.linalg as sla +import scipy.sparse as sp +import scipy.sparse.linalg as spla + +from struphy.polar.homotopy.forms2d import PhysicalFormSpace2D +from struphy.polar.homotopy.layout import CORE +from struphy.polar.homotopy.linalg import SparsePlusLowRank, clean, identity_minus, matmul, max_abs +from struphy.polar.homotopy.polynomials import center_shifted_basis, polynomial_dimension, zernike_basis + +ENRICHMENTS = ("cochain_chain", "cochain", "none") + + +# ---------------------------------------------------------------------------------------------------------- helpers +def cumulative_sum_matrix(n): + """The ``(n + 1) x n`` matrix ``Sigma`` with ``(Sigma v)_i = sum_{i' < i} v_i'`` (``delta Sigma = I``).""" + rows, columns = np.tril_indices(n + 1, -1) + return sp.csr_matrix((np.ones(rows.size), (rows, columns)), shape=(n + 1, n)) + + +def radial_indices(nx, ny, internal_rings, offset=0, layout=CORE): + """Linear indices of the internal (radial index ``< internal_rings``) and external coefficients.""" + return layout.radial_indices(nx, ny, internal_rings, offset) + + +def periodic_primitive_matrix(space, normalization="mean"): + """Periodic primitive ``P`` of a zero-integral periodic 1-form, ``derivative_matrix(space) @ P @ a = a``. + + ``space`` is a periodic B-spline space (forward-difference derivative) or a trigonometric space (``modes`` + attribute). With ``normalization="mean"`` the primitive has zero mean; with a real ``normalization`` it vanishes + there. + """ + n = space.n + primitive = np.zeros((n, n)) + if getattr(space, "is_bspline", False) and space.periodic: + primitive[np.tril_indices(n, -1)] = 1.0 + elif hasattr(space, "modes"): + for k in range(1, space.modes + 1): + sine, cosine = 2 * k - 1, 2 * k + primitive[sine, cosine] = 1.0 + primitive[cosine, sine] = -1.0 + else: + raise TypeError("periodic primitives need a periodic B-spline or trigonometric space") + constant = space.constant_coefficients() + if isinstance(normalization, str): + if normalization != "mean": + raise ValueError("normalization must be 'mean' or a point") + functional = space.coefficient_integrals + else: + functional = space.basis_matrix(np.array([float(normalization)])).toarray().ravel() + return sp.csr_matrix(primitive - np.outer(constant, functional @ primitive)) + + +def identity_columns(n, indices): + indices = np.asarray(indices, dtype=int) + return sp.csr_matrix((np.ones(indices.size), (indices, np.arange(indices.size))), shape=(n, indices.size)) + + +def _rank_tolerance(matrix, rank_tolerance): + if rank_tolerance is not None: + return rank_tolerance + return max(matrix.shape) * np.finfo(float).eps + + +def dual_columns(columns, rank_tolerance=None, context="extraction"): + """Pseudo-inverse columns ``D`` with ``D^T columns = I`` (thin SVD with a rank check).""" + if columns.shape[1] == 0: + return np.zeros((columns.shape[0], 0)) + U, S, Vt = sla.svd(columns, full_matrices=False) + tolerance = _rank_tolerance(columns, rank_tolerance) + rank = int(np.sum(S > tolerance * S[0])) + if rank != columns.shape[1]: + raise ValueError(f"{context} has rank {rank}, but {columns.shape[1]} independent columns are required") + return U @ np.diag(1.0 / S) @ Vt + + +def first_nonzero_rows(A): + """Row of the first nonzero entry of every column of the sparse matrix ``A`` (``-1`` if empty).""" + A = sp.csc_matrix(A) + A.sort_indices() + result = np.full(A.shape[1], -1) + for column in range(A.shape[1]): + start, stop = A.indptr[column], A.indptr[column + 1] + nonzero = np.nonzero(A.data[start:stop])[0] + if nonzero.size: + result[column] = A.indices[start + nonzero[0]] + return result + + +# ---------------------------------------------------------------------------------------------------------- spaces +class PolarFormSpace: + """One degree of a :class:`PolarComplex2D`: extraction ``E`` and duals ``F`` of a physical form space.""" + + def __init__(self, physical_space, extraction, duals, central_count): + self.physical_space = physical_space + self.extraction = clean(extraction).tocsr() + self.duals = clean(duals).tocsr() + self.central_count = central_count + + def __repr__(self): + return ( + f"PolarFormSpace(degree {self.k}, {self.dimension} of {self.physical_space.n} coefficients, " + f"{self.central_count} central)" + ) + + @property + def k(self): + return self.physical_space.k + + @property + def dimension(self): + return self.extraction.shape[1] + + @cached_property + def projector(self): + """The coefficient projector ``E F^T``.""" + return (self.extraction @ self.duals.T).tocsr() + + def project_coefficients(self, coefficients): + return self.extraction @ (self.duals.T @ coefficients) + + def commuting_projection(self, f, additional_order=2, representation="physical"): + """Full coefficients of the polar commuting projection ``E F^T pi f``.""" + full = self.physical_space.commuting_projection(f, additional_order, representation) + return self.project_coefficients(full) + + def __call__(self, f, **kwargs): + return self.commuting_projection(f, **kwargs) + + def central_elements(self): + """Full coefficients of the central polar basis functions (one column each).""" + return self.extraction[:, : self.central_count].toarray() + + +# ---------------------------------------------------------------------------------------------------------- complex +class PolarComplex2D: + """The high-order polar subcomplex ``V0 -> V1 -> V2`` over a disk-like mapping. + + Parameters + ---------- + radial : DerhamSpace1D + Radial (clamped) 0-form space (logical ``x``), e.g. of a Derham complex + (:func:`~struphy.polar.homotopy.derham_spaces.derham_factors`). + angular : DerhamSpace1D (periodic B-splines or trigonometric functions) + Angular (periodic) 0-form space (logical ``y``). + mapping : Mapping2D + Disk-like mapping that collapses ``x = 0`` to the pole (e.g. the cross-section of a Struphy domain, + :class:`~struphy.polar.homotopy.derham_blocks.DomainCrossSection`). + regularity : int + The regularity ``r`` at the pole (polynomial degree of the enrichment). + basis : PolynomialBasis, optional + Polynomial basis of ``T^r``; default ``zernike_basis(r)`` (it is centre-shifted internally). + enrichment : {"cochain_chain", "cochain", "none"} + ``"cochain_chain"`` is the paper's construction; ``"cochain"`` drops the compact 2-forms and curl-like 1-forms + (the reduced spaces of Section 4.2); ``"none"`` keeps only external functions (no constants in ``V0``). + boundary : {"free", "dirichlet"} + ``"dirichlet"`` builds the relative complex (``P(1, y) = 0``, ``A_y(1, y) = 0``). + polynomial_projection : {"commuting", "l2"} + Projection of the polynomials onto the tensor-product spaces. + layout : Layout + Ordering of the tensor-product coefficients (:data:`~struphy.polar.homotopy.layout.STRUPHY` for the Derham + complex, :data:`~struphy.polar.homotopy.layout.CORE` as PolarFEEC.jl). + """ + + def __init__( + self, + radial, + angular, + mapping, + regularity, + basis=None, + enrichment="cochain_chain", + boundary="free", + polynomial_projection="commuting", + additional_quadrature_order=None, + rank_tolerance=None, + check_assumptions=True, + layout=CORE, + ): + if enrichment not in ENRICHMENTS: + raise ValueError(f"enrichment must be one of {ENRICHMENTS}") + if boundary not in ("free", "dirichlet"): + raise ValueError("boundary must be 'free' or 'dirichlet'") + if polynomial_projection not in ("commuting", "l2"): + raise ValueError("polynomial_projection must be 'commuting' or 'l2'") + self.regularity = r = int(regularity) + self.enrichment, self.boundary = enrichment, boundary + self.polynomial_projection = polynomial_projection + self.additional_quadrature_order = additional_quadrature_order + self.rank_tolerance = rank_tolerance + self.mapping = mapping + self.radial, self.angular = radial, angular + self.layout = layout + self._validate(check_assumptions) + + basis = zernike_basis(r) if basis is None else basis + if basis.degree != r: + raise ValueError("the polynomial basis must have the degree of the regularity") + self.polynomial_basis = basis if enrichment == "none" else center_shifted_basis(basis) + + self.spaces = [PhysicalFormSpace2D.from_factors(k, radial, angular, mapping, layout) for k in range(3)] + V0, V1, V2 = self.spaces + nx0, ny0 = radial.n, angular.n + nx1, ny1 = radial.scaled.n, angular.scaled.n + self._shapes = (nx0, ny0, nx1, ny1) + + try: + E0, F0 = self._zero_form_matrices() + E1, E2, central_one, central_two = self._higher_form_extractions(E0) + F2 = self._two_form_duals(E2, central_two) + self._P0 = (E0 @ F0.T).tocsr() + self._P2 = (E2 @ F2.T).tocsr() + F1 = self._one_form_duals(E1, central_one) + except ValueError as error: + if polynomial_projection == "l2": + raise ValueError( + f"{error}. The L2 polynomial projection (a research option of PolarFEEC.jl) is well-defined only " + "when the polynomials are exactly representable in the tensor-product spaces (trigonometric " + "angular spaces and a compatible mapping); use polynomial_projection='commuting'." + ) from error + raise + + if boundary == "dirichlet": + E0, keep0 = _remove_boundary_columns(E0, self.zero_form_dirichlet_rows(), "the zero-form extraction") + F0 = F0[:, keep0] + E1, keep1 = _remove_boundary_columns(E1, self.one_form_dirichlet_rows(), "the one-form extraction") + F1 = F1[:, keep1] + + central_zero = 0 if enrichment == "none" else len(self.polynomial_basis) + self.forms = [ + PolarFormSpace(V0, E0, F0, central_zero), + PolarFormSpace(V1, E1, F1, central_one), + PolarFormSpace(V2, E2, F2, central_two), + ] + self._cache = {} + + # ------------------------------------------------------------------ basics + def __getitem__(self, k): + return self.forms[k] + + def __repr__(self): + dims = tuple(self.dimension(k) for k in range(3)) + return ( + f"PolarComplex2D(regularity {self.regularity}, dimensions {dims}, enrichment={self.enrichment!r}, " + f"boundary={self.boundary!r})" + ) + + top_degree = 2 + + def dimension(self, k): + return self.forms[k].dimension + + def extraction(self, k): + return self.forms[k].extraction + + def duals(self, k): + return self.forms[k].duals + + def _validate(self, check_assumptions): + x, y, r = self.radial, self.angular, self.regularity + if getattr(x, "periodic", False) or not getattr(x, "is_bspline", False): + raise TypeError("the radial factor must be a clamped B-spline space") + if not getattr(y, "periodic", False): + raise TypeError("the angular factor must be periodic or trigonometric") + if r + 2 > x.n: + raise ValueError( + f"the polar construction requires radial dimension >= regularity + 2 (got {x.n} and regularity {r})" + ) + if self.enrichment == "none" or not check_assumptions: + return + if r > x.degree: + raise ValueError( + f"the paper assumes regularity <= radial degree (got regularity {r}, degree {x.degree}); " + "pass check_assumptions=False to construct the spaces anyway" + ) + if y.n < 2 * r + 1: + raise ValueError( + f"the paper assumes an angular dimension >= 2 regularity + 1 (got {y.n} for regularity {r}); " + "pass check_assumptions=False to construct the spaces anyway" + ) + + def _polynomial_columns(self, functions, space): + return space.polynomial_projection_columns( + list(functions), self.polynomial_projection, self.additional_quadrature_order + ) + + # ------------------------------------------------------------------ degree 0 + def _zero_form_matrices(self): + L = self.layout + nx, ny = self.radial.n, self.angular.n + n = nx * ny + rings = self.regularity + 1 + internal, external = L.radial_indices(nx, ny, rings) + external_identity = identity_columns(n, external) + constant = self.angular.constant_coefficients() + denominator = constant @ constant + first_ring = L.index(0, np.arange(ny), nx, ny) + if self.enrichment == "none": + center_dual = np.zeros(n) + center_dual[first_ring] = constant / denominator + external_constant = constant[L.angle(external, nx, ny)] + F = external_identity - sp.csr_matrix(np.outer(center_dual, external_constant)) + return external_identity, F + # project on the full tensor-product space first, then truncate to the internal rings + projected = self._polynomial_columns(self.polynomial_basis.functions, self.spaces[0]) + internal_columns = projected[internal] + duals = dual_columns(internal_columns, self.rank_tolerance, "the zero-form polynomial extraction") + # the first dual reproduces the centre value, which makes the projectors commute with the homotopies + center_dual = np.zeros(internal.size) + center_dual[np.searchsorted(internal, first_ring)] = constant / denominator + duals[:, 0] = center_dual + tolerance = max(1e-10, 1000 * _rank_tolerance(internal_columns, self.rank_tolerance)) * np.linalg.cond( + internal_columns + ) + if np.max(np.abs(duals.T @ internal_columns - np.eye(internal_columns.shape[1]))) > tolerance: + raise ValueError("the projected polynomial basis does not reproduce its prescribed values at the centre") + m = len(self.polynomial_basis) + polynomial = np.zeros((n, m)) + polynomial[internal] = internal_columns + dual = np.zeros((n, m)) + dual[internal] = duals + E = sp.hstack([sp.csr_matrix(polynomial), external_identity]).tocsr() + F = sp.hstack([sp.csr_matrix(dual), external_identity]).tocsr() + return E, F + + # ------------------------------------------------------------------ degrees 1 and 2 + def reduced_two_form_projector(self, regularity): + """``Pi~^{2,r}``: rings ``0..r`` are accumulated in ring ``r`` and later rings are copied (``r = 0``: identity).""" + nx, ny = self.radial.scaled.n, self.angular.scaled.n + i, j = np.meshgrid(np.arange(nx), np.arange(ny), indexing="ij") + columns = self.layout.index(i, j, nx, ny).ravel() + rows = self.layout.index(np.maximum(i, regularity), j, nx, ny).ravel() + return sp.csr_matrix((np.ones(rows.size), (rows, columns)), shape=(nx * ny, nx * ny)) + + @cached_property + def gradient_homotopy_matrix(self): + """Full-space radial homotopy ``G^dagger = [I (x) Sigma, 0]`` from 1-forms to 0-forms.""" + nx0, ny0, nx1, ny1 = self._shapes + radial = self.layout.kron(cumulative_sum_matrix(nx1), sp.identity(ny0)) + return sp.hstack([radial, sp.csr_matrix((radial.shape[0], nx0 * ny1))]).tocsr() + + @cached_property + def curl_homotopy_matrix(self): + """Full-space radial homotopy ``C^dagger = [0; I (x) Sigma]`` from 2-forms to 1-forms.""" + nx0, ny0, nx1, ny1 = self._shapes + radial = self.layout.kron(cumulative_sum_matrix(nx1), sp.identity(ny1)) + return sp.vstack([sp.csr_matrix((nx1 * ny0, radial.shape[1])), radial]).tocsr() + + def _higher_form_extractions(self, E0): + r, L = self.regularity, self.layout + nx0, ny0, nx1, ny1 = self._shapes + V1, V2 = self.spaces[1], self.spaces[2] + G = self.spaces[0].derivative_matrix + m = len(self.polynomial_basis) + if self.enrichment == "none": + gradient_columns = sp.csr_matrix((V1.n, 0)) + else: + gradient_columns = (G @ E0[:, 1:m]).tocsr() + if self.enrichment == "cochain_chain" and r >= 2: + functions = self.polynomial_basis.functions[: polynomial_dimension(r - 2)] + raw = self._polynomial_columns(functions, V2) + conforming = self.reduced_two_form_projector(1) @ raw + compact = np.zeros_like(conforming) + for j in range(ny1): + internal = L.index(np.arange(r), j, nx1, ny1) + compact[internal] = conforming[internal] + # the first external ring cancels the radial sum of each angular column + compact[L.index(r, j, nx1, ny1)] = -conforming[internal].sum(axis=0) + compact = sp.csr_matrix(compact) + else: + compact = sp.csr_matrix((V2.n, 0)) + # curl homotopy lifts; their entries vanish exactly from ring r + 2 on (zero radial sums) + curl_columns = (self.curl_homotopy_matrix @ compact).toarray() + first = nx1 * ny0 + _, outer_y = L.radial_indices(nx0, ny1, r + 1, offset=first) + curl_columns[outer_y] = 0.0 + curl_columns = sp.csr_matrix(curl_columns) + _, external_x = L.radial_indices(nx1, ny0, r) + _, external_y = L.radial_indices(nx0, ny1, r + 1, offset=first) + external_one = identity_columns(V1.n, np.concatenate([external_x, external_y])) + _, external_two = L.radial_indices(nx1, ny1, r) + E1 = sp.hstack([gradient_columns, curl_columns, external_one]).tocsr() + E2 = sp.hstack([compact, identity_columns(V2.n, external_two)]).tocsr() + return E1, E2, gradient_columns.shape[1] + compact.shape[1], compact.shape[1] + + def _two_form_duals(self, E2, central_count): + r = self.regularity + nx, ny = self.radial.scaled.n, self.angular.scaled.n + internal, external = self.layout.radial_indices(nx, ny, r) + central = np.zeros((E2.shape[0], central_count)) + if central_count > 0: + central[internal] = dual_columns( + E2[internal][:, :central_count].toarray(), self.rank_tolerance, "the two-form polynomial extraction" + ) + external_duals = self.reduced_two_form_projector(r)[external].T + F = sp.hstack([sp.csr_matrix(central), external_duals]).tocsr() + tolerance = max(1e-10, 1000 * _rank_tolerance(E2, self.rank_tolerance)) + if max_abs(F.T @ E2 - sp.identity(E2.shape[1])) > tolerance: + raise ValueError("the two-form extraction is incompatible with the derived polar projector") + return F + + @cached_property + def one_form_projector_full(self): + """``Pi^1 = (G Pi^0) G^dagger + (C^dagger Pi^2) C`` on the full tensor-product 1-form coefficients.""" + G = self.spaces[0].derivative_matrix + C = self.spaces[1].derivative_matrix + return ((G @ self._P0) @ self.gradient_homotopy_matrix + (self.curl_homotopy_matrix @ self._P2) @ C).tocsr() + + def _one_form_duals(self, E1, central_count): + n = E1.shape[0] + external = first_nonzero_rows(E1[:, central_count:]) + internal = np.setdiff1d(np.arange(n), external) + E_internal = E1[:, :central_count] + projector = self.one_form_projector_full + internal_rows = projector[internal] + Q = np.zeros((central_count, n)) + conditioning = 1.0 + representation_error = max_abs(internal_rows) + active = np.unique(sp.csc_matrix(internal_rows).nonzero()[1]) + if central_count > 0 and active.size: + A = E_internal[internal].toarray() + if np.linalg.matrix_rank(A) != central_count: + raise ValueError("the internal one-form extraction columns are linearly dependent") + conditioning = np.linalg.cond(A) + rhs = internal_rows[:, active].toarray() + coordinates = np.linalg.lstsq(A, rhs, rcond=None)[0] + Q[:, active] = coordinates + representation_error = np.max(np.abs(A @ coordinates - rhs)) + external_coordinates = projector[external] - sp.csr_matrix(E_internal[external] @ sp.csr_matrix(Q)) + dual_transpose = sp.vstack([sp.csr_matrix(Q), external_coordinates]).tocsr() + dual_transpose.data[np.abs(dual_transpose.data) < 100 * np.finfo(float).eps] = 0.0 + dual_transpose.eliminate_zeros() + F = dual_transpose.T.tocsr() + tolerance = max(1e-10, 1000 * _rank_tolerance(E1, self.rank_tolerance)) * conditioning + biorthogonality = max_abs(F.T @ E1 - sp.identity(E1.shape[1])) + if max(representation_error, biorthogonality) > tolerance: + raise ValueError("the one-form extraction is incompatible with the homotopy projector") + return F + + # ------------------------------------------------------------------ boundary helpers + def zero_form_dirichlet_rows(self): + nx, ny = self.radial.n, self.angular.n + return self.layout.index(nx - 1, np.arange(ny), nx, ny) + + def one_form_dirichlet_rows(self): + nx0, ny0, nx1, ny1 = self._shapes + return nx1 * ny0 + self.layout.index(nx0 - 1, np.arange(ny1), nx0, ny1) + + @cached_property + def center_projection_full(self): + """Rank-one full-space matrix of ``P u = u(0)`` on continuous 0-forms (left, right factors).""" + nx, ny = self.radial.n, self.angular.n + constant = self.angular.constant_coefficients() + functional = np.zeros(nx * ny) + functional[self.layout.index(0, np.arange(ny), nx, ny)] = constant / (constant @ constant) + return self.layout.kron_vector(np.ones(nx), constant), functional + + @cached_property + def two_form_integral_vector(self): + """Integrals of the full tensor-product 2-form basis functions (logical, hence physical).""" + return self.layout.kron_vector( + self.radial.scaled.coefficient_integrals, self.angular.scaled.coefficient_integrals + ) + + # ------------------------------------------------------------------ operators in polar coefficients + def derivative(self, k): + """``D_k = F_{k+1}^T D E_k`` with ``D E_k = E_{k+1} D_k``.""" + if k not in (0, 1): + raise ValueError("degree must be 0 or 1") + key = ("derivative", k) + if key not in self._cache: + D = self.spaces[k].derivative_matrix + self._cache[key] = clean(self.forms[k + 1].duals.T @ D @ self.forms[k].extraction) + return self._cache[key] + + def full_mass_matrix(self, k): + key = ("full_mass", k) + if key not in self._cache: + self._cache[key] = self.spaces[k].mass_matrix() + return self._cache[key] + + def mass_matrix(self, k): + """Polar mass matrix ``E^T M E``.""" + key = ("mass", k) + if key not in self._cache: + E = self.forms[k].extraction + M = (E.T @ self.full_mass_matrix(k) @ E).tocsc() + self._cache[key] = ((M + M.T) / 2).tocsc() + return self._cache[key] + + def mass_solver(self, k): + """Sparse LU factorization of the polar mass matrix (``.solve(b)``).""" + key = ("mass_solver", k) + if key not in self._cache: + self._cache[key] = spla.splu(self.mass_matrix(k)) + return self._cache[key] + + def l2_projection(self, k, f, additional_order=0): + """Full coefficients of the L2 projection into the polar space of degree ``k``.""" + E = self.forms[k].extraction + load = E.T @ self.spaces[k].inner_product(f, additional_order) + return E @ self.mass_solver(k).solve(load) + + def commuting_projection(self, k, f, additional_order=2): + return self.forms[k].commuting_projection(f, additional_order) + + def cohomology_projection(self, k): + """Projection onto the harmonic representatives: centre value (free, ``k = 0``), mean ``B`` (Dirichlet, ``k = 2``).""" + n = self.dimension(k) + if self.boundary == "free": + if k != 0: + return SparsePlusLowRank.zero(n) + left, right = self.center_projection_full + form = self.forms[0] + return SparsePlusLowRank.rank_one(form.duals.T @ left, form.extraction.T @ right) + if k != 2: + return SparsePlusLowRank.zero(n) + key = "mean_projection" + if key not in self._cache: + integral = self.forms[2].extraction.T @ self.two_form_integral_vector + representer = self.mass_solver(2).solve(integral) + normalization = integral @ representer + if not np.isfinite(normalization) or normalization <= 0: + raise ValueError("the polar two-form space does not contain a nonzero mean mode") + self._cache[key] = SparsePlusLowRank.rank_one(representer / normalization, integral) + return self._cache[key] + + def homotopy(self, k): + """Polar homotopy ``V^k -> V^{k-1}``: radial (free) or boundary-compatible (Dirichlet).""" + if k not in (1, 2): + raise ValueError("degree must be 1 or 2") + key = ("homotopy", k) + if key not in self._cache: + if self.boundary == "dirichlet": + if self.dimension(1) != self.dimension(0) + self.dimension(2) - 1: + raise ValueError("the boundary-restricted complex does not admit the boundary homotopy") + value = self._boundary_curl_homotopy() if k == 2 else self._boundary_gradient_homotopy() + else: + H = self.curl_homotopy_matrix if k == 2 else self.gradient_homotopy_matrix + value = clean(self.forms[k - 1].duals.T @ H @ self.forms[k].extraction) + self._cache[key] = value + return self._cache[key] + + def _boundary_curl_homotopy(self): + nx0, ny0, nx1, ny1 = self._shapes + zero_mean = identity_minus(self.cohomology_projection(2)) + radial_lift = matmul(sp.csr_matrix(self.curl_homotopy_matrix @ self.forms[2].extraction), zero_mean) + # periodic primitive of the tangential trace, extended by the outer radial B-spline + trace_rows = self.one_form_dirichlet_rows() + trace = sp.csr_matrix((np.ones(ny1), (np.arange(ny1), trace_rows)), shape=(ny1, self.spaces[1].n)) + extension_rows = self.zero_form_dirichlet_rows() + extension = sp.csr_matrix((np.ones(ny0), (extension_rows, np.arange(ny0))), shape=(self.spaces[0].n, ny0)) + primitive = periodic_primitive_matrix(self.angular, normalization="mean") + potential = matmul((extension @ primitive @ trace).tocsr(), radial_lift) + corrected = radial_lift - matmul(self.spaces[0].derivative_matrix, potential) + return matmul(self.forms[1].duals.T.tocsr(), corrected) + + def _boundary_gradient_homotopy(self): + nx, ny = self.radial.n, self.angular.n + closed_part = identity_minus(self.homotopy(2) @ self.derivative(1)) + radial_potential = matmul(sp.csr_matrix(self.gradient_homotopy_matrix @ self.forms[1].extraction), closed_part) + constant = self.layout.kron_vector(self.radial.constant_coefficients(), self.angular.constant_coefficients()) + boundary_mean = np.zeros(nx * ny) + boundary_mean[self.zero_form_dirichlet_rows()] = self.angular.coefficient_integrals + # the boundary trace of the potential is constant; subtract its mean + mean_part = SparsePlusLowRank.rank_one(constant, boundary_mean) @ radial_potential + corrected = radial_potential - mean_part + return matmul(self.forms[0].duals.T.tocsr(), corrected) + + def homotopy_adjoint(self, k, coefficients): + """``M_k^{-1} h_k^T M_{k-1}`` applied to polar coefficients of degree ``k - 1``.""" + return self.mass_solver(k).solve(self.homotopy(k).T @ (self.mass_matrix(k - 1) @ coefficients)) + + def cohomology_dimensions(self): + """``dim ker d^k / im d^{k-1}`` from dense ranks (for testing).""" + ranks = [np.linalg.matrix_rank(self.derivative(k).toarray()) for k in range(2)] + [0] + return [self.dimension(k) - ranks[k] - (ranks[k - 1] if k > 0 else 0) for k in range(3)] + + def central_counts(self): + return tuple(form.central_count for form in self.forms) + + +def _remove_boundary_columns(extraction, boundary_rows, context): + """Remove the first column with a nonzero in each boundary row (the interpolatory boundary columns).""" + rows = first_nonzero_rows(sp.csr_matrix(extraction).T)[boundary_rows] + if np.any(rows < 0): + raise ValueError(f"{context} does not contain an interpolatory outer-boundary column") + if np.unique(rows).size != rows.size: + raise ValueError(f"{context} has duplicate outer-boundary columns") + keep = np.setdiff1d(np.arange(extraction.shape[1]), rows) + return extraction[:, keep], keep diff --git a/src/struphy/polar/homotopy/complex3d.py b/src/struphy/polar/homotopy/complex3d.py new file mode 100644 index 000000000..ed4e583a4 --- /dev/null +++ b/src/struphy/polar/homotopy/complex3d.py @@ -0,0 +1,168 @@ +"""The 3D polar complex as the total complex of the 2D polar complex and the 1D periodic (toroidal) complex. + +Port of ``src/complexes.jl`` and ``src/polar3d.jl`` of PolarFEEC.jl at the coefficient level:: + + W^m = sum_{p + q = m} A^p (x) B^q, + d(V (x) U) = d_A V (x) U + (-1)^p V (x) d_B U, + h(V (x) U) = h_A V (x) U + (-1)^p P_A V (x) h_B U (ordered homotopy), + P = P_A (x) P_B. + +Within degree ``m`` the blocks are ordered by decreasing ``p``; each block uses the layout ``kron(B, A)`` (coefficients +of ``A`` vary fastest). Operators are returned as dense or sparse matrices (intended for verification and small +problems; the distributed 3D operators live in :class:`~struphy.feec.psydac_derham.Derham`). +""" + +import numpy as np +import scipy.sparse as sp + +from struphy.polar.homotopy.complex2d import periodic_primitive_matrix +from struphy.polar.homotopy.linalg import SparsePlusLowRank, as_dense + + +def _sparse(A): + """Sparse matrix of a sparse, dense or factored matrix (low-rank parts are multiplied out).""" + if isinstance(A, SparsePlusLowRank): + if A.rank == 0: + return A.sparse + return sp.csr_matrix(A.sparse + sp.csr_matrix(A.left @ A.right.T)) + return sp.csr_matrix(A) + + +class PeriodicComplex: + """The 1D periodic complex ``space -> scaled(space)`` with the homotopy normalized at ``base_point``. + + ``h A = int_{base}^{x} (A - int A)``, ``P0 u = u(base)`` and ``P1 A = int A`` (as the paper's 3D construction). + """ + + top_degree = 1 + + def __init__(self, space, base_point=None): + if not getattr(space, "periodic", False): + raise TypeError("the periodic complex needs a periodic B-spline or trigonometric space") + self.space = space + self.base_point = float(space.greville_points()[0] if base_point is None else base_point) + + def dimension(self, k): + return self.space.n + + def derivative(self, k): + assert k == 0 + return self.space.derivative_matrix() + + def extraction(self, k): + return sp.identity(self.space.n, format="csr") + + duals = extraction + + def cohomology_projection(self, k): + if k == 0: + evaluation = self.space.basis_matrix(np.array([self.base_point])).toarray().ravel() + return SparsePlusLowRank.rank_one(self.space.constant_coefficients(), evaluation) + return SparsePlusLowRank.rank_one( + self.space.constant_one_form_coefficients(), self.space.scaled.coefficient_integrals + ) + + def homotopy(self, k): + assert k == 1 + primitive = periodic_primitive_matrix(self.space, normalization=self.base_point) + mean = self.cohomology_projection(1) + return SparsePlusLowRank(primitive) @ SparsePlusLowRank(sp.identity(self.space.n, format="csr") - _sparse(mean)) + + +class TensorProductComplex: + """The total complex ``W^m = sum_{p+q=m} A^p (x) B^q`` of two complexes with contracting homotopies.""" + + def __init__(self, first, second): + self.first, self.second = first, second + self.top_degree = first.top_degree + second.top_degree + + def bidegree_blocks(self, m): + top_a, top_b = self.first.top_degree, self.second.top_degree + return [(p, m - p) for p in range(min(m, top_a), max(0, m - top_b) - 1, -1)] + + def _block_sizes(self, m): + return [self.first.dimension(p) * self.second.dimension(q) for p, q in self.bidegree_blocks(m)] + + def dimension(self, m): + return sum(self._block_sizes(m)) + + def derivative(self, m): + source, target = self.bidegree_blocks(m), self.bidegree_blocks(m + 1) + blocks = [[None] * len(source) for _ in target] + for s, (p, q) in enumerate(source): + if p < self.first.top_degree: + t = target.index((p + 1, q)) + blocks[t][s] = sp.kron(sp.identity(self.second.dimension(q)), _sparse(self.first.derivative(p))) + if q < self.second.top_degree: + t = target.index((p, q + 1)) + block = (-1) ** p * sp.kron(_sparse(self.second.derivative(q)), sp.identity(self.first.dimension(p))) + blocks[t][s] = block if blocks[t][s] is None else blocks[t][s] + block + return _bmat(blocks, self._block_sizes(m + 1), self._block_sizes(m)) + + def homotopy(self, m): + source, target = self.bidegree_blocks(m), self.bidegree_blocks(m - 1) + blocks = [[None] * len(source) for _ in target] + for s, (p, q) in enumerate(source): + if p >= 1: + t = target.index((p - 1, q)) + blocks[t][s] = sp.kron(sp.identity(self.second.dimension(q)), _sparse(self.first.homotopy(p))) + if q >= 1: + projection = self.first.cohomology_projection(p) + if isinstance(projection, SparsePlusLowRank) and projection.is_zero(): + continue + t = target.index((p, q - 1)) + block = (-1) ** p * sp.kron(_sparse(self.second.homotopy(q)), _sparse(projection)) + blocks[t][s] = block if blocks[t][s] is None else blocks[t][s] + block + return _bmat(blocks, self._block_sizes(m - 1), self._block_sizes(m)) + + def cohomology_projection(self, m): + blocks = [ + sp.kron(_sparse(self.second.cohomology_projection(q)), _sparse(self.first.cohomology_projection(p))) + for p, q in self.bidegree_blocks(m) + ] + return sp.block_diag(blocks).tocsr() + + def extraction(self, m): + blocks = [ + sp.kron(_sparse(self.second.extraction(q)), _sparse(self.first.extraction(p))) + for p, q in self.bidegree_blocks(m) + ] + return sp.block_diag(blocks).tocsr() + + def duals(self, m): + blocks = [ + sp.kron(_sparse(self.second.duals(q)), _sparse(self.first.duals(p))) for p, q in self.bidegree_blocks(m) + ] + return sp.block_diag(blocks).tocsr() + + def cohomology_dimensions(self): + ranks = [np.linalg.matrix_rank(as_dense(self.derivative(k))) for k in range(self.top_degree)] + [0] + return [self.dimension(k) - ranks[k] - (ranks[k - 1] if k > 0 else 0) for k in range(self.top_degree + 1)] + + +def _bmat(blocks, row_sizes, column_sizes): + # scipy.sparse.bmat needs at least one block per block row/column to infer sizes: pad with zero blocks + padded = [ + [ + blocks[i][j] if blocks[i][j] is not None else sp.csr_matrix((row_sizes[i], column_sizes[j])) + for j in range(len(column_sizes)) + ] + for i in range(len(row_sizes)) + ] + return sp.bmat(padded, format="csr") + + +class PolarComplex3D(TensorProductComplex): + """The 3D polar complex: the 2D polar complex of the cross-section (x) the periodic toroidal complex. + + Parameters + ---------- + polar2d : PolarComplex2D + Polar complex of the reference cross-section. + toroidal : DerhamSpace1D (periodic B-splines or trigonometric functions) + Periodic toroidal 0-form space. + """ + + def __init__(self, polar2d, toroidal): + super().__init__(polar2d, PeriodicComplex(toroidal)) + self.polar2d, self.toroidal = polar2d, toroidal diff --git a/src/struphy/polar/homotopy/demos/__init__.py b/src/struphy/polar/homotopy/demos/__init__.py new file mode 100644 index 000000000..ad507fd5a --- /dev/null +++ b/src/struphy/polar/homotopy/demos/__init__.py @@ -0,0 +1,8 @@ +"""Studies and figures of the high-order polar FEEC talk and of the PolarFEEC.jl examples, used by the tutorials. + +* :mod:`~struphy.polar.homotopy.demos.basis`: disk mapping, tensor-product and polar basis functions, +* :mod:`~struphy.polar.homotopy.demos.regularity`: smoothness at the centre (PolarFEEC.jl example 04), +* :mod:`~struphy.polar.homotopy.demos.projections`: convergence and commutativity of the polar projections, +* :mod:`~struphy.polar.homotopy.demos.pde`: Poisson and curl-curl on a rounded square (example 05), +* :mod:`~struphy.polar.homotopy.demos.maxwell`: the Maxwell spectrum on the disk (example 06). +""" diff --git a/src/struphy/polar/homotopy/demos/basis.py b/src/struphy/polar/homotopy/demos/basis.py new file mode 100644 index 000000000..4637c6629 --- /dev/null +++ b/src/struphy/polar/homotopy/demos/basis.py @@ -0,0 +1,166 @@ +"""Basis-function figures of the talk (Sections 1, 2 and 5). + +* the disk-like mapping ``R = R1 x (1 + eps x cos(2 pi k y))`` (``eps = 0.2``, ``k = 3``), +* tensor-product 0-forms: cubic clamped B-splines on ``N^x = 5`` cells and cubic periodic B-splines on ``N^y = 12`` + cells (or trigonometric functions in ``y``), +* the polar 0-forms for ``r = 1, 2, 3``: the first ``r + 1`` rings are replaced by ``(r + 1)(r + 2)/2`` projected + polynomials, +* the central polar 0-, 1- and 2-forms for ``r = 3``. + +The spaces are those of a 2D Struphy Derham complex on the domain ``PerturbedDiskCylinder`` (see :mod:`.common`). +""" + +import matplotlib.pyplot as plt +import numpy as np + +from struphy.polar.homotopy.demos import common +from struphy.polar.homotopy.demos.plotting import DIVERGING, clean_axis, draw_mesh + +RS = np.linspace(0, 1, 121) +TS = np.linspace(0, 1, 241) + + +#: cells of the talk's basis figures: cubic splines on N^x = 5 radial cells and N^y = 12 angular cells (B-splines) or +#: 13 trigonometric functions +CELLS = {"bspline": (5, 12), "trig": (5, 13)} + + +def _complex(regularity, angular="bspline", domain=None): + return common.polar_complex(regularity, 3, CELLS[angular], domain=domain, angular=angular) + + +def _panel(ax, mapping, values, title=None, mesh=None): + X, Y = mapping(*np.meshgrid(RS, TS, indexing="ij")) + vmax = np.max(np.abs(values)) or 1.0 + ax.pcolormesh(X, Y, values / vmax, cmap=DIVERGING, vmin=-1, vmax=1, shading="gouraud") + if mesh is not None: + draw_mesh(ax, mapping, *mesh, alpha=0.25) + ax.set_aspect("equal") + clean_axis(ax) + if title: + ax.set_title(title, fontsize=10) + + +def figure_mapping(mapping=None, cells=(5, 12)): + """The disk-like coordinate transformation: logical mesh and its image.""" + mapping = common.cross_section() if mapping is None else mapping + fig, axes = plt.subplots(1, 2, figsize=(9, 4.4)) + rb, tb = np.linspace(0, 1, cells[0] + 1), np.linspace(0, 1, cells[1] + 1) + for r in rb: + axes[0].axvline(r, color="k", lw=0.6, alpha=0.4) + for t in tb: + axes[0].axhline(t, color="k", lw=0.6, alpha=0.4) + axes[0].axvline(0, color="C3", lw=2.5) + axes[0].set_xlabel(r"$\hat x$") + axes[0].set_ylabel(r"$\hat y$") + axes[0].set_title("logical domain") + axes[0].set_aspect("equal") + draw_mesh(axes[1], mapping, rb, tb, alpha=0.4) + axes[1].plot([0], [0], "o", color="C3") + axes[1].set_aspect("equal") + axes[1].set_title(r"$\Phi(\hat x, \hat y) = (R C, R S)$, $\hat x = 0 \mapsto 0$") + fig.tight_layout() + return fig + + +def figure_tensor_basis(angular="bspline", mapping=None, rings=(0, 1, 2, 3), angles=(0, 1, 2, 3)): + """Physical tensor-product 0-form basis functions ``B_i(x) U_j(y)`` (rows: rings, columns: angular index).""" + mapping = common.cross_section() if mapping is None else mapping + space = common.tensor_space(0, 3, CELLS[angular], angular=angular) + radial, ang = space.logical.x, space.logical.y + fig, axes = plt.subplots(len(rings), len(angles), figsize=(2.4 * len(angles), 2.4 * len(rings))) + mesh = (np.linspace(0, 1, radial.cells + 1), np.linspace(0, 1, 13)) + for a, i in enumerate(rings): + for b, j in enumerate(angles): + coefficients = np.zeros(space.n) + coefficients[space.layout.index(i, j, radial.n, ang.n)] = 1.0 + _panel(axes[a, b], mapping, space.evaluate(coefficients, RS, TS), rf"$B_{i} U_{j}$", mesh) + kind = "cubic periodic B-splines" if angular == "bspline" else "trigonometric functions" + fig.suptitle(f"tensor-product 0-forms: cubic clamped B-splines ($N^x = 5$) x {kind}") + fig.tight_layout() + return fig + + +def figure_polar_zero_forms(regularity, angular="bspline", mapping=None): + """The central polar 0-forms of regularity ``r`` (they replace the first ``r + 1`` rings).""" + mapping = common.cross_section() if mapping is None else mapping + c = _complex(regularity, angular) + radial, ang = c.radial, c.angular + central = c[0].central_elements() + n = central.shape[1] + columns = min(n, 5) + rows = int(np.ceil(n / columns)) + fig, axes = plt.subplots(rows, columns, figsize=(2.6 * columns, 2.6 * rows), squeeze=False) + mesh = (np.linspace(0, 1, radial.cells + 1), np.linspace(0, 1, ang.n + 1 if angular == "bspline" else 13)) + for l in range(rows * columns): + ax = axes.ravel()[l] + if l >= n: + ax.axis("off") + continue + _panel( + ax, + mapping, + c.spaces[0].evaluate(central[:, l], RS, TS), + rf"$\bar\Lambda^{{0,{regularity}}}_{{{l + 1}}}$", + mesh, + ) + fig.suptitle(f"polar 0-forms, $r = {regularity}$: the first {regularity + 1} rings replaced by {n} functions") + fig.tight_layout() + return fig + + +def figure_all_forms(regularity=3, angular="bspline", mapping=None): + """Central polar 0-forms, 1-forms (magnitude and arrows) and 2-forms of regularity ``r``.""" + mapping = common.cross_section() if mapping is None else mapping + c = _complex(regularity, angular) + counts = c.central_counts() + columns = max(counts) + fig, axes = plt.subplots(3, columns, figsize=(1.8 * columns, 6), squeeze=False) + rs = np.linspace(1e-3, 1, 121) + X, Y = mapping(*np.meshgrid(rs, TS, indexing="ij")) + for k in range(3): + central = c[k].central_elements() + for l in range(columns): + ax = axes[k, l] + if l >= counts[k]: + ax.axis("off") + continue + values = c.spaces[k].evaluate(central[:, l], rs, TS) + if k == 1: + magnitude = np.hypot(values[0], values[1]) + ax.pcolormesh(X, Y, magnitude / magnitude.max(), cmap="viridis", shading="gouraud") + qr, qt = np.linspace(0.08, 0.95, 7), np.linspace(0, 1, 25)[:-1] + Qr, Qt = np.meshgrid(qr, qt, indexing="ij") + vx, vy = c.spaces[1].evaluate_points(central[:, l], Qr, Qt) + ax.quiver(*mapping(Qr, Qt), vx, vy, color="w", scale=None, width=0.006) + else: + vmax = np.max(np.abs(values)) or 1.0 + ax.pcolormesh(X, Y, values / vmax, cmap=DIVERGING, vmin=-1, vmax=1, shading="gouraud") + ax.set_aspect("equal") + clean_axis(ax) + ax.set_title(rf"$\bar\Lambda^{{{k},{regularity}}}_{{{l + 1}}}$", fontsize=9) + fig.suptitle(f"central polar basis functions for 0-, 1- and 2-forms, $r = {regularity}$") + fig.tight_layout() + return fig + + +def figure_discontinuity(mapping=None): + """A tensor-product 0-form is discontinuous at the origin: ``P(0) = sum_j p_1j U_j(y)`` depends on the angle.""" + mapping = common.cross_section() if mapping is None else mapping + space = common.tensor_space(0, 3, CELLS["bspline"]) + rng = np.random.default_rng(1) + coefficients = 1 + 0.3 * rng.standard_normal(space.n) + fig, axes = plt.subplots(1, 2, figsize=(9, 4)) + values = space.evaluate(coefficients, RS, TS) + X, Y = mapping(*np.meshgrid(RS, TS, indexing="ij")) + m = axes[0].pcolormesh(X, Y, values, cmap="viridis", shading="gouraud") + plt.colorbar(m, ax=axes[0], shrink=0.8) + axes[0].set_aspect("equal") + clean_axis(axes[0]) + axes[0].set_title("random tensor-product 0-form") + axes[1].plot(TS, values[0], label=r"$\lim_{\hat x \to 0} P$") + axes[1].set_xlabel(r"$\hat y$") + axes[1].set_title(r"the value at the origin depends on the angle") + axes[1].legend() + fig.tight_layout() + return fig diff --git a/src/struphy/polar/homotopy/demos/common.py b/src/struphy/polar/homotopy/demos/common.py new file mode 100644 index 000000000..7b9cd0230 --- /dev/null +++ b/src/struphy/polar/homotopy/demos/common.py @@ -0,0 +1,130 @@ +"""Struphy objects of the 2D studies of the talk: the domains, a 2D Derham complex of the cross-section and the polar +complexes built on it. + +All demos use these helpers, so that the figures of the talk are computed with Struphy's spline spaces +(:class:`~struphy.polar.homotopy.derham_spaces.DerhamSpace1D`), Struphy domains +(:class:`~struphy.polar.homotopy.derham_blocks.DomainCrossSection`) and Struphy's coefficient ordering. +""" + +import numpy as np + +from struphy.polar.homotopy.layout import STRUPHY + + +# ---------------------------------------------------------------------------------------------------------- domains +def perturbed_disk(eps=0.2, k=3): + """The perturbed disk of the talk, ``R = x (1 + 0.2 x cos(6 pi y))``.""" + from struphy.geometry import domains + + return domains.PerturbedDiskCylinder(R1=1.0, eps=eps, k=k, Lz=1.0) + + +def star_disk(eps=0.2, k=3): + """The star-shaped disk of the regularity study, ``R = x (1 + 0.2 cos(6 pi y))``.""" + from struphy.geometry import domains + + return domains.PerturbedDiskCylinder(R1=1.0, eps=eps, k=k, Lz=1.0, power=0.0) + + +def unit_disk(): + from struphy.geometry import domains + + return domains.HollowCylinder(a1=0.0, a2=1.0, Lz=1.0) + + +def rounded_square(eps=0.3): + from struphy.geometry import domains + + return domains.RoundedSquareCylinder(eps=eps, Lz=1.0) + + +# ---------------------------------------------------------------------------------------------------------- 1D +def spline_space_1d(degree, cells, periodic=False): + """The 1D B-spline space (0-forms, with its M-spline 1-forms) of a Struphy Derham complex (feectools spaces).""" + from feectools.fem.splines import SplineSpace + + from struphy.polar.homotopy.derham_spaces import DerhamSpace1D + + grid = np.linspace(0, 1, cells + 1) + B = SplineSpace(degree, grid=grid, periodic=periodic) + M = SplineSpace(degree - 1, grid=grid, periodic=periodic, basis="M", pads=degree, parent_multiplicity=1) + return DerhamSpace1D(B, M) + + +def trigonometric_space_1d(n): + """The ``n`` (odd) nodal trigonometric functions of :mod:`struphy.feec.trig_space` (with their 1-forms).""" + from struphy.feec.trig_space import TrigonometricSpace + from struphy.polar.homotopy.derham_spaces import DerhamSpace1D + + return DerhamSpace1D(TrigonometricSpace(n, "B"), TrigonometricSpace(n, "M")) + + +# ---------------------------------------------------------------------------------------------------------- Derham +def _pair(value): + return (value, value) if np.isscalar(value) else tuple(value) + + +def derham_2d(degree=3, cells=(5, 12), domain=None, angular="bspline", nquads=None): + """2D Derham complex of the cross-section (one cell in ``eta3``); ``angular="trig"`` uses ``cells[1]`` + trigonometric functions in ``eta2`` (odd) instead of periodic B-splines.""" + from struphy.polar.homotopy.derham_blocks import poloidal_derham + + if angular not in ("bspline", "trig"): + raise ValueError("angular must be 'bspline' or 'trig'") + if angular == "trig": + return poloidal_derham(_pair(degree), _pair(cells), domain=domain, nquads=nquads, angular="trig") + return poloidal_derham(_pair(degree), _pair(cells), domain=domain, nquads=nquads) + + +def factors(degree=3, cells=(5, 12), angular="bspline"): + """The radial and angular 0-form spaces of the 2D Derham complex.""" + from struphy.polar.homotopy.derham_spaces import derham_factors + + radial, ang, _ = derham_factors(derham_2d(degree, cells, angular=angular)) + return radial, ang + + +def polar_complex( + regularity, + degree=3, + cells=(5, 12), + domain=None, + angular="bspline", + boundary="free", + enrichment="cochain_chain", + cross_section="domain", + polynomial_projection="commuting", + check_assumptions=True, +): + """The polar complex of regularity ``r`` on the 2D Derham complex of ``domain`` (default: the perturbed disk).""" + from struphy.polar.homotopy.derham_blocks import derham_polar_complex + + domain = perturbed_disk() if domain is None else domain + derham = derham_2d(degree, cells, domain=domain, angular=angular) + return derham_polar_complex( + derham, + domain, + regularity, + boundary=boundary, + enrichment=enrichment, + polynomial_projection=polynomial_projection, + cross_section=cross_section, + check_assumptions=check_assumptions, + ) + + +def tensor_space(k, degree=3, cells=(5, 12), domain=None, angular="bspline"): + """The tensor-product ``k``-forms of the 2D Derham complex pushed forward to the cross-section of ``domain``.""" + from struphy.polar.homotopy.derham_blocks import DomainCrossSection + from struphy.polar.homotopy.forms2d import PhysicalFormSpace2D + + domain = perturbed_disk() if domain is None else domain + radial, ang = factors(degree, cells, angular) + return PhysicalFormSpace2D.from_factors(k, radial, ang, DomainCrossSection(domain), STRUPHY) + + +def cross_section(domain=None): + """The 2D map of the cross-section ``eta3 = 0`` of a Struphy domain (for plots).""" + from struphy.polar.homotopy.derham_blocks import DomainCrossSection + + return DomainCrossSection(perturbed_disk() if domain is None else domain) diff --git a/src/struphy/polar/homotopy/demos/fieldlines.py b/src/struphy/polar/homotopy/demos/fieldlines.py new file mode 100644 index 000000000..6638fec0a --- /dev/null +++ b/src/struphy/polar/homotopy/demos/fieldlines.py @@ -0,0 +1,176 @@ +"""A stellarator-like magnetic field in the 3D polar complex of a twisted tube (port of PolarFEEC.jl +``examples/07_torus_3d.jl``, on the Struphy :class:`~struphy.feec.psydac_derham.Derham` with +``polar_splines=True``). + +The tube :class:`~struphy.geometry.domains.HelicalTube` ``("islands")`` has a helical magnetic axis with four field +periods and an elliptical cross-section ``a x b = 0.28 x 0.14`` that rotates by half a turn per period. In logical +coordinates ``(r, theta, zeta)`` the field is the closed 2-form + + B = Phi'(r) dr ^ dtheta + dzeta ^ dpsi, psi = psi_p(r) + sum eps_mn r^m cos 2 pi (m theta - n zeta), + +with toroidal flux ``Phi' = 2 pi a b B0 r`` and poloidal flux ``psi_p' = iota Phi'``, ``iota(r) = 0.6 + 0.3 r^2``. +The modes ``(3, 2)`` and ``(5, 4)`` open island chains on the resonant surfaces ``iota = 2/3`` and ``iota = 4/5``. +Since ``dB = 0`` and the polar commuting projection commutes with ``d``, the projected field has ``div B = 0`` in polar +coefficients up to rounding, on the magnetic axis included. Field lines are traced in the projected field and their +winding is compared with ``iota``. + +The field lines are integrated in logical coordinates with ``zeta`` as the time (``dr/dzeta = b_1 / b_3``, +``dtheta/dzeta = b_2 / b_3`` for the logical 2-form components) by RK4; PolarFEEC.jl traces them in physical space. +""" + +import matplotlib.pyplot as plt +import numpy as np + +from struphy.polar.homotopy.demos.plotting import CATEGORICAL + +B0 = 1.0 +MODES = ((3, 2, 8e-5), (5, 4, 5e-5)) + + +def iota(r): + return 0.6 + 0.3 * r**2 + + +def build(num_elements=(8, 24, 32), degree=(3, 3, 3), regularity=2, domain_resolution=(8, 32, 128)): + """The domain and the Derham complex with high-order polar splines.""" + from struphy.feec.psydac_derham import Derham + from struphy.geometry.domains import HelicalTube + from struphy.io.options import DerhamOptions + from struphy.topology.grids import TensorProductGrid + + domain = HelicalTube("islands", num_elements=domain_resolution, degree=(4, 4, 4)) + options = DerhamOptions( + degree=degree, + bcs=(("free", "free"), None, None), + polar_splines=True, + polar_regularity=regularity, + nquads_proj=tuple(p + 3 for p in degree), + ) + return domain, Derham(TensorProductGrid(num_elements=num_elements), options, domain=domain) + + +def magnetic_field(domain): + """Logical 2-form components ``(b_theta_zeta, b_zeta_r, b_r_theta)`` of the field.""" + a, b = domain.geometry["a"], domain.geometry["b"] + + def phi_prime(r): + return 2 * np.pi * a * b * B0 * r + + def components(r, t, z): + d_theta_psi = np.zeros_like(r * t * z, dtype=float) + d_r_psi = iota(r) * phi_prime(r) + 0 * t * z + for m, n, eps in MODES: + phase = 2 * np.pi * (m * t - n * z) + d_theta_psi = d_theta_psi - 2 * np.pi * m * eps * r**m * np.sin(phase) + d_r_psi = d_r_psi + m * eps * r ** (m - 1) * np.cos(phase) + return -d_theta_psi, d_r_psi, phi_prime(r) + 0 * t * z + + return [lambda r, t, z, i=i: components(r, t, z)[i] for i in range(3)] + + +def project(derham, domain): + """The projected field (polar coefficients) and the relative divergence in polar coefficients.""" + b = derham.P2(magnetic_field(domain)) + div = derham.div.dot(b) + return b, float(np.sqrt(div.inner(div) / b.inner(b))) + + +def spline_field(derham, b): + field = derham.create_spline_function("B", "Hdiv") + field.vector = b + return field + + +def trace(field, seeds, turns, steps_per_turn=100): + """RK4 in zeta for all seeds ``(r, theta, zeta0)``; returns the logical points after every full turn and the + winding numbers ``Delta theta / Delta zeta``.""" + state = np.array(seeds, dtype=float) + h = 1.0 / steps_per_turn + theta_unwrapped = state[:, 1].copy() + + def rhs(s): + markers = np.column_stack([np.clip(s[:, 0], 1e-8, 1.0), s[:, 1] % 1.0, s[:, 2] % 1.0]) + b1, b2, b3 = field(markers) + return np.column_stack([b1 / b3, b2 / b3]) + + sections = [state[:, :2].copy()] + for step in range(turns * steps_per_turn): + k1 = rhs(state) + s2 = state.copy() + s2[:, :2] += 0.5 * h * k1 + s2[:, 2] += 0.5 * h + k2 = rhs(s2) + s3 = state.copy() + s3[:, :2] += 0.5 * h * k2 + s3[:, 2] += 0.5 * h + k3 = rhs(s3) + s4 = state.copy() + s4[:, :2] += h * k3 + s4[:, 2] += h + k4 = rhs(s4) + increment = h / 6 * (k1 + 2 * k2 + 2 * k3 + k4) + state[:, :2] += increment + theta_unwrapped += increment[:, 1] + state[:, 2] += h + # field lines that cross the axis continue on the opposite ray + crossed = state[:, 0] < 0 + state[crossed, 0] *= -1 + state[crossed, 1] += 0.5 + if (step + 1) % steps_per_turn == 0: + sections.append(state[:, :2].copy()) + sections = np.array(sections) + windings = (theta_unwrapped - np.array(seeds)[:, 1]) / turns + return sections, windings + + +def resonant_radius(m, n): + return np.sqrt((n / m - 0.6) / 0.3) + + +def seeds(): + """Seeds on the ray theta = 0 of the cut zeta = 0 and inside both island chains (as PolarFEEC.jl).""" + out = [(r, 0.0, 0.0) for r in np.arange(0.08, 0.97, 0.08)] + for m, n, _ in MODES: + out += [(resonant_radius(m, n) + d, 0.0, 1 - 1 / (2 * n)) for d in (0.008, 0.018)] + return out + + +def study(turns=30, steps_per_turn=100, verbose=True, **kwargs): + domain, derham = build(**kwargs) + b, divergence = project(derham, domain) + field = spline_field(derham, b) + seed_list = seeds() + sections, windings = trace(field, seed_list, turns, steps_per_turn) + if verbose: + print(f"||div b|| / ||b|| in polar coefficients = {divergence:.1e}") + for (r, _, z), w in zip(seed_list, windings): + if z == 0.0 and round(r, 2) in (0.24, 0.40, 0.56): + print(f"r = {r:.2f}: iota of the traced field line {w:.5f}, exact {iota(r):.5f}") + return dict( + domain=domain, derham=derham, b=b, divergence=divergence, seeds=seed_list, sections=sections, windings=windings + ) + + +def figure_poincare(result): + """Poincare section at zeta = 0 (in the logical plane r cos 2 pi theta, r sin 2 pi theta) and iota(r).""" + fig, axes = plt.subplots(1, 2, figsize=(11.5, 4.7)) + sections = result["sections"] + for i, (r0, _, z0) in enumerate(result["seeds"]): + r, t = sections[:, i, 0], sections[:, i, 1] + # sections are taken at zeta = z0 + k; shift the island seeds to zeta = 0 by their phase + color = CATEGORICAL[i % len(CATEGORICAL)] if z0 == 0.0 else "k" + axes[0].plot(r * np.cos(2 * np.pi * t), r * np.sin(2 * np.pi * t), ".", ms=1.5, color=color) + axes[0].set_aspect("equal") + axes[0].set_title(r"Poincare section (logical $r\cos 2\pi\theta$, $r\sin 2\pi\theta$)") + rs = np.linspace(0, 1, 101) + axes[1].plot(rs, iota(rs), "k", label=r"exact $\iota(r) = 0.6 + 0.3 r^2$") + radii = [s[0] for s in result["seeds"] if s[2] == 0.0] + windings = [w for s, w in zip(result["seeds"], result["windings"]) if s[2] == 0.0] + axes[1].plot(radii, windings, "o", color=CATEGORICAL[0], label="traced field lines") + for m, n, _ in MODES: + axes[1].axhline(n / m, color="gray", ls=":") + axes[1].set_xlabel("$r$") + axes[1].set_ylabel(r"$\iota$") + axes[1].legend() + fig.tight_layout() + return fig diff --git a/src/struphy/polar/homotopy/demos/maxwell.py b/src/struphy/polar/homotopy/demos/maxwell.py new file mode 100644 index 000000000..0bae4df39 --- /dev/null +++ b/src/struphy/polar/homotopy/demos/maxwell.py @@ -0,0 +1,183 @@ +"""Spectrally correct Maxwell eigenvalues on the unit disk (port of PolarFEEC.jl ``examples/06_maxwell_spectrum.jl``). + +The polar 1-form space with Dirichlet boundary conditions (zero tangential trace) discretizes the Maxwell eigenproblem +``(curl A, curl B) = lambda (A, B)``, i.e. ``C^T M_2 C a = lambda M_1 a`` in polar coefficients. The nonzero eigenvalues +are ``(j'_{mk})^2``, the squared zeros of the Bessel derivatives ``J_m'`` (double for ``m >= 1``), and +``curl A = J_m(j'_{mk} r) cos(m (theta - theta_0))``. The relative complex has cohomology ``(0, 0, 1)``, so the kernel is +exactly ``grad V0`` with dimension ``dim V0`` and there are ``dim V2 - 1`` nonzero eigenvalues; no spurious modes appear +and the eigenvalues converge like ``h^{2p}``. +""" + +import matplotlib.pyplot as plt +import numpy as np +import scipy.linalg as sla +from scipy.special import jv, jvp + +from struphy.polar.homotopy.demos import common +from struphy.polar.homotopy.demos.plotting import CATEGORICAL, DIVERGING, clean_axis + +DEGREE = 4 +REGULARITY = 3 +MESHES = ((12, 24), (24, 48)) +COMPARED = 60 + + +def dbessel_zeros(m, xmax, step=0.05): + """Zeros of ``J_m'`` in ``(0, xmax)`` by bracketing and bisection.""" + zeros, a = [], max(m, 0.5) + while a < xmax: + b = a + step + fa, fb = jvp(m, a), jvp(m, b) + if np.signbit(fa) != np.signbit(fb): + lo, hi = a, b + for _ in range(200): + c = (lo + hi) / 2 + if not lo < c < hi: + break + if np.signbit(jvp(m, c)) == np.signbit(jvp(m, lo)): + lo = c + else: + hi = c + zeros.append((lo + hi) / 2) + a = b + return zeros + + +def exact_spectrum(lambda_max=2000.0): + """Sorted exact eigenvalues below ``lambda_max`` with multiplicity, as ``(lambda, m, k)`` (zeros of ``J_m'`` from + ``scipy.special.jnp_zeros``).""" + from scipy.special import jnp_zeros + + xmax = np.sqrt(lambda_max) + modes = [] + m = 0 + while True: + count = 1 + zeros = jnp_zeros(m, count) + if zeros[0] >= xmax: + break + while zeros[-1] < xmax: + count *= 2 + zeros = jnp_zeros(m, count) + for k, j in enumerate(zeros[zeros < xmax], start=1): + modes += [(j**2, m, k)] * (1 if m == 0 else 2) + m += 1 + return sorted(modes) + + +def eigenproblem(radial_cells, angular_cells, degree=DEGREE, regularity=REGULARITY): + """Dense generalized eigenproblem ``C^T M_2 C a = lambda M_1 a`` in the Dirichlet polar 1-form space (2D Struphy + Derham complex on the unit disk ``HollowCylinder(a1=0)``).""" + polar = common.polar_complex( + regularity, degree, (radial_cells, angular_cells), domain=common.unit_disk(), boundary="dirichlet" + ) + C = polar.derivative(1) + K = (C.T @ polar.mass_matrix(2) @ C).toarray() + lam, a = sla.eigh((K + K.T) / 2, polar.mass_matrix(1).toarray()) + kernel = int(np.sum(lam < 1e-8 * lam.max())) + return dict(polar=polar, lam=lam[kernel:], a=a[:, kernel:], kernel=kernel) + + +def study(meshes=MESHES, compared=COMPARED, verbose=True): + exact = exact_spectrum() + lam_exact = np.array([e[0] for e in exact[:compared]]) + results = [] + for nr, nt in meshes: + r = eigenproblem(nr, nt) + dims = [r["polar"].dimension(k) for k in range(3)] + error = (r["lam"][:compared] - lam_exact) / lam_exact + r["error"] = np.abs(error) + if verbose: + print( + f"mesh {nr:2d}x{nt:2d}: dim V = {dims}, cohomology {r['polar'].cohomology_dimensions()}, kernel " + f"{r['kernel']} (= dim V0), nonzero {len(r['lam'])} (= dim V2 - 1 = {dims[2] - 1}); " + f"upper bounds: {bool(np.all(error > -1e-12))}" + ) + results.append(r) + if verbose and len(results) == 2: + coarse, fine = results + for i in (1, 5, 10, 20): + rate = np.log2(coarse["error"][i - 1] / fine["error"][i - 1]) + print( + f"lambda_{i:<3d} exact {lam_exact[i - 1]:10.4f} rel. error {coarse['error'][i - 1]:.1e} -> " + f"{fine['error'][i - 1]:.1e} (rate h^{rate:.1f})" + ) + return dict(exact=exact, lam_exact=lam_exact, results=results, meshes=meshes) + + +def figure_spectrum(s): + compared = len(s["lam_exact"]) + idx = np.arange(1, compared + 1) + fine = s["results"][-1] + fig, axes = plt.subplots(1, 2, figsize=(11.5, 4.7)) + axes[0].plot(idx, s["lam_exact"], "o", mfc="none", mec="k", label=r"exact $(j'_{m,k})^2$") + axes[0].plot(idx, fine["lam"][:compared], ".", color=CATEGORICAL[0], label="computed, finest mesh") + axes[0].text( + 0.97, + 0.04, + rf"plus {fine['kernel']} $= \dim\bar V^0$ zero eigenvalues (gradients)", + transform=axes[0].transAxes, + ha="right", + va="bottom", + ) + axes[0].set_xlabel("index $i$") + axes[0].set_ylabel(r"eigenvalue $\lambda_i$") + axes[0].set_title("Nonzero Maxwell eigenvalues on the unit disk") + axes[0].legend(loc="upper left") + for r, (nr, nt), color in zip(s["results"], s["meshes"], (CATEGORICAL[1], CATEGORICAL[0])): + axes[1].semilogy(idx, np.maximum(r["error"], 1e-16), ".", color=color, label=f"{nr}x{nt} mesh") + axes[1].set_xlabel("index $i$") + axes[1].set_ylabel(r"$|\lambda_{h,i} - \lambda_i| / \lambda_i$") + axes[1].set_title(f"Relative error, degree {DEGREE} splines") + axes[1].legend(loc="lower right") + fig.tight_layout() + return fig + + +def curl_mode(r, i, m): + """Polar coefficients of mode ``i`` (degenerate pairs rotated to ``cos(m theta)``) and its curl (full coefficients).""" + polar = r["polar"] + E1, C = polar.extraction(1), polar.spaces[1].derivative_matrix + V2 = polar.spaces[2] + + def curl(a): + return C @ (E1 @ a) + + if m == 0: + return curl(r["a"][:, i]) + c1, c2 = curl(r["a"][:, i]), curl(r["a"][:, i + 1]) + v1 = V2.evaluate(c1, [0.4], [1 / (4 * m)])[0, 0] + v2 = V2.evaluate(c2, [0.4], [1 / (4 * m)])[0, 0] + return curl(v2 * r["a"][:, i] - v1 * r["a"][:, i + 1]) + + +def figure_modes(s): + fine = s["results"][-1] + V2 = fine["polar"].spaces[2] + exact = s["exact"] + distinct, seen = [], set() + for i, (lam, m, k) in enumerate(exact): + if (m, k) not in seen: + seen.add((m, k)) + distinct.append(i) + if len(distinct) == 8: + break + xs = np.linspace(1e-7, 1, 97) + ys = np.linspace(0, 1, 6 * s["meshes"][-1][1] + 1) + X, Y = np.meshgrid(xs, ys, indexing="ij") + Px, Py = X * np.cos(2 * np.pi * Y), X * np.sin(2 * np.pi * Y) + fig, axes = plt.subplots(2, 4, figsize=(12, 6.6)) + for ax, i in zip(axes.ravel(), distinct): + lam, m, k = exact[i] + values = V2.evaluate(curl_mode(fine, i, m), xs, ys) + reference = jv(m, np.sqrt(lam) * X) * np.cos(2 * np.pi * m * Y) + values = values * np.sign(np.sum(values * reference)) / np.max(np.abs(values)) + ax.pcolormesh(Px, Py, values, cmap=DIVERGING, vmin=-1, vmax=1, shading="gouraud") + ax.set_aspect("equal") + clean_axis(ax) + ax.set_title(rf"$(m,k) = ({m},{k}),\ \lambda_h = {fine['lam'][i]:.3f}$", fontsize=11) + fig.suptitle( + f"Maxwell eigenmodes on the unit disk ({s['meshes'][-1][0]}x{s['meshes'][-1][1]} mesh, degree {DEGREE})" + ) + fig.tight_layout() + return fig diff --git a/src/struphy/polar/homotopy/demos/orbits.py b/src/struphy/polar/homotopy/demos/orbits.py new file mode 100644 index 000000000..14e32de69 --- /dev/null +++ b/src/struphy/polar/homotopy/demos/orbits.py @@ -0,0 +1,184 @@ +"""Guiding-centre orbits near the magnetic axis and their energy conservation (talk, Section 9). + +The vacuum field is ``HelicalVacuumField`` on the cylinder ``HollowCylinder(a1=0, a2=0.3, Lz=2 pi)`` (as +``examples/GuidingCenter/vacuum_near_axis``), projected into the polar spline spaces of a Struphy Derham complex of +regularity ``r``: ``|B|_h = P0 |B|``, its discrete gradient, ``b_h = P1 b`` and ``curl b_h``. The guiding centres follow + + dX/dt = (v_par B* + eps mu b x grad |B|_h) / B*_par, dv_par/dt = -mu B* . grad |B|_h / B*_par, + +``B* = |B|_h b + eps v_par curl b``, ``B*_par = b . B*``, integrated with the classical RK4. The exact flow conserves +``E = v_par^2 / 2 + mu |B|_h``; the relative energy error ``max_t |E(t) - E(0)| / (mu |B|)`` after ``20`` transits +measures the time-integration error, which converges like ``dt^4`` only when the projected field is smooth enough at +the axis. ``angular="trig"`` gives the compatible spaces of the talk (trigonometric functions in ``eta2``), the +default B-splines the B-spline panel. (Struphy's particle pushers, used by the example, need B-splines in every +direction; this integrator evaluates the fields with ``SplineFunction``.) +""" + +import matplotlib.pyplot as plt +import numpy as np + +from struphy.polar.homotopy.demos.plotting import CATEGORICAL + +A, LZ = 0.3, 2 * np.pi +MU, EPSILON = 1e-3, 0.01 +MARKERS = tuple( + (r / A, angle, 0.0, sign) + for r, angle in ((0.01, 0.0), (0.02, 0.125), (0.03, 0.25), (0.04, 0.375)) + for sign in (1.0, -1.0) +) + + +class ProjectedVacuumField: + """The helical vacuum field projected into the polar splines of regularity ``r`` (B-splines or trigonometric + functions in ``eta2``), evaluated in Cartesian coordinates.""" + + def __init__(self, regularity, angular="bspline", num_elements=(8, 12, 16), degree=(4, 4, 3)): + from struphy.feec.psydac_derham import Derham + from struphy.fields_background import equils + from struphy.geometry import domains + from struphy.io.options import DerhamOptions + from struphy.topology.grids import TensorProductGrid + + n2 = num_elements[1] + (1 - num_elements[1] % 2 if angular == "trig" else 0) + self.domain = domains.HollowCylinder(a1=0.0, a2=A, Lz=LZ) + options = DerhamOptions( + degree=degree, + bcs=(("free", "free"), None, None), + polar_splines=True, + polar_regularity=regularity, + angular_basis=angular, + nquads_proj=tuple(p + 3 for p in degree), + ) + grid = TensorProductGrid( + num_elements=(num_elements[0], n2, num_elements[2]), mpi_dims_mask=(True, False, False) + ) + self.derham = d = Derham(grid, options, domain=self.domain) + self.equil = equils.HelicalVacuumField(B0=1.0, delta=0.5, m=2, k=3.0) + self.equil.domain = self.domain + self.absB = d.create_spline_function("absB", "H1") + self.absB.vector = d.P0(self.equil.absB0) + self.gradB = d.create_spline_function("gradB", "Hcurl") + self.gradB.vector = d.grad.dot(self.absB.vector) + self.b = d.create_spline_function("b", "Hcurl") + self.b.vector = d.P1([lambda e1, e2, e3, i=i: self.equil.unit_b1(e1, e2, e3)[i] for i in range(3)]) + self.curl_b = d.create_spline_function("curl_b", "Hdiv") + self.curl_b.vector = d.curl.dot(self.b.vector) + + @staticmethod + def logical(x): + """Logical coordinates of Cartesian points of the cylinder (markers ``(N, 3)``).""" + r = np.hypot(x[:, 0], x[:, 1]) + return np.column_stack( + [np.minimum(r / A, 1.0), (np.arctan2(x[:, 1], x[:, 0]) / (2 * np.pi)) % 1.0, (x[:, 2] / LZ) % 1.0] + ) + + def __call__(self, x): + """``|B|_h``, ``grad |B|_h``, ``b_h`` and ``curl b_h`` (Cartesian) at Cartesian points ``x`` (N, 3).""" + eta = self.logical(x) + # the polar fields are conforming: their physical values at the axis are the limits eta1 -> 0 + eta[:, 0] = np.maximum(eta[:, 0], 1e-9) + DF = np.asarray(self.domain.jacobian(eta)).reshape(3, 3, -1).transpose(2, 0, 1) + det = np.linalg.det(DF) + inv_T = np.linalg.inv(DF).transpose(0, 2, 1) + + def one_form(field): + return np.einsum("nij,nj->ni", inv_T, np.column_stack([np.asarray(c) for c in field(eta)])) + + def two_form(field): + return np.einsum("nij,nj->ni", DF, np.column_stack([np.asarray(c) for c in field(eta)])) / det[:, None] + + return np.asarray(self.absB(eta)).ravel(), one_form(self.gradB), one_form(self.b), two_form(self.curl_b) + + def cartesian_markers(self): + eta = np.array([m[:3] for m in MARKERS]) + x = np.asarray(self.domain(eta)).reshape(3, -1).T + return x, np.array([m[3] for m in MARKERS]) + + +def rhs(field, x, v): + B, gradB, b, curl_b = field(x) + b = b / np.linalg.norm(b, axis=1)[:, None] + B_star = B[:, None] * b + EPSILON * v[:, None] * curl_b + B_par = np.einsum("ni,ni->n", b, B_star) + dx = (v[:, None] * B_star + EPSILON * MU * np.cross(b, gradB)) / B_par[:, None] + dv = -MU * np.einsum("ni,ni->n", B_star, gradB) / B_par + return dx, dv, B + + +def integrate(field, dt, periods=20.0, record_every=10): + """RK4 for all markers; returns the recorded positions and the relative energy errors.""" + x, v = field.cartesian_markers() + steps = int(np.ceil(periods * LZ / dt)) + dt = periods * LZ / steps + _, _, B0 = rhs(field, x, v) + energy0 = 0.5 * v**2 + MU * B0 + errors = np.zeros(len(v)) + path = [x.copy()] + for step in range(steps): + k1x, k1v, B = rhs(field, x, v) + errors = np.maximum(errors, np.abs(0.5 * v**2 + MU * B - energy0) / (MU * B0)) + k2x, k2v, _ = rhs(field, x + 0.5 * dt * k1x, v + 0.5 * dt * k1v) + k3x, k3v, _ = rhs(field, x + 0.5 * dt * k2x, v + 0.5 * dt * k2v) + k4x, k4v, _ = rhs(field, x + dt * k3x, v + dt * k3v) + x = x + dt / 6 * (k1x + 2 * k2x + 2 * k3x + k4x) + v = v + dt / 6 * (k1v + 2 * k2v + 2 * k3v + k4v) + if (step + 1) % record_every == 0: + path.append(x.copy()) + _, _, B = rhs(field, x, v) + errors = np.maximum(errors, np.abs(0.5 * v**2 + MU * B - energy0) / (MU * B0)) + return np.array(path), errors + + +def energy_study( + regularities=(0, 1, 2, 3, 4), + dts=(0.04, 0.02, 0.01, 0.005, 0.0025), + angular="bspline", + periods=20.0, + verbose=True, + **kwargs, +): + """Maximum relative energy error over the eight orbits for every regularity and time step.""" + errors = {} + for r in regularities: + field = ProjectedVacuumField(r, angular, **kwargs) + errors[r] = [] + for dt in dts: + _, e = integrate(field, dt, periods) + errors[r].append(float(np.max(e))) + if verbose: + print(f" {angular} r = {r}, dt = {dt}: max relative energy error {errors[r][-1]:.2e}") + return dict(dts=np.array(dts), errors=errors, angular=angular, periods=periods) + + +def figure_energy(study): + fig, ax = plt.subplots(figsize=(5.5, 4.5)) + dts = study["dts"] + for i, (r, e) in enumerate(study["errors"].items()): + ax.loglog(dts, e, "o-", color=CATEGORICAL[i % len(CATEGORICAL)], label=rf"$\mathfrak{{r}} = {r}$") + e = np.array(study["errors"][max(study["errors"])]) + ax.loglog(dts, e[-1] * (dts / dts[-1]) ** 4, "k:", label=r"$O(\Delta t^4)$") + ax.set_xlabel(r"$\Delta t$") + ax.set_ylabel(r"$\max_t |\mathcal{E}(t) - \mathcal{E}(0)| / (\mu |B|)$") + kind = "compatible (trigonometric)" if study["angular"] == "trig" else "B-spline" + ax.set_title(f"energy after {study['periods']:g} transits, {kind} spaces") + ax.legend(fontsize=8) + fig.tight_layout() + return fig + + +def figure_orbits(path, extent=None): + """Orbits projected onto the cross-section (``extent``: half width of the close-up).""" + fig, ax = plt.subplots(figsize=(4.8, 4.8)) + for i in range(path.shape[1]): + ax.plot(path[:, i, 0], path[:, i, 1], ".", ms=1, color=CATEGORICAL[i % len(CATEGORICAL)]) + if extent is None: + ax.add_patch(plt.Circle((0, 0), A, color="k", fill=False)) + extent = A + ax.set_xlim(-extent, extent) + ax.set_ylim(-extent, extent) + ax.set_aspect("equal") + ax.set_xlabel("$x$") + ax.set_ylabel("$y$") + ax.set_title("guiding-centre orbits projected onto a cross-section" if extent == A else "close-up near the axis") + fig.tight_layout() + return fig diff --git a/src/struphy/polar/homotopy/demos/pde.py b/src/struphy/polar/homotopy/demos/pde.py new file mode 100644 index 000000000..d8a75556c --- /dev/null +++ b/src/struphy/polar/homotopy/demos/pde.py @@ -0,0 +1,279 @@ +"""Poisson and curl-curl problems on a rounded square with the polar complex (port of PolarFEEC.jl +``examples/05_polar_pde.jl``). + +A disk-like mapping maps the logical rectangle onto the rounded square ``{g > 0}``, +``g = 1 - x^2 - y^2 - eps Re((x + i y)^4)``. The solvers work in the homotopy gauge of the complex:: + + stiffness_solver(complex, 0)(f) (grad v, grad u) = (v, f) for all v in V0 + stiffness_solver(complex, 1)(J) (curl Phi, curl A) = (Phi, J) for all Phi in V1 + +For the free complex the kernel of the gradient holds the constants, so ``f`` needs zero mean and the gauge fixes +``u(centre) = 0``. For curl-curl it holds the gradients; with ``J = curl* B = (d_y B, -d_x B)`` the free problem needs +``B = 0`` on the boundary, the Dirichlet problem (``A . t = 0``) needs ``B`` to have zero mean. The curl-curl solve is +the "free two-dimensional lunch" of the paper: one polar 2-form mass solve and the homotopy ``h_2``. + +The manufactured solutions are built from ``g``: ``u = g s`` vanishes on the boundary, ``u = g^2 s`` has +``d_n u = 0``. ``B = sin 2x cos y`` is odd in ``x`` (zero mean) and ``g B`` vanishes on the boundary. The expected +rates are ``p + 1`` for ``u`` and ``p`` for ``curl A`` (regularity ``r = p``). +""" + +import matplotlib.pyplot as plt +import numpy as np + +from struphy.polar.homotopy.demos import common +from struphy.polar.homotopy.demos.plotting import CATEGORICAL, DIVERGING, draw_mesh, field_on_disk +from struphy.polar.homotopy.solvers import stiffness_solver + +EPS = 0.15 +QUADRATURE = 8 + + +def boundary_radius(t): + """On the ray at angle ``2 pi t`` the boundary radius solves ``rho^2 + eps cos(8 pi t) rho^4 = 1``.""" + return np.sqrt(2 / (1 + np.sqrt(1 + 4 * EPS * np.cos(8 * np.pi * t)))) + + +def rounded_square(): + """The cross-section of the Struphy domain ``RoundedSquareCylinder`` (``R = x rho_b(y)``).""" + return common.cross_section(rounded_square_domain()) + + +def rounded_square_domain(): + from struphy.geometry import domains + + return domains.RoundedSquareCylinder(a=1.0, eps=EPS, Lz=1.0) + + +# ---------------------------------------------------------------- manufactured solutions as (value, gradient, laplacian) +class Jet: + def __init__(self, value, gradient, laplacian): + self.value, self.gradient, self.laplacian = value, gradient, laplacian + + def __mul__(self, other): + return Jet( + self.value * other.value, + ( + self.value * other.gradient[0] + other.value * self.gradient[0], + self.value * other.gradient[1] + other.value * self.gradient[1], + ), + self.value * other.laplacian + + other.value * self.laplacian + + 2 * (self.gradient[0] * other.gradient[0] + self.gradient[1] * other.gradient[1]), + ) + + +def level(x, y): + return Jet( + 1 - x**2 - y**2 - EPS * (x**4 - 6 * x**2 * y**2 + y**4), + (-2 * x - 4 * EPS * (x**3 - 3 * x * y**2), -2 * y - 4 * EPS * (y**3 - 3 * x**2 * y)), + -4.0 * np.ones_like(x), + ) + + +def shape(x, y): + return Jet(1 + np.sin(2 * x + y), (2 * np.cos(2 * x + y), np.cos(2 * x + y)), -5 * np.sin(2 * x + y)) + + +def u_dirichlet(x, y): + return level(x, y) * shape(x, y) + + +def u_free(x, y): + return level(x, y) * level(x, y) * shape(x, y) + + +def field_curl(x, y): + return Jet( + np.sin(2 * x) * np.cos(y), + (2 * np.cos(2 * x) * np.cos(y), -np.sin(2 * x) * np.sin(y)), + -5 * np.sin(2 * x) * np.cos(y), + ) + + +def curl_dirichlet(x, y): + return field_curl(x, y) + + +def curl_free(x, y): + return level(x, y) * field_curl(x, y) + + +def source(u): + return lambda x, y: -u(x, y).laplacian + + +def current(B): + return lambda x, y: (B(x, y).gradient[1], -B(x, y).gradient[0]) + + +def exact(u, shift=0.0): + return lambda x, y: u(x, y).value - shift + + +EXACT = dict( + u_free=exact(u_free, float(u_free(np.array(0.0), np.array(0.0)).value)), + u_dirichlet=exact(u_dirichlet), + B_free=exact(curl_free), + B_dirichlet=exact(curl_dirichlet), +) + + +def solve_all(p, n, domain=None): + """Free and Dirichlet complexes of degree p with n radial and 6n angular cells (2D Struphy Derham complexes on + ``RoundedSquareCylinder``), and the four solutions.""" + domain = rounded_square_domain() if domain is None else domain + free = common.polar_complex(p, p, (n, 6 * n), domain=domain) + dirichlet = common.polar_complex(p, p, (n, 6 * n), domain=domain, boundary="dirichlet") + q = QUADRATURE + return dict( + free=free, + dirichlet=dirichlet, + u_free=stiffness_solver(free, 0)(source(u_free), additional_order=q), + u_dirichlet=stiffness_solver(dirichlet, 0)(source(u_dirichlet), additional_order=q), + A_free=stiffness_solver(free, 1)(current(curl_free), additional_order=q), + A_dirichlet=stiffness_solver(dirichlet, 1)(current(curl_dirichlet), additional_order=q), + ) + + +def errors(s): + """L2 errors of u (free, Dirichlet) and of curl A (free, Dirichlet).""" + free, q = s["free"], QUADRATURE + V0, V1, V2 = free.spaces + C = V1.derivative_matrix + return np.array( + [ + V0.l2_error(s["u_free"], EXACT["u_free"], q), + V0.l2_error(s["u_dirichlet"], EXACT["u_dirichlet"], q), + V2.l2_error(C @ s["A_free"], EXACT["B_free"], q), + V2.l2_error(C @ s["A_dirichlet"], EXACT["B_dirichlet"], q), + ] + ) + + +def checks(s): + """Compatibility, gauge and boundary checks of the example (returns a dict).""" + free, dirichlet = s["free"], s["dirichlet"] + load = free.spaces[0].inner_product(source(u_free), QUADRATURE) + ts = np.linspace(0, 1, 721) + V0, V1 = free.spaces[0], free.spaces[1] + centre = V0.evaluate(s["u_free"], [0.0], [0.0])[0, 0] + boundary_u = np.max(np.abs(V0.evaluate(s["u_dirichlet"], [1.0], ts))) + A_polar = dirichlet.duals(1).T @ s["A_dirichlet"] + A_free_polar = free.duals(1).T @ s["A_free"] + logical = V1.logical.evaluate(s["A_dirichlet"], [1.0], ts) + return dict( + mean_of_source=float(np.sum(load)), + u_centre=float(centre), + max_u_on_boundary=float(boundary_u), + gauge_free=float(np.max(np.abs(free.homotopy(1) @ A_free_polar))), + gauge_dirichlet=float(np.max(np.abs(dirichlet.homotopy(1) @ A_polar))), + tangential_trace=float(np.max(np.abs(logical[1]))), + ) + + +CELLS = (4, 6, 8, 12, 16, 24) +DEGREES = (2, 3, 4, 5) + + +def convergence(cells=CELLS, degrees=DEGREES, verbose=True): + table = {p: np.array([errors(solve_all(p, n)) for n in cells]) for p in degrees} + if verbose: + rates = lambda e: np.log(e[:-1] / e[1:]) / np.log(np.array(cells[1:]) / np.array(cells[:-1])) # noqa: E731 + print("Observed L2 rates over the last three refinements (expected p + 1 for u, p for curl A)") + print(" p u free u dirichlet curl A free curl A dirichlet") + for p in degrees: + parts = [" ".join(f"{r:4.1f}" for r in rates(table[p][:, j])[-3:]) for j in range(4)] + print(f" {p} " + " ".join(parts)) + return table + + +def _sample(space, coefficients, rs, ts): + return space.evaluate(coefficients, rs, ts) + + +def figure_poisson(s, n=6): + mapping = s["free"].mapping + rs, ts = np.linspace(1e-3, 1, 121), np.linspace(0, 1, 361) + V0 = s["free"].spaces[0] + fig, axes = plt.subplots(1, 2, figsize=(10, 5.2)) + for ax, key, title, symmetric in ( + (axes[0], "u_free", r"$-\Delta u_h = f,\ \partial_n u_h = 0$ (gauge $u_h(0) = 0$)", True), + (axes[1], "u_dirichlet", r"$-\Delta u_h = f,\ u_h = 0$ on $\partial\Omega$", False), + ): + field_on_disk( + ax, + mapping, + rs, + ts, + _sample(V0, s[key], rs, ts), + symmetric=symmetric, + contours=12, + colorbar=True, + label=r"$u_h$", + ) + draw_mesh(ax, mapping, np.linspace(0, 1, n + 1), np.linspace(0, 1, 6 * n + 1)) + ax.set_title(title) + fig.tight_layout() + return fig + + +def figure_curl_curl(s, n=6): + mapping = s["free"].mapping + V1, V2 = s["free"].spaces[1], s["free"].spaces[2] + fig, axes = plt.subplots(1, 2, figsize=(10, 5.2)) + ax = axes[0] + draw_mesh(ax, mapping, np.linspace(0, 1, n + 1), np.linspace(0, 1, 6 * n + 1)) + for r in np.linspace(0.1, 0.95, 9): + ts = np.linspace(0, 1, int(round(44 * r)) + 1)[:-1] + X, Y = mapping(np.full_like(ts, r), ts) + A = V1.evaluate_points(s["A_free"], np.full_like(ts, r), ts) + ax.quiver(X, Y, A[0], A[1], np.hypot(A[0], A[1]), cmap="viridis", scale=12) + ax.set_aspect("equal") + ax.set_title(r"$\mathrm{curl}^\star\mathrm{curl}\,A_h = J$, homotopy gauge $A_h \cdot x = 0$") + rs, ts = np.linspace(1e-3, 1, 121), np.linspace(0, 1, 361) + C = V1.derivative_matrix + field_on_disk( + axes[1], + mapping, + rs, + ts, + _sample(V2, C @ s["A_free"], rs, ts), + contours=12, + colorbar=True, + label=r"$\mathrm{curl}\,A_h$", + ) + draw_mesh(axes[1], mapping, np.linspace(0, 1, n + 1), np.linspace(0, 1, 6 * n + 1)) + axes[1].set_title(r"$\mathrm{curl}\,A_h \approx g(x,y)\sin(2x)\cos(y)$") + fig.tight_layout() + return fig + + +def figure_convergence(table, cells=CELLS, degrees=DEGREES): + hs = 1 / np.array(cells) + fig, axes = plt.subplots(1, 2, figsize=(11, 5)) + panels = ( + (r"$\Vert u - u_h \Vert_{L^2}$", "Poisson: free (o), Dirichlet (s)", (0, 1), 1), + (r"$\Vert \mathrm{curl}\,A - \mathrm{curl}\,A_h\Vert_{L^2}$", "curl-curl: free (o), Dirichlet (s)", (2, 3), 0), + ) + for ax, (ylabel, title, columns, shift) in zip(axes, panels): + for i, p in enumerate(degrees): + for k, j in enumerate(columns): + ax.loglog( + hs, + table[p][:, j], + marker="os"[k], + linestyle=("-", "--")[k], + color=CATEGORICAL[i], + label=f"$p = {p}$" if k == 0 else None, + markersize=5, + ) + segment = hs[-3:] + anchor = 0.3 * np.min(table[p][-1, list(columns)]) + ax.loglog(segment, anchor * (segment / hs[-1]) ** (p + shift), ":", color=CATEGORICAL[i]) + ax.set_xticks(hs, [f"1/{n}" for n in cells]) + ax.set_xlabel("$h = 1/n$") + ax.set_ylabel(ylabel) + ax.set_title(title) + ax.legend(loc="lower right") + fig.tight_layout() + return fig diff --git a/src/struphy/polar/homotopy/demos/performance.py b/src/struphy/polar/homotopy/demos/performance.py new file mode 100644 index 000000000..d3d9c079e --- /dev/null +++ b/src/struphy/polar/homotopy/demos/performance.py @@ -0,0 +1,73 @@ +"""Assembly cost of the polar mass matrices against the spline degree (counterpart of PolarFEEC.jl +``examples/08_assembly_performance.jl`` for the Struphy :class:`~struphy.feec.psydac_derham.Derham` with +``polar_splines=True``). + +PolarFEEC.jl assembles 3D masses by sum factorization (Julia threads); Struphy assembles the tensor-product weighted +mass matrices with its own (pyccel) kernels, and the polar matrices are applied as ``E M E^T`` with the polar +extraction operators. The study measures the assembly time of ``M0 ... M3`` and of one polar matrix-vector product on +the circular torus ``HollowTorus(a1=0)`` for degrees ``p = 1 ... 5``, and checks the volume ``2 pi^2 R0 a^2``. +""" + +import time + +import numpy as np + + +def study(degrees=(1, 2, 3, 4, 5), cells=(12, 12, 12), regularity=1, repeats=3, verbose=True): + from struphy.feec.mass import WeightedMassOperators + from struphy.feec.psydac_derham import Derham + from struphy.geometry.domains import HollowTorus + from struphy.io.options import DerhamOptions + from struphy.topology.grids import TensorProductGrid + + R0, a = 3.0, 1.0 + domain = HollowTorus(a1=0.0, a2=a, R0=R0, tor_period=1) + results = {} + for p in degrees: + options = DerhamOptions( + degree=(p,) * 3, + bcs=(("free", "free"), None, None), + polar_splines=True, + polar_regularity=min(regularity, p), + ) + derham = Derham(TensorProductGrid(num_elements=cells), options, domain=domain) + timings, products = {}, {} + for name in ("M0", "M1", "M2", "M3"): + best = np.inf + for _ in range(repeats): + mass = WeightedMassOperators(derham, domain) + start = time.perf_counter() + M = getattr(mass, name) + x = M.domain.zeros() + M.dot(x) # the matrices are assembled on first use + best = min(best, time.perf_counter() - start) + timings[name] = best + start = time.perf_counter() + M.dot(x) + products[name] = time.perf_counter() - start + # volume (1, 1)_{L2}: the constant 1 has the 3-form proxy det(DF), and M3 has the weight 1 / det(DF) + one = derham.P3(lambda e1, e2, e3: domain.jacobian_det(e1, e2, e3)) + volume = float(WeightedMassOperators(derham, domain).M3.dot(one).inner(one)) + results[p] = dict(assembly=timings, product=products, volume=volume) + if verbose: + print( + f"p = {p}: assembly " + + ", ".join(f"{k} {v:.2f} s" for k, v in timings.items()) + + f"; volume {volume:.6f} (exact {2 * np.pi**2 * R0 * a**2:.6f})" + ) + return results + + +def figure(results): + import matplotlib.pyplot as plt + + degrees = list(results) + fig, ax = plt.subplots(figsize=(5.5, 4.2)) + for name in ("M0", "M1", "M2", "M3"): + ax.semilogy(degrees, [results[p]["assembly"][name] for p in degrees], "o-", label=name) + ax.set_xlabel("degree $p$") + ax.set_ylabel("assembly time and first product [s]") + ax.set_title("polar mass matrices, $12^3$ cells") + ax.legend() + fig.tight_layout() + return fig diff --git a/src/struphy/polar/homotopy/demos/plotting.py b/src/struphy/polar/homotopy/demos/plotting.py new file mode 100644 index 000000000..aed95c511 --- /dev/null +++ b/src/struphy/polar/homotopy/demos/plotting.py @@ -0,0 +1,81 @@ +"""Plotting helpers shared by the demos (matplotlib versions of the conventions of PolarFEEC.jl's ``examples/common.jl``). + +Signed fields use the colour scheme of the paper's figures: MATLAB blue, white, MATLAB orange, with eased transitions +(use a symmetric colour range). +""" + +import matplotlib.pyplot as plt +import numpy as np +from matplotlib.colors import LinearSegmentedColormap + +_BLUE, _WHITE, _ORANGE = (0.0, 0.4470, 0.7410), (1.0, 1.0, 1.0), (0.8500, 0.3250, 0.0980) + + +def _diverging(): + t = np.linspace(0, 1, 256) + eased = 3 * t**2 - 2 * t**3 + colors = [] + for s in eased: + if s < 0.5: + a = s / 0.5 + colors.append(tuple((1 - a) * b + a * w for b, w in zip(_BLUE, _WHITE))) + else: + a = (s - 0.5) / 0.5 + colors.append(tuple((1 - a) * w + a * o for w, o in zip(_WHITE, _ORANGE))) + return LinearSegmentedColormap.from_list("polarfeec_diverging", colors) + + +DIVERGING = _diverging() +SEQUENTIAL = "viridis" +CATEGORICAL = ["#0072B2", "#E69F00", "#009E73", "#CC79A7", "#56B4E9", "#D55E00", "#F0E442", "#000000"] + + +def mapped_grid(mapping, xs, ys): + """Physical coordinates of the logical tensor grid ``xs (x) ys``.""" + X, Y = np.meshgrid(xs, ys, indexing="ij") + return mapping(X, Y) + + +def draw_mesh(ax, mapping, radial_breaks, angular_breaks, color="k", alpha=0.13, linewidth=0.6, boundary=True): + """Mesh lines of a disk-like mapping (rings ``x = const`` and spokes ``y = const``).""" + t = np.linspace(0, 1, 400) + for r in radial_breaks: + if 0 < r < 1: + ax.plot(*mapping(np.full_like(t, r), t), color=color, alpha=alpha, linewidth=linewidth) + for th in angular_breaks: + ax.plot(*mapping(t, np.full_like(t, th)), color=color, alpha=alpha, linewidth=linewidth) + if boundary: + ax.plot(*mapping(np.ones_like(t), t), color=color, linewidth=1.2) + + +def field_on_disk(ax, mapping, xs, ys, values, symmetric=True, cmap=None, contours=0, colorbar=False, label=None): + """``pcolormesh`` of values on the mapped logical grid ``xs (x) ys`` (values of shape ``(len(xs), len(ys))``).""" + X, Y = mapped_grid(mapping, xs, ys) + vmax = np.max(np.abs(values)) or 1.0 + vmin = -vmax if symmetric else np.min(values) + cmap = (DIVERGING if symmetric else SEQUENTIAL) if cmap is None else cmap + mesh = ax.pcolormesh(X, Y, values, cmap=cmap, vmin=vmin, vmax=vmax, shading="gouraud") + if contours: + ax.contour( + X, Y, values, levels=np.linspace(vmin, vmax, contours + 2)[1:-1], colors="k", linewidths=0.6, alpha=0.5 + ) + ax.set_aspect("equal") + if colorbar: + plt.colorbar(mesh, ax=ax, label=label, shrink=0.8) + return mesh + + +def clean_axis(ax): + ax.set_xticks([]) + ax.set_yticks([]) + for spine in ax.spines.values(): + spine.set_visible(False) + + +def loglog_convergence(ax, hs, errors, rate, color, label=None, marker="o", linestyle="-", reference=True): + """Error curve against ``h`` with a dotted ``h^rate`` reference below the finest points.""" + ax.loglog(hs, errors, marker=marker, linestyle=linestyle, color=color, label=label, markersize=5) + if reference: + segment = np.asarray(hs[-3:]) + anchor = 0.3 * errors[-1] + ax.loglog(segment, anchor * (segment / segment[-1]) ** rate, ":", color=color, linewidth=1.3) diff --git a/src/struphy/polar/homotopy/demos/projections.py b/src/struphy/polar/homotopy/demos/projections.py new file mode 100644 index 000000000..e0ed6bc67 --- /dev/null +++ b/src/struphy/polar/homotopy/demos/projections.py @@ -0,0 +1,173 @@ +"""Convergence and commutativity of the polar projection operators (talk, Section 6). + +The polar projections are ``pi_bar = Pi_bar o pi`` with the commuting tensor-product projection ``pi`` (geometric +degrees of freedom), and the polar L2 projections for comparison. With ``r_x = r_y = r`` (B-spline degree equal to +the regularity) the 0-form projections converge like ``h^{r+1}`` and the 1- and 2-form projections like ``h^r``. The +commuting projections commute with ``grad``, ``curl`` and with the homotopies ``grad^dagger``, ``curl^dagger`` up to +round-off (and quadrature), the L2 projections only up to the approximation error. + +The test functions, the mesh (``N`` radial and ``4 N`` angular cells) and the perturbed disk are choices of this port; +the talk does not state them. The spaces are those of 2D Struphy Derham complexes on the domain +``PerturbedDiskCylinder`` (see :mod:`.common`). +""" + +import matplotlib.pyplot as plt +import numpy as np + +from struphy.polar.homotopy.demos import common +from struphy.polar.homotopy.demos.plotting import CATEGORICAL + + +def f0(x, y): + return np.sin(3 * x + 1) * np.cos(2 * y - 0.5) + 0.5 * np.cos(4 * x * y) + + +def grad_f0(x, y): + return ( + 3 * np.cos(3 * x + 1) * np.cos(2 * y - 0.5) - 2 * y * np.sin(4 * x * y), + -2 * np.sin(3 * x + 1) * np.sin(2 * y - 0.5) - 2 * x * np.sin(4 * x * y), + ) + + +def A1(x, y): + return (np.cos(x - y), x * np.exp(y)) + + +def curl_A1(x, y): + return np.exp(y) - np.sin(x - y) + + +def B2(x, y): + return np.sin(x + 2 * y) + x**2 + + +def _gauss01(n=20): + t, w = np.polynomial.legendre.leggauss(n) + return (t + 1) / 2, w / 2 + + +def poincare_gradient(A): + """``(grad^dagger A)(x) = int_0^1 A(t x) . x dt``.""" + t, w = _gauss01() + + def h(x, y): + total = 0.0 + for tq, wq in zip(t, w): + ax, ay = A(tq * x, tq * y) + total = total + wq * (ax * x + ay * y) + return total + + return h + + +def poincare_curl(B): + """``(curl^dagger B)(x) = (-y, x) int_0^1 t B(t x) dt``.""" + t, w = _gauss01() + + def h(x, y): + s = sum(wq * tq * B(tq * x, tq * y) for tq, wq in zip(t, w)) + return (-y * s, x * s) + + return h + + +def complexes(regularity, cells=(4, 8, 16, 32), angular_factor=4, domain=None): + """Polar complexes with B-splines of degree ``r`` on ``N`` radial and ``4 N`` angular cells.""" + domain = common.perturbed_disk() if domain is None else domain + return [ + common.polar_complex(regularity, regularity, (n, angular_factor * n), domain=domain, check_assumptions=False) + for n in cells + ] + + +def _norm(c, k, coefficients): + M = c.full_mass_matrix(k) + return float(np.sqrt(max(coefficients @ (M @ coefficients), 0.0))) + + +def project(c, k, f, method): + return c[k](f, additional_order=6) if method == "commuting" else c.l2_projection(k, f, additional_order=0) + + +def convergence(regularities=(1, 2, 3, 4, 5, 6), cells=(4, 8, 16, 32), method="commuting"): + """Relative L2 errors of the 0-, 1-, 2-form projections, ``errors[r]`` of shape ``(len(cells), 3)``.""" + targets = (f0, A1, B2) + errors = {} + for r in regularities: + rows = [] + for c in complexes(r, cells): + row = [] + for k, f in enumerate(targets): + u = project(c, k, f, method) + zero = np.zeros_like(u) + row.append(c.spaces[k].l2_error(u, f, 4) / c.spaces[k].l2_error(zero, f, 4)) + rows.append(row) + errors[r] = np.array(rows) + return errors + + +def commutation(regularities=(1, 2, 3, 4, 5, 6), cells=(4, 8, 16, 32), method="commuting"): + """Relative defects of ``grad``, ``curl``, ``grad^dagger`` and ``curl^dagger`` with the projections.""" + defects = {} + for r in regularities: + rows = [] + for c in complexes(r, cells): + G, C = c.spaces[0].derivative_matrix, c.spaces[1].derivative_matrix + HG, HC = c.gradient_homotopy_matrix, c.curl_homotopy_matrix + P = lambda k, f: project(c, k, f, method) # noqa: E731 + pairs = ( + (1, P(1, grad_f0), G @ P(0, f0)), + (2, P(2, curl_A1), C @ P(1, A1)), + (0, P(0, poincare_gradient(A1)), HG @ P(1, A1)), + (1, P(1, poincare_curl(B2)), HC @ P(2, B2)), + ) + rows.append([_norm(c, k, a - b) / _norm(c, k, a) for k, a, b in pairs]) + defects[r] = np.array(rows) + return defects + + +def figure_convergence(errors, cells=(4, 8, 16, 32), title=None): + hs = 1 / np.array(cells) + fig, axes = plt.subplots(1, 3, figsize=(12, 4.2), sharey=True) + for k, ax in enumerate(axes): + for i, (r, e) in enumerate(errors.items()): + ax.loglog(hs, e[:, k], "o-", color=CATEGORICAL[i % len(CATEGORICAL)], markersize=4, label=f"$r = {r}$") + r_max = max(errors) + rate = r_max + 1 if k == 0 else r_max + e = errors[r_max][:, k] + ax.loglog(hs[-3:], 0.3 * e[-1] * (hs[-3:] / hs[-1]) ** rate, "k:", label=f"$O(h^{rate})$") + ax.set_xlabel("$h$") + ax.set_title(f"{k}-forms") + if k == 0: + ax.set_ylabel("relative $L^2$ error") + ax.legend(fontsize=8) + if title: + fig.suptitle(title) + fig.tight_layout() + return fig + + +def figure_commutation(defects, cells=(4, 8, 16, 32), title=None): + hs = 1 / np.array(cells) + names = (r"grad", r"curl", r"grad$^\dagger$", r"curl$^\dagger$") + fig, axes = plt.subplots(1, 4, figsize=(13, 3.8), sharey=True) + for j, ax in enumerate(axes): + for i, (r, d) in enumerate(defects.items()): + ax.loglog( + hs, + np.maximum(d[:, j], 1e-17), + "o-", + color=CATEGORICAL[i % len(CATEGORICAL)], + markersize=4, + label=f"$r = {r}$", + ) + ax.set_ylim(1e-17, 2) + ax.set_xlabel("$h$") + ax.set_title(names[j]) + if j == 0: + ax.set_ylabel("defect") + axes[-1].legend(fontsize=8) + if title: + fig.suptitle(title) + fig.tight_layout() + return fig diff --git a/src/struphy/polar/homotopy/demos/regularity.py b/src/struphy/polar/homotopy/demos/regularity.py new file mode 100644 index 000000000..92bbe857e --- /dev/null +++ b/src/struphy/polar/homotopy/demos/regularity.py @@ -0,0 +1,397 @@ +"""Smooth polar splines at the centre of a disk (port of PolarFEEC.jl ``examples/04_polar_regularity.jl``). + +A disk-like mapping collapses the edge ``x = 0`` of the logical square onto the centre, so tensor-product splines are +in general not smooth there. A polar complex of regularity ``r`` replaces the ``r + 1`` innermost rings by projected +polynomials of degree ``<= r``, so its 0-forms are ``C^r`` at the centre. A function is ``C^k`` at the centre exactly +when, for each ``j <= k``, its ``j``-th derivative along the ray at angle ``alpha`` contains only the angular harmonics +``j, j - 2, ...``. The size ``delta_j`` of the remaining harmonics is measured by polynomial fits along 48 rays. + +The defaults are those of the Julia example: the star-shaped disk ``R = x (1 + 0.2 cos(6 pi y))`` (the Struphy +domain ``PerturbedDiskCylinder(power=0)``), B-splines of degree 5 on 8 cells in ``x`` and 24 Fourier modes in ``y`` +(``angular="trig"``; ``angular="bspline"``: 24 periodic B-spline cells), regularities ``r = 0, ..., 3``. The spaces are +those of a 2D Struphy Derham complex (see :mod:`.common`). +""" + +import matplotlib.pyplot as plt +import numpy as np + +from struphy.polar.homotopy.demos import common +from struphy.polar.homotopy.demos.plotting import CATEGORICAL, DIVERGING, clean_axis +from struphy.polar.homotopy.forms2d import PhysicalFormSpace2D + +#: the default angular space; the talk uses 24 Fourier modes (``"trig"``) +ANGULAR = "trig" + + +def angular_cells(angular): + """24 periodic B-spline cells, or 24 Fourier modes (49 trigonometric functions).""" + return 24 if angular == "bspline" else 49 + + +def star_disk(): + """The cross-section of the star-shaped disk of the Julia example (radius linear in ``x``).""" + return common.cross_section(common.star_disk()) + + +def f(x, y): + return np.sin(3 * x + 1) * np.cos(2 * y - 0.5) + 0.5 * np.cos(4 * x * y) + + +def grad_f(x, y): + return ( + 3 * np.cos(3 * x + 1) * np.cos(2 * y - 0.5) - 2 * y * np.sin(4 * x * y), + -2 * np.sin(3 * x + 1) * np.sin(2 * y - 0.5) - 2 * x * np.sin(4 * x * y), + ) + + +def rays(n=48): + return 2 * np.pi * (np.arange(1, n + 1) - 0.5) / n + + +def ray_derivatives(u, alpha, a=0.06, degree=10, samples=40): + """One-sided derivatives ``d^k u(0+)``, ``k = 0..3``, along the ray at angle ``alpha`` (least-squares fit).""" + s = a * (1 - np.cos(np.pi * np.arange(1, samples + 1) / (2 * samples))) + V = (s[:, None] / a) ** np.arange(degree + 1)[None, :] + c = np.linalg.lstsq(V, u(s * np.cos(alpha), s * np.sin(alpha)), rcond=None)[0] + return np.array([_factorial(k) * c[k] / a**k for k in range(4)]) + + +def _factorial(k): + out = 1 + for i in range(2, k + 1): + out *= i + return out + + +def nonpolynomial_part(d, alpha, k): + """Max norm of the part of ``d(alpha)`` outside the harmonics ``k, k - 2, ...`` (least squares).""" + columns = [] + for m in range(k, -1, -2): + columns.append(np.cos(m * alpha)) + if m > 0: + columns.append(np.sin(m * alpha)) + H = np.column_stack(columns) + residual = d - H @ np.linalg.lstsq(H, d, rcond=None)[0] + return np.max(np.abs(residual)) + + +def smoothness_defects(u): + """Relative defects ``delta_k`` of ``d^k u(0+; alpha)`` over 48 rays, ``k = 1, 2, 3``.""" + alpha = rays() + D = np.column_stack([ray_derivatives(u, a) for a in alpha]) + return np.array([nonpolynomial_part(D[k], alpha, k) / np.max(np.abs(D[k])) for k in range(1, 4)]) + + +def _inverse_series(x, order): + """Taylor coefficients of ``eta(s)`` for the polynomial ``s = x(eta)`` with ``x[0] = 0``, up to ``s^order``.""" + from numpy.polynomial import polynomial as P + + identity = np.zeros(order + 1) + identity[1] = 1.0 + eta = identity / x[1] + for _ in range(order + 1): + eta = eta + (identity - _compose(x, eta, order)) / x[1] + return eta + + +def _compose(u, eta, order): + """Taylor coefficients of ``u(eta(s))`` up to ``s^order`` (``u``, ``eta`` coefficient arrays, ``eta[0] = 0``).""" + from numpy.polynomial import polynomial as P + + out, power = np.zeros(order + 1), np.zeros(order + 1) + power[0] = 1.0 + for c in u: + out[: power.size] += c * power[: order + 1] + power = P.polymul(power, eta)[: order + 1] + return out + + +def line_derivative_jumps(field, domain, num_cells, degree, eta3=0.0, orders=5): + """Jumps ``|d^k u/dx^k(0+) - d^k u/dx^k(0-)|``, ``k = 0, ..., orders - 1``, of a Struphy 0-form across the pole + along the line through the rays ``eta2 = 0`` and ``eta2 = 1/2``. + + On the first radial cell the field and the mapping are polynomials of degree ``degree`` in ``eta1`` along each + ray, so the one-sided derivatives in the distance ``x`` to the pole are exact (polynomial fit in ``eta1`` and + series inversion of ``x(eta1)``).""" + import math + + import cunumpy as xp + from numpy.polynomial import polynomial as P + + e1 = np.linspace(0, 1 / num_cells, 3 * degree + 3) + pole = np.array([xp.to_numpy(c) for c in domain(xp.zeros(1), xp.zeros(1), xp.asarray([eta3]))]).reshape(3) + derivatives = [] + for e2 in (0.0, 0.5): + etas = (xp.asarray(e1), xp.asarray([e2]), xp.asarray([eta3])) + points = np.array([xp.to_numpy(c).ravel() for c in domain(*etas)]) + distance = np.hypot(points[0] - pole[0], points[1] - pole[1]) * (1 if e2 == 0.0 else -1) + u = P.polyfit(e1, xp.to_numpy(field(*etas)).ravel(), degree) + x = P.polyfit(e1, distance, degree) + x[0] = 0.0 + series = _compose(u, _inverse_series(x, degree), degree) + derivatives.append([math.factorial(k) * series[k] for k in range(orders)]) + return np.abs(np.subtract(*derivatives)) + + +def field_smoothness_defects(field, domain, eta3=0.0, a=0.06, degree=10, samples=40): + """The defects ``delta_k`` of a Struphy 0-form ``field`` (a :class:`~struphy.feec.psydac_derham.SplineFunction`) + on a domain whose cross-section is parametrized by the polar angle ``2 pi eta2`` (e.g. + ``PerturbedDiskCylinder``, ``HollowCylinder(a1=0)``): the field is sampled along each ray ``eta2 = alpha / 2 pi`` + and fitted by a polynomial in the distance ``s`` to the pole.""" + import cunumpy as xp + + alpha = rays() + t = 1 - np.cos(np.pi * np.arange(1, 4 * samples + 1) / (8 * samples)) + pole = np.array([xp.to_numpy(c) for c in domain(xp.zeros(1), xp.zeros(1), xp.asarray([eta3]))]).reshape(3) + D = [] + for angle in alpha: + e1, e2, e3 = xp.asarray(t), xp.asarray([angle / (2 * np.pi)]), xp.asarray([eta3]) + points = np.array([xp.to_numpy(c).ravel() for c in domain(e1, e2, e3)]) + s = np.hypot(points[0] - pole[0], points[1] - pole[1]) + assert np.allclose(np.arctan2(points[1] - pole[1], points[0] - pole[0]) % (2 * np.pi), angle % (2 * np.pi)), ( + "the domain is not parametrized by the polar angle 2 pi eta2" + ) + values = xp.to_numpy(field(e1, e2, e3)).ravel() + inside = s <= a + V = (s[inside, None] / a) ** np.arange(degree + 1)[None, :] + c = np.linalg.lstsq(V, values[inside], rcond=None)[0] + D.append([_factorial(k) * c[k] / a**k for k in range(4)]) + D = np.array(D).T + return np.array([nonpolynomial_part(D[k], alpha, k) / np.max(np.abs(D[k])) for k in range(1, 4)]) + + +def build(angular=None, regularities=(0, 1, 2, 3), degree=5, radial_cells=8): + """Tensor-product 0-form space and the polar complexes of the study (Struphy Derham spaces on the star disk).""" + angular = ANGULAR if angular is None else angular + domain = common.star_disk() + cells = (radial_cells, angular_cells(angular)) + tensor = common.tensor_space(0, degree, cells, domain=domain, angular=angular) + polar = { + r: common.polar_complex(r, degree, cells, domain=domain, angular=angular, check_assumptions=False) + for r in regularities + } + return tensor.mapping, tensor, polar + + +def study(angular=None, regularities=(0, 1, 2, 3), verbose=True, **kwargs): + """Dimensions, cohomology and the defects of the commuting projections of ``f``; returns a dict of results.""" + mapping, tensor, polar = build(angular, regularities, **kwargs) + lines = [" r dim V0 dim V1 dim V2 central 0/1/2 cohomology"] + for r, c in polar.items(): + dims = tuple(c.dimension(k) for k in range(3)) + lines.append( + f" {r} {dims[0]:6d} {dims[1]:6d} {dims[2]:6d} {c.central_counts()} {c.cohomology_dimensions()}" + ) + spaces = [("tensor product", tensor, tensor.commuting_projection(f))] + [ + (f"r = {r}", c.spaces[0], c[0](f)) for r, c in polar.items() + ] + + def physical(space, coefficients): + return lambda X, Y: space.evaluate_physical(coefficients, X, Y) + + defects = {name: smoothness_defects(physical(space, u)) for name, space, u in spaces} + errors = {name: space.l2_error(u, f) for name, space, u in spaces} + lines.append("\nrelative defects delta_k of d^k u(0+; alpha) (C^k at the centre <=> delta_1 = ... = delta_k = 0)") + lines.append(f" {'exact f':15s} " + " ".join(f"{d:9.1e}" for d in smoothness_defects(f))) + for name in defects: + lines.append(f" {name:15s} " + " ".join(f"{d:9.1e}" for d in defects[name]) + f" {errors[name]:9.1e}") + if verbose: + print("\n".join(lines)) + return dict(mapping=mapping, tensor=tensor, polar=polar, spaces=spaces, defects=defects, errors=errors) + + +STENCILS = (([0], [1.0]), ([-1, 1], [-0.5, 0.5]), ([-1, 0, 1], [1.0, -2.0, 1.0])) + + +def ring_defects(space1, gradient_coefficients, rho): + """Non-polynomial part of ``d_s^k (u_h - f)`` on the circle ``|x| = rho`` (relative to ``d_s^k f``), ``k = 1..3``.""" + alpha = rays() + h = rho / 5 + e = np.stack([np.cos(alpha), np.sin(alpha)]) + + def radial_derivative(gradient, k): + total = 0.0 + for j, w in zip(*STENCILS[k - 1]): + X, Y = (rho + j * h) * e[0], (rho + j * h) * e[1] + gx, gy = gradient(X, Y) + total = total + w * (gx * e[0] + gy * e[1]) + return total / h ** (k - 1) + + def grad_uh(X, Y): + return space1.evaluate_physical(gradient_coefficients, X, Y) + + out = [] + for k in (1, 2, 3): + dh, df = radial_derivative(grad_uh, k), radial_derivative(grad_f, k) + out.append(nonpolynomial_part(dh - df, alpha, k) / np.max(np.abs(df))) + return np.array(out) + + +def figure_central_basis(regularity=3, results=None): + """The central 0-forms of the regularity-3 complex (ordered like the Zernike polynomials).""" + mapping = star_disk() if results is None else results["mapping"] + c = results["polar"][regularity] if results is not None else build(regularities=(regularity,))[2][regularity] + space = c.spaces[0] + central = c[0].central_elements() + ts = np.linspace(0, 1, 241) + fig = plt.figure(figsize=(12, 6)) + grid = fig.add_gridspec(regularity + 1, 2 * (regularity + 1) + 2, width_ratios=[4] * 2 + [1] * (2 * regularity + 2)) + + def panel(ax, column, rs): + values = space.evaluate(central[:, column], rs, ts) + values = values / np.max(np.abs(values)) + X, Y = np.meshgrid(rs, ts, indexing="ij") + ax.pcolormesh(*mapping(X, Y), values, cmap=DIVERGING, vmin=-1, vmax=1, shading="gouraud") + ax.plot(*mapping(np.full_like(ts, 0.5), ts), "k--", linewidth=0.8, alpha=0.6) + ax.set_aspect("equal") + clean_axis(ax) + + ax = fig.add_subplot(grid[:, :2]) + panel(ax, 0, np.linspace(0, 1, 161)) + for r in np.linspace(0, 1, c.radial.cells + 1)[1:]: + ax.plot(*mapping(np.full_like(ts, r), ts), "k", linewidth=0.6 if r < 1 else 1.2, alpha=0.25 if r < 1 else 1) + ax.set_title("$n = 0$: the central function of $T_1 = 1$") + index = 1 + for n in range(1, regularity + 1): + for j in range(n + 1): + start = 2 + (regularity - n) + 2 * j + axn = fig.add_subplot(grid[n - 1 if n > 0 else 0, start : start + 2]) + panel(axn, index, np.linspace(0, 0.5, 81)) + axn.set_title(f"$n={n}$, $m={2 * j - n}$", fontsize=10) + index += 1 + fig.suptitle(rf"central 0-forms of $\bar V^{{0,{regularity}}}$ (zoom to $\hat x < 1/2$)") + return fig + + +def figure_regularity(results, radii=None): + """``d_s^k u_h`` is continuous at the centre (defect -> 0) exactly when ``k <= r``.""" + radii = np.logspace(-4, -1.1, 30) if radii is None else radii + curves = {} + for name, space, u in results["spaces"]: + polar = results["polar"] + space1 = ( + polar[int(name.split("=")[1])].spaces[1] + if name.startswith("r") + else PhysicalFormSpace2D(space.logical.with_degree(1), space.mapping) + ) + gradient = space.logical.derivative_matrix @ u + curves[name] = np.column_stack([ring_defects(space1, gradient, rho) for rho in radii]) + styles = [ + (CATEGORICAL[0], "-", 4), + ("k", "--", 1.5), + (CATEGORICAL[1], "-", 2), + (CATEGORICAL[2], "-", 2), + (CATEGORICAL[5], "-", 2), + ] + fig, axes = plt.subplots(1, 4, figsize=(12, 4.5), gridspec_kw=dict(width_ratios=[1, 1, 1, 0.7])) + for k in range(3): + ax = axes[k] + for (name, curve), (color, linestyle, width) in zip(curves.items(), styles): + ax.loglog(radii, curve[k], color=color, linestyle=linestyle, linewidth=width, label=name) + ax.set_ylim(1e-9, 1e-2) + ax.set_xlabel(r"$\rho$") + ax.set_title(f"$k = {k + 1}$") + if k == 0: + ax.set_ylabel(r"defect of $\partial_s^k (u_h - f)$ on $|x| = \rho$") + if k == 2: + ax.legend(loc="lower right", fontsize=9) + names = list(results["defects"]) + delta = np.array([results["defects"][n] for n in names]) + ax = axes[3] + ax.imshow(delta < 1e-6, cmap="Reds_r", alpha=0.25, aspect="auto") + for i in range(delta.shape[0]): + for k in range(3): + ax.text(k, i, f"{delta[i, k]:.0e}", ha="center", va="center", fontsize=9) + ax.set_xticks(range(3), [f"{k}" for k in (1, 2, 3)]) + ax.set_yticks(range(len(names)), names) + ax.set_xlabel("$k$") + ax.set_title(r"$\delta_k$ at $\rho = 0$") + fig.suptitle(r"$\partial_s^k u_h$ is continuous at the centre (defect $\to 0$) exactly when $k \leq \mathfrak{r}$") + fig.tight_layout() + return fig + + +def bspline_angular_defects(regularity=3): + """The defects with a B-spline angular factor: regularity up to the spline approximation of the polynomials.""" + c = common.polar_complex(regularity, 5, (8, 24), domain=common.star_disk(), check_assumptions=False) + u = c[0](f) + return smoothness_defects(lambda X, Y: c.spaces[0].evaluate_physical(u, X, Y)) + + +# ---------------------------------------------------------------------------------------------------- line figures +def line_derivatives(complex_, form, coefficients, k_max, cells=2, samples=24, fit_degree=10): + """``d^k u / dx^k`` (``k <= k_max``) of the 0-form (``form=0``) or the 2-form density (``form=2``) along the line + ``y = 0`` through the pole (the rays ``eta2 = 0`` and ``1/2``), on the first ``cells`` radial cells each side. + + On each radial cell the field is fitted by a polynomial in ``x`` (exact for 0-forms on the star disk, where ``x`` + is linear in ``eta1`` along these rays). Returns ``x``, the derivative curves ``(k_max + 1, len(x))`` and the + one-sided limits ``(k_max + 1, 2)`` at the pole.""" + space = complex_.spaces[form] + mapping = space.mapping + breaks = np.linspace(0, 1, complex_.radial.cells + 1)[: cells + 1] + xs, curves, limits = [], [], np.zeros((k_max + 1, 2)) + for side, eta2 in ((0, 0.5), (1, 0.0)): + pieces_x, pieces_d = [], [] + for c in range(cells): + t = breaks[c] + (breaks[c + 1] - breaks[c]) * (1 - np.cos(np.pi * (np.arange(samples) + 0.5) / samples)) / 2 + X, Y = mapping(t, np.full_like(t, eta2)) + x = X if side else X # signed distance along y = 0 (X < 0 on the ray eta2 = 1/2) + u = np.asarray(space.evaluate_points(coefficients, t, np.full_like(t, eta2))) + fit = np.polynomial.Polynomial.fit(x, u, min(fit_degree, samples - 1)) + grid = np.linspace(x.min(), x.max(), 60) + pieces_x.append(grid) + pieces_d.append(np.array([fit.deriv(k)(grid) if k else fit(grid) for k in range(k_max + 1)])) + if c == 0: + x0 = 0.0 + limits[:, side] = [fit.deriv(k)(x0) if k else fit(x0) for k in range(k_max + 1)] + order = np.argsort(np.concatenate(pieces_x)) + xs.append(np.concatenate(pieces_x)[order]) + curves.append(np.concatenate(pieces_d, axis=1)[:, order]) + x = np.concatenate([xs[0], [np.nan], xs[1]]) + curve = np.concatenate([curves[0], np.full((k_max + 1, 1), np.nan), curves[1]], axis=1) + return x, curve, limits + + +def random_polar_element(complex_, form, seed=3): + """Full coefficients of a random element of the polar space of degree ``form``.""" + rng = np.random.default_rng(seed) + E = complex_.extraction(form) + return E @ rng.standard_normal(E.shape[1]) + + +def figure_regularity_lines(form=0, regularities=None, orders=None, angular=None, degree=5, radial_cells=8, seed=3): + """The talk's regularity figures: ``d^k`` of a random polar ``form``-form (0-forms: ``r = 0..3``, ``k = 1..4``; + 2-form densities: ``r = 1..4``, ``k = 0..3``) along a line through the pole; grey panels jump at the pole.""" + angular = ANGULAR if angular is None else angular + regularities = (0, 1, 2, 3) if form == 0 else (1, 2, 3, 4) if regularities is None else regularities + regularities = regularities if regularities is not None else (0, 1, 2, 3) + orders = (1, 2, 3, 4) if form == 0 else (0, 1, 2, 3) if orders is None else orders + cells = (radial_cells, angular_cells(angular)) + smooth_up_to = (lambda r: r) if form == 0 else (lambda r: r - 2) + fig, axes = plt.subplots(len(orders), len(regularities), figsize=(2.3 * len(regularities), 1.6 * len(orders))) + for col, r in enumerate(regularities): + c = common.polar_complex(r, degree, cells, domain=common.star_disk(), angular=angular, check_assumptions=False) + x, curves, limits = line_derivatives(c, form, random_polar_element(c, form, seed), max(orders)) + color = CATEGORICAL[col % len(CATEGORICAL)] + for row, k in enumerate(orders): + ax = axes[row, col] + jump = k > smooth_up_to(r) + ax.set_facecolor("0.94" if jump else "white") + ax.plot(x, curves[k], color=color, linewidth=1.6) + ax.plot([0, 0], limits[k], "o", color=color, markersize=4) + ax.axvline(0, color="0.7", linewidth=0.8) + for b in np.linspace(0, 1, radial_cells + 1)[1:3]: + for sign_x in (np.nanmax(x) / (2 / radial_cells), np.nanmin(x) / (2 / radial_cells)): + ax.axvline(b * sign_x, color="0.7", linestyle=":", linewidth=0.8) + ax.set_xticks([]) + ax.set_yticks([]) + if row == 0: + ax.set_title(f"$r = {r}$", color=color, fontsize=13) + if col == 0: + ax.set_ylabel(f"$k = {k}$", rotation=0, labelpad=22, fontsize=12) + if row == len(orders) - 1: + ax.set_xlabel("$x = 0$") + name = "0-form $u$" if form == 0 else "2-form density" + fig.suptitle(rf"$d^k/dx^k$ of a random polar {name} along $y = 0$ (grey: discontinuous at the pole)", fontsize=11) + fig.tight_layout() + return fig diff --git a/src/struphy/polar/homotopy/demos/schematics.py b/src/struphy/polar/homotopy/demos/schematics.py new file mode 100644 index 000000000..2c3540f24 --- /dev/null +++ b/src/struphy/polar/homotopy/demos/schematics.py @@ -0,0 +1,221 @@ +"""The schematic diagrams of the talk: de Rham diagrams (drawn with ``tikzfigure``, needs LaTeX), the 3D coordinate +transformation of the helical tube and the steps of the boundary-compatible homotopy (computed with the Struphy +domains and polar complexes, drawn with matplotlib).""" + +import matplotlib.pyplot as plt +import numpy as np + +from struphy.polar.homotopy.demos import common +from struphy.polar.homotopy.demos.plotting import DIVERGING, clean_axis, draw_mesh + +# ---------------------------------------------------------------------------------------------------------- diagrams +ARROW = "-{Latex}" +LABEL = r"font=\scriptsize" + + +def _figure(): + from tikzfigure import TikzFigure + from tikzfigure.core.matrix import Matrix + + if not hasattr(Matrix, "arrow"): + raise ImportError("the diagrams need tikzfigure >= 0.5.0 (Matrix.arrow): pip install -U tikzfigure") + fig = TikzFigure() + fig.add_package("amssymb") + fig.usetikzlibrary("arrows.meta") + fig.colorlet("struphyblue", "rgb,1:red,0;green,0.447;blue,0.741") + return fig + + +def _diagram(fig, rows, maps, projections, label, x=0, reverse=False): + """A commutative diagram (a ``tikzfigure`` matrix of math nodes): ``maps`` label the arrows between neighbouring + columns in every row (right to left with ``reverse``), ``projections`` the arrows from the last to the first row.""" + matrix = fig.add_matrix( + rows, + x=x, + y=0, + label=label, + options="matrix of math nodes", + row_sep="1.25cm", + column_sep="1.7cm", + cell_style="text height=1.6ex, text depth=0.4ex", + ) + for row in range(1, len(rows) + 1): + for col, name in enumerate(maps, start=1): + ends = ((row, col + 1), (row, col)) if reverse else ((row, col), (row, col + 1)) + matrix.arrow(*ends, label=name, label_options=LABEL, arrows=ARROW) + for col, name in enumerate(projections, start=1): + matrix.arrow( + (len(rows), col), + (1, col), + label=f"${name}$", + label_position="left", + label_options=LABEL, + arrows=ARROW, + color="struphyblue", + ) + return matrix + + +def derham_logical_and_physical(): + """The commuting projections of the logical and of the physical 2D de Rham complex.""" + fig = _figure() + _diagram( + fig, + [[r"\hat V^0_h", r"\hat V^1_h", r"\hat V^2_h"], [r"\hat H^1", r"\hat H(\mathrm{curl})", r"\hat L^2"]], + ("grad", "curl"), + (r"\hat\pi^0", r"\hat\pi^1", r"\hat\pi^2"), + "l", + ) + physical = _diagram( + fig, + [["V^0_h", "V^1_h", "V^2_h"], ["H^1", r"H(\mathrm{curl})", "L^2"]], + ("grad", "curl"), + (r"\pi^0", r"\pi^1", r"\pi^2"), + "p", + x=9, + ) + # an arrow is drawn after its matrix: the arrow between the two belongs to the second + physical.arrow( + "l.east", + "p.west", + label="pushforward $F_*$", + label_options=LABEL, + arrows=ARROW, + options=["dashed", "shorten >=8mm"], + ) + return fig + + +def derham_polar_cochain_chain(): + """The polar subspaces form a cochain complex (grad, curl) and a chain complex (the homotopies) with commuting + projections ``pi_bar^{k, r}``.""" + rows = [ + [r"\bar V^{0,\mathfrak r}", r"\bar V^{1,\mathfrak r}", r"\bar V^{2,\mathfrak r}"], + ["H^1", r"H(\mathrm{curl})", "L^2"], + ] + projections = (r"\bar\pi^{0,\mathfrak r}", r"\bar\pi^{1,\mathfrak r}", r"\bar\pi^{2,\mathfrak r}") + fig = _figure() + _diagram(fig, rows, ("grad", "curl"), projections, "c") + homotopies = (r"$\mathrm{grad}^\dagger$", r"$\mathrm{curl}^\dagger$") + _diagram(fig, rows, homotopies, projections, "h", x=9, reverse=True) + return fig + + +# ---------------------------------------------------------------------------------------------------------- geometry +def coordinate_transformation_3d(n=(8, 16, 48)): + """The logical cube (the face ``eta1 = 0`` collapses to the magnetic axis) and the helical tube with two + cross-sections, the mesh lines and the axis.""" + from struphy.geometry import domains + + domain = domains.HelicalTube("stellarator", num_elements=(8, 24, 64), degree=(4, 4, 4)) + fig = plt.figure(figsize=(11, 4.2)) + ax = fig.add_subplot(1, 2, 1, projection="3d") + for t in np.linspace(0, 1, 5): + for s in np.linspace(0, 1, 5): + ax.plot([t, t], [s, s], [0, 1], color="0.6", linewidth=0.5) + ax.plot([t, t], [0, 1], [s, s], color="0.6", linewidth=0.5) + ax.plot([0, 1], [t, t], [s, s], color="0.6", linewidth=0.5) + face = np.linspace(0, 1, 2) + Y, Z = np.meshgrid(face, face) + ax.plot_surface(np.zeros_like(Y), Y, Z, color="#E69F00", alpha=0.5) + ax.set_xlabel(r"$\eta_1$") + ax.set_ylabel(r"$\eta_2$") + ax.set_zlabel(r"$\eta_3$") + ax.set_title(r"logical cube ($\eta_1 = 0$: the axis)") + ax = fig.add_subplot(1, 2, 2, projection="3d") + e1, e2, e3 = np.array([1.0]), np.linspace(0, 1, 49), np.linspace(0, 0.8, 120) + X, Y, Z = (np.asarray(c)[0] for c in domain(e1, e2, e3)) + ax.plot_surface(X, Y, Z, color="0.85", alpha=0.35, linewidth=0) + for z in np.linspace(0, 0.8, 9): + ring = np.asarray(domain(e1, e2, np.array([z])))[:, 0, :, 0] + ax.plot(*ring, color="0.4", linewidth=0.5) + for z in (0.0, 0.8): + r, t = np.linspace(0, 1, 9), np.linspace(0, 1, 49) + S = np.asarray(domain(r, t, np.array([z])))[:, :, :, 0] + ax.plot_surface(*S, color="0.95", edgecolor="0.4", linewidth=0.3) + axis = np.asarray(domain(np.array([0.0]), np.array([0.0]), e3))[:, 0, 0, :] + ax.plot(*axis, color="#E69F00", linewidth=2) + ax.set_box_aspect((1, 1, 0.6)) + ax.set_axis_off() + ax.set_title("HelicalTube (magnetic axis in orange)") + fig.tight_layout() + return fig + + +def boundary_homotopy(regularity=3): + """The boundary-compatible homotopy ``curl^dagger_partial`` of a zero-mean 2-form ``B``: the radial homotopy + ``curl^dagger B`` has a tangential trace on the boundary; its periodic primitive, extended into the domain, is + subtracted as a gradient, which leaves a potential ``A`` with ``curl A = B`` and ``A x n = 0``.""" + c = common.polar_complex(regularity, 4, (6, 18), boundary="dirichlet") + mapping = c.spaces[0].mapping + + def B(x, y): + return np.sin(3 * x + 1) * np.cos(2 * y) + x * y + + b = c[2](B) + # zero mean in polar coefficients, then back to full coefficients + polar = c.duals(2).T @ b + polar = polar - np.asarray(c.cohomology_projection(2) @ polar).ravel() + b0 = c.extraction(2) @ polar + radial = c.curl_homotopy_matrix @ b0 + bounded = c.extraction(1) @ np.asarray(c.homotopy(2) @ polar).ravel() + rs, ts = np.linspace(1e-3, 1, 81), np.linspace(0, 1, 181) + X, Y = mapping(*np.meshgrid(rs, ts, indexing="ij")) + panels = [ + (c.spaces[2].evaluate(b0, rs, ts), r"$B$ (zero mean)"), + ( + np.hypot(*c.spaces[1].evaluate(radial, rs, ts)), + r"$|\mathrm{curl}^\dagger B|$: tangential trace on $\partial\Omega$", + ), + ( + np.hypot(*c.spaces[1].evaluate(radial - bounded, rs, ts)), + r"$|\mathrm{grad}\,\phi|$: boundary primitive (outer ring)", + ), + (np.hypot(*c.spaces[1].evaluate(bounded, rs, ts)), r"$|\mathrm{curl}^\dagger_\partial B|$: $A \times n = 0$"), + ] + fig, axes = plt.subplots(1, 4, figsize=(14, 3.8)) + # one colour range for the potentials (the extension of the boundary primitive lives in the outer ring and + # saturates there) + common_limit = np.max(panels[1][0]) or 1.0 + for index, (ax, (values, title)) in enumerate(zip(axes, panels)): + signed = np.min(values) < 0 + limit = (np.max(np.abs(values)) or 1.0) if index == 0 else common_limit + ax.pcolormesh( + X, + Y, + values, + cmap=DIVERGING if signed else "viridis", + vmin=-limit if signed else 0, + vmax=limit, + shading="gouraud", + ) + draw_mesh(ax, mapping, np.linspace(0, 1, 7), np.linspace(0, 1, 19), alpha=0.15) + ax.set_aspect("equal") + clean_axis(ax) + ax.set_title(title, fontsize=10) + fig.tight_layout() + return fig + + +def write_all(folder): + """Write all schematics into ``folder`` (``.pdf`` for the TikZ diagrams and ``.png``).""" + from pathlib import Path + + folder = Path(folder) + folder.mkdir(parents=True, exist_ok=True) + for name, make in ( + ("derham_logical_and_physical_2d", derham_logical_and_physical), + ("derham_polar_cochain_chain_2d", derham_polar_cochain_chain), + ): + fig = make() + for ext in ("pdf", "png"): + fig.savefig(folder / f"{name}.{ext}") + print("wrote", folder / f"{name}.png") + for name, make in ( + ("coordinate_transformation_3d", coordinate_transformation_3d), + ("bdy_homotopy", boundary_homotopy), + ): + fig = make() + fig.savefig(folder / f"{name}.png", dpi=150, bbox_inches="tight") + plt.close(fig) + print("wrote", folder / f"{name}.png") diff --git a/src/struphy/polar/homotopy/demos/spaces.py b/src/struphy/polar/homotopy/demos/spaces.py new file mode 100644 index 000000000..f7305ae40 --- /dev/null +++ b/src/struphy/polar/homotopy/demos/spaces.py @@ -0,0 +1,266 @@ +"""1D spaces, tensor-product forms and mapped forms (ports of PolarFEEC.jl ``examples/01_spaces_1d.jl``, +``02_commuting_projections.jl`` and ``03_mapped_forms.jl``). + +* 01: clamped and periodic B-splines and trigonometric functions with their derivative-scaled companions; the exact + coefficient derivative ``d/dx B_i = C_{i-1} - C_i`` and the commuting 1D diagram ``d pi^0 f = pi^1 f'``. +* 02: 0-, 1- and 2-forms on the unit square (clamped x periodic), ``curl o grad = 0`` exactly, projections that + commute with the derivatives to round-off; commuting and L2 projections converge at order ``p + 1`` for 0-forms and + ``p`` for 1- and 2-forms. +* 03: a quarter annulus (``HollowCylinder(a1=1, a2=2, poc=4)``; the Julia example warps it); physical 0-, 1- and + 2-forms (scalar, covariant Piola and density pushforwards) with commuting projections + ``grad pi^0 phi = pi^1 grad phi`` and ``curl pi^1 u = pi^2 curl u``. + +The 1D spaces are the spline spaces of Struphy's Derham complex (feectools ``SplineSpace``) and the trigonometric +spaces of :mod:`struphy.feec.trig_space`; the square and the annulus are Struphy domains. +""" + +import matplotlib.pyplot as plt +import numpy as np + +from struphy.polar.homotopy.demos import common +from struphy.polar.homotopy.demos.plotting import CATEGORICAL, DIVERGING +from struphy.polar.homotopy.forms2d import PhysicalFormSpace2D +from struphy.polar.homotopy.layout import STRUPHY + + +def complex_step(f, x, *args, argnum=0, h=1e-30): + """Derivative of the (analytic) function ``f`` in its argument ``argnum`` by the complex step.""" + values = [np.asarray(x, dtype=complex), *[np.asarray(a, dtype=complex) for a in args]] + values[argnum] = values[argnum] + 1j * h + return np.imag(f(*values)) / h + + +# ---------------------------------------------------------------------------------------------------- 01 +def spaces_1d(verbose=True): + """Commuting-diagram residuals ``|d pi^0 f - pi^1 f'|`` of the three 1D families.""" + families = { + "clamped B-splines (3, 6)": (common.spline_space_1d(3, 6), lambda x: np.exp(-x) * np.sin(4 * np.pi * x)), + "periodic B-splines (3, 6)": (common.spline_space_1d(3, 6, True), lambda x: np.exp(np.sin(2 * np.pi * x))), + "trigonometric (2 modes)": (common.trigonometric_space_1d(5), lambda x: np.exp(np.sin(2 * np.pi * x))), + } + residuals = {} + for name, (space, f) in families.items(): + u = space.commuting_projection(f) + derivative = space.scaled.commuting_projection(lambda x, f=f: complex_step(f, x), additional_order=8) + residuals[name] = float(np.max(np.abs(space.derivative_matrix() @ u - derivative))) + if verbose: + print(f"{name:28s} |d pi0 f - pi1 f'| = {residuals[name]:.1e}") + return families, residuals + + +def figure_spaces_1d(): + families, _ = spaces_1d(verbose=False) + x = np.linspace(0, 1, 801) + fig, axes = plt.subplots(2, 3, figsize=(12, 6), sharex=True) + for j, (name, (space, f)) in enumerate(families.items()): + for k, s in enumerate((space, space.scaled)): + values = s.basis_matrix(x).toarray() + for i in range(values.shape[1]): + axes[k, j].plot(x, values[:, i], color=CATEGORICAL[i % len(CATEGORICAL)], lw=1.2) + for p in s.greville_points(): + axes[k, j].axvline(p, color="gray", lw=0.5, ls=":") + axes[k, j].set_title(("$V^0$: " if k == 0 else "$V^1$ (scaled): ") + name, fontsize=10) + for ax in axes[1]: + ax.set_xlabel("$x$") + fig.tight_layout() + return fig + + +# ---------------------------------------------------------------------------------------------------- 02 +def identity(): + """The unit square (the Struphy domain ``Cuboid``).""" + from struphy.geometry import domains + from struphy.polar.homotopy.derham_blocks import DomainCrossSection + + return DomainCrossSection(domains.Cuboid(), centered=False) + + +def _g(x): + return np.exp(-(((x - 0.5) / 0.3) ** 2)) + + +def psi(x, y): + return _g(x) * np.sin(2 * np.pi * (y - x / 2)) + + +def A_field(x, y): + """``A = (-d_y psi, d_x psi)``.""" + return (-complex_step(psi, x, y, argnum=1), complex_step(psi, x, y, argnum=0)) + + +def curl_A(x, y): + """``curl A = Laplace psi`` (analytic).""" + g = _g(x) + dg = -2 * (x - 0.5) / 0.09 * g + ddg = (4 * (x - 0.5) ** 2 / 0.09**2 - 2 / 0.09) * g + phase = 2 * np.pi * (y - x / 2) + S, C = np.sin(phase), np.cos(phase) + return ddg * S + 2 * dg * (-np.pi * C) + g * (-(np.pi**2) * S) + g * (-4 * np.pi**2 * S) + + +def tensor_spaces(p, n): + return [ + PhysicalFormSpace2D.from_factors( + k, common.spline_space_1d(p, n), common.spline_space_1d(p, n, True), identity(), STRUPHY + ) + for k in range(3) + ] + + +def identities(p=3, n=8, verbose=True): + """``curl o grad = 0`` and the commutation residuals of the projections on the square.""" + V0, V1, V2 = tensor_spaces(p, n) + G, C = V0.derivative_matrix, V1.derivative_matrix + curl_grad = float(abs(C @ G).max()) + grad_residual = float( + np.max( + np.abs( + G @ V0.commuting_projection(psi) + - V1.commuting_projection( + lambda x, y: (complex_step(psi, x, y, argnum=0), complex_step(psi, x, y, argnum=1)), + additional_order=4, + ) + ) + ) + ) + curl_residual = float( + np.max( + np.abs( + C @ V1.commuting_projection(A_field, additional_order=4) + - V2.commuting_projection(curl_A, additional_order=4) + ) + ) + ) + if verbose: + print( + f"|curl grad| = {curl_grad:.1e}, |grad pi0 psi - pi1 grad psi| = {grad_residual:.1e}, " + f"|curl pi1 A - pi2 curl A| = {curl_residual:.1e}" + ) + return curl_grad, grad_residual, curl_residual + + +def convergence(degrees=(1, 2, 3, 4, 5), cells=(8, 16, 32, 64, 128), verbose=True): + """L2 errors of the commuting and L2 projections of psi, A and curl A, ``errors[method][p]`` (cells x 3).""" + targets = (psi, A_field, curl_A) + errors = {"commuting": {}, "l2": {}} + for p in degrees: + rows = {"commuting": [], "l2": []} + for n in cells: + spaces = tensor_spaces(p, n) + for method in rows: + row = [] + for k, (space, f) in enumerate(zip(spaces, targets)): + u = space.commuting_projection(f, 4) if method == "commuting" else space.l2_projection(f) + row.append(space.l2_error(u, f, 4)) + rows[method].append(row) + for method in rows: + errors[method][p] = np.array(rows[method]) + if verbose: + e = errors["commuting"][p] + print(f"p = {p}: observed rates (commuting) {np.log2(e[-2] / e[-1]).round(2)}") + return errors + + +def figure_convergence(errors, cells=(8, 16, 32, 64, 128)): + hs = 1 / np.array(cells) + fig, axes = plt.subplots(1, 3, figsize=(12, 4.7), sharey=True) + for k, ax in enumerate(axes): + for i, p in enumerate(errors["commuting"]): + color = CATEGORICAL[i % len(CATEGORICAL)] + ax.loglog(hs, errors["commuting"][p][:, k], "o-", color=color, label=f"$p = {p}$", ms=4) + ax.loglog(hs, errors["l2"][p][:, k], "s", mfc="none", color=color, ms=5) + rate = p + 1 if k == 0 else p + e = errors["commuting"][p][:, k] + ax.loglog(hs[-3:], 0.3 * e[-1] * (hs[-3:] / hs[-1]) ** rate, ":", color=color) + ax.set_ylim(1e-13, 10) + ax.set_xlabel("$h$") + ax.set_title([r"0-form $\psi$", "1-form $A$", r"2-form curl $A$"][k]) + axes[0].set_ylabel(r"$L^2$ error (circles: commuting, squares: $L^2$ projection)") + axes[0].legend(fontsize=8) + fig.tight_layout() + return fig + + +# ---------------------------------------------------------------------------------------------------- 03 +def quarter_annulus(): + """The quarter annulus ``1 < r < 2`` (the Struphy domain ``HollowCylinder(a1=1, a2=2, poc=4)``).""" + from struphy.geometry import domains + from struphy.polar.homotopy.derham_blocks import DomainCrossSection + + return DomainCrossSection(domains.HollowCylinder(a1=1.0, a2=2.0, Lz=1.0, poc=4), centered=False) + + +def vortices(x, y): + """Two Gaussian vortices (strengths 1 and -0.8, width 0.4).""" + out = 0.0 + for (cx, cy), s in (((1.45, 0.55), 1.0), ((0.6, 1.45), -0.8)): + out = out + s * np.exp(-((x - cx) ** 2 + (y - cy) ** 2) / 0.4**2) + return out + + +def mapped_spaces(n=8): + mapping = quarter_annulus() + return [ + PhysicalFormSpace2D.from_factors( + k, common.spline_space_1d(3, n), common.spline_space_1d(3, 2 * n), mapping, STRUPHY + ) + for k in range(3) + ] + + +def mapped_identities(n=8, verbose=True): + V0, V1, V2 = mapped_spaces(n) + + def grad_phi(x, y): + return (complex_step(vortices, x, y, argnum=0), complex_step(vortices, x, y, argnum=1)) + + def u(x, y): + gx, gy = grad_phi(x, y) + return (-gy + 0.3 * y, gx + 0.2 * x**2) + + def curl_u(x, y): + h = 1e-5 + dvx = (u(x + h, y)[1] - u(x - h, y)[1]) / (2 * h) + duy = (u(x, y + h)[0] - u(x, y - h)[0]) / (2 * h) + return dvx - duy + + o = 8 + grad_residual = float( + np.max( + np.abs(V0.derivative_matrix @ V0.commuting_projection(vortices, o) - V1.commuting_projection(grad_phi, o)) + ) + ) + curl_residual = float( + np.max(np.abs(V1.derivative_matrix @ V1.commuting_projection(u, o) - V2.commuting_projection(curl_u, o))) + ) + if verbose: + print( + f"mapped: |grad pi0 phi - pi1 grad phi| = {grad_residual:.1e}, |curl pi1 u - pi2 curl u| = " + f"{curl_residual:.1e} (finite differences in curl u)" + ) + return dict(spaces=(V0, V1, V2), phi=V0.commuting_projection(vortices, o), u=V1.commuting_projection(u, o)) + + +def figure_mapped_forms(n=8): + result = mapped_identities(n, verbose=False) + V0, V1, V2 = result["spaces"] + mapping = V0.mapping + xs, ys = np.linspace(0, 1, 81), np.linspace(0, 1, 161) + X, Y = mapping(*np.meshgrid(xs, ys, indexing="ij")) + fig, axes = plt.subplots(1, 3, figsize=(12, 4.6)) + phi = V0.evaluate(result["phi"], xs, ys) + axes[0].pcolormesh(X, Y, phi, cmap=DIVERGING, vmin=-np.abs(phi).max(), vmax=np.abs(phi).max(), shading="gouraud") + axes[0].contour(X, Y, phi, 12, colors="k", linewidths=0.5) + axes[0].set_title(r"$\phi_h$ (0-form)") + qx, qy = np.linspace(0.05, 0.95, 9), np.linspace(0.02, 0.98, 17) + Qx, Qy = np.meshgrid(qx, qy, indexing="ij") + ux, uy = V1.evaluate_points(result["u"], Qx, Qy) + axes[1].quiver(*mapping(Qx, Qy), ux, uy, np.hypot(ux, uy), cmap="viridis") + axes[1].set_title(r"$u_h$ (1-form, covariant Piola)") + curl = V2.evaluate(V1.derivative_matrix @ result["u"], xs, ys) + axes[2].pcolormesh(X, Y, curl, cmap=DIVERGING, vmin=-np.abs(curl).max(), vmax=np.abs(curl).max(), shading="gouraud") + axes[2].set_title(r"curl $u_h$ (2-form)") + for ax in axes: + ax.set_aspect("equal") + fig.tight_layout() + return fig diff --git a/src/struphy/polar/homotopy/demos/spectrum.py b/src/struphy/polar/homotopy/demos/spectrum.py new file mode 100644 index 000000000..904e5d4ca --- /dev/null +++ b/src/struphy/polar/homotopy/demos/spectrum.py @@ -0,0 +1,472 @@ +"""Spectrum explorer of the 2D curl-curl operator on the unit disk (the interactive slide of the talk). + +For the polar 1-form spaces with Dirichlet conditions of regularity ``r = 0, ..., 5`` and the enrichments +``"cochain_chain"`` (the paper's spaces) and ``"cochain"`` (the reduced spaces), the eigenproblem +``(curl A, curl B) = lambda (A, B)`` is solved on a 2D Struphy Derham complex (unit disk ``HollowCylinder(a1=0)``, +B-splines of degree 6 on 10 x 14 cells, as the talk). For every eigenpair the explorer shows + +* the discrete eigenvalue ``lambda_h`` and the exact one ``(j'_{m,k})^2`` (squared zeros of ``J_m'``), +* the relative eigenvalue error and the *regularity deficit* ``||x - Pi^6 x|| / ||x||`` of the eigenfunction ``x`` + (``Pi^6`` the L2 projection onto the polar 1-forms of regularity 6, L2 norms), +* images of ``A_x``, ``A_y`` and ``curl A`` on the disk and around the centre. + +``write_explorer(folder)`` writes ``spectrum_explorer.html`` with ``spectrum_data.js`` and the images into ``folder``; +open the HTML file in a browser. +""" + +import json +from pathlib import Path + +import numpy as np +import scipy.linalg as sla + +from struphy.polar.homotopy.demos import common +from struphy.polar.homotopy.demos.maxwell import exact_spectrum +from struphy.polar.homotopy.demos.plotting import DIVERGING + +DEGREE = 6 +CELLS = (10, 14) +REGULARITIES = (0, 1, 2, 3, 4, 5) +REFERENCE_REGULARITY = 6 +ENRICHMENTS = ("cochain_chain", "cochain") +ZOOM = 0.25 +FIELDS = ("A_x", "A_y", "curl A") + + +def complex_(regularity, enrichment="cochain_chain", degree=DEGREE, cells=CELLS): + return common.polar_complex( + regularity, + degree, + cells, + domain=common.unit_disk(), + boundary="dirichlet", + enrichment=enrichment, + check_assumptions=False, + ) + + +def series(regularity, enrichment="cochain_chain", reference=None, count=None, degree=DEGREE, cells=CELLS): + """Nonzero eigenpairs of ``C^T M_2 C a = lambda M_1 a`` (ascending), with the matched exact eigenvalues and the + regularity deficits; the eigenvectors are returned as full tensor-product coefficients.""" + c = complex_(regularity, enrichment, degree, cells) + C = c.derivative(1) + K = (C.T @ c.mass_matrix(2) @ C).toarray() + lam, a = sla.eigh((K + K.T) / 2, c.mass_matrix(1).toarray()) + kernel = int(np.sum(lam < 1e-8 * lam.max())) + lam, a = lam[kernel:], a[:, kernel:] + x = c.extraction(1) @ a + count = len(lam) if count is None else min(count, len(lam)) + exact = exact_spectrum(lambda_max=1.5 * lam[count - 1]) + deficit = np.full(count, np.nan) + if reference is not None: + # L2 projection onto the polar space of the reference regularity, relative L2 norm (as the talk) + import scipy.sparse.linalg as spla + + M = c.full_mass_matrix(1) + E = reference.extraction(1) + xs = x[:, :count] + projected = E @ spla.splu((E.T @ M @ E).tocsc()).solve(E.T @ (M @ xs)) + difference = xs - projected + deficit = np.sqrt(np.einsum("ij,ij->j", difference, M @ difference) / np.einsum("ij,ij->j", xs, M @ xs)) + return dict( + complex=c, + regularity=regularity, + enrichment=enrichment, + dimension=c.dimension(1), + kernel=kernel, + lam=lam[:count], + vectors=x[:, :count], + exact=[e[0] if n < len(exact) else None for n, e in enumerate(exact[:count])] + [None] * (count - len(exact)), + m=[e[1] for e in exact[:count]] + [None] * (count - len(exact)), + k=[e[2] for e in exact[:count]] + [None] * (count - len(exact)), + deficit=deficit, + ) + + +# ---------------------------------------------------------------------------------------------------------- images +class ModeImages: + """Rasters of ``A_x``, ``A_y`` (physical components) and ``curl A`` (density) of 1-form coefficients on the full + disk and around the centre, by precomputed evaluation matrices.""" + + def __init__(self, complex_, pixels=96): + self.pixels = pixels + mapping = complex_.spaces[1].mapping + logical1, logical2 = complex_.spaces[1].logical, complex_.spaces[2].logical + self.views = [] + for extent in (1.0, ZOOM): + t = (np.arange(pixels) + 0.5) / pixels * 2 - 1 + X, Y = np.meshgrid(extent * t, -extent * t, indexing="xy") + inside = np.hypot(X, Y) < 1.0 + x, y = mapping.inverse(X[inside], Y[inside]) + J = mapping.jacobian_matrix(np.maximum(x, 1e-12), y) + det = J[..., 0, 0] * J[..., 1, 1] - J[..., 0, 1] * J[..., 1, 0] + G1 = [self._evaluation(logical1, c, x, y) for c in range(2)] + G2 = self._evaluation(logical2, 0, x, y) + self.views.append(dict(inside=inside, J=J, det=det, G1=G1, G2=G2)) + self.curl = complex_.spaces[1].derivative_matrix + + @staticmethod + def _evaluation(logical, component, x, y): + c = logical.components[component] + Bx, By = c.fx.basis_matrix(x).toarray(), c.fy.basis_matrix(y).toarray() + G = (Bx[:, :, None] * By[:, None, :]).reshape(len(x), -1) # Struphy ordering i * ny + j + offset = logical.offsets[component] + full = np.zeros((len(x), logical.n)) + full[:, offset : offset + G.shape[1]] = G + return full + + def fields(self, coefficients): + """Per view, the three fields on the pixels (NaN outside the disk).""" + curl = self.curl @ coefficients + out = [] + for v in self.views: + ax, ay = v["G1"][0] @ coefficients, v["G1"][1] @ coefficients + J, det = v["J"], v["det"] + # covariant pushforward J^{-T} (a_x, a_y) + px = (J[..., 1, 1] * ax - J[..., 1, 0] * ay) / det + py = (-J[..., 0, 1] * ax + J[..., 0, 0] * ay) / det + b = (v["G2"] @ curl) / det + images = [] + for values in (px, py, b): + image = np.full(v["inside"].shape, np.nan) + image[v["inside"]] = values + images.append(image) + out.append(images) + return out + + def sprite(self, coefficients): + """RGB sprite: columns ``A_x, A_y, curl A``, rows full disk and centre zoom (own symmetric scale per tile).""" + tiles = self.fields(coefficients) + p = self.pixels + sprite = np.full((2 * p, 3 * p, 3), 255, dtype=np.uint8) + for row, images in enumerate(tiles): + for col, image in enumerate(images): + limit = np.nanmax(np.abs(image)) or 1.0 + rgba = DIVERGING(0.5 + 0.5 * np.nan_to_num(image / limit)) + rgb = (255 * rgba[..., :3]).astype(np.uint8) + rgb[np.isnan(image)] = 255 + sprite[row * p : (row + 1) * p, col * p : (col + 1) * p] = rgb + return sprite + + +# ---------------------------------------------------------------------------------------------------------- output +def compute(count=None, regularities=REGULARITIES, enrichments=ENRICHMENTS, degree=DEGREE, cells=CELLS, verbose=True): + """The spectra of all spaces (``count`` eigenpairs each, default: all).""" + data = {} + for enrichment in enrichments: + reference = complex_(REFERENCE_REGULARITY, "cochain_chain", degree, cells) + data[enrichment] = [] + for r in regularities: + s = series(r, enrichment, reference, count, degree, cells) + data[enrichment].append(s) + if verbose: + lam_exact = np.array([e if e is not None else np.nan for e in s["exact"]]) + error = np.abs(s["lam"] - lam_exact) / lam_exact + print( + f"{enrichment:13s} r = {r}: dim V1 = {s['dimension']}, kernel {s['kernel']}, " + f"lambda_1 error {error[0]:.1e}, median deficit {np.nanmedian(s['deficit']):.1e}" + ) + return data + + +def write_explorer(folder, count=120, pixels=96, regularities=REGULARITIES, degree=DEGREE, cells=CELLS, verbose=True): + """Compute the spectra and write the explorer (HTML, data and images) into ``folder``.""" + from PIL import Image + + folder = Path(folder) + images = folder / "spectrum" + images.mkdir(parents=True, exist_ok=True) + data = compute(count, regularities, degree=degree, cells=cells, verbose=verbose) + spaces = [] + for enrichment, list_ in data.items(): + series_out = [] + for s in list_: + raster = ModeImages(s["complex"], pixels) + names = [] + for n in range(len(s["lam"])): + name = f"{enrichment}_r{s['regularity']}_{n + 1:03d}.png" + Image.fromarray(raster.sprite(s["vectors"][:, n])).convert("P", palette=Image.ADAPTIVE).save( + images / name, optimize=True + ) + names.append(name) + series_out.append( + dict( + r=s["regularity"], + dimension=s["dimension"], + **{"lambda": [float(v) for v in s["lam"]]}, + exact=[None if e is None else float(e) for e in s["exact"]], + m=s["m"], + k=s["k"], + deficit=[None if not np.isfinite(d) else float(d) for d in s["deficit"]], + image=names, + ) + ) + spaces.append(dict(enrichment=enrichment, series=series_out)) + meta = dict( + degree=degree, + radialCells=cells[0], + angularCells=cells[1], + boundary="dirichlet", + referenceRegularity=REFERENCE_REGULARITY, + zoom=ZOOM, + tile=pixels, + fields=list(FIELDS), + domain="HollowCylinder(a1=0, a2=1) (unit disk)", + ) + exact = [dict(k=k, m=m, **{"lambda": float(lam)}) for lam, m, k in exact_spectrum(lambda_max=2000.0)] + payload = json.dumps(dict(meta=meta, exact=exact, spaces=spaces), separators=(",", ":")) + (folder / "spectrum_data.js").write_text( + "// Generated by struphy.polar.homotopy.demos.spectrum\nwindow.SPECTRUM = " + payload + ";\n" + ) + (folder / "spectrum_explorer.html").write_text(_HTML) + if verbose: + print(f"wrote {folder / 'spectrum_explorer.html'}") + return folder / "spectrum_explorer.html" + + +_HTML = r""" + + + + +Curl-curl spectrum explorer + + + +
+

Spectrum of the 2D curl–curl operator on the unit disk

+
+
+
+
+
+ + +
+
+ vertical axis + + + +
+ +
← → eigenpair · ↑ ↓ regularity · f field · z centre · g mesh · hover to preview, click to pin · drag to zoom
+
+
+
+ + + +
+
+ + + +
+
+
+
same eigenpair index, all regularities
+
+
+
+
+ + + + +""" diff --git a/src/struphy/polar/homotopy/demos/talk_figures.py b/src/struphy/polar/homotopy/demos/talk_figures.py new file mode 100644 index 000000000..e35ae7a38 --- /dev/null +++ b/src/struphy/polar/homotopy/demos/talk_figures.py @@ -0,0 +1,210 @@ +"""All figures (and videos) of R. Remmerswaal's NMPP seminar talk on high-order polar splines, computed with Struphy. + +Run as ``python -m struphy.polar.homotopy.demos.talk_figures OUT_FOLDER [--full] [--videos] [--only NAME ...]``. +Without ``--full`` the studies run at a coarse resolution (minutes); ``--full`` uses the resolution of the talk +(the wave, 3D wave and orbit studies then take hours). The file names are those of the talk (``figures/tikz`` and +``media``); the schematic diagrams of the talk are drawn by :mod:`struphy.polar.homotopy.demos.schematics`. + +Every figure is computed with Struphy objects: the polar splines of Struphy Derham complexes (B-splines or +``angular_basis="trig"``), Struphy domains (``PerturbedDiskCylinder``, ``HollowCylinder(a1=0)``, ``HelicalTube``) and +Struphy's mass operators, projectors and field evaluation. +""" + +import argparse +from pathlib import Path + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt # noqa: E402 +import numpy as np # noqa: E402 + + +def _save(fig, folder, name): + path = Path(folder) / f"{name}.png" + fig.savefig(path, dpi=150, bbox_inches="tight") + plt.close(fig) + print("wrote", path) + + +# ---------------------------------------------------------------------------------------------------------- figures +def basis_figures(folder, full=False): + from struphy.polar.homotopy.demos import basis + + _save(basis.figure_mapping(), folder, "disk_mapping") + _save(basis.figure_tensor_basis("bspline"), folder, "disc_basis_functions_tensor_bspline") + _save(basis.figure_tensor_basis("trig"), folder, "disc_basis_functions_tensor_fourier") + for r in (1, 2, 3): + _save(basis.figure_polar_zero_forms(r), folder, f"disc_basis_functions_r{r}") + _save(basis.figure_all_forms(3), folder, "allforms") + + +def regularity_figures(folder, full=False): + from struphy.polar.homotopy.demos import regularity + + _save(regularity.figure_regularity_lines(0), folder, "result_regularity_zeroform") + _save(regularity.figure_regularity_lines(2), folder, "result_regularity_twoform") + + +def projection_figures(folder, full=False): + from struphy.polar.homotopy.demos import projections + + regularities = (1, 2, 3, 4, 5, 6) if full else (1, 2, 3) + cells = (8, 12, 16, 24, 32, 48) if full else (4, 8, 16) + for method, label in (("commuting", "commuting projection"), ("l2", "$L^2$ projection")): + errors = projections.convergence(regularities, cells, method=method) + _save(projections.figure_convergence(errors, cells, title=label), folder, f"polar_projection_{method}") + defects = projections.commutation(regularities, cells, method=method) + _save( + projections.figure_commutation(defects, cells, title=label), + folder, + f"polar_projection_commutation_{method}", + ) + + +def _wave_settings(full): + from struphy.polar.homotopy.demos import waves + + if full: + return dict(regularities=(0, 1, 2, 3, 4, 5), radial_cells=(8, 16, 32, 48), reference_cells=64, degree=5), 1.0 + initial = lambda x, y: waves.gaussian(x, y, width=0.15) # noqa: E731 (the talk's pulse is not resolved coarse) + return dict(regularities=(0, 1, 2), radial_cells=(4, 8, 16), reference_cells=32, degree=3, initial=initial), 0.3 + + +def wave_figures(folder, full=False): + from struphy.polar.homotopy.demos import waves + + settings, final_time = _wave_settings(full) + names = [ + ("polar_0form_wave", dict(form=0)), + ("polar_1form_wave", dict(form=1)), + ("polar_1form_wave_convergence_nonmatching", dict(form=1, non_matching=True)), + ] + for name, extra in names: + form = extra.pop("form") + study = waves.convergence(form, final_time=final_time, verbose=True, **settings, **extra) + stem = name if "nonmatching" in name else f"{name}_convergence" + _save(waves.figure_convergence(study, degree=settings["degree"]), folder, f"{stem}_linf") + for index, n in enumerate(study["cells"]): + if n in (8, 48) or (not full and index in (0, len(study["cells"]) - 1)): + tag = f"nr{n}" + _save(waves.figure_time_traces(study, index), folder, f"{stem}_linf_time_{tag}") + + +def wave3d_figures(folder, full=False): + from struphy.polar.homotopy.demos import waves3d + + if full: + study = waves3d.convergence((1, 2, 3, 4), (5, 8, 12), reference_mesh=16, final_time=0.4) + else: + study = waves3d.convergence((1, 2), (3, 4), reference_mesh=5, final_time=0.1, degree=3) + _save(waves3d.figure_convergence(study), folder, "polar_3d_wave_convergence_linf") + from struphy.polar.homotopy.demos.waves import figure_time_traces + + _save(figure_time_traces(study, 0), folder, f"polar_3d_wave_convergence_linf_time_nr{study['cells'][0]}") + + +def orbit_figures(folder, full=False): + from struphy.polar.homotopy.demos import orbits + + field = orbits.ProjectedVacuumField(2) + path, _ = orbits.integrate(field, 0.04, periods=20.0 if full else 2.0) + _save(orbits.figure_orbits(path), folder, "result_vacuum_orbits_overview") + _save(orbits.figure_orbits(path, extent=0.05), folder, "result_vacuum_orbits_closeup") + kwargs = ( + dict(regularities=(0, 1, 2, 3, 4), dts=(0.04, 0.02, 0.01, 0.005, 0.0025), periods=20.0) + if full + else dict(regularities=(0, 2, 4), dts=(0.08, 0.04, 0.02), periods=1.0) + ) + for angular, name in (("bspline", "bspline"), ("trig", "compatible")): + study = orbits.energy_study(angular=angular, **kwargs) + _save(orbits.figure_energy(study), folder, f"result_vacuum_orbits_energy_{name}") + + +def extra_figures(folder, full=False): + """Figures of PolarFEEC.jl that the talk does not show: Poisson/curl-curl (example 05), the Maxwell spectrum + (example 06), field lines in the twisted tube (example 07).""" + from struphy.polar.homotopy.demos import fieldlines, maxwell, pde + + solution = pde.solve_all(3, 6) + _save(pde.figure_poisson(solution), folder, "polarfeec_05_poisson") + _save(pde.figure_curl_curl(solution), folder, "polarfeec_05_curl_curl") + spectrum = maxwell.study(meshes=maxwell.MESHES if full else ((8, 16), (16, 32)), compared=40, verbose=False) + _save(maxwell.figure_spectrum(spectrum), folder, "polarfeec_06_maxwell_spectrum") + result = fieldlines.study(turns=150 if full else 25, verbose=False) + _save(fieldlines.figure_poincare(result), folder, "polarfeec_07_poincare") + + +# ---------------------------------------------------------------------------------------------------------- videos +def videos(folder, full=False): + """The wave videos of the talk (mp4, needs ffmpeg).""" + from struphy.polar.homotopy.demos import waves + + folder = Path(folder) + final_time = 1.0 if full else 0.3 + kwargs = dict(degree=5, radial_cells=8) if full else dict(degree=3, radial_cells=6) + series = [ + ("polar_0form_wave_p5_cochain_bspline16", 0, dict(enrichment="cochain", angular_cells=16)), + ( + "polar_0form_wave_p5_cochain_trigonometric15", + 0, + dict(enrichment="cochain", angular="trig", angular_cells=15), + ), + ("polar_0form_wave_p5_none_bspline16", 0, dict(enrichment="none", angular_cells=16)), + ("polar_1form_wave_p5_cochain_chain_bspline16", 1, dict(angular_cells=16)), + ("polar_1form_wave_p5_cochain_chain_bspline16_nonmatching", 1, dict(angular_cells=16, non_matching=True)), + ( + "polar_1form_wave_p5_cochain_trigonometric15", + 1, + dict(enrichment="cochain", angular="trig", angular_cells=15), + ), + ] + regularities = (0, 1, 2, 3, 4, 5) if full else (0, 2) + for name, form, extra in series: + for r in regularities: + if extra.get("enrichment") == "none" and r > 0 and form == 1: + continue + result = waves.run_family(form, (r,), final_time=final_time, verbose=False, **kwargs, **extra) + path = folder / f"{name}_r{r}.mp4" + waves.animate(result, path, regularity=r) + print("wrote", path) + + +def schematic_figures(folder, full=False): + """The schematic diagrams (de Rham diagrams with tikzfigure and LaTeX, the 3D coordinate transformation, the + boundary-compatible homotopy).""" + from struphy.polar.homotopy.demos import schematics + + schematics.write_all(folder) + + +GROUPS = { + "basis": basis_figures, + "schematics": schematic_figures, + "regularity": regularity_figures, + "projections": projection_figures, + "waves": wave_figures, + "waves3d": wave3d_figures, + "orbits": orbit_figures, + "extra": extra_figures, + "videos": videos, +} + + +def main(argv=None): + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("folder", help="output folder") + parser.add_argument("--full", action="store_true", help="the resolution of the talk (hours)") + parser.add_argument("--videos", action="store_true", help="also write the wave videos (mp4)") + parser.add_argument("--only", nargs="+", choices=list(GROUPS), help="only these groups of figures") + args = parser.parse_args(argv) + folder = Path(args.folder) + folder.mkdir(parents=True, exist_ok=True) + groups = args.only or [g for g in GROUPS if g != "videos" or args.videos] + for group in groups: + print(f"--- {group}") + GROUPS[group](folder, args.full) + + +if __name__ == "__main__": + main() diff --git a/src/struphy/polar/homotopy/demos/waves.py b/src/struphy/polar/homotopy/demos/waves.py new file mode 100644 index 000000000..65739eaad --- /dev/null +++ b/src/struphy/polar/homotopy/demos/waves.py @@ -0,0 +1,331 @@ +"""0- and 1-form wave equations on the perturbed disk (talk, Sections 1, 3 and 5). + +The semidiscretization in polar coefficients is (form degree ``k = 0, 1``, ``D_k`` the tensor-product gradient or curl):: + + M_p v'' + K_p v = 0, M_p = E_k^T M_k E_k, K_p = (D_k E_k)^T M_{k+1} (D_k E_k), + +with Dirichlet conditions (``P = 0``, resp. ``A x n = 0``), integrated by the explicit staggered leapfrog with +``pi = M_p v'`` (one stiffness product and one mass solve per step):: + + pi^{n+1/2} = pi^{n-1/2} - dt K_p v^n, v^{n+1} = v^n + dt M_p^{-1} pi^{n+1/2}, + +and each regularity uses its own CFL step ``dt = 0.9 * 2 / sqrt(lambda_max)``. This is the method of the wave script of +PolarFEEC.jl (``examples/polar_1form_wave.jl``, moved to the seminar repository). The observable is ``u`` for 0-forms +and ``curl A`` for 1-forms. + +Initial data. The talk uses the Gaussian ``exp(-((x - 0.05)^2 + y^2) / (2 * 0.04^2))`` with zero velocity, for ``P`` +and for ``curl A``; the Julia script uses a compactly supported bump. For 1-forms the initial ``A`` is the homotopy +lift ``A = h_2 (I - B) b`` of the projected (zero-mean) 2-form ``b``, so ``curl A = b - B b``. + +Not stated in the talk and chosen here: the reference solution of the convergence study (the finest mesh with the +same degree), the sampling of the L-infinity norm (a physical grid) and the final time ``T = 1``. +""" + +import matplotlib.pyplot as plt +import numpy as np +import scipy.sparse.linalg as spla + +from struphy.polar.homotopy.demos import common +from struphy.polar.homotopy.demos.plotting import CATEGORICAL, DIVERGING, clean_axis, draw_mesh +from struphy.polar.homotopy.linalg import as_dense + +AMPLITUDE, PETALS = 0.2, 3 +CENTER, WIDTH = (0.05, 0.0), 0.04 + + +def perturbed_disk(): + """The cross-section of the Struphy domain ``PerturbedDiskCylinder`` of the talk.""" + return common.cross_section(common.perturbed_disk(AMPLITUDE, PETALS)) + + +def gaussian(x, y, center=CENTER, width=WIDTH): + return np.exp(-((x - center[0]) ** 2 + (y - center[1]) ** 2) / (2 * width**2)) + + +def bump(x, y, center=(0.38, 0.0), radius=0.24): + """Laplacian of the compact bump ``psi = exp(-1/(1 - s))`` of the Julia script (its 0-form initial condition).""" + dx, dy = x - center[0], y - center[1] + s = (dx**2 + dy**2) / radius**2 + out = np.zeros_like(s) + inside = s < 1 + d = 1 - s[inside] + out[inside] = -4 * (1 - s[inside] - s[inside] ** 2) * np.exp(-1 / d - 4 * np.log(d)) / radius**2 + return out + + +class WaveProblem: + """Polar wave semidiscretization of degree ``form`` and regularity ``r`` on a 2D Struphy Derham complex. + + ``domain`` (default: the perturbed disk) defines the pushforward. With ``non_matching=True`` the polar spaces are + constructed on the unit disk (``polar_cross_section="unit_disk"``) and pushed forward with ``domain`` (the + "non-matching" study of the talk). ``angular`` is ``"bspline"`` (``angular_cells`` periodic B-spline cells) or + ``"trig"`` (``angular_cells`` trigonometric functions, odd). + """ + + def __init__( + self, + form, + regularity, + degree=5, + radial_cells=8, + angular="bspline", + angular_cells=16, + enrichment="cochain_chain", + domain=None, + non_matching=False, + initial=gaussian, + ): + from struphy.polar.homotopy.derham_blocks import DomainCrossSection + + self.form, self.regularity = form, regularity + domain = common.perturbed_disk(AMPLITUDE, PETALS) if domain is None else domain + complex_ = common.polar_complex( + regularity, + degree, + (radial_cells, angular_cells), + domain=domain, + angular=angular, + boundary="dirichlet", + enrichment=enrichment, + cross_section="unit_disk" if non_matching else "domain", + check_assumptions=False, + ) + if non_matching: + # the physical spaces (masses, observables) use the pushforward of the domain + mapping = DomainCrossSection(domain) + for space in complex_.spaces: + space.mapping = mapping + complex_._cache.clear() + self.complex = c = complex_ + E = c.extraction(form) + D = c.spaces[form].derivative_matrix + self.E, self.DE = E, (D @ E).tocsr() + M, Mc = c.full_mass_matrix(form), c.full_mass_matrix(form + 1) + self.mass = ((E.T @ M @ E + (E.T @ M @ E).T) / 2).tocsc() + self.stiffness = (self.DE.T @ Mc @ self.DE).tocsr() + self.mass_lu = spla.splu(self.mass) + lam = spla.eigsh(self.stiffness, k=1, M=self.mass, which="LM", return_eigenvectors=False)[0] + self.lambda_max = float(lam) + self.dt = 0.9 * 2 / np.sqrt(self.lambda_max) + self.initial = self._initial_coefficients(initial) + + def _initial_coefficients(self, initial): + c = self.complex + if self.form == 0: + return c.mass_solver(0).solve(c.extraction(0).T @ c.spaces[0].inner_product(initial, 2)) + b = c.mass_solver(2).solve(c.extraction(2).T @ c.spaces[2].inner_product(initial, 2)) + return as_dense(c.homotopy(2)) @ b if c.boundary == "dirichlet" else c.homotopy(2) @ b + + def observable(self, v): + """Full coefficients of ``u`` (0-forms) or ``curl A`` (1-forms).""" + return self.E @ v if self.form == 0 else self.DE @ v + + @property + def observable_space(self): + return self.complex.spaces[0 if self.form == 0 else 2] + + def evolve(self, final_time, record_every=1): + """Staggered leapfrog; returns times, observable snapshots and the relative staggered-energy error.""" + dt = self.dt + steps = int(np.ceil(final_time / dt)) + v = self.initial.copy() + force = self.stiffness @ v + momentum = 0.5 * dt * force # pi^{-1/2} for zero initial velocity + times, snapshots = [0.0], [self.observable(v)] + energies = [] + for step in range(1, steps + 1): + momentum -= dt * force + velocity = self.mass_lu.solve(momentum) + v += dt * velocity + force = self.stiffness @ v + energies.append(0.5 * momentum @ velocity + 0.5 * (v @ force - dt * velocity @ force)) + if step % record_every == 0 or step == steps: + times.append(step * dt) + snapshots.append(self.observable(v)) + energies = np.array(energies) + energy_error = np.max(np.abs(energies - energies[0])) / abs(energies[0]) if energies.size else 0.0 + return np.array(times), np.array(snapshots), energy_error + + +def interpolate(times, snapshots, t): + i = np.clip(np.searchsorted(times, t, side="right") - 1, 0, len(times) - 2) + w = (t - times[i]) / (times[i + 1] - times[i]) + return (1 - w) * snapshots[i] + w * snapshots[i + 1] + + +def sample_points(n=121, radius=1.0): + """A physical sampling grid inside the domain (logical points, uniform in a box around the disk).""" + mapping = perturbed_disk() + x = np.linspace(-1.2 * radius, 1.2 * radius, n) + X, Y = np.meshgrid(x, x, indexing="ij") + logical_x, logical_y = mapping.inverse(X.ravel(), Y.ravel()) + # the pole itself is left out: the pushforward of 2-forms divides by det J = 0 there (and the 2-forms of r = 0 + # have no limit at the pole) + inside = (logical_x < 1 - 1e-12) & (logical_x > 1e-12) + return logical_x[inside], logical_y[inside] + + +def sampled(problem, coefficients, points): + return problem.observable_space.evaluate_points(coefficients, *points) + + +def run_family(form, regularities=(0, 1, 2, 3, 4, 5), final_time=1.0, output_times=None, verbose=True, **kwargs): + """Solutions for every regularity, interpolated to common output times.""" + output_times = np.linspace(0, final_time, 101) if output_times is None else output_times + family = {} + for r in regularities: + problem = WaveProblem(form, r, **kwargs) + times, snapshots, energy_error = problem.evolve(final_time) + family[r] = dict( + problem=problem, + frames=np.array([interpolate(times, snapshots, t) for t in output_times]), + energy_error=energy_error, + ) + if verbose: + print( + f" {form}-form r = {r}: {problem.mass.shape[0]} polar modes, dt = {problem.dt:.4g} " + f"({len(times) - 1} steps), max relative staggered-energy error {energy_error:.2e}" + ) + return dict(form=form, output_times=output_times, family=family) + + +def figure_snapshots(result, frames=(0, 50, 100), radial_cells=8): + """The observable at the first, middle and last output time, one row per regularity.""" + family = result["family"] + rs, ts = np.linspace(1e-3, 1, 101), np.linspace(0, 1, 241) + fig, axes = plt.subplots(len(family), len(frames), figsize=(3 * len(frames), 3 * len(family)), squeeze=False) + reference = family[max(family)] + for i, (r, entry) in enumerate(family.items()): + problem = entry["problem"] + mapping = problem.observable_space.mapping + for j, frame in enumerate(frames): + values = problem.observable_space.evaluate(entry["frames"][frame], rs, ts) + limit = np.max(np.abs(reference["problem"].observable_space.evaluate(reference["frames"][frame], rs, ts))) + X, Y = mapping(*np.meshgrid(rs, ts, indexing="ij")) + axes[i, j].pcolormesh(X, Y, values, cmap=DIVERGING, vmin=-limit, vmax=limit, shading="gouraud") + draw_mesh( + axes[i, j], + mapping, + np.linspace(0, 1, problem.complex.radial.cells + 1), + np.linspace(0, 1, problem.complex.angular.n + 1), + alpha=0.2, + ) + axes[i, j].set_aspect("equal") + clean_axis(axes[i, j]) + axes[i, j].set_title(rf"$\mathfrak{{r}} = {r}$, $t = {result['output_times'][frame]:.2f}$", fontsize=10) + name = "u" if result["form"] == 0 else r"\mathrm{curl}\,A" + fig.suptitle(f"${name}$ of the {result['form']}-form wave equation") + fig.tight_layout() + return fig + + +def convergence( + form, + regularities=(0, 1, 2, 3, 4, 5), + radial_cells=(8, 16, 32), + reference_cells=64, + final_time=1.0, + angular="bspline", + angular_factor=2, + degree=5, + points=None, + verbose=True, + **kwargs, +): + """Relative L-infinity errors against a fine-mesh reference: time traces and the maximum over the final 10%.""" + output_times = np.linspace(0, final_time, 51) + points = sample_points(81) if points is None else points + + def solve(r, n): + cells = angular_factor * n if angular == "bspline" else 2 * (angular_factor * n // 2) + 1 + problem = WaveProblem(form, r, degree=degree, radial_cells=n, angular=angular, angular_cells=cells, **kwargs) + times, snapshots, _ = problem.evolve(final_time) + frames = np.array([sampled(problem, interpolate(times, snapshots, t), points) for t in output_times]) + return frames + + reference = solve(max(regularities), reference_cells) + scale = np.max(np.abs(reference), axis=1) + traces, final = {}, {} + for r in regularities: + traces[r], final[r] = [], [] + for n in radial_cells: + error = np.max(np.abs(solve(r, n) - reference), axis=1) / scale + traces[r].append(error) + final[r].append(np.max(error[output_times >= 0.9 * final_time])) + if verbose: + print(f" {form}-form r = {r}, N^x = {n}: max error over the final 10% {final[r][-1]:.3e}") + return dict(form=form, times=output_times, cells=np.array(radial_cells), traces=traces, final=final) + + +def figure_convergence(study, degree=5): + hs = 1 / study["cells"] + fig, ax = plt.subplots(figsize=(5.5, 4.5)) + for i, (r, e) in enumerate(study["final"].items()): + ax.loglog(hs, e, "o-", color=CATEGORICAL[i % len(CATEGORICAL)], label=rf"$\mathfrak{{r}} = {r}$") + rate = degree + 1 if study["form"] == 0 else degree + e = np.array(study["final"][max(study["final"])]) + ax.loglog(hs, 0.5 * e[-1] * (hs / hs[-1]) ** rate, "k:", label=f"$O(h^{rate})$") + name = "P" if study["form"] == 0 else "B" + ax.set_xlabel("$h$") + ax.set_ylabel( + rf"$\Vert {name}_h - {name}_\mathrm{{ref}}\Vert_{{L^\infty}} / \Vert {name}_\mathrm{{ref}}\Vert_{{L^\infty}}$" + ) + ax.set_title("maximum error during the final 10%") + ax.legend(fontsize=8) + fig.tight_layout() + return fig + + +def figure_time_traces(study, index=0): + fig, ax = plt.subplots(figsize=(5.5, 4.5)) + for i, (r, traces) in enumerate(study["traces"].items()): + ax.semilogy( + study["times"], traces[index], color=CATEGORICAL[i % len(CATEGORICAL)], label=rf"$\mathfrak{{r}} = {r}$" + ) + name = "P" if study["form"] == 0 else "B" + ax.set_xlabel("$t$") + ax.set_ylabel( + rf"$\Vert {name}_h - {name}_\mathrm{{ref}}\Vert_{{L^\infty}} / \Vert {name}_\mathrm{{ref}}\Vert_{{L^\infty}}$" + ) + ax.set_title(f"$N^x = {study['cells'][index]}$") + ax.legend(fontsize=8) + fig.tight_layout() + return fig + + +def animate(result, path, fps=25, regularity=None): + """Write an animation (gif or mp4) of the observable for one regularity (as the videos of the talk).""" + from matplotlib.animation import FuncAnimation + + r = max(result["family"]) if regularity is None else regularity + entry = result["family"][r] + problem = entry["problem"] + mapping = problem.observable_space.mapping + rs, ts = np.linspace(1e-3, 1, 81), np.linspace(0, 1, 181) + X, Y = mapping(*np.meshgrid(rs, ts, indexing="ij")) + fig, ax = plt.subplots(figsize=(4.5, 4.5)) + values = problem.observable_space.evaluate(entry["frames"][0], rs, ts) + limit = np.max(np.abs(values)) + mesh = ax.pcolormesh(X, Y, values, cmap=DIVERGING, vmin=-limit, vmax=limit, shading="gouraud") + draw_mesh( + ax, + mapping, + np.linspace(0, 1, problem.complex.radial.cells + 1), + np.linspace(0, 1, problem.complex.angular.n + 1), + alpha=0.2, + ) + ax.set_aspect("equal") + clean_axis(ax) + + def update(frame): + values = problem.observable_space.evaluate(entry["frames"][frame], rs, ts) + limit = np.max(np.abs(values)) or 1.0 + mesh.set_array(values.ravel()) + mesh.set_clim(-limit, limit) + ax.set_title(rf"$\mathfrak{{r}} = {r}$, $t = {result['output_times'][frame]:.2f}$") + return (mesh,) + + animation = FuncAnimation(fig, update, frames=len(entry["frames"])) + animation.save(path, fps=fps) + plt.close(fig) + return path diff --git a/src/struphy/polar/homotopy/demos/waves3d.py b/src/struphy/polar/homotopy/demos/waves3d.py new file mode 100644 index 000000000..2c00823ad --- /dev/null +++ b/src/struphy/polar/homotopy/demos/waves3d.py @@ -0,0 +1,207 @@ +"""The 3D 1-form wave equation in the stellarator-like tube (talk, Section 8), on the Struphy Derham complex with polar +splines. + +Maxwell's equations ``dE/dt = curl B``, ``dB/dt = -curl E`` with ``E x n = 0`` on the boundary, semidiscretized with the +polar 1- and 2-forms of a 3D Struphy ``Derham`` complex (masses ``WeightedMassOperators.M1, M2``, the polar curl, +``PolarMassMatrixPreconditioner``), integrated by the leapfrog of ``MaxwellWeakAmpere(algo="leapfrog")`` with +``dt = 0.9 * 2 / sqrt(lambda_max)``. The domain is ``HelicalTube("stellarator")``, the mesh ``n x 2n x 8n`` cells of +degree 4 (as the talk). The initial field ``B_0 = curl A_0`` is the magnetic pulse of +``examples/Maxwell/gaussian_helical_tube`` (``A_0 = (0, 0, a)`` in logical components, a Gaussian next to the axis); +``B_0`` is the discrete curl of the projection of ``A_0``, so ``div B = 0`` exactly. + +The error measure is the talk's: the maximum relative L-infinity error during the final 10% of the run against a +finer mesh of the highest regularity, sampled at logical points of the cross-sections ``eta3 = 0`` and ``1/2``. +""" + +import matplotlib.pyplot as plt +import numpy as np + +from struphy.polar.homotopy.demos.plotting import CATEGORICAL, DIVERGING, clean_axis + +WIDTH_FACTOR, AXIAL_WIDTH = 0.04, 0.1 + + +def helical_tube(): + from struphy.geometry import domains + + return domains.HelicalTube("stellarator", num_elements=(8, 24, 64), degree=(4, 4, 4)) + + +def initial_potential(domain): + """The logical component ``a_3`` of ``A_0`` of the example (a Gaussian in the cross-section and along the axis).""" + g = domain.geometry + width = WIDTH_FACTOR * g["a"] + + def section(e1, e2): + radius = e1 * (1 + g["eps"] * e1 * np.cos(2 * np.pi * g["k"] * e2)) + return radius * g["a"] * np.cos(2 * np.pi * e2), radius * g["b"] * np.sin(2 * np.pi * e2) + + center = section(0.05, 0.0) + speed = 2 * np.pi * np.sqrt((g["R0"] + g["delta"]) ** 2 + (g["periods"] * g["delta_z"]) ** 2) + + def a3(e1, e2, e3): + x2, y2 = section(e1, e2) + arclength = speed * (np.mod(e3 + 0.5, 1.0) - 0.5) + return np.exp( + -((x2 - center[0]) ** 2 + (y2 - center[1]) ** 2) / (2 * width**2) - arclength**2 / (2 * AXIAL_WIDTH**2) + ) + + return a3 + + +class Wave3D: + """Polar 1-form wave (Maxwell) in the helical tube with regularity ``r`` on ``n x 2n x 8n`` cells.""" + + def __init__(self, regularity, n, degree=4, domain=None, tol=1e-12): + from feectools.linalg.solvers import inverse + + from struphy.feec.mass import WeightedMassOperators + from struphy.feec.preconditioner import PolarMassMatrixPreconditioner + from struphy.feec.psydac_derham import Derham + from struphy.io.options import DerhamOptions + from struphy.topology.grids import TensorProductGrid + + self.domain = helical_tube() if domain is None else domain + self.regularity, self.n = regularity, n + options = DerhamOptions( + degree=(degree,) * 3, + bcs=(("free", "dirichlet"), None, None), + polar_splines=True, + polar_regularity=regularity, + nquads_proj=(degree + 2,) * 3, + ) + grid = TensorProductGrid(num_elements=(n, 2 * n, 8 * n), mpi_dims_mask=(False, True, True)) + import warnings + + with warnings.catch_warnings(): + # the cross-sections of the tube rotate: the regularity r >= 2 holds near eta3 = 0 (see derham_blocks) + warnings.simplefilter("ignore", UserWarning) + self.derham = derham = Derham(grid, options, domain=self.domain) + mass = WeightedMassOperators(derham, self.domain) + self.M1, self.M2, self.curl = mass.M1, mass.M2, derham.curl + self.M1_inv = inverse(self.M1, "pcg", pc=PolarMassMatrixPreconditioner(self.M1), tol=tol, maxiter=1000) + self.weak_curl = self.M1_inv @ self.curl.T @ self.M2 + self.lambda_max = self._largest_eigenvalue() + self.dt = 0.9 * 2 / np.sqrt(self.lambda_max) + a3 = initial_potential(self.domain) + A0 = derham.P1([lambda e1, e2, e3: 0 * e1, lambda e1, e2, e3: 0 * e1, a3]) + self.b0 = derham.boundary_ops["2"].dot(self.curl.dot(A0)) + + def _largest_eigenvalue(self, iterations=25): + """Power iteration on ``M1^{-1} C^T M2 C`` (the square of the largest discrete frequency).""" + x = self.M1.domain.zeros() + rng = np.random.default_rng(0) + for block in getattr(x, "pol", []): + block[:] = rng.standard_normal(block.shape) + tp = x.tp if hasattr(x, "tp") else x + for sub in getattr(tp, "blocks", [tp]): + sub._data[:] = rng.standard_normal(sub._data.shape) + x.update_ghost_regions() + x = self.derham.boundary_ops["1"].dot(x) + value = 0.0 + for _ in range(iterations): + y = self.weak_curl.dot(self.curl.dot(x)) + value = float(np.sqrt(y.inner(self.M1.dot(y)) / x.inner(self.M1.dot(x)))) + x = y * (1.0 / np.sqrt(y.inner(y))) + return value + + def evolve(self, final_time, output_times, sample): + """Leapfrog ``b <- b - dt/2 C e``, ``e <- e + dt M1^{-1} C^T M2 b``, ``b <- b - dt/2 C e``; returns the sampled + 2-form proxies (``len(output_times)`` arrays) and the relative energy error.""" + steps = int(np.ceil(final_time / self.dt)) + dt = final_time / steps + e = self.M1.domain.zeros() + b = self.b0.copy() + field = self.derham.create_spline_function("b", "Hdiv") + frames, times, t = [], [], 0.0 + energy0 = 0.5 * b.inner(self.M2.dot(b)) + energies = [] + + def record(): + field.vector = b + frames.append(sample(field)) + times.append(t) + + record() + for _ in range(steps): + b -= 0.5 * dt * self.curl.dot(e) + e += dt * self.weak_curl.dot(b) + b -= 0.5 * dt * self.curl.dot(e) + t += dt + energies.append(0.5 * (b.inner(self.M2.dot(b)) + e.inner(self.M1.dot(e)))) + record() + frames = np.array(frames) + values = np.array([np.interp(output_times, times, frames[:, i]) for i in range(frames.shape[1])]).T + energy_error = np.max(np.abs(np.array(energies) - energy0)) / energy0 if energies else 0.0 + return values, energy_error + + +def sample_points(n=24): + """Logical points in the cross-sections ``eta3 = 0`` and ``1/2`` (markers).""" + e1 = (np.arange(n) + 0.5) / n + e2 = np.arange(2 * n) / (2 * n) + E1, E2 = np.meshgrid(e1, e2, indexing="ij") + return np.concatenate([np.column_stack([E1.ravel(), E2.ravel(), np.full(E1.size, z)]) for z in (0.0, 0.5)], axis=0) + + +def convergence(regularities=(1, 2, 3, 4), meshes=(4, 6, 8), reference_mesh=12, final_time=0.4, degree=4, verbose=True): + """Maximum relative L-infinity errors during the final 10% against the finest mesh of the highest regularity.""" + markers = sample_points() + output_times = np.linspace(0, final_time, 41) + + def sample(field): + values = field(markers) + return np.concatenate([np.asarray(v) for v in values]) + + def solve(r, n): + wave = Wave3D(r, n, degree) + values, energy_error = wave.evolve(final_time, output_times, sample) + if verbose: + print(f" 3D wave r = {r}, n = {n}: dt = {wave.dt:.3g}, relative energy error {energy_error:.1e}") + return values + + reference = solve(max(regularities), reference_mesh) + scale = np.max(np.abs(reference), axis=1) + traces, final = {}, {} + for r in regularities: + traces[r], final[r] = [], [] + for n in meshes: + error = np.max(np.abs(solve(r, n) - reference), axis=1) / scale + traces[r].append(error) + final[r].append(np.max(error[output_times >= 0.9 * final_time])) + return dict(form=1, times=output_times, cells=np.array(meshes), traces=traces, final=final, degree=degree) + + +def figure_cross_section(wave, coefficients=None, eta3=0.0, title=None): + """The eta3-component of B (physical density of the 2-form) in a cross-section of the tube.""" + field = wave.derham.create_spline_function("b", "Hdiv") + field.vector = wave.b0 if coefficients is None else coefficients + e1, e2 = np.linspace(1e-3, 1, 81), np.linspace(0, 1, 161) + b3 = np.asarray(field(e1, e2, np.array([eta3]))[2])[:, :, 0] + det = np.asarray(wave.domain.jacobian_det(e1, e2, np.array([eta3])))[:, :, 0] + X, Y, Z = (np.asarray(c)[:, :, 0] for c in wave.domain(e1, e2, np.array([eta3]))) + R = np.hypot(X, Y) + values = b3 / det + fig, ax = plt.subplots(figsize=(4.5, 4.5)) + limit = np.max(np.abs(values)) or 1.0 + ax.pcolormesh(R, Z, values, cmap=DIVERGING, vmin=-limit, vmax=limit, shading="gouraud") + ax.set_aspect("equal") + clean_axis(ax) + ax.set_title(title or rf"$B_\zeta$ at $\zeta = {eta3}$, $\mathfrak{{r}} = {wave.regularity}$") + return fig + + +def figure_convergence(study): + hs = 1 / study["cells"] + fig, ax = plt.subplots(figsize=(5.5, 4.5)) + for i, (r, e) in enumerate(study["final"].items()): + ax.loglog(hs, e, "o-", color=CATEGORICAL[i % len(CATEGORICAL)], label=rf"$\mathfrak{{r}} = {r}$") + rate = study["degree"] + e = np.array(study["final"][max(study["final"])]) + ax.loglog(hs, 0.5 * e[-1] * (hs / hs[-1]) ** rate, "k:", label=f"$O(h^{rate})$") + ax.set_xlabel("$h = 1/n$") + ax.set_ylabel(r"$\Vert B_h - B_\mathrm{ref}\Vert_{L^\infty} / \Vert B_\mathrm{ref}\Vert_{L^\infty}$") + ax.set_title("3D 1-form wave: maximum error during the final 10%") + ax.legend(fontsize=8) + fig.tight_layout() + return fig diff --git a/src/struphy/polar/homotopy/derham_blocks.py b/src/struphy/polar/homotopy/derham_blocks.py new file mode 100644 index 000000000..3c0f3cf2d --- /dev/null +++ b/src/struphy/polar/homotopy/derham_blocks.py @@ -0,0 +1,569 @@ +"""High-order polar extraction blocks of the 3D :class:`~struphy.feec.psydac_derham.Derham` complex. + +The 3D polar complex is the total complex of the 2D polar complex on the poloidal cross-section (logical +``eta1``-``eta2``, built with :class:`~struphy.polar.homotopy.complex2d.PolarComplex2D`) and the 1D complex in +``eta3`` (paper Section 5):: + + V^m = sum_{p + q = m} V2D^p (x) V1D^q. + +The blocks are applied by the polar operators of :mod:`struphy.polar.linear_operators`. Regularity 1 spans the C1 +polar splines of Toshniwal et al. (2017) and Holderied & Possanner (2022), which they replace. Per tensor-product component, the first +``n_rings`` rings hold the polar coefficients (the central polar basis functions and the external basis functions of +these rings); all later rings are standard tensor-product coefficients. The ring counts ``r + 2`` (0-forms), +``(r + 1, r + 2)`` (1-forms) and ``r + 1`` (2-forms) of regularity ``r`` are chosen such that + +* the extraction ``E`` and the duals ``F`` (``F^T E = I``) are the identity outside the polar rings, +* the exterior derivative maps no tensor-product coefficient into the polar coefficients. + +The duals are those of the paper with the non-conforming ring-0 ``y``-component removed, ``F~ = Z F``: on the +coefficients of physical (conforming) fields both coincide, and ``F~^T E = I``, ``D_bar = F~^T D E`` and the homotopies +are unchanged. +""" + +import cunumpy as xp +import numpy as np +import scipy.sparse as sp + +from struphy.polar.homotopy.complex2d import PolarComplex2D +from struphy.polar.homotopy.derham_spaces import derham_factors +from struphy.polar.homotopy.layout import STRUPHY +from struphy.polar.homotopy.polynomials import PolynomialBasis, zernike_basis + + +# ------------------------------------------------------------------------------------------------- cross-sections +def _host(values): + """Domain evaluations (arrays or tuples of arrays of the active backend) as one NumPy array.""" + if isinstance(values, (tuple, list)): + return np.stack([xp.to_numpy(v) for v in values]) + return xp.to_numpy(values) + + +class CrossSection: + """A 2D map ``(x, y) -> (X, Y)`` of the logical square onto a cross-section (``__call__``, ``jacobian_matrix``, + ``jacobian``), as used by the polar construction.""" + + has_polar_axis = True + + def jacobian(self, x, y): + J = self.jacobian_matrix(x, y) + return J[..., 0, 0] * J[..., 1, 1] - J[..., 0, 1] * J[..., 1, 0] + + +class UnitDiskCrossSection(CrossSection): + """The unit disk ``(X, Y) = x (cos 2 pi y, sin 2 pi y)``, the reference cross-section ``"unit_disk"``.""" + + def __call__(self, x, y): + x, y = np.broadcast_arrays(np.asarray(x, dtype=float), np.asarray(y, dtype=float)) + return x * np.cos(2 * np.pi * y), x * np.sin(2 * np.pi * y) + + def jacobian_matrix(self, x, y): + x, y = np.broadcast_arrays(np.asarray(x, dtype=float), np.asarray(y, dtype=float)) + c, s = np.cos(2 * np.pi * y), np.sin(2 * np.pi * y) + J = np.empty(x.shape + (2, 2)) + J[..., 0, 0], J[..., 0, 1] = c, -2 * np.pi * x * s + J[..., 1, 0], J[..., 1, 1] = s, 2 * np.pi * x * c + return J + + +def _tensor_axes(x, y): + """The 1D axes ``(x[:, 0], y[0, :])`` if ``x, y`` are an ``"ij"`` meshgrid, else ``None``.""" + if x.ndim == 2 and x.shape[1] > 0 and x.shape[0] > 0: + if np.all(x == x[:, :1]) and np.all(y == y[:1, :]): + return x[:, 0], y[0, :] + return None + + +class DomainCrossSection(CrossSection): + """The poloidal cross-section ``eta3 = const`` of a struphy :class:`~struphy.geometry.base.Domain` as a 2D map. + + Physical points are expressed in an orthonormal basis ``(q1, q2)`` of the (planar) cross-section, centred at the + pole (``centered=False``: the absolute ``x, y`` coordinates of a cross-section ``z = const``); the Jacobian is + ``[q1 q2]^T DF[:, :2]`` from the domain's analytic Jacobian. Tensor grids (``"ij"`` meshgrids) are evaluated with + one call of the domain. + """ + + def __init__(self, domain, eta3=0.0, centered=True): + self.domain, self.eta3 = domain, float(eta3) + e3 = np.array([self.eta3]) + if not centered: + # absolute coordinates of a cross-section in the plane z = const (domains without a pole) + self.pole, self.basis = np.zeros(3), np.eye(3)[:, :2] + return + pole = self._map(np.array([0.0]), np.array([0.0]), e3).reshape(3) + a = self._map(np.array([1.0]), np.array([0.0]), e3).reshape(3) - pole + b = self._map(np.array([1.0]), np.array([0.25]), e3).reshape(3) - pole + q1 = a / np.linalg.norm(a) + b = b - (b @ q1) * q1 + q2 = b / np.linalg.norm(b) + self.pole, self.basis = pole, np.stack([q1, q2], axis=1) + J = self.jacobian_matrix(np.array([0.5]), np.array([0.3])) + if np.linalg.det(J[0]) < 0: + self.basis[:, 1] *= -1 + # the cross-section must be planar + e1, e2 = np.linspace(0, 1, 5), np.linspace(0, 1, 7) + P = self._map(e1, e2, e3).reshape(3, -1) - pole[:, None] + normal = np.cross(self.basis[:, 0], self.basis[:, 1]) + size = np.max(np.linalg.norm(P, axis=0)) + if np.max(np.abs(normal @ P)) > 1e-8 * size: + raise ValueError("the polar splines need a planar cross-section eta3 = const of the domain") + + def _map(self, e1, e2, e3): + """Physical points of the domain on a tensor grid, on the host (the domain evaluates on any backend).""" + return _host(self.domain(*(xp.asarray(e) for e in (e1, e2, e3)))) + + def _domain_jacobian(self, e1, e2): + return _host(self.domain.jacobian(*(xp.asarray(e) for e in (e1, e2, [self.eta3])))) + + def __call__(self, x, y): + x, y = np.broadcast_arrays(np.asarray(x, dtype=float), np.asarray(y, dtype=float)) + axes = _tensor_axes(x, y) + if axes is not None: + P = self._map(axes[0], axes[1], np.array([self.eta3])).reshape(3, -1) + Q = (self.basis.T @ (P - self.pole[:, None])).reshape(2, *x.shape) + return Q[0], Q[1] + P = _host(self.domain(self._markers(x, y))).reshape(3, -1) + out = self.basis.T @ (P - self.pole[:, None]) + return out[0].reshape(x.shape), out[1].reshape(x.shape) + + def jacobian_matrix(self, x, y): + x, y = np.broadcast_arrays(np.asarray(x, dtype=float), np.asarray(y, dtype=float)) + axes = _tensor_axes(x, y) + if axes is not None: + DF = self._domain_jacobian(axes[0], axes[1]).reshape(3, 3, -1)[:, :2, :] + return np.einsum("ka,kbn->nab", self.basis, DF).reshape(x.shape + (2, 2)) + DF = _host(self.domain.jacobian(self._markers(x, y))).reshape(3, 3, -1)[:, :2, :] + return np.einsum("ka,kbn->nab", self.basis, DF).reshape(x.shape + (2, 2)) + + def _markers(self, x, y): + """Scattered logical points as Struphy markers ``(N, 3)`` on the active backend.""" + flat_x, flat_y = x.ravel(), y.ravel() + return xp.asarray(np.column_stack([flat_x, flat_y % 1.0, np.full(flat_x.size, self.eta3)])) + + def inverse(self, X, Y, tol=1e-13, maxiter=50): + """Logical coordinates of physical points of the cross-section (Newton's method started at the polar angle, + for cross-sections parametrized by ``eta2`` as an angle around the pole).""" + X, Y = np.broadcast_arrays(np.asarray(X, dtype=float), np.asarray(Y, dtype=float)) + target = np.stack([X.ravel(), Y.ravel()]) + size = float(np.mean(np.hypot(*self(np.ones(16), np.arange(16) / 16)))) + x = np.clip(np.hypot(*target) / size, 0.0, 1.0) + y = (np.arctan2(target[1], target[0]) / (2 * np.pi)) % 1.0 + for _ in range(maxiter): + residual = np.stack(self(x, y)) - target + if np.max(np.abs(residual)) < tol * size: + break + J = self.jacobian_matrix(np.maximum(x, 1e-12), y) + step = np.linalg.solve(J, -residual.T[..., None])[..., 0] + x = np.clip(x + step[:, 0], 0.0, 1.0) + y = (y + step[:, 1]) % 1.0 + return x.reshape(X.shape), y.reshape(X.shape) + + +def cross_section_mapping(domain, cross_section="domain"): + """The 2D reference cross-section of the polar complex: ``"domain"``, ``"unit_disk"`` or a 2D map (an object + with ``__call__`` and ``jacobian_matrix``).""" + if not isinstance(cross_section, str): + return cross_section + if cross_section == "unit_disk": + return UnitDiskCrossSection() + if cross_section == "domain": + if domain is None: + raise ValueError("the cross-section 'domain' needs a domain") + return DomainCrossSection(domain) + raise ValueError("polar_cross_section must be 'domain', 'unit_disk' or a 2D map") + + +def _warn_if_cross_section_varies(mapping): + """The polynomials are projected on the cross-section ``eta3 = 0``; on other cross-sections of a different shape + (e.g. a rotating ellipse) the polar splines of regularity ``r >= 2`` are only approximately ``C^r``.""" + import warnings + + e1, e2 = np.linspace(0, 1, 5), np.linspace(0, 1, 9) + X0, Y0 = mapping(*np.meshgrid(e1, e2, indexing="ij")) + size = np.max(np.hypot(X0, Y0)) + for eta3 in (0.29, 0.61): + other = DomainCrossSection(mapping.domain, eta3) + X, Y = other(*np.meshgrid(e1, e2, indexing="ij")) + if max(np.max(np.abs(X - X0)), np.max(np.abs(Y - Y0))) > 1e-8 * size: + warnings.warn( + "the cross-sections eta3 = const of the domain differ in shape: the polar splines of regularity >= 2 " + "are built on the cross-section eta3 = 0 and are only approximately C^r on the others", + stacklevel=3, + ) + return + + +def polar_radius(mapping, eta1): + """Mean distance of the circle ``eta1 = const`` of a 2D cross-section from its pole (its physical size).""" + eta2 = (np.arange(16) + 0.5) / 16 + x, y = mapping(np.full_like(eta2, eta1), eta2) + x0, y0 = mapping(np.zeros(1), np.zeros(1)) + return float(np.mean(np.hypot(x - x0[0], y - y0[0]))) + + +def scaled_zernike_basis(regularity, radius): + """The Zernike polynomials of the disk of ``radius`` scaled to values of order one on that disk. + + The central polar basis functions are the projections of these polynomials onto the polar rings; with the radius + of the polar rings (about ``p / N1`` of the cross-section) they are as large as the tensor-product splines, which + keeps the polar mass and stiffness matrices as well conditioned as the tensor-product ones under refinement.""" + basis = zernike_basis(regularity, radius=radius) + return PolynomialBasis([lambda x, y, f=f: radius * f(x, y) for f in basis.functions], regularity) + + +# ------------------------------------------------------------------------------------------------- blocks +def poloidal_derham(degree, num_elements, domain=None, nquads=None, comm=None, angular="bspline"): + """A 2D Derham complex of the cross-section: ``num_elements = (N1, N2)`` cells and ``degree = (p1, p2)`` in the + poloidal plane, one periodic cell of degree one in ``eta3`` (tensor-product splines; the polar complexes of the 2D + studies are built on it with :func:`derham_polar_complex`). With ``angular="trig"``, ``N2`` (odd) trigonometric + functions in ``eta2``.""" + from struphy.feec.psydac_derham import Derham + from struphy.io.options import DerhamOptions + from struphy.topology.grids import TensorProductGrid + + p1, p2 = degree + options = DerhamOptions( + degree=(p1, p2, 1), + bcs=(("free", "free"), None, None), + nquads=None if nquads is None else (nquads[0], nquads[1], 2), + nquads_proj=(p1 + 3, p2 + 3, 2), + angular_basis=angular, + ) + # (eta2 is never decomposed with trigonometric functions) + mask = (True, False, False) if angular == "trig" else (False, True, False) + grid = TensorProductGrid(num_elements=(num_elements[0], num_elements[1], 1), mpi_dims_mask=mask) + return Derham(grid, options, comm=comm, domain=domain) + + +def derham_polar_complex( + derham, + domain, + regularity, + boundary="free", + enrichment="cochain_chain", + polynomial_projection="commuting", + cross_section="domain", + check_assumptions=True, +): + """The 2D polar complex of a Struphy Derham complex: built from the Derham's 1D spline spaces (eta1, eta2), on the + cross-section of the Struphy ``domain``, in Struphy's coefficient ordering. + + This is the complex of the polar splines of the Derham complex (``boundary="free"``); with + ``boundary="dirichlet"`` it is the relative complex of the 2D studies (``P = 0``, ``A x n = 0`` at ``eta1 = 1``). + The polynomials are the Zernike polynomials scaled to the polar rings (:func:`scaled_zernike_basis`).""" + assert not derham.spl_kind[0], "polar splines need clamped splines in eta1" + assert derham.spl_kind[1], "polar splines need periodic splines in eta2" + r = int(regularity) + radial, angular, _ = derham_factors(derham) + mapping = cross_section_mapping(domain, cross_section) + if isinstance(mapping, DomainCrossSection) and r >= 2: + _warn_if_cross_section_varies(mapping) + radius = polar_radius(mapping, min(1.0, derham.degree[0] / derham.num_elements[0])) + c = PolarComplex2D( + radial, + angular, + mapping, + r, + enrichment=enrichment, + boundary=boundary, + polynomial_projection=polynomial_projection, + basis=scaled_zernike_basis(r, radius), + check_assumptions=check_assumptions, + layout=STRUPHY, + ) + c.polar_radius = radius + return c + + +class PolarExtractionBlocksHomotopy: + """High-order polar extraction blocks for :class:`~struphy.feec.psydac_derham.Derham`. + + Parameters + ---------- + domain : struphy.geometry.base.Domain + Mapping with a pole at ``eta1 = 0``. + derham : struphy.feec.psydac_derham.Derham + The (partially constructed) tensor-product de Rham complex. + regularity : int + Regularity ``r`` at the pole. + enrichment : str + ``"cochain_chain"`` (paper), ``"cochain"`` or ``"none"``. + polynomial_projection : str + ``"commuting"`` (paper) or ``"l2"``. + cross_section : str or CrossSection + Reference cross-section of the 2D polar complex: ``"domain"`` (the domain at ``eta3 = 0``), ``"unit_disk"`` + (as PolarFEEC.jl for general 3D maps) or a 2D mapping. + """ + + method = "homotopy" + + def __init__( + self, + domain, + derham, + regularity=1, + enrichment="cochain_chain", + polynomial_projection="commuting", + cross_section="domain", + check_assumptions=True, + ): + assert not derham.spl_kind[0], "polar splines need clamped splines in eta1" + assert derham.spl_kind[1], "polar splines need periodic splines in eta2" + N2 = derham.num_elements[1] + self.regularity = r = int(regularity) + self.complex2d = c = derham_polar_complex( + derham, + domain, + r, + enrichment=enrichment, + polynomial_projection=polynomial_projection, + cross_section=cross_section, + check_assumptions=check_assumptions, + ) + self.polar_radius = c.polar_radius + nx0, ny0, nx1, ny1 = c._shapes + # shapes (rings, angular coefficients) of the 2D tensor-product components + self.shapes = {"0": (nx0, ny0), "1x": (nx1, ny0), "1y": (nx0, ny1), "2": (nx1, ny1)} + self.n2 = N2 + self.n3 = {"B": derham.H1_1d_serial[2].spaces[0].nbasis, "M": derham.L2_1d_serial[2].spaces[0].nbasis} + self.spl_kind3 = derham.spl_kind[2] + + # 2D extraction and duals per 2D tensor-product component (Struphy ordering) + first = nx1 * ny0 + E, F = [c.extraction(k) for k in range(3)], [c.duals(k).tolil() for k in range(3)] + # F~ = Z F: remove the non-conforming ring-0 y-component from the 1-form duals + F[1][first + STRUPHY.index(0, np.arange(ny1), nx0, ny1), :] = 0.0 + F = [f.tocsr() for f in F] + rows = {"0": slice(0, nx0 * ny0), "1x": slice(0, first), "1y": slice(first, None), "2": slice(0, nx1 * ny1)} + degree = {"0": 0, "1x": 1, "1y": 1, "2": 2} + self._E2D = {key: E[degree[key]][rows[key]].tocsc() for key in rows} + self._F2D = {key: F[degree[key]][rows[key]].tocsc() for key in rows} + + # polar rings per 2D component and the polar coefficient sets (columns of the 2D extractions) + self.rings = {"0": r + 2, "1x": r + 1, "1y": r + 2, "2": r + 1} + central = [c[k].central_count for k in range(3)] + self._sets = { + "0": self._polar_columns("0", 0, central[0], ["0"]), + "1x": self._polar_columns("1x", None, None, ["1x"]), + "1y": self._polar_columns("1y", 1, central[1], ["1y"]), + "2": self._polar_columns("2", 2, central[2], ["2"]), + } + self._check_structure() + + # 3D spaces: tensor-product components (2D component, sign, eta3 basis) and polar sets per component + # (2D component, sign, eta3 basis, coupling group) per tp component; (polar set, coupling group) per component. + # Components of one group belong to the same 2D form (e.g. the x- and y-components of 2D 1-forms). + self.spaces = { + "0": dict(tp=[("0", 1, "B", 0)], polar=[("0", 0)]), + "1": dict( + tp=[("1x", 1, "B", 0), ("1y", 1, "B", 0), ("0", 1, "M", 1)], polar=[("1x", 0), ("1y", 0), ("0", 1)] + ), + "2": dict( + tp=[("1y", 1, "M", 0), ("1x", -1, "M", 0), ("2", 1, "B", 1)], polar=[("1y", 0), ("1x", 0), ("2", 1)] + ), + "3": dict(tp=[("2", 1, "M", 0)], polar=[("2", 0)]), + "v": dict(tp=[("0", 1, "B", g) for g in range(3)], polar=[("0", g) for g in range(3)]), + } + self._e_ten_to_pol, self._p_ten_to_pol, self._p_ten_to_ten = {}, {}, {} + for form in self.spaces: + self._e_ten_to_pol[form] = self._extraction_blocks(form, self._E2D) + self._p_ten_to_pol[form] = self._extraction_blocks(form, self._F2D) + self._p_ten_to_ten[form] = None + self._build_derivatives() + + # ------------------------------------------------------------------ layout + def _polar_columns(self, key, degree, central_count, components): + """Central columns plus the external identity columns whose coefficient lies in the polar rings.""" + c = self.complex2d + nx0, ny0, nx1, ny1 = c._shapes + first = nx1 * ny0 + if key == "1x": + E = c.extraction(1) + lead = _first_rows(E) + ring = STRUPHY.ring(lead, nx1, ny0) + mask = (lead < first) & (ring < self.rings["1x"]) & (np.arange(E.shape[1]) >= c[1].central_count) + return np.flatnonzero(mask) + E = c.extraction(degree) + lead = _first_rows(E) + if key == "1y": + ring = np.where(lead >= first, STRUPHY.ring(lead - first, nx0, ny1), -1) + external = (np.arange(E.shape[1]) >= central_count) & (lead >= first) & (ring < self.rings["1y"]) + else: + nx, ny = self.shapes[key] + external = (np.arange(E.shape[1]) >= central_count) & (STRUPHY.ring(lead, nx, ny) < self.rings[key]) + return np.concatenate([np.arange(central_count), np.flatnonzero(external)]) + + def _check_structure(self): + """E and F~ of polar columns live in the polar rings; all other columns are the identity there.""" + for degree, keys in ((0, ["0"]), (1, ["1x", "1y"]), (2, ["2"])): + polar = np.unique(np.concatenate([self._sets[k] for k in keys])) + for key in keys: + limit = self.rings[key] * self.shapes[key][1] + for M in (self._E2D[key], self._F2D[key]): + block = M[:, polar] + if block.nnz and block.nonzero()[0].max() >= limit: + raise AssertionError(f"polar basis/duals leave the polar rings of component {key}") + others = np.setdiff1d(np.arange(self.complex2d.dimension(degree)), polar) + for M in (self._E2D, self._F2D): + stacked = sp.vstack([M[k][:, others] for k in keys]).tocsc() + stacked.eliminate_zeros() + if np.any(np.diff(stacked.indptr) != 1) or not np.allclose(stacked.data, 1.0): + raise AssertionError(f"external {degree}-form basis/duals are not the identity") + + def space_layout(self, form): + """``(n_polar, n_rings)`` per component of the 3D space ``form`` (``"0"``, ``"1"``, ``"2"``, ``"3"``, ``"v"``).""" + info = self.spaces[form] + n_polar = tuple(len(self._sets[key]) for key, _ in info["polar"]) + n_rings = tuple(self.rings[key] for key, _, _, _ in info["tp"]) + return n_polar, n_rings + + def _extraction_blocks(self, form, matrices): + """Blocks ``[m][n]`` (polar of component m <- first rings of tp component n): ``(M[rings, set_m])^T sign``.""" + info = self.spaces[form] + blocks = [] + for m, (polar_key, group_m) in enumerate(info["polar"]): + row = [] + type_m = info["tp"][m][2] + for n, (key, sign, type_n, group_n) in enumerate(info["tp"]): + if type_m != type_n: + row.append(None) + continue + limit = self.rings[key] * self.shapes[key][1] + if group_m == group_n: + block = sign * matrices[key][:limit, self._sets[polar_key]].T + else: + block = sp.csr_matrix((len(self._sets[polar_key]), limit)) + row.append(sp.csr_matrix(block)) + blocks.append(row) + return blocks + + # ------------------------------------------------------------------ derivatives + def _derivative_terms(self, name): + """Terms ``(tp_out, tp_in, sign, A2D, eta3 factor)`` of the 3D derivative ``name``.""" + c = self.complex2d + nx0, ny0, nx1, ny1 = c._shapes + first = nx1 * ny0 + G, C = c.spaces[0].derivative_matrix, c.spaces[1].derivative_matrix + Gx, Gy = G[:first], G[first:] + Cx, Cy = C[:, :first], C[:, first:] + I0, Ix, Iy, I2 = (sp.identity(n, format="csr") for n in (nx0 * ny0, first, nx0 * ny1, nx1 * ny1)) + if name == "grad": + return [(0, 0, 1, Gx, "I"), (1, 0, 1, Gy, "I"), (2, 0, 1, I0, "dz")] + if name == "curl": + return [ + (0, 2, 1, Gy, "I"), + (0, 1, -1, Iy, "dz"), + (1, 0, 1, Ix, "dz"), + (1, 2, -1, Gx, "I"), + (2, 0, 1, Cx, "I"), + (2, 1, 1, Cy, "I"), + ] + if name == "div": + return [(0, 0, 1, Cy, "I"), (0, 1, -1, Cx, "I"), (0, 2, 1, I2, "dz")] + raise ValueError(name) + + def _e3(self, kind, type_out, type_in): + import cunumpy as xp + + from struphy.eigenvalue_solvers.derivatives import grad_1d_matrix + + if kind == "I": + assert type_out == type_in + return sp.identity(self.n3[type_in], format="csr") + return sp.csr_matrix(xp.to_numpy(grad_1d_matrix(self.spl_kind3, self.n3["B"]))) + + def _build_derivatives(self): + forms = {"grad": ("0", "1"), "curl": ("1", "2"), "div": ("2", "3")} + for name, (source, target) in forms.items(): + V, W = self.spaces[source], self.spaces[target] + terms = self._derivative_terms(name) + pol_to_pol, pol_to_ten, e3 = [], [], [] + for m, (polar_out, group_out) in enumerate(W["polar"]): + rows_pp, rows_pt, rows_e3 = [], [], [] + out_key, _, type_out, _ = W["tp"][m] + n_rings_out = self.rings[out_key] + n2 = self.shapes[out_key][1] + limit = (n_rings_out + 1) * n2 + for n, (polar_in, group_in) in enumerate(V["polar"]): + kinds = set() + pp = sp.csr_matrix((len(self._sets[polar_out]), len(self._sets[polar_in]))) + pt = sp.csr_matrix((limit, len(self._sets[polar_in]))) + for k_out, k_in, sign, A, kind in terms: + key_out, sign_out, _, tp_group_out = W["tp"][k_out] + key_in, sign_in, _, tp_group_in = V["tp"][k_in] + if tp_group_in != group_in: + continue + # 3D tensor-product coefficients of component k_out of the derivative of the polar basis + image = sign * sign_in * (A @ self._E2D[key_in][:, self._sets[polar_in]]) + image = _clean(image) + if image.nnz == 0: + continue + if tp_group_out == group_out: + kinds.add(kind) + pp = pp + sign_out * (self._F2D[key_out][:, self._sets[polar_out]].T @ image) + if k_out == m: + kinds.add(kind) + if image.nonzero()[0].max() >= limit: + raise AssertionError(f"{name}: polar image leaves the first tensor-product ring") + first_ring = sp.lil_matrix(image[:limit]) + first_ring[: n_rings_out * n2] = 0.0 + pt = pt + first_ring.tocsr() + pp, pt = _clean(pp), _clean(pt) + if not kinds or (pp.nnz == 0 and pt.nnz == 0): + rows_pp.append(None), rows_pt.append(None), rows_e3.append(None) + continue + if len(kinds) != 1: + raise AssertionError(f"{name}: mixed eta3 factors in block [{m}][{n}]") + type_in = V["tp"][n][2] + rows_pp.append(pp) + rows_pt.append(pt) + rows_e3.append(self._e3(kinds.pop(), type_out, type_in)) + pol_to_pol.append(rows_pp) + pol_to_ten.append(rows_pt) + e3.append(rows_e3) + setattr(self, f"_{name}_pol_to_pol", pol_to_pol) + setattr(self, f"_{name}_pol_to_ten", pol_to_ten) + setattr(self, f"_{name}_e3", e3) + + # ------------------------------------------------------------------ interface of the polar operators + @property + def n_polar(self): + return [self.space_layout(form)[0] for form in ("0", "1", "2", "3")] + + @property + def n_rings(self): + return [self.space_layout(form)[1] for form in ("0", "1", "2", "3")] + + @property + def e_ten_to_pol(self): + return self._e_ten_to_pol + + @property + def p_ten_to_pol(self): + return self._p_ten_to_pol + + @property + def p_ten_to_ten(self): + return self._p_ten_to_ten + + grad_pol_to_pol = property(lambda self: self._grad_pol_to_pol) + grad_pol_to_ten = property(lambda self: self._grad_pol_to_ten) + grad_e3 = property(lambda self: self._grad_e3) + curl_pol_to_pol = property(lambda self: self._curl_pol_to_pol) + curl_pol_to_ten = property(lambda self: self._curl_pol_to_ten) + curl_e3 = property(lambda self: self._curl_e3) + div_pol_to_pol = property(lambda self: self._div_pol_to_pol) + div_pol_to_ten = property(lambda self: self._div_pol_to_ten) + div_e3 = property(lambda self: self._div_e3) + + +def _first_rows(E): + from struphy.polar.homotopy.complex2d import first_nonzero_rows + + return first_nonzero_rows(E) + + +def _clean(A): + A = sp.csr_matrix(A) + A.data[np.abs(A.data) < 100 * np.finfo(float).eps] = 0.0 + A.eliminate_zeros() + return A diff --git a/src/struphy/polar/homotopy/derham_spaces.py b/src/struphy/polar/homotopy/derham_spaces.py new file mode 100644 index 000000000..a0c13e31c --- /dev/null +++ b/src/struphy/polar/homotopy/derham_spaces.py @@ -0,0 +1,208 @@ +"""The 1D spline spaces of a Struphy :class:`~struphy.feec.psydac_derham.Derham` complex, with the data of the polar +construction. + +:class:`DerhamSpace1D` wraps a feectools ``SplineSpace`` of the Derham complex (B-splines for 0-forms, M-splines for +1-forms, in Struphy's ordering) and provides what the polar algebra needs: basis values +(``feectools.core.bsplines.collocation_matrix``), the coefficient derivative +(:func:`~struphy.eigenvalue_solvers.derivatives.grad_1d_matrix`), the geometric degrees of freedom on Struphy's +projection grids (:func:`~struphy.feec.psydac_derham.get_pts_and_wts`: values at the Greville points, integrals over +the Greville cells split at the breakpoints), integrals and quadrature on the elements. +""" + +from functools import cached_property + +import numpy as np +import scipy.sparse as sp + + +def basis_values_1d(space, x): + """Values ``(len(x), nbasis)`` of the basis functions of a 1D space of the Derham complex: B-/M-splines (periodic + splines with all their periodic copies, also when there are fewer functions than ``degree + 1``) or the nodal + trigonometric functions of :mod:`struphy.feec.trig_space`.""" + from feectools.core.bsplines import collocation_matrix + + x = np.atleast_1d(np.asarray(x, dtype=float)).ravel() + if getattr(space, "is_trigonometric", False): + return space.basis_values(x)[0] + periodic, n = bool(space.periodic), int(space.nbasis) + points = x % 1.0 if periodic else x + # (feectools returns a 1 x 1 matrix for a single point: evaluate at least two) + padded = points if points.size > 1 else np.concatenate([points, points]) + # the unrolled basis of the (extended) knot vector; periodic copies are folded onto their index + values = np.asarray(collocation_matrix(space.knots, space.degree, False, space.basis, padded), dtype=float) + values = values[: points.size] + if not periodic: + values[(x < 0) | (x > 1)] = 0.0 + return values + folded = np.zeros((points.size, n)) + for j in range(values.shape[1]): + folded[:, j % n] += values[:, j] + return folded + + +def _gauss(n): + return np.polynomial.legendre.leggauss(max(int(n), 1)) + + +def _drop_small(matrix): + matrix = sp.csr_matrix(matrix) + matrix.data[np.abs(matrix.data) < 100 * np.finfo(float).eps * max(np.max(np.abs(matrix.data), initial=0.0), 1)] = 0 + matrix.eliminate_zeros() + return matrix + + +class DerhamSpace1D: + """A uni-variate space of the Derham complex (``space.basis`` ``"B"``: 0-forms, ``"M"``: 1-forms): B-/M-splines, + or the nodal trigonometric functions of :mod:`struphy.feec.trig_space` (``angular_basis="trig"``). + + Parameters + ---------- + space : feectools.fem.splines.SplineSpace + The 1D space, e.g. ``derham.H1_1d_serial[d].spaces[0]`` or ``derham.L2_1d_serial[d].spaces[0]``. + partner : feectools.fem.splines.SplineSpace, optional + The other space of the 1D complex (the M-splines of a B-spline space and vice versa). + """ + + is_bspline = True + + def __init__(self, space, partner=None): + self.space = space + self._partner = partner + self.n = int(space.nbasis) + self.degree = int(space.degree) + self.periodic = bool(space.periodic) + self.is_scaled = space.basis == "M" + self.breaks = np.asarray(space.breaks, dtype=float) + self.cells = len(self.breaks) - 1 + + def __repr__(self): + kind = "periodic" if self.periodic else "clamped" + return f"DerhamSpace1D({kind} {self.space.basis}-splines of degree {self.degree}, n = {self.n})" + + def __len__(self): + return self.n + + @property + def local_degree(self): + return self.degree + + @cached_property + def scaled(self): + """The M-spline (1-form) companion of a B-spline space.""" + if self.is_scaled or self._partner is None: + raise ValueError("only B-spline spaces with a partner have a scaled (1-form) companion") + return DerhamSpace1D(self._partner, self.space) + + @cached_property + def parent(self): + """The B-spline space of an M-spline space (``self`` for B-splines).""" + if not self.is_scaled: + return self + return DerhamSpace1D(self._partner, self.space) + + # ------------------------------------------------------------------ basis and derivative + def basis_matrix(self, x): + """Sparse matrix of basis values with rows for the points ``x`` and columns for the basis functions.""" + return sp.csr_matrix(basis_values_1d(self.space, x)) + + def derivative_matrix(self): + """Coefficient derivative into the M-splines, ``(D u)_i = u_{i+1} - u_i`` (periodically wrapped).""" + import cunumpy as xp + + from struphy.eigenvalue_solvers.derivatives import grad_1d_matrix + + if self.is_scaled: + raise ValueError("the derivative of 1-forms is zero in one dimension") + return sp.csr_matrix(xp.to_numpy(grad_1d_matrix(self.periodic, self.n))) + + def constant_coefficients(self): + """Coefficients of the constant function one (partition of unity of the B-splines).""" + if self.is_scaled: + raise ValueError("use constant_one_form_coefficients for M-splines") + return np.ones(self.n) + + def constant_one_form_coefficients(self): + """M-spline coefficients of the constant 1-form ``dx`` (histopolation of one).""" + target = self if self.is_scaled else self.scaled + points, R = target.gdof_rule_matrix(0) + return sp.linalg.spsolve(sp.csc_matrix(target.histopolation_matrix), R @ np.ones(points.size)) + + # ------------------------------------------------------------------ degrees of freedom + def greville_points(self): + return np.asarray(self.parent.space.greville, dtype=float) + + def gdof_rule_matrix(self, additional_order=2): + """Points of the geometric degrees of freedom and the sparse matrix ``R`` with ``dofs = R @ f(points)``. + + Values at the Greville points for B-splines; for M-splines integrals over the Greville cells with + ``degree + 1 + additional_order`` Gauss points on each part of a cell split at the breakpoints (Struphy's + projection grid of :func:`~struphy.feec.psydac_derham.get_pts_and_wts`).""" + from struphy.feec.psydac_derham import get_pts_and_wts + + if not self.is_scaled: + points = self.greville_points() + return points, sp.identity(self.n, format="csr") + n_quad = self.degree + 1 + int(additional_order) + pts, wts, subs = get_pts_and_wts(self.space, 0, self.n - 1, n_quad=n_quad) + pts, wts, subs = np.asarray(pts), np.asarray(wts), np.asarray(subs) + cell = np.cumsum(subs == 0) - 1 + rows = np.repeat(cell, pts.shape[1]) + R = sp.csr_matrix((wts.ravel(), (rows, np.arange(pts.size))), shape=(self.n, pts.size)) + return pts.ravel(), R + + @cached_property + def histopolation_matrix(self): + """Square matrix from coefficients to geometric degrees of freedom (exact for the spline integrands; the + identity for the nodal trigonometric functions, which interpolate and histopolate on the nodes).""" + from struphy.feec.trig_space import is_trigonometric + + if is_trigonometric(self.space): + return sp.identity(self.n, format="csr") + points, R = self.gdof_rule_matrix(0) + return _drop_small(R @ self.basis_matrix(points)) + + def geometric_dofs(self, f, additional_order=2): + points, R = self.gdof_rule_matrix(additional_order) + return R @ np.asarray(f(points), dtype=float) + + def commuting_projection(self, f, additional_order=2): + return sp.linalg.spsolve(sp.csc_matrix(self.histopolation_matrix), self.geometric_dofs(f, additional_order)) + + # ------------------------------------------------------------------ quadrature, integrals and masses + def quadrature_order(self, additional_order): + return self.degree + 1 + int(additional_order) + + def quadrature(self, order): + """Gauss points and weights with ``order`` points on each element.""" + values, weights = _gauss(order) + a, b = self.breaks[:-1, None], self.breaks[1:, None] + points = (a + b) / 2 + (b - a) / 2 * values[None, :] + return points.ravel(), ((b - a) / 2 * weights[None, :]).ravel() + + @cached_property + def coefficient_integrals(self): + """Integrals of the basis functions over ``[0, 1]``.""" + points, weights = self.quadrature(self.degree + 2) + values = np.asarray(self.basis_matrix(points).T @ weights).ravel() + values[np.abs(values) < 100 * np.finfo(float).eps] = 0.0 + return values + + def mass_matrix(self, coefficient=None, additional_quadrature_order=None): + additional = ( + 0 if coefficient is None else (2 if additional_quadrature_order is None else additional_quadrature_order) + ) + points, weights = self.quadrature(self.quadrature_order(additional)) + if coefficient is not None: + weights = weights * coefficient(points) + B = self.basis_matrix(points) + return _drop_small(B.T @ sp.diags(weights) @ B) + + def inner_product(self, f, additional_order=2): + points, weights = self.quadrature(self.quadrature_order(additional_order)) + return np.asarray(self.basis_matrix(points).T @ (weights * f(points))).ravel() + + +def derham_factors(derham): + """The radial (``eta1``), angular (``eta2``) and axial (``eta3``) 0-form spaces of a Derham complex.""" + H1, L2 = derham.H1_1d_serial, derham.L2_1d_serial + return tuple(DerhamSpace1D(H1[d].spaces[0], L2[d].spaces[0]) for d in range(3)) diff --git a/src/struphy/polar/homotopy/forms2d.py b/src/struphy/polar/homotopy/forms2d.py new file mode 100644 index 000000000..c0f61d4ff --- /dev/null +++ b/src/struphy/polar/homotopy/forms2d.py @@ -0,0 +1,436 @@ +"""Tensor-product form spaces on the logical unit square and their push-forwards to a physical domain. + +Port of ``src/tensor_forms.jl`` and ``src/physical_forms.jl`` of PolarFEEC.jl, restricted to two dimensions. + +Coefficient layout: components are stacked (``0``-forms: one, ``1``-forms: ``x``- then ``y``-component, ``2``-forms: +one); within a component the coefficients are ordered by a :class:`~struphy.polar.homotopy.layout.Layout`: Struphy's +``i * n_y + j`` (:data:`~struphy.polar.homotopy.layout.STRUPHY`) or PolarFEEC.jl's ``i + j * n_x`` +(:data:`~struphy.polar.homotopy.layout.CORE`, the default). The derivatives are ``G = [dx (x) I; I (x) dy]`` and +``C = [-I (x) dy, dx (x) I]`` (radial factor first). The 1D factors are +:class:`~struphy.polar.homotopy.derham_spaces.DerhamSpace1D` (the spaces of a Derham complex) or the test-only +reference spaces of ``struphy.polar.tests.reference_core``. +""" + +from functools import cached_property + +import numpy as np +import scipy.sparse as sp +import scipy.sparse.linalg as spla + +from struphy.polar.homotopy.layout import CORE + + +def _drop_small(M): + M = sp.csr_matrix(M) + if M.nnz: + scale = np.max(np.abs(M.data)) + M.data[np.abs(M.data) < 100 * np.finfo(float).eps * scale] = 0.0 + M.eliminate_zeros() + return M + + +def mass_additional_order(coefficient, requested): + """Additional quadrature order of a mass matrix: 0 without and 2 (default) with a coefficient.""" + if coefficient is None: + if requested not in (None, 0): + raise ValueError("additional_quadrature_order must be zero when no coefficient is supplied") + return 0 + if requested is None: + return 2 + if int(requested) != requested or requested < 0: + raise ValueError("additional quadrature order must be a nonnegative integer") + return int(requested) + + +class FormComponent: + """One vector-proxy component: its two 1D factors and the mask of scaled directions.""" + + def __init__(self, fx, fy, mask): + self.fx, self.fy, self.mask = fx, fy, mask + + @property + def factors(self): + return (self.fx, self.fy) + + @property + def shape(self): + return (self.fx.n, self.fy.n) + + @property + def size(self): + return self.fx.n * self.fy.n + + +class FormSpace2D: + """Logical tensor-product ``k``-form space ``(k = 0, 1, 2)`` built from the 1D 0-form spaces ``x`` and ``y``.""" + + def __init__(self, k, x, y, layout=CORE): + if k not in (0, 1, 2): + raise ValueError("form degree must be 0, 1 or 2") + self.k, self.x, self.y, self.layout = k, x, y, layout + if k == 0: + self.components = [FormComponent(x, y, (False, False))] + elif k == 1: + self.components = [FormComponent(x.scaled, y, (True, False)), FormComponent(x, y.scaled, (False, True))] + else: + self.components = [FormComponent(x.scaled, y.scaled, (True, True))] + self.sizes = [c.size for c in self.components] + self.offsets = list(np.cumsum([0] + self.sizes[:-1])) + self.n = int(sum(self.sizes)) + + def __repr__(self): + return f"FormSpace2D({self.k}, {self.x!r}, {self.y!r})" + + def __len__(self): + return self.n + + def with_degree(self, k): + return FormSpace2D(k, self.x, self.y, self.layout) + + @property + def shape(self): + """``(n_x, n_y)`` of a single-component space.""" + if len(self.components) != 1: + raise ValueError("shape is defined for single-component spaces") + return self.components[0].shape + + def component_range(self, index): + return slice(self.offsets[index], self.offsets[index] + self.sizes[index]) + + # ------------------------------------------------------------------ exterior derivative + @cached_property + def derivative_matrix(self): + """Coefficient matrix of the exterior derivative ``d: V^k -> V^{k+1}``.""" + x, y, kron = self.x, self.y, self.layout.kron + if self.k == 0: + Ix, Iy = sp.identity(x.n), sp.identity(y.n) + return sp.vstack([kron(x.derivative_matrix(), Iy), kron(Ix, y.derivative_matrix())]).tocsr() + if self.k == 1: + Ix1, Iy1 = sp.identity(x.scaled.n), sp.identity(y.scaled.n) + return sp.hstack([-kron(Ix1, y.derivative_matrix()), kron(x.derivative_matrix(), Iy1)]).tocsr() + raise ValueError("the top-degree form space has no exterior derivative") + + # ------------------------------------------------------------------ geometric degrees of freedom + def _component_rules(self, index, additional_order): + c = self.components[index] + return c.fx.gdof_rule_matrix(additional_order), c.fy.gdof_rule_matrix(additional_order) + + @cached_property + def _histopolation_factorizations(self): + result = [] + for c in self.components: + result.append(tuple(spla.splu(sp.csc_matrix(f.histopolation_matrix)) for f in c.factors)) + return result + + def histopolation_matrix(self): + blocks = [self.layout.kron(c.fx.histopolation_matrix, c.fy.histopolation_matrix) for c in self.components] + return sp.block_diag(blocks).tocsr() + + def solve_histopolation(self, dofs): + """Solve ``histopolation_matrix @ u = dofs`` direction by direction (``dofs`` may have several columns).""" + dofs = np.asarray(dofs, dtype=float) + vector = dofs.ndim == 1 + dofs = dofs.reshape(self.n, -1) + result = np.empty_like(dofs) + order = self.layout.order + for index, c in enumerate(self.components): + lu_x, lu_y = self._histopolation_factorizations[index] + nx, ny = c.shape + block = dofs[self.component_range(index)].reshape(nx, ny, -1, order=order) + block = lu_x.solve(np.ascontiguousarray(block).reshape(nx, -1)).reshape(nx, ny, -1) + block = np.ascontiguousarray(np.moveaxis(block, 1, 0)) + block = lu_y.solve(block.reshape(ny, -1)).reshape(ny, nx, -1) + block = np.moveaxis(block, 0, 1) + result[self.component_range(index)] = block.reshape(nx * ny, -1, order=order) + return result[:, 0] if vector else result + + def geometric_dofs(self, f, additional_order=2): + """GDOFs of the logical proxy ``f(x, y)`` (scalar, or a tuple of components for 1-forms).""" + result = np.empty(self.n) + for index in range(len(self.components)): + (px, Rx), (py, Ry) = self._component_rules(index, additional_order) + X, Y = np.meshgrid(px, py, indexing="ij") + values = f(X, Y) + if self.k == 1: + values = values[index] + values = np.broadcast_to(np.asarray(values, dtype=float), X.shape) + result[self.component_range(index)] = (Rx @ values @ Ry.T).ravel(order=self.layout.order) + return result + + def commuting_projection(self, f, additional_order=2): + return self.solve_histopolation(self.geometric_dofs(f, additional_order)) + + # ------------------------------------------------------------------ quadrature + def quadrature_points(self, additional_order, component=None): + """Tensor quadrature shared by the components (or of one ``component``): 1D points and weights per direction.""" + components = self.components if component is None else [self.components[component]] + rules = [] + for d in range(2): + factor = components[0].factors[d] + order = max(c.factors[d].quadrature_order(additional_order) for c in components) + rules.append(factor.quadrature(order)) + return rules + + def basis_at(self, px, py, component): + """Sparse matrix of the component's basis values at the tensor points (rows in the order of the layout).""" + c = self.components[component] + return self.layout.kron(c.fx.basis_matrix(px), c.fy.basis_matrix(py)).tocsr() + + def evaluate(self, coefficients, x, y): + """Logical proxy of the element with the given coefficients at the tensor grid ``x (x) y``. + + Returns arrays of shape ``(len(x), len(y))`` (a list of two for 1-forms). + """ + x, y = np.atleast_1d(x), np.atleast_1d(y) + values = [] + for index in range(len(self.components)): + B = self.basis_at(x, y, index) + values.append( + (B @ coefficients[self.component_range(index)]).reshape(len(x), len(y), order=self.layout.order) + ) + return values if self.k == 1 else values[0] + + def evaluate_points(self, coefficients, x, y): + """Logical proxy at the scattered logical points ``(x[k], y[k])`` (a list of two arrays for 1-forms).""" + x, y = np.broadcast_arrays(np.atleast_1d(np.asarray(x, dtype=float)), np.atleast_1d(np.asarray(y, dtype=float))) + shape = x.shape + x, y = x.ravel(), y.ravel() + values = [] + for index, c in enumerate(self.components): + C = np.asarray(coefficients[self.component_range(index)]).reshape(c.fx.n, c.fy.n, order=self.layout.order) + Bx, By = c.fx.basis_matrix(x), c.fy.basis_matrix(y) + values.append(np.asarray(By.multiply(Bx @ C).sum(axis=1)).reshape(shape)) + return values if self.k == 1 else values[0] + + def mass_matrix(self, coefficient=None, additional_quadrature_order=None): + """Logical mass matrix (Kronecker products of 1D masses without a coefficient).""" + if coefficient is None: + mass_additional_order(None, additional_quadrature_order) + return sp.block_diag( + [self.layout.kron(c.fx.mass_matrix(), c.fy.mass_matrix()) for c in self.components] + ).tocsr() + additional = mass_additional_order(coefficient, additional_quadrature_order) + blocks = [] + for index in range(len(self.components)): + (px, wx), (py, wy) = self.quadrature_points(additional, index) + X, Y = np.meshgrid(px, py, indexing="ij") + w = (np.outer(wx, wy) * coefficient(X, Y)).ravel(order=self.layout.order) + B = self.basis_at(px, py, index) + blocks.append(B.T @ sp.diags(w) @ B) + return _drop_small(sp.block_diag(blocks)) + + +class PhysicalFormSpace2D: + """Push-forward of a logical form space by a 2D mapping. + + ====== ====================================== ==================== + degree push-forward of a logical proxy ``v`` metric weight ``W`` + ====== ====================================== ==================== + 0 ``v o F^{-1}`` ``|J|`` + 1 ``(J^{-T} v) o F^{-1}`` (covariant) ``|J| G^{-1}`` + 2 ``(v / det J) o F^{-1}`` ``1 / |J|`` + ====== ====================================== ==================== + + Physical targets are functions of the physical coordinates ``(X, Y)``; elements are evaluated at logical points. + """ + + def __init__(self, logical, mapping): + self.logical = logical + self.mapping = mapping + + @classmethod + def from_factors(cls, k, x, y, mapping, layout=CORE): + return cls(FormSpace2D(k, x, y, layout), mapping) + + def __repr__(self): + return f"PhysicalFormSpace2D({self.logical!r}, {type(self.mapping).__name__})" + + @property + def k(self): + return self.logical.k + + @property + def layout(self): + return self.logical.layout + + @property + def n(self): + return self.logical.n + + def __len__(self): + return self.n + + def with_degree(self, k): + return PhysicalFormSpace2D(self.logical.with_degree(k), self.mapping) + + @property + def derivative_matrix(self): + return self.logical.derivative_matrix + + # ------------------------------------------------------------------ Piola maps + def pullback_values(self, values, x, y): + """Pull back physical proxy values given at ``F(x, y)`` to logical proxies at ``(x, y)``.""" + if self.k == 0: + return values + if self.k == 2: + return self.mapping.jacobian(x, y) * values + J = self.mapping.jacobian_matrix(x, y) + vx, vy = values + return [J[..., 0, 0] * vx + J[..., 1, 0] * vy, J[..., 0, 1] * vx + J[..., 1, 1] * vy] + + def pullback(self, f): + """Logical proxy ``(x, y) -> pullback of f`` of a physical target ``f(X, Y)``.""" + + def target(x, y): + X, Y = self.mapping(x, y) + return self.pullback_values(f(X, Y), x, y) + + return target + + def pushforward_values(self, values, x, y): + """Push logical proxy values at ``(x, y)`` forward to the physical proxy at ``F(x, y)``.""" + if self.k == 0: + return values + if self.k == 2: + return values / self.mapping.jacobian(x, y) + J = self.mapping.jacobian_matrix(x, y) + det = J[..., 0, 0] * J[..., 1, 1] - J[..., 0, 1] * J[..., 1, 0] + ax, ay = values + # J^{-T} = adj(J)^T / det + return [(J[..., 1, 1] * ax - J[..., 1, 0] * ay) / det, (-J[..., 0, 1] * ax + J[..., 0, 0] * ay) / det] + + def metric_weight(self, x, y): + """``|J|``, ``|J| G^{-1}`` (array ``(..., 2, 2)``) or ``1/|J|``.""" + if self.k == 0: + return np.abs(self.mapping.jacobian(x, y)) + if self.k == 2: + return 1.0 / np.abs(self.mapping.jacobian(x, y)) + J = self.mapping.jacobian_matrix(x, y) + det = J[..., 0, 0] * J[..., 1, 1] - J[..., 0, 1] * J[..., 1, 0] + G = np.einsum("...ki,...kj->...ij", J, J) + Ginv = np.empty_like(G) + Ginv[..., 0, 0] = G[..., 1, 1] + Ginv[..., 1, 1] = G[..., 0, 0] + Ginv[..., 0, 1] = -G[..., 0, 1] + Ginv[..., 1, 0] = -G[..., 1, 0] + Ginv /= (det**2)[..., None, None] + W = np.abs(det)[..., None, None] * Ginv + return (W + np.swapaxes(W, -1, -2)) / 2 + + # ------------------------------------------------------------------ quadrature-based operations + def _quadrature(self, additional_order): + """Shared quadrature (coupled 1-form components) or the single component's quadrature.""" + return self.logical.quadrature_points(additional_order) + + def mass_matrix(self, coefficient=None, additional_quadrature_order=None): + """Physical mass matrix ``int u . v dX`` (``coefficient(X, Y)`` optional).""" + additional = mass_additional_order(coefficient, additional_quadrature_order) + (px, wx), (py, wy) = self._quadrature(additional) + X, Y = np.meshgrid(px, py, indexing="ij") + w = np.outer(wx, wy) + if coefficient is not None: + w = w * coefficient(*self.mapping(X, Y)) + W = self.metric_weight(X, Y) + bases = [self.logical.basis_at(px, py, c) for c in range(len(self.logical.components))] + if self.k != 1: + weights = (w * W).ravel(order=self.layout.order) + return _drop_small(bases[0].T @ sp.diags(weights) @ bases[0]) + blocks = [[None, None], [None, None]] + for a in range(2): + for b in range(2): + weights = (w * W[..., a, b]).ravel(order=self.layout.order) + blocks[a][b] = bases[a].T @ sp.diags(weights) @ bases[b] + return _drop_small(sp.bmat(blocks)) + + def inner_product(self, f, additional_order=2, representation="physical"): + """Load vector ``(Lambda_i, f)`` of a physical (or logical) target.""" + target = self.pullback(f) if representation == "physical" else f + (px, wx), (py, wy) = self._quadrature(additional_order) + X, Y = np.meshgrid(px, py, indexing="ij") + w = np.outer(wx, wy) + W = self.metric_weight(X, Y) + values = target(X, Y) + result = np.empty(self.n) + if self.k != 1: + weighted = (w * W * np.broadcast_to(values, X.shape)).ravel(order=self.layout.order) + return self.logical.basis_at(px, py, 0).T @ weighted + vx, vy = (np.broadcast_to(v, X.shape) for v in values) + weighted = [w * (W[..., 0, 0] * vx + W[..., 0, 1] * vy), w * (W[..., 1, 0] * vx + W[..., 1, 1] * vy)] + for c in range(2): + result[self.logical.component_range(c)] = self.logical.basis_at(px, py, c).T @ weighted[c].ravel( + order=self.layout.order + ) + return result + + def geometric_dofs(self, f, additional_order=2, representation="physical"): + target = self.pullback(f) if representation == "physical" else f + return self.logical.geometric_dofs(target, additional_order) + + def commuting_projection(self, f, additional_order=2, representation="physical"): + return self.logical.solve_histopolation(self.geometric_dofs(f, additional_order, representation)) + + def l2_projection(self, f, additional_quadrature_order=None, mass=None): + additional = mass_additional_order(None, additional_quadrature_order) + M = self.mass_matrix() if mass is None else mass + return spla.spsolve(sp.csc_matrix(M), self.inner_product(f, additional)) + + def l2_error(self, coefficients, f, additional_order=2, representation="physical"): + """``|| u_h - f ||_{L2}`` by Gauss quadrature on the logical cells.""" + target = self.pullback(f) if representation == "physical" else f + (px, wx), (py, wy) = self._quadrature(additional_order) + X, Y = np.meshgrid(px, py, indexing="ij") + w = np.outer(wx, wy) + W = self.metric_weight(X, Y) + values = self.logical.evaluate(coefficients, px, py) + exact = target(X, Y) + if self.k != 1: + difference = values - exact + return np.sqrt(max(np.sum(w * W * difference**2), 0.0)) + dx, dy = values[0] - exact[0], values[1] - exact[1] + quadratic = W[..., 0, 0] * dx**2 + 2 * W[..., 0, 1] * dx * dy + W[..., 1, 1] * dy**2 + return np.sqrt(max(np.sum(w * quadratic), 0.0)) + + def evaluate(self, coefficients, x, y): + """Physical proxy of the element at the logical tensor grid ``x (x) y``.""" + x, y = np.atleast_1d(x), np.atleast_1d(y) + X, Y = np.meshgrid(x, y, indexing="ij") + return self.pushforward_values(self.logical.evaluate(coefficients, x, y), X, Y) + + def evaluate_points(self, coefficients, x, y): + """Physical proxy of the element at the scattered logical points ``(x[k], y[k])``.""" + x, y = np.broadcast_arrays(np.atleast_1d(np.asarray(x, dtype=float)), np.atleast_1d(np.asarray(y, dtype=float))) + return self.pushforward_values(self.logical.evaluate_points(coefficients, x, y), x, y) + + def evaluate_physical(self, coefficients, X, Y): + """Physical proxy of the element at physical points (requires ``mapping.inverse``).""" + x, y = self.mapping.inverse(X, Y) + return self.evaluate_points(coefficients, x, y) + + def polynomial_projection_columns(self, functions, method="commuting", additional_order=None): + """Projections (one column each) of polynomials into a single-component (0- or 2-form) space. + + ``method="commuting"``: the GDOFs of every polynomial (pulled back with the factor 1 or ``det J``) followed by + the histopolation solve, as in the paper; ``"l2"``: the full tensor-product L2 projection. + """ + if not functions: + return np.zeros((self.n, 0)) + if method == "l2": + if additional_order not in (None, 0): + raise ValueError("additional_quadrature_order must be zero for the L2 polynomial projection") + loads = np.column_stack([self.inner_product(f, 0) for f in functions]) + lu = spla.splu(sp.csc_matrix(self.mass_matrix())) + return lu.solve(loads) + if method != "commuting": + raise ValueError("polynomial_projection must be 'l2' or 'commuting'") + additional = 2 if additional_order is None else additional_order + (px, Rx), (py, Ry) = self.logical._component_rules(0, additional) + Xl, Yl = np.meshgrid(px, py, indexing="ij") + X, Y = self.mapping(Xl, Yl) + factor = 1.0 if self.k == 0 else self.mapping.jacobian(Xl, Yl) + dofs = np.column_stack( + [ + (Rx @ (factor * np.broadcast_to(f(X, Y), X.shape)) @ Ry.T).ravel(order=self.layout.order) + for f in functions + ] + ) + return self.logical.solve_histopolation(dofs) diff --git a/src/struphy/polar/homotopy/layout.py b/src/struphy/polar/homotopy/layout.py new file mode 100644 index 000000000..4c8913859 --- /dev/null +++ b/src/struphy/polar/homotopy/layout.py @@ -0,0 +1,64 @@ +"""Ordering of the coefficients of a 2D tensor-product component of shape ``(nx, ny)`` (``x`` radial, ``y`` +angular). + +* :data:`STRUPHY`: the ordering of Struphy's ``StencilVector``s, ``i * ny + j`` (the radial index ``i`` varies + slowest, a ring ``i`` is a contiguous block), Kronecker products ``kron(A_x, A_y)``; +* :data:`CORE`: the ordering of PolarFEEC.jl, ``i + j * nx`` (the radial index varies fastest), + Kronecker products ``kron(A_y, A_x)``. + +The polar construction is written for any layout; the Derham complex uses :data:`STRUPHY`. +""" + +import numpy as np +import scipy.sparse as sp + + +class Layout: + """Linear indices, reshapes and Kronecker products of one coefficient ordering.""" + + def __init__(self, name, order): + self.name = name + #: the NumPy order (``"C"`` or ``"F"``) of ``reshape(nx, ny)`` / ``ravel`` of a component + self.order = order + + def __repr__(self): + return f"Layout({self.name!r})" + + def index(self, i, j, nx, ny): + """Linear index of the coefficient ``(i, j)`` (broadcasting).""" + i, j = np.asarray(i), np.asarray(j) + return i * ny + j if self.order == "C" else i + j * nx + + def ring(self, linear, nx, ny): + """Radial index ``i`` of linear indices.""" + linear = np.asarray(linear) + return linear // ny if self.order == "C" else linear % nx + + def angle(self, linear, nx, ny): + """Angular index ``j`` of linear indices.""" + linear = np.asarray(linear) + return linear % ny if self.order == "C" else linear // nx + + def kron(self, radial, angular): + """Sparse matrix of ``radial (x) angular`` acting on coefficients in this layout.""" + if self.order == "C": + return sp.kron(radial, angular, format="csr") + return sp.kron(angular, radial, format="csr") + + def kron_vector(self, radial, angular): + """Coefficients ``radial[i] * angular[j]`` in this layout.""" + if self.order == "C": + return np.kron(radial, angular) + return np.kron(angular, radial) + + def radial_indices(self, nx, ny, internal_rings, offset=0): + """Sorted linear indices of the internal (radial index ``< internal_rings``) and external coefficients.""" + if not 0 <= internal_rings <= nx: + raise ValueError(f"the number of internal rings must lie between 0 and {nx}") + linear = np.arange(nx * ny) + ring = self.ring(linear, nx, ny) + return offset + linear[ring < internal_rings], offset + linear[ring >= internal_rings] + + +STRUPHY = Layout("struphy", "C") +CORE = Layout("core", "F") diff --git a/src/struphy/polar/homotopy/linalg.py b/src/struphy/polar/homotopy/linalg.py new file mode 100644 index 000000000..d214e9e08 --- /dev/null +++ b/src/struphy/polar/homotopy/linalg.py @@ -0,0 +1,153 @@ +"""Small linear-algebra helpers: matrices kept as a sparse part plus a factored low-rank part.""" + +import numpy as np +import scipy.sparse as sp + + +class SparsePlusLowRank: + """The matrix ``S + U V^T`` with sparse ``S`` and dense ``U`` (n x r), ``V`` (m x r).""" + + __array_priority__ = 100 + + def __init__(self, sparse, left=None, right=None): + self.sparse = sp.csr_matrix(sparse) + n, m = self.sparse.shape + self.left = np.zeros((n, 0)) if left is None else np.asarray(left, dtype=float).reshape(n, -1) + self.right = np.zeros((m, 0)) if right is None else np.asarray(right, dtype=float).reshape(m, -1) + + @classmethod + def rank_one(cls, left, right): + left, right = np.asarray(left, dtype=float), np.asarray(right, dtype=float) + return cls(sp.csr_matrix((left.size, right.size)), left[:, None], right[:, None]) + + @classmethod + def zero(cls, n, m=None): + return cls(sp.csr_matrix((n, n if m is None else m))) + + @property + def shape(self): + return self.sparse.shape + + @property + def rank(self): + return self.left.shape[1] + + @property + def T(self): + return SparsePlusLowRank(self.sparse.T, self.right, self.left) + + def toarray(self): + return self.sparse.toarray() + self.left @ self.right.T + + def is_zero(self): + return self.sparse.nnz == 0 and (self.rank == 0 or not np.any(self.left) or not np.any(self.right)) + + def __matmul__(self, other): + if isinstance(other, SparsePlusLowRank): + # (S1 + U1 V1^T)(S2 + U2 V2^T) = S1 S2 + [S1 U2 + U1 (V1^T U2)] V2^T + U1 (S2^T V1)^T + sparse = self.sparse @ other.sparse + left = np.hstack([self.sparse @ other.left + self.left @ (self.right.T @ other.left), self.left]) + right = np.hstack([other.right, other.sparse.T @ self.right]) + return SparsePlusLowRank(sparse, left, right) + if sp.issparse(other): + return SparsePlusLowRank(self.sparse @ other, self.left, other.T @ self.right) + other = np.asarray(other) + return self.sparse @ other + self.left @ (self.right.T @ other) + + def __rmatmul__(self, other): + if sp.issparse(other): + return SparsePlusLowRank(other @ self.sparse, other @ self.left, self.right) + other = np.asarray(other) + return other @ self.sparse + (other @ self.left) @ self.right.T + + def __add__(self, other): + if isinstance(other, SparsePlusLowRank): + return SparsePlusLowRank( + self.sparse + other.sparse, np.hstack([self.left, other.left]), np.hstack([self.right, other.right]) + ) + if sp.issparse(other): + return SparsePlusLowRank(self.sparse + other, self.left, self.right) + return self.toarray() + other + + __radd__ = __add__ + + def __neg__(self): + return SparsePlusLowRank(-self.sparse, -self.left, self.right) + + def __sub__(self, other): + return self + (-other) + + def __rsub__(self, other): + return (-self) + other + + def __mul__(self, scalar): + return SparsePlusLowRank(scalar * self.sparse, scalar * self.left, self.right) + + __rmul__ = __mul__ + + +def identity_minus(P): + """``I - P`` for a square :class:`SparsePlusLowRank` ``P``.""" + n = P.shape[0] + return SparsePlusLowRank(sp.identity(n, format="csr") - P.sparse, -P.left, P.right) + + +def as_dense(A): + if isinstance(A, SparsePlusLowRank) or sp.issparse(A): + return A.toarray() + return np.asarray(A) + + +def clean(A, tolerance=100 * np.finfo(float).eps): + """Drop entries below ``tolerance`` of a sparse matrix (absolute, as in PolarFEEC.jl).""" + A = sp.csr_matrix(A, copy=True) + A.data[np.abs(A.data) < tolerance] = 0.0 + A.eliminate_zeros() + return A + + +def max_abs(A): + if sp.issparse(A): + return float(np.max(np.abs(A.data))) if A.nnz else 0.0 + A = as_dense(A) + return float(np.max(np.abs(A))) if A.size else 0.0 + + +def kron_factored(A, B): + """``kron(A, B)`` where ``A`` and/or ``B`` may be :class:`SparsePlusLowRank`.""" + if not isinstance(A, SparsePlusLowRank) and not isinstance(B, SparsePlusLowRank): + return sp.kron(A, B, format="csr") + A = A if isinstance(A, SparsePlusLowRank) else SparsePlusLowRank(A) + B = B if isinstance(B, SparsePlusLowRank) else SparsePlusLowRank(B) + # (SA + UA VA^T) (x) (SB + UB VB^T): the mixed terms kron(UA VA^T, SB) are not low rank and are densified + sparse = sp.kron(A.sparse, B.sparse, format="csr") + dense = np.zeros(sparse.shape) + any_dense = False + if A.rank and B.sparse.nnz: + dense += np.kron(A.left @ A.right.T, B.sparse.toarray()) + any_dense = True + if B.rank and A.sparse.nnz: + dense += np.kron(A.sparse.toarray(), B.left @ B.right.T) + any_dense = True + left = np.zeros((sparse.shape[0], A.rank * B.rank)) + right = np.zeros((sparse.shape[1], A.rank * B.rank)) + column = 0 + for i in range(A.rank): + for j in range(B.rank): + left[:, column] = np.kron(A.left[:, i], B.left[:, j]) + right[:, column] = np.kron(A.right[:, i], B.right[:, j]) + column += 1 + if any_dense: + sparse = sp.csr_matrix(sparse.toarray() + dense) + return SparsePlusLowRank(sparse, left, right) + + +def matmul(A, B): + """``A @ B`` for any combination of sparse, dense and :class:`SparsePlusLowRank` operands.""" + if isinstance(A, SparsePlusLowRank): + return A @ B + if isinstance(B, SparsePlusLowRank): + if sp.issparse(A): + return SparsePlusLowRank(A @ B.sparse, A @ B.left, B.right) + return np.asarray(A) @ B.sparse + (np.asarray(A) @ B.left) @ B.right.T + return A @ B diff --git a/src/struphy/polar/homotopy/polynomials.py b/src/struphy/polar/homotopy/polynomials.py new file mode 100644 index 000000000..033feea70 --- /dev/null +++ b/src/struphy/polar/homotopy/polynomials.py @@ -0,0 +1,117 @@ +"""Graded bases of the Cartesian polynomials of total degree at most ``degree``. + +Port of ``src/polynomials.jl`` of PolarFEEC.jl. The functions are vectorized: ``f(x, y)`` accepts arrays of physical +coordinates. +""" + +from math import comb, factorial + +import numpy as np + + +def polynomial_dimension(degree): + """Dimension ``(r + 1)(r + 2)/2`` of the polynomials of total degree at most ``r`` (0 for ``r < 0``).""" + return (degree + 1) * (degree + 2) // 2 if degree >= 0 else 0 + + +class PolynomialBasis: + """A graded basis: the first ``(q + 1)(q + 2)/2`` functions span the polynomials of degree ``q``, for every + ``q <= degree``, and the first function is a nonzero constant.""" + + def __init__(self, functions, degree): + if len(functions) != polynomial_dimension(degree): + raise ValueError(f"a total-degree-{degree} polynomial basis must contain {polynomial_dimension(degree)}") + self.functions = list(functions) + self.degree = int(degree) + + def __len__(self): + return len(self.functions) + + def __getitem__(self, index): + return self.functions[index] + + def __iter__(self): + return iter(self.functions) + + +def _constant_one(x, y): + return np.ones(np.broadcast(np.asarray(x), np.asarray(y)).shape) + + +def center_shifted_basis(basis): + """The basis of the polar construction: ``T_1 = 1`` and ``T_l = P_l - P_l(0, 0)``.""" + samples = [basis[0](np.array(a), np.array(b)) for a, b in ((0.0, 0.0), (1.0, 0.0), (0.0, 1.0), (0.37, 0.29))] + if not np.allclose(samples, samples[0]) or samples[0] == 0: + raise ValueError("the first polynomial basis function must be a nonzero constant") + functions = [_constant_one] + for f in basis.functions[1:]: + center = float(f(np.array(0.0), np.array(0.0))) + functions.append(lambda x, y, f=f, c=center: f(x, y) - c) + return PolynomialBasis(functions, basis.degree) + + +def monomial_basis(order): + """Monomials ``1, x, y, x^2, x y, y^2, ...`` (by total degree, then decreasing power of ``x``).""" + functions = [] + for total in range(order + 1): + for ydeg in range(total + 1): + xdeg = total - ydeg + functions.append(lambda x, y, a=xdeg, b=ydeg: np.asarray(x, dtype=float) ** a * np.asarray(y) ** b) + return PolynomialBasis(functions, order) + + +def _zernike_radial_coefficients(n, m): + half = (n - m) // 2 + return [ + ((-1) ** s * factorial(n - s) / (factorial(s) * factorial((n + m) // 2 - s) * factorial(half - s)), half - s) + for s in range(half + 1) + ] + + +def _zernike_function(n, index, radius): + m = abs(index) + coefficients = _zernike_radial_coefficients(n, m) + use_sine = index < 0 + normalization = np.sqrt((1 if m == 0 else 2) * (n + 1) / np.pi) / radius + + def zernike(x, y): + sx, sy = np.asarray(x) / radius, np.asarray(y) / radius + if m == 0: + harmonic = np.ones(np.broadcast(sx, sy).shape) + else: + # real/imaginary parts of (sx + i sy)^m, written with binomials (complex-step safe) + terms = [comb(m, j) * sx ** (m - j) * sy**j for j in range(m + 1)] + if use_sine: + harmonic = sum(terms[j] * (-1) ** ((j - 1) // 2) for j in range(1, m + 1, 2)) + else: + harmonic = sum(terms[j] * (-1) ** (j // 2) for j in range(0, m + 1, 2)) + r2 = sx**2 + sy**2 + result = sum(c * r2**power for c, power in coefficients) * harmonic + return normalization * result + + return zernike + + +def zernike_basis(order, radius=1.0, shifted=False): + """Real Zernike polynomials through radial degree ``order``, ordered by ``n``, then ``m = -n, -n + 2, ..., n``. + + Negative ``m`` is the sine mode. Unnormalised, the ordering is ``1, y, x, 2xy, 2(x^2 + y^2) - 1, x^2 - y^2, ...``. + Each function has unit L2 norm on the disk of the given ``radius``. With ``shifted=True`` the constant mode is + ``1`` and the value at the origin is subtracted from the other modes. + """ + if order < 0: + raise ValueError("order must be nonnegative") + if not np.isfinite(radius) or radius <= 0: + raise ValueError("radius must be positive and finite") + functions = [] + for n in range(order + 1): + for index in range(-n, n + 1, 2): + f = _zernike_function(n, index, float(radius)) + if shifted: + if n == 0: + f = _constant_one + else: + center = float(f(np.array(0.0), np.array(0.0))) + f = (lambda g, c: lambda x, y: g(x, y) - c)(f, center) + functions.append(f) + return PolynomialBasis(functions, order) diff --git a/src/struphy/polar/homotopy/solvers.py b/src/struphy/polar/homotopy/solvers.py new file mode 100644 index 000000000..8f8784df4 --- /dev/null +++ b/src/struphy/polar/homotopy/solvers.py @@ -0,0 +1,168 @@ +"""Solvers for ``d* d A = J`` based on the homotopy identities (port of ``src/solvers.jl`` of PolarFEEC.jl). + +In the homotopy gauge the solution is ``A = h b``, where ``b = dA`` is the ``M``-orthogonal projection of +``M^{-1} h^T l`` onto ``range(d_k)``. For the 2D curl-curl problem (``k = 1``) this is a single polar 2-form mass solve +(the "free two-dimensional lunch"), ``A = h_2 (I - B) M_2^{-1} h_2^T l`` with the mean projection ``B`` for Dirichlet +boundary conditions. For ``k = 0`` (Poisson) the projection is a sparse saddle-point solve. +""" + +import numpy as np +import scipy.linalg as sla +import scipy.sparse as sp +import scipy.sparse.linalg as spla + +from struphy.polar.homotopy.linalg import SparsePlusLowRank, as_dense + + +def factor_basis(P, side): + """Orthonormal basis of the span of the nontrivial ``left`` or ``right`` factor columns of ``P = L R^T``.""" + n = P.shape[0] if side == "left" else P.shape[1] + if not isinstance(P, SparsePlusLowRank) or P.rank == 0: + return np.zeros((n, 0)) + L, R = P.left, P.right + active = [j for j in range(L.shape[1]) if np.linalg.norm(L[:, j]) > 0 and np.linalg.norm(R[:, j]) > 0] + if not active: + return np.zeros((n, 0)) + factors = (L if side == "left" else R)[:, active] + Q, Rq, _ = sla.qr(factors, mode="economic", pivoting=True) + diagonal = np.abs(np.diag(Rq)) + rank = int(np.sum(diagonal > 1e-12 * diagonal.max())) + return Q[:, :rank] + + +def source_load(complex_, k, source, additional_order=None, representation="physical"): + """Full tensor-product load vector ``(Lambda_i, J)`` of a callable, a full or a polar coefficient vector.""" + space = complex_.spaces[k] + if callable(source): + return space.inner_product(source, 0 if additional_order is None else additional_order, representation) + source = np.asarray(source, dtype=float) + if source.size == space.n: + return complex_.full_mass_matrix(k) @ source + if source.size == complex_.dimension(k): + return complex_.full_mass_matrix(k) @ (complex_.extraction(k) @ source) + raise ValueError("the source has neither full nor polar coefficient length") + + +class StiffnessSolver: + """Solver of ``(d_k A, d_k Phi) + shift (A, Phi) = (J, Phi)`` on the polar space of degree ``k`` (2D). + + ``gauge="homotopy"`` (requires ``shift = 0``) returns the solution ``A = h b`` in the homotopy gauge + (``h A = 0``); ``gauge="none"`` solves the (shifted) stiffness system directly. + """ + + def __init__(self, complex_, k, gauge="homotopy", shift=0.0): + top = complex_.top_degree + if not 0 <= k < top: + raise ValueError(f"k must lie between 0 and {top - 1}") + if gauge not in ("homotopy", "none"): + raise ValueError("gauge must be 'homotopy' or 'none'") + if gauge == "homotopy" and shift != 0: + raise ValueError("the homotopy method solves the unshifted problem; use gauge='none' for shift > 0") + if gauge == "none" and shift == 0: + raise ValueError("the direct method needs an invertible operator: use shift > 0 or gauge='homotopy'") + self.complex, self.k, self.gauge, self.shift = complex_, k, gauge, shift + if gauge == "none": + D = complex_.derivative(k) + A = D.T @ complex_.mass_matrix(k + 1) @ D + shift * complex_.mass_matrix(k) + self._lu = spla.splu(sp.csc_matrix((A + A.T) / 2)) + else: + self._factorize_range_projection() + + # ------------------------------------------------------------------ range projection + def _factorize_range_projection(self): + """Sparse LU of the saddle system ``[M d^T 0 V ...; d 0 d^T 0 V ...; ...]`` over degrees ``k+1..top``.""" + c, k = self.complex, self.k + degrees = list(range(k + 1, c.top_degree + 1)) + sizes = [c.dimension(j) for j in degrees] + constraints = [factor_basis(c.cohomology_projection(j), "right") for j in degrees] + ranks = [V.shape[1] for V in constraints] + nblocks = len(degrees) + len(degrees) + blocks = [[None] * nblocks for _ in range(nblocks)] + all_sizes = sizes + ranks + blocks[0][0] = c.mass_matrix(k + 1) + for index, j in enumerate(degrees): + if j < c.top_degree: + D = c.derivative(j) + blocks[index + 1][index] = D + blocks[index][index + 1] = D.T + if ranks[index] > 0: + constraint = len(degrees) + index + blocks[index][constraint] = sp.csr_matrix(constraints[index]) + blocks[constraint][index] = sp.csr_matrix(constraints[index].T) + for i, size in enumerate(all_sizes): + if blocks[i][i] is None: + blocks[i][i] = sp.csr_matrix((size, size)) + system = sp.bmat(blocks, format="csc") + self._range_lu = spla.splu(system) + self._range_size = system.shape[0] + + def range_projection(self, load): + """``M``-orthogonal projection of ``M^{-1} load`` onto ``range(d_k)`` (polar coefficients of degree k+1).""" + rhs = np.zeros(self._range_size) + rhs[: load.size] = load + return self._range_lu.solve(rhs)[: load.size] + + # ------------------------------------------------------------------ solves + def solve(self, load): + """Polar coefficients of the solution for a polar load vector ``E^T (Lambda, J)``.""" + load = np.asarray(load, dtype=float) + if self.gauge == "none": + return self._lu.solve(load) + h = self.complex.homotopy(self.k + 1) + b = self.range_projection(h.T @ load) + A = h @ b + defect = np.max(np.abs(self.complex.derivative(self.k) @ A - b)) if b.size else 0.0 + if defect > 1e-8 * max(1.0, np.max(np.abs(b))): + raise ArithmeticError("the polar spaces do not satisfy the homotopy identity required by the solver") + return A + + def compatibility_residual(self, load): + """``max |(D_{k-1}^T l, V^T l)|``: the source must be orthogonal to ``ker d_k``.""" + c, k = self.complex, self.k + parts = [] + if k >= 1: + parts.append(c.derivative(k - 1).T @ load) + parts.append(factor_basis(c.cohomology_projection(k), "left").T @ load) + parts = np.concatenate(parts) + return float(np.max(np.abs(parts))) if parts.size else 0.0 + + def __call__(self, source, additional_order=None, compatibility_tolerance=None, representation="physical"): + """Full tensor-product coefficients of the solution for a callable or coefficient source.""" + c, k = self.complex, self.k + load = c.extraction(k).T @ source_load(c, k, source, additional_order, representation) + if self.shift == 0: + residual = self.compatibility_residual(load) + tolerance = ( + 1e-10 * max(np.max(np.abs(load)), 1.0) if compatibility_tolerance is None else (compatibility_tolerance) + ) + if residual > tolerance: + raise ValueError( + f"the source is not L2-orthogonal to the kernel of the derivative (residual {residual:.3e}, " + f"tolerance {tolerance:.3e})" + ) + return c.extraction(k) @ self.solve(load) + + +def stiffness_solver(complex_, k, gauge="homotopy", shift=0.0): + """Cached :class:`StiffnessSolver` of ``complex_``.""" + key = ("stiffness_solver", k, gauge, float(shift)) + if key not in complex_._cache: + complex_._cache[key] = StiffnessSolver(complex_, k, gauge, shift) + return complex_._cache[key] + + +def range_projection(complex_, k, load): + """The projection ``b = dA`` of the homotopy method for polar coefficients ``load`` of degree ``k + 1``.""" + return stiffness_solver(complex_, k).range_projection(np.asarray(load, dtype=float)) + + +def curl_curl_free_lunch(complex_, load): + """2D curl-curl in the homotopy gauge with one 2-form mass solve: ``A = h_2 (I - B) M_2^{-1} h_2^T l``. + + Identical to ``stiffness_solver(complex_, 1).solve(load)``; written out to show the algorithm of the paper. + """ + h = complex_.homotopy(2) + x = complex_.mass_solver(2).solve(h.T @ load) + if complex_.boundary == "dirichlet": + x = x - as_dense(complex_.cohomology_projection(2)) @ x + return h @ x diff --git a/src/struphy/polar/kron_2dx1d.py b/src/struphy/polar/kron_2dx1d.py index 6df8a7bec..e4f3fff9b 100644 --- a/src/struphy/polar/kron_2dx1d.py +++ b/src/struphy/polar/kron_2dx1d.py @@ -2,8 +2,7 @@ 2D x 1D Kronecker structure of linear operators on polar spline spaces. Polar extraction operators only act in the poloidal plane :math:`(\eta_1, \eta_2)` and are the identity -in :math:`\eta_3`. Inter-/histopolation matrices of the commuting projectors and (axisymmetric) weighted -mass matrices thus have the structure +in :math:`\eta_3`. (Axisymmetric) weighted mass matrices thus have the structure .. math:: @@ -792,34 +791,6 @@ def _apply_3(self, L: np.ndarray, c: int, op_3) -> np.ndarray: return xp.ascontiguousarray(T3.backward(op_3(X)).T).reshape(l1, l2, l3) -class PolarKron2Dx1DOperator(_PolarKron2Dx1DBase): - r""" - Linear operator :math:`\bigoplus_G \mathbb A^G_\textnormal{pol} \otimes \mathbb A^G_3` on a PolarDerhamSpace. - - Parameters - ---------- - space : PolarDerhamSpace - Domain and codomain. - - groups : list[Kron2Dx1DGroup] - The factors of each group of components. - """ - - def dot(self, v: PolarVector, out: PolarVector | None = None) -> PolarVector: - if xp.is_gpu(v.pol[0]): - if not hasattr(self, "_device_factors"): - self._device_factors = [(xp.asarray(g.A_3), DeviceSparseMatrix(g.A_pol)) for g in self._groups] - factors = self._device_factors - else: - factors = [(g.A_3, g.A_pol) for g in self._groups] - ops_3 = [lambda X, A=A: A @ X for A, _ in factors] - ops_pol = [lambda Y, modes, A=A: A.dot(Y) for _, A in factors] - return self._apply(v, out, ops_3, ops_pol) - - def transpose(self, conjugate=False) -> "PolarKron2Dx1DOperator": - return PolarKron2Dx1DOperator(self._space, [g.transpose() for g in self._groups]) - - class PolarKron2Dx1DSolver(_PolarKron2Dx1DBase): r""" (Approximate) inverse :math:`\bigoplus_G (\mathbb A^G_\textnormal{pol})^{-1} \otimes (\mathbb A^G_3)^{-1}` on a PolarDerhamSpace. @@ -1103,54 +1074,6 @@ def spline_bases(fem_space) -> list[list[str]]: return [[s.basis for s in comp.spaces] for comp in _scalar_spaces(fem_space)] -def _spline_space_1d(fem_space, c: int, d: int): - return _scalar_spaces(fem_space)[c].spaces[d] - - -def collocation_matrix_1d(fem_space, c: int, d: int) -> np.ndarray: - """1d interpolation (B-splines) or histopolation (M-splines) matrix of component c in direction d.""" - S = _spline_space_1d(fem_space, c, d) - return np.asarray(S.imat if S.basis == "B" else S.hmat, dtype=float) - - -def projector_groups( - space: PolarDerhamSpace, fem_space, dofs_extraction_op, base_extraction_op -) -> list[Kron2Dx1DGroup]: - r""" - Factors of the polar inter-/histopolation matrix :math:`\mathbb P \mathcal I \mathbb E^T = \bigoplus_G (\mathbb P_\textnormal{pol} (\mathcal I_1 \otimes \mathcal I_2) \mathbb E_\textnormal{pol}^T) \otimes \mathcal I_3`. - - Parameters - ---------- - space : PolarDerhamSpace - The polar space. - - fem_space : TensorFemSpace | VectorFemSpace - The tensor-product FE space of the projector. - - dofs_extraction_op, base_extraction_op : PolarExtractionOperator - DOF and basis extraction operators. - """ - groups = [] - for comps in component_groups(space): - P = poloidal_extraction_matrix( - space, comps, dofs_extraction_op.blocks_ten_to_pol, dofs_extraction_op.blocks_ten_to_ten - ) - E = poloidal_extraction_matrix(space, comps, base_extraction_op.blocks_ten_to_pol, None) - tp_blocks = [] - outer = [] - for c in comps: - A1 = collocation_matrix_1d(fem_space, c, 0) - A2 = collocation_matrix_1d(fem_space, c, 1) - tp_blocks.append(spa.kron(spa.csr_matrix(A1), spa.csr_matrix(A2))) - nr = space.n_rings[c] - outer.append((A1[nr + 1 :, nr + 1 :], A2)) - A_pol = (P @ spa.block_diag(tp_blocks, format="csr") @ E.T).tocsr() - A_pol.eliminate_zeros() - A_3 = collocation_matrix_1d(fem_space, comps[0], 2) - groups.append(Kron2Dx1DGroup(comps, A_pol, A_3, outer, pole_mask(space, comps))) - return groups - - # -------------------------------------------------------------------------------------- # Weighted mass matrices # -------------------------------------------------------------------------------------- diff --git a/src/struphy/polar/linear_operators.py b/src/struphy/polar/linear_operators.py index 1da50c261..204521e54 100644 --- a/src/struphy/polar/linear_operators.py +++ b/src/struphy/polar/linear_operators.py @@ -42,7 +42,7 @@ def device_sparse_module(): class DeviceSparseMatrix: """Device CSR copy of a host polar matrix, for products with device arrays. - The polar blocks (:class:`~struphy.polar.extraction_operators.PolarExtractionBlocksC1`) are setup data built + The polar blocks (:class:`~struphy.polar.homotopy.derham_blocks.PolarExtractionBlocksHomotopy`) are setup data built with SciPy on the host. On the CuPy backend the polar operators apply them to coefficients that live on the device; this class holds a CSR copy of a block on the device (``cupyx.scipy.sparse``) and computes ``A @ x`` there (SpMV/SpMM on the GPU). Blocks without non-zero entries (or with an empty dimension) are not copied, diff --git a/src/struphy/polar/tests/data/c1_reference.npz b/src/struphy/polar/tests/data/c1_reference.npz new file mode 100644 index 000000000..20fbde8ab Binary files /dev/null and b/src/struphy/polar/tests/data/c1_reference.npz differ diff --git a/src/struphy/polar/tests/reference_core/__init__.py b/src/struphy/polar/tests/reference_core/__init__.py new file mode 100644 index 000000000..bdbefe6a0 --- /dev/null +++ b/src/struphy/polar/tests/reference_core/__init__.py @@ -0,0 +1,29 @@ +"""Frozen reference of the PolarFEEC.jl port (test-only). + +The 1D B-spline and trigonometric spaces and the disk mappings of the original port. The library builds the polar +complexes on Struphy's Derham spaces and domains (:mod:`struphy.polar.homotopy.derham_blocks`); the tests use this +reference to check them against an independent implementation and to test the homotopy algebra on arbitrary +analytic mappings (``DiskMapping(radius=...)``). +""" + +from struphy.polar.tests.reference_core.mappings import ( + CoordinateTransformation2D, + DiskMapping, + complex_step, + perturbed_disk, +) +from struphy.polar.tests.reference_core.spaces1d import ( + ClampedBSplineSpace, + PeriodicBSplineSpace, + TrigonometricSpace, +) + +__all__ = [ + "ClampedBSplineSpace", + "CoordinateTransformation2D", + "DiskMapping", + "PeriodicBSplineSpace", + "TrigonometricSpace", + "complex_step", + "perturbed_disk", +] diff --git a/src/struphy/polar/tests/reference_core/mappings.py b/src/struphy/polar/tests/reference_core/mappings.py new file mode 100644 index 000000000..b875effab --- /dev/null +++ b/src/struphy/polar/tests/reference_core/mappings.py @@ -0,0 +1,139 @@ +"""Two-dimensional (disk-like) coordinate transformations of the logical unit square. + +Port of the disk-like mappings of ``src/mappings.jl`` of PolarFEEC.jl. Derivatives are computed with the complex-step +method, which is exact to rounding for analytic functions written with NumPy ufuncs (the Julia code uses +ForwardDiff). All methods are vectorized in the logical coordinates ``(x, y)``: ``x`` is the radial and ``y`` the +periodic angular coordinate. +""" + +import numpy as np + +_STEP = 1e-30 + + +def complex_step(f, x, *args, argnum=0): + """Derivative of the analytic function ``f`` in its argument ``argnum`` by the complex-step method.""" + arguments = [np.asarray(x, dtype=float)] + [np.asarray(a, dtype=float) for a in args] + arguments[argnum] = arguments[argnum] + 1j * _STEP + return np.imag(f(*arguments)) / _STEP + + +class Mapping2D: + """Base class of 2D coordinate transformations ``F(x, y) = (X, Y)``.""" + + has_polar_axis = False + + def __call__(self, x, y): + raise NotImplementedError + + def jacobian_matrix(self, x, y): + """Jacobian ``J[..., i, j] = dF_i/dx_j`` of shape ``broadcast(x, y).shape + (2, 2)``.""" + x, y = np.broadcast_arrays(np.asarray(x, dtype=float), np.asarray(y, dtype=float)) + J = np.empty(x.shape + (2, 2)) + for j in range(2): + columns = complex_step(lambda a, b: np.stack(self(a, b), axis=-1), x, y, argnum=j) + J[..., :, j] = columns + return J + + def jacobian(self, x, y): + J = self.jacobian_matrix(x, y) + return J[..., 0, 0] * J[..., 1, 1] - J[..., 0, 1] * J[..., 1, 0] + + +class DiskMapping(Mapping2D): + """The disk-like map ``F(x, y) = (R(x, y) C(y), R(x, y) S(y))`` of the paper. + + Parameters + ---------- + radius : float or callable + Either the disk radius ``R1`` (``R = R1 x``) or a function ``R(x, y)`` with ``R(0, y) = 0`` and ``dR/dx > 0``. + angular : tuple of callables, optional + Periodic functions ``(C, S)`` with ``C S' - S C' > 0``; the default is ``(cos(2 pi y), sin(2 pi y))``. + + Example: the perturbed disk of the paper is ``DiskMapping(radius=lambda x, y: x * (1 + 0.2 * x * cos(6 pi y)))``. + """ + + has_polar_axis = True + + def __init__(self, radius=1.0, angular=None): + if callable(radius): + self.radius = radius + self.R1 = None + else: + if radius <= 0: + raise ValueError("radius must be positive") + R1 = float(radius) + self.R1 = R1 + self.radius = lambda x, y: R1 * x + if angular is None: + self.cosine = lambda y: np.cos(2 * np.pi * y) + self.sine = lambda y: np.sin(2 * np.pi * y) + self.circular = True + else: + self.cosine, self.sine = angular + self.circular = False + + def __call__(self, x, y): + x, y = np.broadcast_arrays(np.asarray(x, dtype=float), np.asarray(y, dtype=float)) + R = self.radius(x, y) + return R * self.cosine(y), R * self.sine(y) + + def _derivatives(self, x, y): + x, y = np.broadcast_arrays(np.asarray(x, dtype=float), np.asarray(y, dtype=float)) + R = self.radius(x, y) * np.ones_like(x) + Rx = complex_step(self.radius, x, y, argnum=0) * np.ones_like(x) + Ry = complex_step(self.radius, x, y, argnum=1) * np.ones_like(x) + C, S = self.cosine(y), self.sine(y) + Cy, Sy = complex_step(self.cosine, y), complex_step(self.sine, y) + return R, Rx, Ry, C, S, Cy, Sy + + def jacobian_matrix(self, x, y): + R, Rx, Ry, C, S, Cy, Sy = self._derivatives(x, y) + J = np.empty(R.shape + (2, 2)) + J[..., 0, 0] = Rx * C + J[..., 0, 1] = Ry * C + R * Cy + J[..., 1, 0] = Rx * S + J[..., 1, 1] = Ry * S + R * Sy + return J + + def jacobian(self, x, y): + R, Rx, _, C, S, Cy, Sy = self._derivatives(x, y) + return R * Rx * (C * Sy - S * Cy) + + def inverse(self, X, Y, iterations=60): + """Logical coordinates of physical points (polar angle and a safeguarded radial Newton/bisection).""" + X, Y = np.broadcast_arrays(np.asarray(X, dtype=float), np.asarray(Y, dtype=float)) + if not self.circular: + raise NotImplementedError("the inverse is implemented for the default angular functions") + rho = np.hypot(X, Y) + y = (np.arctan2(Y, X) / (2 * np.pi)) % 1.0 + outer = self.radius(np.ones_like(y), y) + lower, upper = np.zeros_like(rho), np.ones_like(rho) + x = np.clip(rho / outer, 0.0, 1.0) + for _ in range(iterations): + residual = self.radius(x, y) - rho + upper = np.where(residual > 0, x, upper) + lower = np.where(residual <= 0, x, lower) + slope = complex_step(self.radius, x, y, argnum=0) + with np.errstate(divide="ignore", invalid="ignore"): + candidate = x - residual / slope + x = np.where((candidate > lower) & (candidate < upper), candidate, (lower + upper) / 2) + return x, y + + +class CoordinateTransformation2D(Mapping2D): + """A general map ``F(x, y) -> (X, Y)`` given by an analytic, vectorized function.""" + + def __init__(self, function, has_polar_axis=False): + self.function = function + self.has_polar_axis = has_polar_axis + + def __call__(self, x, y): + x, y = np.broadcast_arrays(np.asarray(x), np.asarray(y)) + X, Y = self.function(x, y) + return X * np.ones_like(x), Y * np.ones_like(x) + + +def perturbed_disk(R1=1.0, eps=0.2, k=3): + """The paper's disk ``R = R1 x (1 + eps x cos(2 pi k y))`` (three petals for ``k = 3``).""" + return DiskMapping(radius=lambda x, y: R1 * x * (1 + eps * x * np.cos(2 * np.pi * k * y))) diff --git a/src/struphy/polar/tests/reference_core/spaces1d.py b/src/struphy/polar/tests/reference_core/spaces1d.py new file mode 100644 index 000000000..f60664f7d --- /dev/null +++ b/src/struphy/polar/tests/reference_core/spaces1d.py @@ -0,0 +1,573 @@ +"""One-dimensional 0-form spaces and their derivative-scaled (1-form) companions. + +Port of ``src/spaces1d.jl`` and ``src/quadrature.jl`` of PolarFEEC.jl (R. Remmerswaal). All spaces live on the +logical unit interval and all indices are zero-based (the Julia code is one-based): + +* clamped and periodic B-splines satisfy ``d/dx B_i = C_{i-1} - C_i``, so the coefficient derivative is the forward + difference ``(du)_i = u_{i+1} - u_i`` (periodically wrapped for periodic splines); +* trigonometric spaces are ordered ``1, sin(2 pi x), cos(2 pi x), sin(4 pi x), ...``; +* the geometric degrees of freedom are values at Greville points (0-forms) and integrals over Greville edges + (1-forms). +""" + +from functools import cached_property + +import numpy as np +import scipy.sparse as sp + +GAUSS_RULES = {} + + +def gauss_rule(order): + """Gauss-Legendre rule with ``max(order, 2)`` points on ``[-1, 1]``.""" + n = max(int(order), 2) + if n not in GAUSS_RULES: + GAUSS_RULES[n] = np.polynomial.legendre.leggauss(n) + return GAUSS_RULES[n] + + +def _check_additional_order(order): + if int(order) != order or order < 0: + raise ValueError("additional quadrature order must be a nonnegative integer") + return int(order) + + +def mass_additional_order(coefficient, requested): + """Additional quadrature order of a mass matrix: 0 without and 2 (default) with a coefficient.""" + if coefficient is None: + if requested not in (None, 0): + raise ValueError("additional_quadrature_order must be zero when no coefficient is supplied") + return 0 + return 2 if requested is None else _check_additional_order(requested) + + +class Space1D: + """Base class of the 1D spaces on the logical unit interval.""" + + #: whether the space is periodic + periodic = False + #: whether the space is a (scaled) B-spline space + is_bspline = True + #: whether the space is a scaled (1-form) space + is_scaled = False + + def __len__(self): + return self.n + + @property + def parent(self): + """The unscaled parent space (``self`` for unscaled spaces).""" + return self + + # ------------------------------------------------------------------ quadrature + def integration_breaks(self): + """Breakpoints at which quadrature rules are split.""" + return np.linspace(0.0, 1.0, self.parent.cells + 1) + + def minimal_quadrature_order(self): + """Number of Gauss points per cell that integrates products of two basis functions exactly.""" + return self.local_degree + 1 + + def quadrature_order(self, additional_order): + return self.minimal_quadrature_order() + _check_additional_order(additional_order) + + def quadrature(self, order): + """Global quadrature points and weights on ``[0, 1]`` with ``order`` Gauss points per cell.""" + values, weights = gauss_rule(order) + breaks = self.integration_breaks() + a, b = breaks[:-1, None], breaks[1:, None] + points = (a + b) / 2 + (b - a) / 2 * values[None, :] + return points.ravel(), ((b - a) / 2 * weights[None, :]).ravel() + + def basis_matrix(self, x): + """Sparse matrix of basis values with rows for the points ``x`` and columns for the basis functions.""" + raise NotImplementedError + + # ------------------------------------------------------------------ degrees of freedom + def greville_points(self): + raise NotImplementedError + + def edge_bounds(self): + """Greville edges ``(a, b)`` of the 1-form degrees of freedom (wrapped by one period if periodic).""" + points = self.greville_points() + if self.periodic: + nxt = np.roll(points, -1) + nxt = np.where(nxt <= points, nxt + 1.0, nxt) + return np.stack([points, nxt], axis=1) + return np.stack([points[:-1], points[1:]], axis=1) + + def integration_segments(self, a, b): + """Split ``[a, b]`` into segments inside ``[0, 1]`` (periodic spaces wrap around).""" + if not self.periodic: + return [(a, b)] + width = b - a + left = a % 1.0 + right = left + width + if right <= 1 + 100 * np.finfo(float).eps: + return [(left, min(right, 1.0))] + return [(left, 1.0), (0.0, right - 1.0)] + + def edge_rule(self, a, b, values, weights): + """Quadrature points and weights on ``[a, b]``, split at the element boundaries.""" + breaks = self.integration_breaks() + points, rule_weights = [], [] + for left_segment, right_segment in self.integration_segments(a, b): + first = np.clip(np.searchsorted(breaks, left_segment, side="right") - 1, 0, len(breaks) - 2) + last = np.clip(np.searchsorted(breaks, right_segment, side="left") - 1, first, len(breaks) - 2) + for cell in range(first, last + 1): + left = max(left_segment, breaks[cell]) + right = min(right_segment, breaks[cell + 1]) + if left >= right: + continue + scale, midpoint = (right - left) / 2, (right + left) / 2 + points.append(midpoint + scale * values) + rule_weights.append(scale * weights) + if not points: + return np.zeros(0), np.zeros(0) + return np.concatenate(points), np.concatenate(rule_weights) + + def gdof_rules(self, additional_order=2): + """For every geometric DOF, the quadrature points and weights that realize it. + + A single unit-weight point at a Greville point for 0-form spaces, and a Gauss rule on the Greville edge (split + at the element boundaries) for scaled spaces. + """ + if not self.is_scaled: + return [(np.array([x]), np.array([1.0])) for x in self.greville_points()] + values, weights = gauss_rule(self.quadrature_order(additional_order)) + return [self.edge_rule(a, b, values, weights) for a, b in self.parent_edges()] + + def gdof_rule_matrix(self, additional_order=2): + """Points of all GDOF rules and the sparse matrix ``R`` with ``dofs = R @ f(points)``.""" + rules = self.gdof_rules(additional_order) + points = np.concatenate([r[0] for r in rules]) + rows = np.concatenate([np.full(len(r[0]), i) for i, r in enumerate(rules)]) + weights = np.concatenate([r[1] for r in rules]) + R = sp.csr_matrix((weights, (rows, np.arange(len(points)))), shape=(len(rules), len(points))) + return points, R + + def parent_edges(self): + return self.parent.edge_bounds() + + def geometric_dofs(self, f, additional_order=2): + """Geometric degrees of freedom of the function ``f`` (vectorized in its argument).""" + points, R = self.gdof_rule_matrix(additional_order) + return R @ np.asarray(f(points), dtype=float) + + @cached_property + def histopolation_matrix(self): + """Square matrix mapping coefficients to geometric degrees of freedom.""" + if not self.is_scaled: + return sp.csr_matrix(self.basis_matrix(self.greville_points())) + return self._edge_integral_matrix() + + def _edge_integral_matrix(self): + values, weights = gauss_rule((self.local_degree + 2) // 2) + rows = [] + for a, b in self.parent_edges(): + points, w = self.edge_rule(a, b, values, weights) + rows.append(w @ self.basis_matrix(points)) + H = sp.csr_matrix(np.vstack(rows)) + H.eliminate_zeros() + return H + + def commuting_projection(self, f, additional_order=2): + return sp.linalg.spsolve(sp.csc_matrix(self.histopolation_matrix), self.geometric_dofs(f, additional_order)) + + # ------------------------------------------------------------------ integrals and masses + def mass_matrix(self, coefficient=None, additional_quadrature_order=None): + """Sparse mass matrix, exactly integrated without a coefficient.""" + order = self.quadrature_order(mass_additional_order(coefficient, additional_quadrature_order)) + points, weights = self.quadrature(order) + if coefficient is not None: + weights = weights * coefficient(points) + B = sp.csr_matrix(self.basis_matrix(points)) + M = (B.T @ sp.diags(weights) @ B).tocsr() + return _drop_small(M) + + def inner_product(self, f, additional_order=2): + points, weights = self.quadrature(self.quadrature_order(additional_order)) + return np.asarray(sp.csr_matrix(self.basis_matrix(points)).T @ (weights * f(points))).ravel() + + @cached_property + def coefficient_integrals(self): + """Integrals of the basis functions over ``[0, 1]``.""" + points, weights = self.quadrature(self.local_degree + 2) + values = np.asarray(sp.csr_matrix(self.basis_matrix(points)).T @ weights).ravel() + values[np.abs(values) < 100 * np.finfo(float).eps] = 0.0 + return values + + +def _drop_small(M): + M = sp.csr_matrix(M) + if M.nnz: + scale = np.max(np.abs(M.data)) + M.data[np.abs(M.data) < 100 * np.finfo(float).eps * scale] = 0.0 + M.eliminate_zeros() + return M + + +# ====================================================================== B-spline evaluation +def _de_boor(knots, degree, span, x): + """Values of the ``degree + 1`` B-splines active on knot span ``span`` (zero-based, vectorized in ``x``). + + Returns an array of shape ``(len(x), degree + 1)`` whose column ``r`` belongs to basis function ``span - degree + r``. + """ + x = np.asarray(x, dtype=float) + N = np.zeros((x.size, degree + 1)) + N[:, 0] = 1.0 + left = np.zeros((x.size, degree + 1)) + right = np.zeros((x.size, degree + 1)) + for j in range(1, degree + 1): + left[:, j] = x - knots[span + 1 - j] + right[:, j] = knots[span + j] - x + saved = np.zeros(x.size) + for r in range(j): + temp = N[:, r] / (right[:, r + 1] + left[:, j - r]) + N[:, r] = saved + right[:, r + 1] * temp + saved = left[:, j - r] * temp + N[:, j] = saved + return N + + +class ClampedBSplineSpace(Space1D): + """Uniform open (clamped) B-spline space of ``degree`` on ``cells`` cells in ``[0, 1]``.""" + + def __init__(self, degree, cells): + if degree < 0 or cells <= 0: + raise ValueError("degree must be nonnegative and cells positive") + self.degree = int(degree) + self.cells = int(cells) + h = 1.0 / cells + self.knots = np.concatenate([np.zeros(degree + 1), h * np.arange(1, cells), np.ones(degree + 1)]) + self.n = cells + degree + + def __repr__(self): + return f"ClampedBSplineSpace({self.degree}, {self.cells})" + + def __eq__(self, other): + return type(other) is type(self) and (other.degree, other.cells) == (self.degree, self.cells) + + def __hash__(self): + return hash((type(self), self.degree, self.cells)) + + @property + def local_degree(self): + return self.degree + + def active_cells(self, x): + x = np.asarray(x, dtype=float) + return np.clip(np.floor(x * self.cells).astype(int), 0, self.cells - 1) + + def basis_matrix(self, x): + x = np.atleast_1d(np.asarray(x, dtype=float)) + cells = self.active_cells(x) + p = self.degree + values = np.zeros((x.size, p + 1)) + for cell in np.unique(cells): + mask = cells == cell + values[mask] = _de_boor(self.knots, p, cell + p, x[mask]) + values[(x < 0) | (x > 1)] = 0.0 + rows = np.repeat(np.arange(x.size), p + 1) + columns = (cells[:, None] + np.arange(p + 1)[None, :]).ravel() + return sp.csr_matrix((values.ravel(), (rows, columns)), shape=(x.size, self.n)) + + def greville_points(self): + p = self.degree + if p == 0: + return np.linspace(0.0, 1.0, self.n) + return np.array([self.knots[j + 1 : j + p + 1].sum() / p for j in range(self.n)]) + + def derivative_matrix(self): + m = self.n - 1 + rows = np.repeat(np.arange(m), 2) + columns = np.stack([np.arange(m), np.arange(1, m + 1)], axis=1).ravel() + values = np.tile([-1.0, 1.0], m) + return sp.csr_matrix((values, (rows, columns)), shape=(m, self.n)) + + @cached_property + def scaled(self): + if self.degree == 0: + raise ValueError("a degree-zero B-spline space has no scaled space") + return ScaledClampedSpace(self) + + def constant_coefficients(self): + return np.ones(self.n) + + +class ScaledClampedSpace(Space1D): + """``C_i = p / (t_{i+p+1} - t_{i+1}) B^{p-1}_{i+1}``: the 1-form companion of a clamped B-spline space.""" + + is_scaled = True + + def __init__(self, parent): + self._parent = parent + self.reduced = ClampedBSplineSpace(parent.degree - 1, parent.cells) + self.n = parent.n - 1 + t, p = parent.knots, parent.degree + self.scales = np.array([p / (t[i + p + 1] - t[i + 1]) for i in range(self.n)]) + + def __repr__(self): + return f"scaled({self._parent!r})" + + @property + def parent(self): + return self._parent + + @property + def local_degree(self): + return self.reduced.degree + + def basis_matrix(self, x): + return sp.csr_matrix(self.reduced.basis_matrix(x) @ sp.diags(self.scales)) + + def greville_points(self): + return self._parent.greville_points() + + +class PeriodicBSplineSpace(Space1D): + """Uniform periodic B-spline space of ``degree`` with ``cells`` basis functions. + + Basis function ``i`` (zero-based) is supported on ``[i, i + p + 1] / n``. + """ + + periodic = True + + def __init__(self, degree, cells): + if degree < 0 or cells <= degree: + raise ValueError("degree must be nonnegative and periodic cells must exceed the degree") + self.degree = int(degree) + self.cells = int(cells) + self.n = cells + # extended uniform knots t_k = (k - p) / n + self.knots = (np.arange(cells + 2 * degree + 1) - degree) / cells + + def __repr__(self): + return f"PeriodicBSplineSpace({self.degree}, {self.cells})" + + def __eq__(self, other): + return type(other) is type(self) and (other.degree, other.cells) == (self.degree, self.cells) + + def __hash__(self): + return hash((type(self), self.degree, self.cells)) + + @property + def local_degree(self): + return self.degree + + def active_cells(self, x): + u = np.asarray(x, dtype=float) % 1.0 + return np.clip(np.floor(u * self.cells).astype(int), 0, self.cells - 1) + + def _active(self, x): + """Cells, values ``(npts, p + 1)`` and zero-based indices of the active B-splines.""" + x = np.atleast_1d(np.asarray(x, dtype=float)) + u = x - np.floor(x) + cells = self.active_cells(u) + p = self.degree + values = np.zeros((x.size, p + 1)) + for cell in np.unique(cells): + mask = cells == cell + values[mask] = _de_boor(self.knots, p, cell + p, u[mask]) + indices = (cells[:, None] - p + np.arange(p + 1)[None, :]) % self.cells + return values, indices + + def basis_matrix(self, x): + values, indices = self._active(x) + rows = np.repeat(np.arange(values.shape[0]), self.degree + 1) + return sp.csr_matrix((values.ravel(), (rows, indices.ravel())), shape=(values.shape[0], self.n)) + + def greville_points(self): + n, p = self.n, self.degree + return ((np.arange(n) + (p + 1) / 2) % n) / n + + def derivative_matrix(self): + n = self.n + rows = np.repeat(np.arange(n), 2) + columns = np.stack([np.arange(n), (np.arange(n) + 1) % n], axis=1).ravel() + return sp.csr_matrix((np.tile([-1.0, 1.0], n), (rows, columns)), shape=(n, n)) + + @cached_property + def scaled(self): + if self.degree == 0: + raise ValueError("a degree-zero periodic B-spline space has no scaled space") + return ScaledPeriodicSpace(self) + + def constant_coefficients(self): + return np.ones(self.n) + + def constant_one_form_coefficients(self): + return np.full(self.n, 1.0 / self.n) + + +class ScaledPeriodicSpace(Space1D): + """``C_i = n B^{p-1}_{i+1}``: the 1-form companion of a periodic B-spline space (supported on ``[i+1, i+p+1]/n``).""" + + periodic = True + is_scaled = True + + def __init__(self, parent): + self._parent = parent + self.reduced = PeriodicBSplineSpace(parent.degree - 1, parent.cells) + self.n = parent.n + + def __repr__(self): + return f"scaled({self._parent!r})" + + @property + def parent(self): + return self._parent + + @property + def local_degree(self): + return self.reduced.degree + + def basis_matrix(self, x): + values, reduced_indices = self.reduced._active(x) + indices = (reduced_indices - 1) % self.n + rows = np.repeat(np.arange(values.shape[0]), self.reduced.degree + 1) + return sp.csr_matrix( + (self.n * values.ravel(), (rows, indices.ravel())), + shape=(values.shape[0], self.n), + ) + + def greville_points(self): + return self._parent.greville_points() + + +# ====================================================================== trigonometric spaces +class TrigonometricSpace(Space1D): + """Real Fourier space with ``2 modes + 1`` functions ``1, sin(2 pi x), cos(2 pi x), sin(4 pi x), ...``. + + Zero-based index ``2k - 1`` is ``sin(2 pi k x)`` and ``2k`` is ``cos(2 pi k x)``. + """ + + periodic = True + is_bspline = False + + def __init__(self, modes): + if modes < 0: + raise ValueError("modes must be nonnegative") + self.modes = int(modes) + self.n = 2 * modes + 1 + + def __repr__(self): + return f"TrigonometricSpace({self.modes})" + + def __eq__(self, other): + return type(other) is type(self) and other.modes == self.modes + + def __hash__(self): + return hash((type(self), self.modes)) + + @property + def local_degree(self): + return 2 * self.modes + + def integration_breaks(self): + return np.array([0.0, 1.0]) + + def minimal_quadrature_order(self): + return 2 * self.modes + 1 + + def quadrature(self, order): + """Midpoint rule with ``max(order, 2 modes + 1)`` points (exact for the trigonometric products).""" + n = max(order, 2 * self.modes + 1) + return (np.arange(n) + 0.5) / n, np.full(n, 1.0 / n) + + def _values(self, x, scale): + x = np.atleast_1d(np.asarray(x, dtype=float)) + values = np.empty((x.size, self.n)) + values[:, 0] = 1.0 + for k in range(1, self.modes + 1): + factor = 2 * np.pi * k if scale else 1.0 + values[:, 2 * k - 1] = factor * np.sin(2 * np.pi * k * x) + values[:, 2 * k] = factor * np.cos(2 * np.pi * k * x) + return values + + def basis_matrix(self, x): + return sp.csr_matrix(self._values(x, False)) + + def greville_points(self): + return np.arange(self.n) / self.n + + def derivative_matrix(self): + rows, columns, values = [], [], [] + for k in range(1, self.modes + 1): + sine, cosine = 2 * k - 1, 2 * k + rows += [cosine, sine] + columns += [sine, cosine] + values += [1.0, -1.0] + return sp.csr_matrix((values, (rows, columns)), shape=(self.n, self.n)) + + def mass_matrix(self, coefficient=None, additional_quadrature_order=None): + if coefficient is None: + mass_additional_order(None, additional_quadrature_order) + diagonal = np.full(self.n, 0.5) + diagonal[0] = 1.0 + return sp.diags(diagonal).tocsr() + return super().mass_matrix(coefficient, additional_quadrature_order) + + @cached_property + def scaled(self): + return ScaledTrigonometricSpace(self) + + def constant_coefficients(self): + e = np.zeros(self.n) + e[0] = 1.0 + return e + + def constant_one_form_coefficients(self): + return self.constant_coefficients() + + +class ScaledTrigonometricSpace(TrigonometricSpace): + """``1, 2 pi k sin(2 pi k x), 2 pi k cos(2 pi k x), ...``: the 1-form companion of a trigonometric space.""" + + is_scaled = True + + def __init__(self, parent): + super().__init__(parent.modes) + self._parent = parent + + def __repr__(self): + return f"scaled({self._parent!r})" + + @property + def parent(self): + return self._parent + + def parent_edges(self): + return self.edge_bounds() + + def basis_matrix(self, x): + return sp.csr_matrix(self._values(x, True)) + + def mass_matrix(self, coefficient=None, additional_quadrature_order=None): + if coefficient is None: + mass_additional_order(None, additional_quadrature_order) + diagonal = np.ones(self.n) + for index in range(1, self.n): + frequency = 2 * np.pi * ((index + 1) // 2) + diagonal[index] = frequency**2 / 2 + return sp.diags(diagonal).tocsr() + return Space1D.mass_matrix(self, coefficient, additional_quadrature_order) + + def _edge_integral_matrix(self): + rows = [] + for a, b in self.edge_bounds(): + row = np.empty(self.n) + segments = self.integration_segments(a, b) + row[0] = sum(right - left for left, right in segments) + for k in range(1, self.modes + 1): + w = 2 * np.pi * k + row[2 * k - 1] = sum(np.cos(w * left) - np.cos(w * right) for left, right in segments) + row[2 * k] = sum(np.sin(w * right) - np.sin(w * left) for left, right in segments) + rows.append(row) + H = np.vstack(rows) + H[np.abs(H) < 100 * np.finfo(float).eps] = 0.0 + return sp.csr_matrix(H) + + @cached_property + def scaled(self): + raise ValueError("scaled spaces have no scaled companion") diff --git a/src/struphy/polar/tests/test_factorization_cache.py b/src/struphy/polar/tests/test_factorization_cache.py index 31db5b8b2..eb4dba8d8 100644 --- a/src/struphy/polar/tests/test_factorization_cache.py +++ b/src/struphy/polar/tests/test_factorization_cache.py @@ -79,10 +79,11 @@ def factorize(matrix): assert len(cache) == 0 -def test_projectors_share_factors_across_spaces_and_solver_instances(monkeypatch): - _, derham, _ = make_polar_derham("IGAPolarCylinder", with_mass=False) - projectors = [getattr(derham, "P" + form) for form in ("0", "1", "2", "3")] - solvers = [P.pc for P in projectors] +def test_preconditioners_share_factors_across_spaces_and_solver_instances(monkeypatch): + from struphy.feec.preconditioner import PolarMassMatrixPreconditioner + + _, derham, mass = make_polar_derham("IGAPolarCylinder") + solvers = [PolarMassMatrixPreconditioner(getattr(mass, "M" + form)).solver for form in ("0", "1", "2", "3")] assert all(s.domain._factorization_cache is solvers[0].domain._factorization_cache for s in solvers) seen = {} reused = 0 @@ -104,7 +105,7 @@ def unexpected_factorization(*args, **kwargs): assert duplicate._lu_3[0] is solvers[1]._lu_3[0] assert duplicate._pol_solvers[0] is solvers[1]._pol_solvers[0] _, other_derham, _ = make_polar_derham("IGAPolarCylinder", with_mass=False) - assert other_derham.P0.I.domain._factorization_cache is not solvers[0].domain._factorization_cache + assert other_derham.V0pol._factorization_cache is not solvers[0].domain._factorization_cache def test_mass_weight_update_reuses_1d_factors_and_rebuilds_poloidal_factors(): @@ -145,8 +146,10 @@ def check_device_factor_reuse(fake=False): assert device_factor is not host_factor assert kron._factorized_1d(matrix.copy(), cache) is device_factor np.testing.assert_allclose(xp.to_numpy(device_factor.solve(xp.asarray(matrix))), np.eye(2), atol=1e-12) - _, derham, _ = make_polar_derham("IGAPolarCylinder", with_mass=False) - first = derham.P1.pc + from struphy.feec.preconditioner import PolarMassMatrixPreconditioner + + _, derham, mass = make_polar_derham("IGAPolarCylinder") + first = PolarMassMatrixPreconditioner(mass.M1).solver x = fill(first.domain.zeros(), seed=91) expected = first.dot(x) duplicate = kron.PolarKron2Dx1DSolver(first.domain, first.groups) diff --git a/src/struphy/polar/tests/test_homotopy_3d.py b/src/struphy/polar/tests/test_homotopy_3d.py new file mode 100644 index 000000000..5671cad8b --- /dev/null +++ b/src/struphy/polar/tests/test_homotopy_3d.py @@ -0,0 +1,52 @@ +"""Tests of the 3D polar complex (total complex of the 2D polar complex and the toroidal complex), ported from +``test/test_polar3d.jl`` and ``test/test_complexes.jl`` of PolarFEEC.jl.""" + +import numpy as np +import pytest + +from struphy.polar.homotopy import PolarComplex2D +from struphy.polar.homotopy.complex3d import PeriodicComplex, PolarComplex3D +from struphy.polar.homotopy.linalg import as_dense +from struphy.polar.tests.reference_core import ( + ClampedBSplineSpace, + DiskMapping, + PeriodicBSplineSpace, + TrigonometricSpace, +) + +SECTION = DiskMapping(radius=lambda r, t: 0.5 * r * (1 + 0.1 * r * np.cos(4 * np.pi * t))) + + +@pytest.mark.parametrize("space", [PeriodicBSplineSpace(3, 8), TrigonometricSpace(3)]) +def test_periodic_complex(space): + c = PeriodicComplex(space) + d, h = c.derivative(0).toarray(), as_dense(c.homotopy(1)) + P0, P1 = as_dense(c.cohomology_projection(0)), as_dense(c.cohomology_projection(1)) + n = space.n + assert np.allclose(h @ d, np.eye(n) - P0, atol=1e-12) + assert np.allclose(d @ h, np.eye(n) - P1, atol=1e-12) + assert np.allclose(h @ P1, 0, atol=1e-12) and np.allclose(P0 @ h, 0, atol=1e-12) + + +@pytest.mark.parametrize("boundary, cohomology", [("free", [1, 1, 0, 0]), ("dirichlet", [0, 0, 1, 1])]) +@pytest.mark.parametrize("toroidal", [PeriodicBSplineSpace(3, 8), TrigonometricSpace(3)]) +def test_polar_complex_3d(boundary, cohomology, toroidal): + c2 = PolarComplex2D(ClampedBSplineSpace(3, 3), PeriodicBSplineSpace(3, 8), SECTION, regularity=2, boundary=boundary) + c3 = PolarComplex3D(c2, toroidal) + nz = toroidal.n + for m in range(4): + expected = sum(c2.dimension(p) * nz for p, q in c3.bidegree_blocks(m)) + assert c3.dimension(m) == expected + E, F = c3.extraction(m), c3.duals(m) + assert np.allclose((F.T @ E).toarray(), np.eye(c3.dimension(m)), atol=1e-12) + assert c3.cohomology_dimensions() == cohomology + for m in range(2): + assert np.max(np.abs(as_dense(c3.derivative(m + 1) @ c3.derivative(m)))) < 1e-13 + for m in range(4): + n = c3.dimension(m) + contraction = np.zeros((n, n)) + if m < 3: + contraction += as_dense(c3.homotopy(m + 1) @ c3.derivative(m)) + if m > 0: + contraction += as_dense(c3.derivative(m - 1) @ c3.homotopy(m)) + assert np.allclose(contraction, np.eye(n) - as_dense(c3.cohomology_projection(m)), atol=1e-11) diff --git a/src/struphy/polar/tests/test_homotopy_applications.py b/src/struphy/polar/tests/test_homotopy_applications.py new file mode 100644 index 000000000..ecce3d9f8 --- /dev/null +++ b/src/struphy/polar/tests/test_homotopy_applications.py @@ -0,0 +1,156 @@ +"""Tests of the applications of the high-order polar splines: wave equations (core and Struphy models with the +leapfrog), the helical tube, the helical vacuum field (finite on the axis), field lines and the 1D/2D space demos.""" + +import numpy as np +import pytest + + +# ---------------------------------------------------------------------------------------------------- core demos +@pytest.mark.parametrize("form", [0, 1]) +def test_wave_leapfrog_conserves_energy(form): + from struphy.polar.homotopy.demos import waves + + problem = waves.WaveProblem(form, 2, degree=3, radial_cells=4, angular_cells=12) + _, snapshots, energy_error = problem.evolve(0.2) + assert energy_error < 1e-10 + assert np.all(np.isfinite(snapshots)) + # the CFL step is set by the largest eigenvalue of M^{-1} K + assert problem.dt == pytest.approx(0.9 * 2 / np.sqrt(problem.lambda_max)) + + +def test_spaces_demos_commute(): + from struphy.polar.homotopy.demos import spaces + + _, residuals = spaces.spaces_1d(verbose=False) + assert max(residuals.values()) < 1e-12 + curl_grad, grad_residual, curl_residual = spaces.identities(verbose=False) + assert curl_grad == 0 and grad_residual < 1e-12 and curl_residual < 1e-12 + # the quarter annulus HollowCylinder(a1=1, a2=2, poc=4) + V0, V1, _ = spaces.mapped_identities(verbose=False)["spaces"] + assert abs(V1.derivative_matrix @ V0.derivative_matrix).max() == 0 + + +# ---------------------------------------------------------------------------------------------------- geometry +def test_helical_tube(): + from struphy.geometry.domains import HelicalTube + from struphy.geometry.domains.helical_tube.helical_tube import helical_axis + + for preset in ("stellarator", "islands"): + domain = HelicalTube(preset, num_elements=(6, 24, 48), degree=(3, 3, 3)) + assert domain.pole + e1, e2, e3 = np.linspace(0.05, 1, 5), np.linspace(0, 1, 7), np.linspace(0, 1, 9) + assert np.all(domain.jacobian_det(e1, e2, e3) > 0) + g = domain.geometry + axis = np.array(domain(np.array([0.0]), np.array([0.3]), e3)).reshape(3, -1) + exact = helical_axis(e3, g["R0"], g["delta"], g["delta_z"], g["periods"]) + assert np.max(np.abs(axis - exact)) < 1e-3 + + +def test_helical_vacuum_field_is_vacuum_and_finite_on_axis(): + from struphy import equils + from struphy.geometry.domains import HollowCylinder + + eq = equils.HelicalVacuumField(delta=0.5, k=3.0) + eq.domain = HollowCylinder(a1=0.0, a2=0.3, Lz=2 * np.pi) + point, h = np.array([0.1, 0.05, 0.3]), 1e-5 + + def B(p): + return np.array(eq.b_xyz(*(np.array([c]) for c in p)), dtype=float).ravel() + + d = [(B(point + h * e) - B(point - h * e)) / (2 * h) for e in np.eye(3)] + assert abs(d[0][0] + d[1][1] + d[2][2]) < 1e-9 + assert max(abs(d[1][2] - d[2][1]), abs(d[2][0] - d[0][2]), abs(d[0][1] - d[1][0])) < 1e-9 + etas = (np.array([0.0, 0.5]), np.array([0.2]), np.array([0.4])) + for name in ("absB0", "b1", "b2", "unit_b1", "gradB1", "curl_unit_b_dot_b0", "curl_unit_b2"): + assert np.all(np.isfinite(np.array(getattr(eq, name)(*etas), dtype=float))), name + + +def test_cartesian_equilibria_unchanged_away_from_axis(): + """The direct Cartesian evaluation (finite on a polar axis) equals the former pull-back/push-forward.""" + from struphy import equils + from struphy.fields_background.base import FluidEquilibriumWithB, MHDequilibrium + from struphy.geometry.domains import HollowCylinder + + eq = equils.ScrewPinch(a=0.3, R0=1.0) + eq.domain = HollowCylinder(a1=0.05, a2=0.3, Lz=2 * np.pi) + etas = (np.linspace(0.1, 1, 4), np.linspace(0, 1, 5), np.linspace(0, 1, 3)) + for name, base in ( + ("b_cart", FluidEquilibriumWithB), + ("b1", FluidEquilibriumWithB), + ("b2", FluidEquilibriumWithB), + ("j_cart", MHDequilibrium), + ("j2", MHDequilibrium), + ("gradB_cart", FluidEquilibriumWithB), + ): + new, old = getattr(eq, name)(*etas), getattr(base, name)(eq, *etas) + if isinstance(new, tuple): + new, old = new[0], old[0] + assert np.allclose(np.array(new), np.array(old), atol=1e-13), name + + +# ---------------------------------------------------------------------------------------------------- Struphy +def test_field_lines_divergence_free_and_iota(): + from struphy.polar.homotopy.demos import fieldlines + + result = fieldlines.study(turns=2, steps_per_turn=40, verbose=False, num_elements=(6, 18, 16)) + assert result["divergence"] < 1e-12 + for (r, _, z), w in zip(result["seeds"], result["windings"]): + if z == 0.0 and 0.2 < r < 0.6: + assert w == pytest.approx(fieldlines.iota(r), abs=5e-3) + + +def _run_simulation(model, domain, derham_opts, dt, tend, folder, setup): + from struphy import EnvironmentOptions, Simulation, Time, grids + + sim = Simulation( + model=model, + env=EnvironmentOptions(out_folders=str(folder), sim_folder="sim"), + time_opts=Time(dt=dt, Tend=tend), + domain=domain, + grid=grids.TensorProductGrid(num_elements=(4, 12, 1)), + derham_opts=derham_opts, + ) + setup(model) + sim.run() + from struphy import Output + + return Output(folder / "sim") + + +def test_scalar_wave_matches_core_leapfrog(tmp_path): + """The ScalarWave model with polar splines integrates exactly the core leapfrog (same spaces and dt).""" + from struphy import DerhamOptions, domains + from struphy.initial.base import GenericPerturbation + from struphy.linear_algebra.solver import SolverParameters + from struphy.models import ScalarWave + from struphy.polar.homotopy.demos.waves import WaveProblem + + def pulse(x, y): + return np.exp(-((x - 0.2) ** 2 + y**2) / (2 * 0.2**2)) + + problem = WaveProblem(0, 2, degree=3, radial_cells=4, angular_cells=12) + dt = problem.dt + c = problem.complex + problem.initial = c.duals(0).T @ c.spaces[0].commuting_projection(pulse) + + def setup(model): + model.em_fields.phi.save_data = True + model.propagators.wave.options = model.propagators.wave.Options( + precond="PolarMassMatrixPreconditioner", solver_params=SolverParameters(tol=1e-13, maxiter=500) + ) + model.em_fields.phi.add_perturbation( + GenericPerturbation(lambda x, y, z: pulse(x, y) + 0 * z, given_in_basis="physical") + ) + + derham_opts = DerhamOptions( + degree=(3, 3, 1), + bcs=(("free", "dirichlet"), None, None), + polar_splines=True, + polar_regularity=2, + ) + out = _run_simulation(ScalarWave(), domains.PerturbedDiskCylinder(), derham_opts, dt, 3 * dt, tmp_path, setup) + times, snapshots, _ = problem.evolve(3 * dt + 1e-12) + e1, e2 = np.linspace(0.0, 1, 9), np.linspace(0, 1, 13) + phi = out.evaluate("em_fields/phi", t=-1, eta1=e1, eta2=e2, eta3=np.array([0.5])).values.reshape(9, 13) + core = problem.observable_space.evaluate(snapshots[3], e1, e2) + assert np.max(np.abs(phi - core)) < 1e-9 * np.max(np.abs(core)) diff --git a/src/struphy/polar/tests/test_homotopy_convergence.py b/src/struphy/polar/tests/test_homotopy_convergence.py new file mode 100644 index 000000000..e9b57dbc8 --- /dev/null +++ b/src/struphy/polar/tests/test_homotopy_convergence.py @@ -0,0 +1,65 @@ +"""Convergence of the polar projections under uniform mesh refinement (port of ``test/test_convergence.jl``). + +With cubic splines in both directions, the 0-form projections converge with order 4 and the 1- and 2-form projections +with order 3 in the physical L2 norm, also on the perturbed disk where the spaces are not r-compatible with the +mapping. +""" + +import numpy as np +import pytest + +from struphy.polar.homotopy import PolarComplex2D +from struphy.polar.tests.reference_core import ClampedBSplineSpace, PeriodicBSplineSpace, perturbed_disk + + +def f(x, y): + return np.exp(x) * np.sin(2 * y + 1) + + +def grad_f(x, y): + return (np.exp(x) * np.sin(2 * y + 1), 2 * np.exp(x) * np.cos(2 * y + 1)) + + +def A(x, y): + return (np.cos(x - y), x * np.exp(y)) + + +def B(x, y): + return np.sin(x + 2 * y) + x**2 + + +CELLS = (4, 8, 16) + + +def _complexes(): + return [ + PolarComplex2D(ClampedBSplineSpace(3, n), PeriodicBSplineSpace(3, 2 * n), perturbed_disk(), regularity=2) + for n in CELLS + ] + + +@pytest.mark.convergence +@pytest.mark.parametrize("method", ["commuting", "l2"]) +def test_polar_projection_convergence(method): + targets = (f, A, B) + errors = [[], [], []] + for polar in _complexes(): + for k in range(3): + if method == "commuting": + projection = polar[k](targets[k]) + else: + projection = polar.l2_projection(k, targets[k]) + errors[k].append(polar.spaces[k].l2_error(projection, targets[k], additional_order=4)) + for k, rate in enumerate((4, 3, 3)): + e = np.array(errors[k]) + assert np.all(e[1:] / e[:-1] < 1) + assert np.log2(e[-2] / e[-1]) > rate - 0.5 + + +@pytest.mark.convergence +def test_gradient_of_commuting_projection_converges(): + errors = [] + for polar in _complexes(): + gradient = polar.spaces[0].derivative_matrix @ polar[0](f) + errors.append(polar.spaces[1].l2_error(gradient, grad_f, additional_order=4)) + assert np.log2(errors[-2] / errors[-1]) > 2.5 diff --git a/src/struphy/polar/tests/test_homotopy_core.py b/src/struphy/polar/tests/test_homotopy_core.py new file mode 100644 index 000000000..a6a1b3718 --- /dev/null +++ b/src/struphy/polar/tests/test_homotopy_core.py @@ -0,0 +1,412 @@ +"""Tests of the serial high-order polar FEEC core, ported from ``test/test_paper_theorems.jl``, +``test/test_boundary_homotopy.jl`` and ``test/test_curl_curl.jl`` of PolarFEEC.jl (same parameters and tolerances).""" + +import numpy as np +import pytest +import scipy.sparse as sp + +from struphy.polar.homotopy import PolarComplex2D, monomial_basis, zernike_basis +from struphy.polar.homotopy.linalg import as_dense +from struphy.polar.homotopy.solvers import curl_curl_free_lunch, range_projection +from struphy.polar.tests.reference_core import ( + ClampedBSplineSpace, + CoordinateTransformation2D, + DiskMapping, + PeriodicBSplineSpace, + TrigonometricSpace, + complex_step, + perturbed_disk, +) + +DISC = DiskMapping() + + +def gauss01(n): + t, w = np.polynomial.legendre.leggauss(n) + return (t + 1) / 2, w / 2 + + +def integrate(f, a, b, n=40): + t, w = gauss01(n) + return (b - a) * sum(wi * f(a + (b - a) * ti) for ti, wi in zip(t, w)) + + +def range_coordinates(form, A): + """Coordinates of the columns of ``A`` in the extraction, or ``None`` if not in its range.""" + A = as_dense(A) + coordinates = form.duals.T @ A + if np.max(np.abs(form.extraction @ coordinates - A)) <= 1e-10 * max(1.0, np.max(np.abs(A))): + return coordinates + return None + + +def continuous_zero_domain(c): + """Coefficient vectors of continuous 0-forms (constant first ring).""" + nx, ny = c.radial.n, c.angular.n + centre = np.zeros(nx * ny) + centre[::nx] = c.angular.constant_coefficients() + others = [i + j * nx for j in range(ny) for i in range(1, nx)] + return np.column_stack([centre, np.eye(nx * ny)[:, others]]) + + +def poincare_gradient(A): + t, w = gauss01(20) + + def h(x, y): + return sum(wq * (A(tq * x, tq * y)[0] * x + A(tq * x, tq * y)[1] * y) for tq, wq in zip(t, w)) + + return h + + +def poincare_curl(B): + t, w = gauss01(20) + + def h(x, y): + s = sum(wq * tq * B(tq * x, tq * y) for tq, wq in zip(t, w)) + return (-y * s, x * s) + + return h + + +# ------------------------------------------------------------------------------------------------- 1D lemmas +@pytest.mark.parametrize( + "space", + [ + ClampedBSplineSpace(1, 3), + ClampedBSplineSpace(4, 6), + PeriodicBSplineSpace(1, 5), + PeriodicBSplineSpace(4, 8), + PeriodicBSplineSpace(3, 6), + TrigonometricSpace(3), + ], +) +def test_unisolvence_of_1d_gdofs(space): + for s in (space, space.scaled): + assert np.linalg.matrix_rank(s.histopolation_matrix.toarray()) == s.n + + +def test_vanishing_physical_derivative(): + radial = ClampedBSplineSpace(4, 6) + for i in range(1, 5): + value = radial.basis_matrix(np.array([1e-6, 5e-7])).toarray()[:, i] + assert np.log2(value[0] / value[1]) == pytest.approx(i, abs=1e-3) + + +def test_derivative_identities_1d(): + for space in (ClampedBSplineSpace(3, 6), PeriodicBSplineSpace(3, 6), TrigonometricSpace(3)): + x = np.linspace(0.013, 0.987, 23) + u = np.cos(np.arange(space.n) + 0.3) + h = 1e-6 + derivative = (space.basis_matrix(x + h) @ u - space.basis_matrix(x - h) @ u) / (2 * h) + assert np.allclose(space.scaled.basis_matrix(x) @ (space.derivative_matrix() @ u), derivative, atol=1e-6) + + +# ------------------------------------------------------------------------------------------------- polar subspaces +CASES = [(r, kind) for r in (1, 2, 3) for kind in ("trig", "bspline")] + + +def _angular(r, kind): + return TrigonometricSpace(r + 1) if kind == "trig" else PeriodicBSplineSpace(3, 2 * r + 3) + + +@pytest.fixture(scope="module", params=CASES, ids=[f"r{r}-{k}" for r, k in CASES]) +def polar(request): + r, kind = request.param + return PolarComplex2D(ClampedBSplineSpace(3, 7), _angular(r, kind), DISC, regularity=r) + + +def test_dimensions(polar): + r = polar.regularity + nx0, ny0, nx1, ny1 = polar._shapes + n0, n1G, n1C, n2 = (r + 1) * (r + 2) // 2, (r + 3) * r // 2, (r - 1) * r // 2, (r - 1) * r // 2 + assert polar.central_counts() == (n0, n1G + n1C, n2) + assert polar.dimension(0) == n0 + (nx0 - r - 1) * ny0 + assert polar.dimension(1) == n1G + n1C + (nx1 - r) * ny0 + (nx0 - r - 1) * ny1 + assert polar.dimension(2) == n2 + (nx1 - r) * ny1 + + +def test_conformity_conditions(polar): + nx0, ny0, nx1, ny1 = polar._shapes + E1, E2 = polar.extraction(1).toarray(), polar.extraction(2).toarray() + first = nx1 * ny0 + a_x_first = E1[0:first:nx1] + a_y = E1[first:] + assert np.max(np.abs(a_y[0::nx0])) == 0 + assert np.max(np.abs(a_y[1::nx0] - polar.angular.derivative_matrix() @ a_x_first)) <= 1e-13 + assert np.max(np.abs(E2[0::nx1])) == 0 + + +def test_conforming_two_form_projector(polar): + r = polar.regularity + nx1, ny1 = polar._shapes[2:] + Pt = polar.reduced_two_form_projector(r) + assert (Pt @ Pt != Pt).nnz == 0 + image = Pt @ np.sin(np.arange(1, nx1 * ny1 + 1)) + assert np.all(image[[i + j * nx1 for j in range(ny1) for i in range(r)]] == 0) + w = polar.two_form_integral_vector + assert np.allclose(Pt.T @ w, w) + + +def test_cochain_and_chain_complex(polar): + G, C = polar.spaces[0].derivative_matrix, polar.spaces[1].derivative_matrix + E0, E1, E2 = (polar.extraction(k) for k in range(3)) + assert range_coordinates(polar[1], G @ E0) is not None + assert range_coordinates(polar[2], C @ E1) is not None + assert range_coordinates(polar[0], polar.gradient_homotopy_matrix @ E1) is not None + assert range_coordinates(polar[1], polar.curl_homotopy_matrix @ E2) is not None + assert np.max(np.abs(as_dense(polar.homotopy(1) @ polar.homotopy(2)))) <= 1e-12 + ranks = np.linalg.matrix_rank(as_dense(polar.derivative(0))) + np.linalg.matrix_rank(as_dense(polar.homotopy(2))) + assert ranks == polar.dimension(1) + + +def test_cohomology(polar): + assert polar.cohomology_dimensions() == [1, 0, 0] + dirichlet = PolarComplex2D(polar.radial, polar.angular, DISC, regularity=polar.regularity, boundary="dirichlet") + assert dirichlet.cohomology_dimensions() == [0, 0, 1] + + +def test_discrete_homotopy_operators(polar): + for k in range(3): + n = polar.dimension(k) + contraction = np.zeros((n, n)) + if k < 2: + contraction += as_dense(polar.homotopy(k + 1) @ polar.derivative(k)) + if k > 0: + contraction += as_dense(polar.derivative(k - 1) @ polar.homotopy(k)) + assert np.allclose(contraction, np.eye(n) - as_dense(polar.cohomology_projection(k)), atol=1e-11) + + +def test_one_form_projection_fixes_external_inputs(polar): + E1 = polar.extraction(1) + central = polar[1].central_count + external = [E1[:, c].nonzero()[0].min() for c in range(central, E1.shape[1])] + P1 = polar[1].projector.toarray() + assert np.allclose(P1[:, external], np.eye(E1.shape[0])[:, external], atol=1e-12) + + +def test_commuting_projection_matrices(polar): + P = [polar[k].projector.toarray() for k in range(3)] + G, C = polar.spaces[0].derivative_matrix.toarray(), polar.spaces[1].derivative_matrix.toarray() + HG, HC = polar.gradient_homotopy_matrix.toarray(), polar.curl_homotopy_matrix.toarray() + continuous = continuous_zero_domain(polar) + assert np.allclose(P[1] @ G @ continuous, G @ P[0] @ continuous, atol=1e-11) + assert np.allclose(C @ P[1], P[2] @ C, atol=1e-11) + assert np.allclose(HG @ P[1], P[0] @ HG, atol=1e-11) + assert np.allclose(HC @ P[2], P[1] @ HC, atol=1e-11) + assert np.allclose(P[1], G @ P[0] @ HG + HC @ P[2] @ C, atol=1e-11) + + +def f0(x, y): + return np.exp(x) * np.cos(y) + x * y**2 + + +def grad_f0(x, y): + return (np.exp(x) * np.cos(y) + y**2, -np.exp(x) * np.sin(y) + 2 * x * y) + + +def A1(x, y): + return (np.sin(y) + x**2, x * y + np.cos(x)) + + +def curl_A1(x, y): + return y - np.sin(x) - np.cos(y) + + +def test_commuting_projections_on_fields(polar): + G, C = polar.spaces[0].derivative_matrix, polar.spaces[1].derivative_matrix + o = dict(additional_order=6) + assert np.allclose(G @ polar[0](f0, **o), polar[1](grad_f0, **o), atol=1e-9) + assert np.allclose(C @ polar[1](A1, **o), polar[2](curl_A1, **o), atol=1e-9) + assert np.allclose( + polar.gradient_homotopy_matrix @ polar[1](A1, **o), polar[0](poincare_gradient(A1), **o), atol=1e-8 + ) + assert np.allclose( + polar.curl_homotopy_matrix @ polar[2](curl_A1, **o), polar[1](poincare_curl(curl_A1), **o), atol=1e-8 + ) + + +def test_conservation_properties(polar): + def B(x, y): + return np.exp(x - y) + x**2 + + projected = polar[2](B, additional_order=6) + w = polar.two_form_integral_vector + exact = ( + 2 + * np.pi + * integrate( + lambda rho: rho * integrate(lambda t: B(rho * np.cos(2 * np.pi * t), rho * np.sin(2 * np.pi * t)), 0, 1), + 0, + 1, + ) + ) + assert w @ projected == pytest.approx(exact, rel=1e-9) + + def A(x, y): + return (-y * np.exp(x), x + y**3) + + nx0, ny0, nx1, ny1 = polar._shapes + circulation = np.zeros(polar.spaces[1].n) + first = nx1 * ny0 + integrals = polar.angular.scaled.coefficient_integrals + for j in range(ny1): + circulation[first + nx0 - 1 + j * nx0] = integrals[j] + + def tangential(t): + a = A(np.cos(2 * np.pi * t), np.sin(2 * np.pi * t)) + return a[0] * (-2 * np.pi * np.sin(2 * np.pi * t)) + a[1] * 2 * np.pi * np.cos(2 * np.pi * t) + + boundary = integrate(tangential, 0, 1) + assert circulation @ polar[1](A, additional_order=6) == pytest.approx(boundary, rel=1e-9) + + +def test_polynomial_reproduction_under_compatibility(polar): + if not isinstance(polar.angular, TrigonometricSpace): + pytest.skip("r-compatibility needs trigonometric angular functions") + r = polar.regularity + points = [(rho, t) for rho in (0.0, 1e-3, 0.05, 0.3, 0.9) for t in (0.0, 0.21, 0.66)] + for T in zernike_basis(r): + projected = polar[0](T) + for rho, t in points: + X, Y = DISC(rho, t) + assert polar.spaces[0].evaluate(projected, rho, t)[0, 0] == pytest.approx(T(X, Y), abs=1e-11) + assert np.allclose(projected, polar.spaces[0].commuting_projection(T), atol=1e-11) + if r >= 2: + for T in zernike_basis(r - 2): + projected = polar[2](T) + for rho, t in points: + if rho > 0: + X, Y = DISC(rho, t) + assert polar.spaces[2].evaluate(projected, rho, t)[0, 0] == pytest.approx(T(X, Y), abs=1e-10) + G = polar.spaces[0].derivative_matrix + for T in zernike_basis(r): + + def gradient(x, y, T=T): + return (complex_step(T, x, y, argnum=0), complex_step(T, x, y, argnum=1)) + + assert np.allclose(polar[1](gradient), G @ polar[0](T), atol=1e-9) + + +# ------------------------------------------------------------------------------------------------- boundary +def test_boundary_compatible_homotopy_identities(): + polar = PolarComplex2D( + ClampedBSplineSpace(3, 6), TrigonometricSpace(3), perturbed_disk(), regularity=3, boundary="dirichlet" + ) + hG, hC, Bm = (as_dense(a) for a in (polar.homotopy(1), polar.homotopy(2), polar.cohomology_projection(2))) + D0, D1 = polar.derivative(0).toarray(), polar.derivative(1).toarray() + assert np.allclose(hG @ D0, np.eye(D0.shape[1]), atol=1e-10) + assert np.allclose(D0 @ hG + hC @ D1, np.eye(D0.shape[0]), atol=1e-10) + assert np.allclose(D1 @ hC, np.eye(D1.shape[0]) - Bm, atol=1e-11) + assert np.max(np.abs(hG @ hC)) <= 1e-11 + M2 = polar.mass_matrix(2).toarray() + assert np.allclose(Bm.T @ M2, M2 @ Bm, atol=1e-11) + w = polar.extraction(2).T @ polar.two_form_integral_vector + assert np.allclose(w @ Bm, w, atol=1e-12) + + +def generic_disc_map(): + return CoordinateTransformation2D( + lambda r, t: (r * np.cos(2 * np.pi * t), r * np.sin(2 * np.pi * t)), has_polar_axis=True + ) + + +def scaled_residual(actual, expected): + actual, expected = as_dense(actual), as_dense(expected) + return np.max(np.abs(actual - expected)) / max(1.0, np.max(np.abs(expected))) + + +@pytest.mark.parametrize("regularity", [1, 4]) +@pytest.mark.parametrize("radial_cells", [2, 4]) +@pytest.mark.parametrize("geometry", ["disc", "perturbed disc"]) +@pytest.mark.parametrize("angular_kind", ["trig", "bspline"]) +def test_boundary_homotopy_regression(regularity, radial_cells, geometry, angular_kind): + transformation = generic_disc_map() if geometry == "disc" else perturbed_disk() + angular = ( + TrigonometricSpace(regularity) + if angular_kind == "trig" + else PeriodicBSplineSpace(regularity, 4 if regularity == 1 else 10) + ) + radial = ClampedBSplineSpace(regularity, radial_cells) + polar = PolarComplex2D( + radial, angular, transformation, regularity, basis=monomial_basis(regularity), boundary="dirichlet" + ) + D0, D1 = polar.derivative(0), polar.derivative(1) + hG, hC, Bm = polar.homotopy(1), polar.homotopy(2), polar.cohomology_projection(2) + I0, I1, I2 = (np.eye(polar.dimension(k)) for k in range(3)) + assert scaled_residual(hG @ D0, I0) < 8e-11 + assert scaled_residual(as_dense(D0) @ as_dense(hG) + as_dense(hC) @ as_dense(D1), I1) < 8e-11 + assert scaled_residual(as_dense(D1) @ as_dense(hC), I2 - as_dense(Bm)) < 3e-11 + norm_inf = lambda A: np.max(np.sum(np.abs(as_dense(A)), axis=1)) # noqa: E731 + assert norm_inf(as_dense(hG) @ as_dense(hC)) / max(1.0, norm_inf(hG) * norm_inf(hC)) < 1e-11 + full_curl = polar.spaces[1].derivative_matrix + E1, E2 = polar.extraction(1), polar.extraction(2) + assert scaled_residual(full_curl @ E1 @ as_dense(hC), E2 @ (I2 - as_dense(Bm))) < 3e-11 + assert np.all((polar.extraction(0) @ as_dense(hG))[polar.zero_form_dirichlet_rows()] == 0) + assert np.all((E1 @ as_dense(hC))[polar.one_form_dirichlet_rows()] == 0) + B = as_dense(Bm) + assert scaled_residual(B @ B, B) < 3e-12 + integral = E2.T @ polar.two_form_integral_vector + assert scaled_residual(integral @ B, integral) < 3e-12 + M2 = polar.mass_matrix(2).toarray() + assert scaled_residual(B.T @ M2, M2 @ B) < 3e-11 + # the projected constant, computed independently through the L2 projection + projected_constant = polar.l2_projection(2, lambda x, y: np.ones_like(x)) + constant_reduced = polar.duals(2).T @ projected_constant + constant_integral = integral @ constant_reduced + assert constant_integral > 0 + assert constant_integral == pytest.approx(constant_reduced @ M2 @ constant_reduced, abs=3e-11) + unit = constant_reduced / constant_integral + assert np.max(np.abs(B @ unit - unit)) < 3e-11 + # lifted fields have vanishing tangential trace on the boundary + data = polar.duals(2).T @ np.kron( + polar.angular.constant_one_form_coefficients(), np.eye(radial.scaled.n)[regularity] + ) + mean_free = data - B @ data + lifted = as_dense(hC) @ data + assert np.max(np.abs(as_dense(D1) @ lifted - mean_free)) / max(1.0, np.max(np.abs(mean_free))) < 3e-11 + field = E1 @ lifted + assert np.all(field[polar.one_form_dirichlet_rows()] == 0) + for t in (0.0, 0.137, 0.41, 0.79): + values = polar.spaces[1].logical.evaluate(field, [1.0], [t]) + tangent = transformation.jacobian_matrix(1.0, t)[:, 1] + # logical 1-form proxy: A . dF/dy = (J^T A)_y + assert abs(values[1][0, 0]) < 3e-11 * max(1.0, np.linalg.norm(tangent)) + potential = polar.extraction(0) @ (as_dense(hG) @ np.cos(np.arange(1, polar.dimension(1) + 1))) + assert np.all(potential[polar.zero_form_dirichlet_rows()] == 0) + + +# ------------------------------------------------------------------------------------------------- 2D curl-curl +def test_curl_curl_free_lunch(): + """C1: one mass solve ``a = h_2 (I - B) M_2^{-1} l`` equals the direct range projection.""" + petals = DiskMapping(radius=lambda r, t: 0.3 * r * (1 + 0.2 * r * np.cos(6 * np.pi * t))) + section = PolarComplex2D( + ClampedBSplineSpace(3, 4), PeriodicBSplineSpace(3, 8), petals, regularity=2, boundary="dirichlet" + ) + D = section.derivative(1) + M2 = section.mass_matrix(2) + load = M2 @ np.cos(np.arange(1, section.dimension(2) + 1)) + mean = as_dense(section.cohomology_projection(2)) + x = section.mass_solver(2).solve(load) + a = as_dense(section.homotopy(2)) @ (x - mean @ x) + b = range_projection(section, 1, load) + assert np.linalg.norm(D @ a - b) <= 1e-10 * np.linalg.norm(b) + assert np.linalg.norm(D.T @ (load - M2 @ (D @ a))) <= 1e-10 * np.linalg.norm(D.T @ load) + x = np.sin(np.arange(1, section.dimension(2) + 1)) + assert np.linalg.norm(D @ (as_dense(section.homotopy(2)) @ x) - (x - mean @ x)) <= 1e-12 * np.linalg.norm(x) + # the written-out algorithm of the paper, for a 1-form load + load1 = section.mass_matrix(1) @ np.cos(np.arange(1, section.dimension(1) + 1)) + expected = range_projection(section, 1, section.homotopy(2).T @ load1) + assert np.allclose(D @ curl_curl_free_lunch(section, load1), expected, atol=1e-10 * np.max(np.abs(expected))) + + +def test_mass_and_l2_projection_reproduce_extraction(): + c = PolarComplex2D(ClampedBSplineSpace(3, 8), TrigonometricSpace(3), perturbed_disk(), regularity=3) + for k in range(3): + E = c.extraction(k) + assert np.allclose((c.duals(k).T @ E).toarray(), np.eye(E.shape[1]), atol=5e-12) + P = c[k].projector + assert abs((P @ P - P)).max() < 5e-12 + assert abs((P @ E - E)).max() < 5e-12 + M = c.mass_matrix(k).toarray() + assert np.allclose(M, (E.T @ c.full_mass_matrix(k) @ E).toarray(), atol=1e-12) + assert sp.issparse(c.derivative(0)) diff --git a/src/struphy/polar/tests/test_homotopy_demos.py b/src/struphy/polar/tests/test_homotopy_demos.py new file mode 100644 index 000000000..abdce3ca9 --- /dev/null +++ b/src/struphy/polar/tests/test_homotopy_demos.py @@ -0,0 +1,168 @@ +"""Regression tests of the polar-FEEC demos (ports of PolarFEEC.jl examples 04-06 and the talk studies).""" + +import numpy as np +import pytest + + +def test_regularity_study(): + """The k-th derivative is continuous at the centre exactly when k <= r (PolarFEEC.jl example 04).""" + from struphy.polar.homotopy.demos import regularity + + results = regularity.study(verbose=False) + for r in range(4): + delta = results["defects"][f"r = {r}"] + for k in range(1, 4): + if k <= r: + assert delta[k - 1] < 1e-7, (r, k, delta) + else: + assert delta[k - 1] > 1e-6, (r, k, delta) + assert np.allclose(results["defects"]["tensor product"], results["defects"]["r = 0"]) + + +def test_pde_checks_and_reference_errors(): + """Gauges, boundary conditions and the L2 errors of PolarFEEC.jl example 05 (p = 3, n = 6).""" + from struphy.polar.homotopy.demos import pde + + solution = pde.solve_all(3, 6) + checks = pde.checks(solution) + assert abs(checks["mean_of_source"]) < 1e-12 + assert abs(checks["u_centre"]) < 1e-12 + assert checks["max_u_on_boundary"] < 1e-12 and checks["tangential_trace"] < 1e-12 + assert checks["gauge_free"] < 1e-10 and checks["gauge_dirichlet"] < 1e-10 + # values of PolarFEEC.jl + reference = [3.4128468787e-03, 4.2082992517e-04, 1.5702930010e-03, 1.2889397521e-03] + assert np.allclose(pde.errors(solution), reference, rtol=1e-8) + + +@pytest.mark.convergence +def test_pde_convergence_rates(): + from struphy.polar.homotopy.demos import pde + + cells = (6, 8, 12) + table = pde.convergence(cells, (2, 3), verbose=False) + for p, errors in table.items(): + rates = np.log(errors[:-1] / errors[1:]) / np.log(np.array(cells[1:]) / np.array(cells[:-1]))[:, None] + assert np.all(rates[-1, :2] > p + 1 - 0.3), (p, rates) + assert np.all(rates[-1, 2:] > p - 0.3), (p, rates) + + +def test_maxwell_spectrum(): + """No spurious modes: kernel = dim V0, dim V2 - 1 nonzero eigenvalues, upper bounds of the exact ones.""" + from struphy.polar.homotopy.demos import maxwell + + s = maxwell.study(meshes=((8, 16),), compared=20, verbose=False) + r = s["results"][0] + polar = r["polar"] + assert r["kernel"] == polar.dimension(0) + assert len(r["lam"]) == polar.dimension(2) - 1 + assert np.max(r["error"]) < 2e-2 + + +def test_projection_commutation(): + from struphy.polar.homotopy.demos import projections + + defects = projections.commutation((2,), (4, 8), method="commuting") + assert np.max(defects[2]) < 1e-12 + errors = projections.convergence((2,), (4, 8, 16)) + rates = np.log2(errors[2][:-1] / errors[2][1:])[-1] + assert rates[0] > 2.5 and rates[1] > 1.6 and rates[2] > 1.6 + + +def test_spectrum_explorer(tmp_path): + """Curl-curl eigenvalues on the unit disk: no spurious modes, upper bounds of the exact ones; the deficit with + respect to a higher regularity vanishes for the modes that already are in the more regular space.""" + import json + + from struphy.polar.homotopy.demos import spectrum + + kwargs = dict(degree=4, cells=(4, 14)) + reference = spectrum.complex_(3, **kwargs) + deficits = [] + for r in (0, 1, 2): + s = spectrum.series(r, reference=reference, count=8, **kwargs) + assert s["kernel"] == s["complex"].dimension(0) + exact = np.array(s["exact"], dtype=float) + assert np.all(s["lam"] >= exact * (1 - 1e-12)) + assert np.max(np.abs(s["lam"] - exact) / exact) < 2e-3 + deficits.append(s["deficit"][0]) + assert deficits[2] < 1e-10 < deficits[0] + + html = spectrum.write_explorer(tmp_path, count=3, pixels=16, regularities=(0, 1), verbose=False, **kwargs) + assert html.exists() + text = (tmp_path / "spectrum_data.js").read_text() + data = json.loads(text[text.index("=") + 1 :].strip().rstrip(";")) + assert [s["enrichment"] for s in data["spaces"]] == ["cochain_chain", "cochain"] + for s in data["spaces"]: + for entry in s["series"]: + assert len(entry["lambda"]) == 3 + assert all((tmp_path / "spectrum" / name).exists() for name in entry["image"]) + + +def test_wave_3d_helical_tube(): + """Leapfrog Maxwell in the helical tube with polar splines: a stable time step and a small energy error.""" + import warnings + + from struphy.polar.homotopy.demos import waves3d + + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + wave = waves3d.Wave3D(1, 3, degree=2) + markers = waves3d.sample_points(3) + values, energy_error = wave.evolve( + 0.05, np.linspace(0, 0.05, 3), lambda f: np.concatenate([np.asarray(v) for v in f(markers)]) + ) + assert np.all(np.isfinite(values)) and np.max(np.abs(values[0])) > 0 + assert energy_error < 0.1 + + +@pytest.mark.parametrize("angular", ["bspline", "trig"]) +def test_guiding_centre_orbits(angular): + """Guiding centres in the projected vacuum field near the axis: finite orbits, conserved energy.""" + from struphy.polar.homotopy.demos import orbits + + field = orbits.ProjectedVacuumField(2, angular, num_elements=(4, 8, 8), degree=(3, 3, 3)) + path, errors = orbits.integrate(field, 0.05, periods=0.2) + assert np.all(np.isfinite(path)) + assert np.max(np.hypot(path[..., 0], path[..., 1])) < orbits.A + assert np.max(errors) < 1e-6 + + +def test_talk_figures_cli(tmp_path): + """The figure script runs (one quick group).""" + import matplotlib + + matplotlib.use("Agg") + from struphy.polar.homotopy.demos import talk_figures + + talk_figures.main([str(tmp_path), "--only", "regularity"]) + assert (tmp_path / "result_regularity_zeroform.png").exists() + assert (tmp_path / "result_regularity_twoform.png").exists() + + +def test_wave_convergence_is_finite(): + """The sampling grid of the wave studies leaves out the pole, where the 2-form pushforward is singular.""" + from struphy.polar.homotopy.demos import waves + + x, _ = waves.sample_points(21) + assert np.all(x > 0) + study = waves.convergence( + 1, regularities=(0, 1), radial_cells=(3, 4), reference_cells=6, degree=2, final_time=0.05, verbose=False + ) + for errors in study["final"].values(): + assert np.all(np.isfinite(errors)) + + +def test_schematic_diagrams(): + """The de Rham diagrams are tikzfigure matrices with labelled arrows (TikZ code only, no LaTeX needed).""" + matrix = pytest.importorskip("tikzfigure.core.matrix") + if not hasattr(matrix.Matrix, "arrow"): + pytest.skip("needs tikzfigure >= 0.5.0") + from struphy.polar.homotopy.demos import schematics + + code = schematics.derham_logical_and_physical().generate_tikz() + assert code.count("\\matrix") == 2 + assert "(l-1-1) to node[above, font=\\scriptsize] {grad} (l-1-2)" in code + assert "(l.east) to node[above, font=\\scriptsize] {pushforward $F_*$} (p.west)" in code + code = schematics.derham_polar_cochain_chain().generate_tikz() + assert "(h-1-2) to node[above, font=\\scriptsize] {$\\mathrm{grad}^\\dagger$} (h-1-1)" in code + assert code.count("color=struphyblue") == 6 diff --git a/src/struphy/polar/tests/test_homotopy_derham.py b/src/struphy/polar/tests/test_homotopy_derham.py new file mode 100644 index 000000000..fc3e476a6 --- /dev/null +++ b/src/struphy/polar/tests/test_homotopy_derham.py @@ -0,0 +1,160 @@ +"""Tests of the high-order (homotopy) polar splines in :class:`~struphy.feec.psydac_derham.Derham`.""" + +import numpy as np +import pytest + + +def _derham(domain, num_elements=(8, 12, 1), degree=(3, 3, 1), regularity=2, bcs=(("free", "free"), None, None), **kw): + from maybempi import MPI + + from struphy.feec.psydac_derham import Derham + from struphy.io.options import DerhamOptions + from struphy.topology.grids import TensorProductGrid + + opts = DerhamOptions(degree=degree, bcs=bcs, polar_splines=True, polar_regularity=regularity, **kw) + # eta1 is not decomposed: the process at the pole holds all polar rings + grid = TensorProductGrid(num_elements=num_elements, mpi_dims_mask=(False, True, True)) + return Derham(grid, opts, comm=MPI.COMM_WORLD, domain=domain) + + +def _domain(name): + from struphy.geometry import domains + + if name == "PerturbedDiskCylinder": + return domains.PerturbedDiskCylinder(R1=1.0, eps=0.2, k=3, Lz=1.0) + if name == "HollowCylinder": + return domains.HollowCylinder(a1=0.0, a2=1.0, Lz=1.0) + if name == "HollowTorus": + return domains.HollowTorus(a1=0.0, a2=0.3, R0=1.0, tor_period=1) + return getattr(domains, name)(num_elements=(8, 12), degree=(3, 3)) + + +@pytest.mark.parametrize("domain_name", ["PerturbedDiskCylinder", "HollowCylinder", "IGAPolarCylinder"]) +@pytest.mark.parametrize("regularity", [1, 2, 3]) +def test_dimensions_match_total_complex(domain_name, regularity): + derham = _derham(_domain(domain_name), regularity=regularity, num_elements=(8, 12, 4), degree=(3, 3, 2)) + c = derham.ck_blocks.complex2d + n3B, n3M = derham.ck_blocks.n3["B"], derham.ck_blocks.n3["M"] + d0, d1, d2 = (c.dimension(k) for k in range(3)) + expected = {"H1": d0 * n3B, "Hcurl": d1 * n3B + d0 * n3M, "Hdiv": d1 * n3M + d2 * n3B, "L2": d2 * n3M} + for space_id, dim in expected.items(): + assert derham.polar_coeff_spaces[derham.space_to_form[space_id]].dimension == dim, space_id + + +@pytest.mark.parametrize("domain_name", ["PerturbedDiskCylinder", "IGAPolarCylinder"]) +@pytest.mark.parametrize("regularity", [1, 3]) +def test_poloidal_blocks_equal_core_derivatives(domain_name, regularity): + """With one cell in eta3 the poloidal parts of grad/curl are the core matrices D_bar = F~^T D E.""" + derham = _derham(_domain(domain_name), regularity=regularity) + blocks = derham.ck_blocks + c = blocks.complex2d + sets = blocks._sets + D0, D1 = c.derivative(0).toarray(), c.derivative(1).toarray() + s0, s1x, s1y, s2 = sets["0"], sets["1x"], sets["1y"], sets["2"] + # grad: Hcurl x <- H1, Hcurl y <- H1 + assert np.allclose(blocks.grad_pol_to_pol[0][0].toarray(), D0[np.ix_(s1x, s0)], atol=1e-12) + assert np.allclose(blocks.grad_pol_to_pol[1][0].toarray(), D0[np.ix_(s1y, s0)], atol=1e-12) + # curl: Hdiv z <- Hcurl x, y + assert np.allclose(blocks.curl_pol_to_pol[2][0].toarray(), D1[np.ix_(s2, s1x)], atol=1e-12) + assert np.allclose(blocks.curl_pol_to_pol[2][1].toarray(), D1[np.ix_(s2, s1y)], atol=1e-12) + # div: the polar Hdiv basis functions are Lambda^{1,1} E_l (B_x = a_y, B_y = -a_x), whose divergence is the 2D + # curl of E_l + assert np.allclose(blocks.div_pol_to_pol[0][0].toarray(), D1[np.ix_(s2, s1y)], atol=1e-12) + assert np.allclose(blocks.div_pol_to_pol[0][1].toarray(), D1[np.ix_(s2, s1x)], atol=1e-12) + + +@pytest.mark.parametrize("domain_name", ["PerturbedDiskCylinder", "HollowTorus"]) +def test_exactness(domain_name): + """d o d = 0 for the polar operators (polar and tensor-product parts).""" + from struphy.polar.tests.test_polar import _rand_polar_vector + + derham = _derham(_domain(domain_name), regularity=2, num_elements=(6, 12, 4), degree=(3, 3, 1)) + spaces = [derham.polar_coeff_spaces[f] for f in ("0", "1", "2", "3")] + ops = [derham.grad_bcfree, derham.curl_bcfree, derham.div_bcfree] + for k in range(2): + x = _rand_polar_vector(derham, spaces[k], seed=7) + y = ops[k + 1].dot(ops[k].dot(x)) + assert np.sqrt(y.inner(y)) < 1e-12 * np.sqrt(x.inner(x)) + + +def _logical_function(a, b, c): + return ( + np.sin(2 * np.pi * b) * a**2 * np.cos(2 * np.pi * c) + a**3 * np.cos(4 * np.pi * b) + a * np.sin(2 * np.pi * b) + ) + + +@pytest.mark.parametrize("domain_name", ["PerturbedDiskCylinder", "HollowTorus"]) +@pytest.mark.parametrize("regularity", [1, 3]) +@pytest.mark.parametrize("eta3", [(4, 2), (1, 1)]) +def test_commuting_projectors(domain_name, regularity, eta3): + """grad P0 = P1 grad, curl P1 = P2 curl, div P2 = P3 div (up to the quadrature of the GDOFs); also with a single + periodic cell in eta3 (the usual setting of 2D runs).""" + derham = _derham( + _domain(domain_name), + regularity=regularity, + num_elements=(8, 12, eta3[0]), + degree=(3, 3, eta3[1]), + nquads_proj=(8, 10, 10), + ) + h = 1e-6 + + def d(fun, i): + def g(a, b, c): + ep, em = [a, b, c], [a, b, c] + ep[i] = ep[i] + h + em[i] = em[i] - h + return (fun(*ep) - fun(*em)) / (2 * h) + + return g + + f = _logical_function + u, g = derham.P0(f), derham.P1([d(f, 0), d(f, 1), d(f, 2)]) + r = derham.grad.dot(u) - g + assert np.sqrt(r.inner(r)) < 1e-9 * max(np.sqrt(g.inner(g)), 1.0) + + A = [ + lambda a, b, c: np.cos(2 * np.pi * b) * a + 0 * c, + lambda a, b, c: a**2 * np.sin(2 * np.pi * b) * np.cos(2 * np.pi * c), + lambda a, b, c: a**3 * np.sin(2 * np.pi * b), + ] + curlA = [ + lambda a, b, c: d(A[2], 1)(a, b, c) - d(A[1], 2)(a, b, c), + lambda a, b, c: d(A[0], 2)(a, b, c) - d(A[2], 0)(a, b, c), + lambda a, b, c: d(A[1], 0)(a, b, c) - d(A[0], 1)(a, b, c), + ] + a1, b2 = derham.P1(A), derham.P2(curlA) + r = derham.curl.dot(a1) - b2 + assert np.sqrt(r.inner(r)) < 1e-9 * max(np.sqrt(b2.inner(b2)), 1.0) + B = [lambda a, b, c: a * np.cos(2 * np.pi * b) * np.cos(2 * np.pi * c), lambda a, b, c: a**2 + 0 * b + 0 * c, A[2]] + + def divB(a, b, c): + return d(B[0], 0)(a, b, c) + d(B[1], 1)(a, b, c) + d(B[2], 2)(a, b, c) + + b2, r3 = derham.P2(B), derham.P3(divB) + r = derham.div.dot(b2) - r3 + assert np.sqrt(r.inner(r)) < 1e-9 * max(np.sqrt(r3.inner(r3)), 1.0) + + +@pytest.mark.parametrize("domain_name", ["HollowCylinder", "IGAPolarTorus"]) +@pytest.mark.parametrize( + "bcs", [(("free", "free"), None, None), (("free", "dirichlet"), None, ("dirichlet", "dirichlet"))] +) +def test_mass_matrices_and_preconditioner(domain_name, bcs): + from feectools.linalg.solvers import inverse + + from struphy.feec.mass import WeightedMassOperators + from struphy.feec.preconditioner import PolarMassMatrixPreconditioner + from struphy.polar.tests.test_polar import _rand_polar_vector, _rel_err + + domain = _domain(domain_name) + derham = _derham(domain, num_elements=(8, 12, 4), degree=(3, 3, 2), regularity=3, bcs=bcs) + mass = WeightedMassOperators(derham, domain) + for name in ("M0", "M1", "M2", "M3", "Mv"): + M = getattr(mass, name) + pc = PolarMassMatrixPreconditioner(M) + x = M._V_boundary_op.dot(_rand_polar_vector(derham, M.domain, seed=4321)) + b = M.dot(x) + solver = inverse(M, "pcg", pc=pc, tol=1e-10, maxiter=300) + y = solver.dot(b) + assert solver.get_info()["success"], name + assert _rel_err(y, x) < 1e-8, name diff --git a/src/struphy/polar/tests/test_homotopy_layout.py b/src/struphy/polar/tests/test_homotopy_layout.py new file mode 100644 index 000000000..a853d57d7 --- /dev/null +++ b/src/struphy/polar/tests/test_homotopy_layout.py @@ -0,0 +1,178 @@ +"""The polar complex built from the 1D spaces of a Struphy Derham complex in Struphy's coefficient ordering equals the +reference construction (the ported PolarFEEC.jl spaces in PolarFEEC.jl's ordering). + +The bases differ only by the ordering of the tensor-product coefficients (Struphy's ``i * ny + j`` with periodic +splines shifted by the degree, against ``i + j * nx``) and of the external polar basis functions, so the comparison +uses the full-space operators ``E F^T`` (commuting projections), ``E D F^T`` (derivatives) and ``E h F^T`` +(homotopies).""" + +import numpy as np +import pytest +import scipy.sparse as sp + + +def _permutation(shapes, shift): + """Matrix mapping reference (PolarFEEC.jl ordering) to Struphy coefficients of stacked components.""" + blocks = [] + for nx, ny in shapes: + i, j = np.meshgrid(np.arange(nx), np.arange(ny), indexing="ij") + struphy = (i * ny + j).ravel() + reference = (i + ((j - shift) % ny) * nx).ravel() + blocks.append(sp.csr_matrix((np.ones(nx * ny), (struphy, reference)), shape=(nx * ny, nx * ny))) + return sp.block_diag(blocks).tocsr() + + +def _dense(A): + from struphy.polar.homotopy.linalg import as_dense + + return as_dense(A) + + +def _complexes(degree, cells, regularity, boundary, enrichment="cochain_chain"): + from struphy.feec.psydac_derham import Derham + from struphy.io.options import DerhamOptions + from struphy.polar.homotopy import PolarComplex2D + from struphy.polar.homotopy.derham_spaces import derham_factors + from struphy.polar.homotopy.layout import STRUPHY + from struphy.polar.tests.reference_core import ClampedBSplineSpace, PeriodicBSplineSpace, perturbed_disk + from struphy.topology.grids import TensorProductGrid + + derham = Derham( + TensorProductGrid(num_elements=(*cells, 1)), + DerhamOptions(degree=(degree, degree, 1), bcs=(("free", "free"), None, None)), + ) + radial, angular, _ = derham_factors(derham) + mapping = perturbed_disk() + kwargs = dict(boundary=boundary, enrichment=enrichment) + struphy = PolarComplex2D(radial, angular, mapping, regularity, layout=STRUPHY, **kwargs) + reference = PolarComplex2D( + ClampedBSplineSpace(degree, cells[0]), PeriodicBSplineSpace(degree, cells[1]), mapping, regularity, **kwargs + ) + return struphy, reference + + +@pytest.mark.parametrize("degree, cells", [(3, (6, 12)), (4, (5, 9))]) +@pytest.mark.parametrize("regularity", [1, 2, 3]) +@pytest.mark.parametrize("boundary", ["free", "dirichlet"]) +def test_struphy_layout_equals_reference(degree, cells, regularity, boundary): + struphy, reference = _complexes(degree, cells, regularity, boundary) + nx0, ny0, nx1, ny1 = reference._shapes + shapes = {0: [(nx0, ny0)], 1: [(nx1, ny0), (nx0, ny1)], 2: [(nx1, ny1)]} + P = {k: _permutation(shapes[k], degree) for k in range(3)} + assert [struphy.dimension(k) for k in range(3)] == [reference.dimension(k) for k in range(3)] + assert struphy.central_counts() == reference.central_counts() + + def full(c, k_out, k_in, A): + return _dense(c.extraction(k_out)) @ _dense(A) @ _dense(c.duals(k_in)).T + + for k in range(3): + a = full(struphy, k, k, sp.identity(struphy.dimension(k))) + b = P[k] @ full(reference, k, k, sp.identity(reference.dimension(k))) @ P[k].T + assert np.max(np.abs(a - b)) < 1e-12 * max(1.0, np.max(np.abs(b))), f"projector of degree {k}" + for k in range(2): + a = full(struphy, k + 1, k, struphy.derivative(k)) + b = P[k + 1] @ full(reference, k + 1, k, reference.derivative(k)) @ P[k].T + assert np.max(np.abs(a - b)) < 1e-11 * max(1.0, np.max(np.abs(b))), f"derivative of degree {k}" + for k in (1, 2): + a = full(struphy, k - 1, k, struphy.homotopy(k)) + b = P[k - 1] @ full(reference, k - 1, k, reference.homotopy(k)) @ P[k].T + assert np.max(np.abs(a - b)) < 1e-10 * max(1.0, np.max(np.abs(b))), f"homotopy of degree {k}" + assert struphy.cohomology_dimensions() == reference.cohomology_dimensions() + + +def test_derivatives_equal_the_derham_derivatives(): + """The 2D derivatives of the Struphy-ordered forms are the poloidal blocks of Derham's grad and curl.""" + from struphy.feec.psydac_derham import Derham + from struphy.io.options import DerhamOptions + from struphy.polar.homotopy.derham_spaces import derham_factors + from struphy.polar.homotopy.forms2d import FormSpace2D + from struphy.polar.homotopy.layout import STRUPHY + from struphy.topology.grids import TensorProductGrid + + derham = Derham( + TensorProductGrid(num_elements=(6, 9, 1)), + DerhamOptions(degree=(3, 2, 1), bcs=(("free", "free"), None, None)), + ) + radial, angular, _ = derham_factors(derham) + G = FormSpace2D(0, radial, angular, STRUPHY).derivative_matrix + C = FormSpace2D(1, radial, angular, STRUPHY).derivative_matrix + n3 = derham.H1_1d_serial[2].spaces[0].nbasis + assert n3 == 1 + grad = derham.grad_tp.tosparse() + curl = derham.curl_tp.tosparse() + n0 = radial.n * angular.n + n1x, n1y = radial.scaled.n * angular.n, radial.n * angular.scaled.n + assert np.allclose(grad[: n1x + n1y, :n0].toarray(), G.toarray()) + n2z = radial.scaled.n * angular.scaled.n + offset = curl.shape[0] - n2z + assert np.allclose(curl[offset:, : n1x + n1y].toarray(), C.toarray()) + + +def test_cross_sections(): + """The reference cross-sections of the polar splines: the domain at eta3 = 0 (one tensor-grid evaluation of the + domain), the unit disk, and a warning for regularity >= 2 when the cross-sections change shape along eta3.""" + import warnings + + from struphy.feec.psydac_derham import Derham + from struphy.geometry import domains + from struphy.io.options import DerhamOptions + from struphy.polar.homotopy.derham_blocks import DomainCrossSection, UnitDiskCrossSection + from struphy.topology.grids import TensorProductGrid + + disk = DomainCrossSection(domains.HollowCylinder(a1=0.0, a2=1.0, Lz=1.0)) + e1, e2 = np.meshgrid(np.linspace(0, 1, 5), np.linspace(0, 1, 7), indexing="ij") + X, Y = disk(e1, e2) + X_ref, Y_ref = UnitDiskCrossSection()(e1, e2) + assert np.allclose(X, X_ref) and np.allclose(Y, Y_ref) + assert np.allclose(disk.jacobian_matrix(e1, e2), UnitDiskCrossSection().jacobian_matrix(e1, e2)) + # scattered points give the same values as the tensor grid + Xs, Ys = disk(e1.ravel()[::-1], e2.ravel()[::-1]) + assert np.allclose(Xs, X.ravel()[::-1]) and np.allclose(Ys, Y.ravel()[::-1]) + + def build(domain, regularity, **kwargs): + options = DerhamOptions( + degree=(3, 3, 2), + bcs=(("free", "free"), None, None), + polar_splines=True, + polar_regularity=regularity, + **kwargs, + ) + return Derham(TensorProductGrid(num_elements=(6, 12, 4)), options, domain=domain) + + with warnings.catch_warnings(): + warnings.simplefilter("error") + build(domains.HollowTorus(a1=0.0, a2=0.3, R0=1.0), 2) # the same cross-section everywhere + build(domains.HelicalTube("islands", num_elements=(6, 24, 32), degree=(3, 3, 3)), 1) + with pytest.warns(UserWarning, match="cross-sections eta3 = const of the domain differ"): + build(domains.HelicalTube("islands", num_elements=(6, 24, 32), degree=(3, 3, 3)), 2) + unit = build( + domains.HelicalTube("islands", num_elements=(6, 24, 32), degree=(3, 3, 3)), 2, polar_cross_section="unit_disk" + ) + assert unit.ck_blocks.regularity == 2 + + +@pytest.mark.parametrize("domain_name", ["PerturbedDiskCylinder", "HollowCylinder"]) +def test_complex_masses_equal_struphy_masses(domain_name): + """The polar masses ``E^T M E`` of the Derham-based polar complex equal those with Struphy's weighted mass operators + on a cylinder of length one (2D 0-forms: V0, 1-forms: the eta1-eta2 components of V1, 2-forms: the eta3 component + of V2). (The tensor-product 1-form masses themselves are only formal at the pole: the ring-0 eta2-components have + divergent integrals, which the polar functions do not use.)""" + from struphy.feec.mass import WeightedMassOperators + from struphy.geometry import domains + from struphy.polar.homotopy.derham_blocks import derham_polar_complex, poloidal_derham + + kwargs = {"Lz": 1.0} if domain_name == "PerturbedDiskCylinder" else {"a1": 0.0, "a2": 1.0, "Lz": 1.0} + domain = getattr(domains, domain_name)(**kwargs) + nquads = 12 + derham = poloidal_derham((3, 3), (6, 12), domain=domain, nquads=(nquads, nquads)) + c = derham_polar_complex(derham, domain, 2) + mass = WeightedMassOperators(derham, domain) + M0, M1, M2 = (getattr(mass, name).matrix.tosparse().toarray() for name in ("M0", "M1", "M2")) + n1 = c.spaces[1].n + M2z = M2[-c.spaces[2].n :, -c.spaces[2].n :] + for k, struphy in ((0, M0), (1, M1[:n1, :n1]), (2, M2z)): + E = c.extraction(k).toarray() + # the same Gauss rule (nquads points per cell) as Struphy's mass operators + full = c.spaces[k].mass_matrix(coefficient=lambda X, Y: 1.0 + 0 * X, additional_quadrature_order=nquads - 4) + own, reference = E.T @ full.toarray() @ E, E.T @ struphy @ E + assert np.max(np.abs(own - reference)) < 1e-10 * np.max(np.abs(reference)), k diff --git a/src/struphy/polar/tests/test_polar.py b/src/struphy/polar/tests/test_polar.py index 157b20d0d..4e5a8c58e 100644 --- a/src/struphy/polar/tests/test_polar.py +++ b/src/struphy/polar/tests/test_polar.py @@ -5,6 +5,14 @@ logger = logging.getLogger("struphy") +def _pole_safe(num_elements): + """At least five elements in eta1 per process: the process at the pole holds the polar rings of regularity 1 + (three rings of 0-forms) and two tensor-product rings, also when eta1 is decomposed.""" + from maybempi import MPI + + return [max(num_elements[0], 5 * MPI.COMM_WORLD.Get_size()), *num_elements[1:]] + + @pytest.mark.parametrize("num_elements", [[8, 9, 6]]) @pytest.mark.parametrize("degree", [[3, 2, 4]]) @pytest.mark.parametrize( @@ -16,13 +24,14 @@ ) def test_spaces(num_elements, degree, bcs): from struphy.feec.psydac_derham import Derham + from struphy.geometry.domains import HollowCylinder from struphy.io.options import DerhamOptions from struphy.polar.basic import PolarDerhamSpace, PolarVector from struphy.topology.grids import TensorProductGrid grid = TensorProductGrid(num_elements=num_elements) - derham_opts = DerhamOptions(degree=degree, bcs=bcs) - derham = Derham(grid, derham_opts) + derham_opts = DerhamOptions(degree=degree, bcs=bcs, polar_splines=True) + derham = Derham(grid, derham_opts, domain=HollowCylinder(a1=0.0, a2=1.0, Lz=3.0)) logger.info("polar V0:") V = PolarDerhamSpace(derham, "H1") @@ -195,7 +204,6 @@ def test_extraction_ops_and_derivatives(num_elements, degree, bcs): from struphy.geometry.domains import IGAPolarCylinder from struphy.io.options import DerhamOptions from struphy.polar.basic import PolarDerhamSpace, PolarVector - from struphy.polar.extraction_operators import PolarExtractionBlocksC1 from struphy.polar.linear_operators import PolarExtractionOperator, PolarLinearOperator from struphy.topology.grids import TensorProductGrid @@ -203,6 +211,8 @@ def test_extraction_ops_and_derivatives(num_elements, degree, bcs): rank = comm.Get_rank() size = comm.Get_size() + num_elements = _pole_safe(num_elements) + # create control points params_map = {"num_elements": num_elements[:2], "degree": degree[:2], "Lz": 3.0, "a": 1.0} domain = IGAPolarCylinder(**params_map) @@ -317,8 +327,8 @@ def test_extraction_ops_and_derivatives(num_elements, degree, bcs): @pytest.mark.parametrize("num_elements", [[8, 3, 2], [8, 6, 2]]) def test_polar_adjoints_small_nel2(num_elements): - """ = for polar extraction operators and derivatives, - including Nel2 = 3 where n_polar == n2 for H1 and H1vec.""" + """ = for polar extraction operators and derivatives, including the smallest number of + angular elements Nel2 = 3 of regularity 1.""" import cunumpy as xp from maybempi import MPI @@ -330,6 +340,7 @@ def test_polar_adjoints_small_nel2(num_elements): from struphy.topology.grids import TensorProductGrid degree = [2, 2, 1] + num_elements = _pole_safe(num_elements) domain = IGAPolarCylinder(num_elements=num_elements[:2], degree=degree[:2], Lz=1.0, a=1.0) grid = TensorProductGrid(num_elements=num_elements) derham_opts = DerhamOptions(degree=degree, bcs=(("free", "free"), None, None), polar_splines=True) @@ -379,7 +390,7 @@ def test_polar_ops_eta1_decomposition(): comm = MPI.COMM_WORLD degree = [2, 2, 1] - num_elements = [16, 6, 4] + num_elements = _pole_safe([16, 6, 4]) domain = IGAPolarCylinder(num_elements=num_elements[:2], degree=degree[:2], Lz=1.0, a=1.0) grid = TensorProductGrid(num_elements=num_elements, mpi_dims_mask=(True, False, False)) derham_opts = DerhamOptions(degree=degree, bcs=(("free", "free"), None, None), polar_splines=True) @@ -443,7 +454,7 @@ def test_polar_vector_ghost_sync(): comm = MPI.COMM_WORLD degree = [2, 2, 1] - num_elements = [16, 6, 4] + num_elements = _pole_safe([16, 6, 4]) domain = IGAPolarCylinder(num_elements=num_elements[:2], degree=degree[:2], Lz=1.0, a=1.0) grid = TensorProductGrid(num_elements=num_elements, mpi_dims_mask=(True, False, False)) derham_opts = DerhamOptions(degree=degree, bcs=(("free", "free"), None, None), polar_splines=True) @@ -524,11 +535,10 @@ def _rel_err(a, b): (("free", "dirichlet"), None, ("dirichlet", "dirichlet")), ], ) -@pytest.mark.parametrize("poloidal_solver", ["schur", "lu", "ilu"]) -def test_projectors(num_elements, degree, bcs, poloidal_solver): - """Polar commuting projectors: the 2D x 1D Kronecker inter-/histopolation matrices agree with - P * I * E^T, projection of splines is the identity (with and without boundary conditions), - grad P0 = P1 grad, and exactness of the polar sequence.""" +def test_projectors(num_elements, degree, bcs): + """Polar commuting projectors: the polar inter-/histopolation matrix P * I * E^T is the identity, projection of + splines is the identity (with and without boundary conditions), grad P0 = P1 grad, and exactness of the polar + sequence.""" import cunumpy as xp from maybempi import MPI @@ -539,11 +549,10 @@ def test_projectors(num_elements, degree, bcs, poloidal_solver): comm = MPI.COMM_WORLD + num_elements = _pole_safe(num_elements) domain = IGAPolarCylinder(num_elements=num_elements[:2], degree=degree[:2], Lz=3.0, a=1.0) grid = TensorProductGrid(num_elements=num_elements) - derham_opts = DerhamOptions( - degree=degree, bcs=bcs, nquads_proj=[6, 6, 6], polar_splines=True, polar_poloidal_solver=poloidal_solver - ) + derham_opts = DerhamOptions(degree=degree, bcs=bcs, nquads_proj=[6, 6, 6], polar_splines=True) derham = Derham(grid, derham_opts, comm=comm, domain=domain) # projection of a spline returns the spline, with and without boundary conditions @@ -553,11 +562,9 @@ def test_projectors(num_elements, degree, bcs, poloidal_solver): x = _rand_polar_vector(derham, P.I.domain, seed=1234) y = _rand_polar_vector(derham, P.I.domain, seed=5678) - # 2D x 1D Kronecker operators: P * I * E^T and its transpose + # polar inter-/histopolation matrix P * I * E^T (the identity) and its transpose Pd, Eb, B = P.dofs_extraction_op, P.base_extraction_op, P.boundary_op - if poloidal_solver == "schur": - # (the operators do not depend on the poloidal solver; the Kronecker dot product is slow) - assert _rel_err(P.I.dot(x), Pd.dot(P.projector_tensor.imat_kronecker.dot(Eb.T.dot(x)))) < 1e-12, form + assert _rel_err(P.I.dot(x), Pd.dot(P.projector_tensor.imat_kronecker.dot(Eb.T.dot(x)))) < 1e-12, form assert xp.isclose(P.IT.dot(x).dot(y), x.dot(P.I.dot(y)), rtol=1e-12), form assert xp.isclose(P.I0T.dot(x).dot(y), x.dot(P.I0.dot(y)), rtol=1e-12), form assert _rel_err(P.I0.dot(x), B.dot(P.I.dot(B.dot(x)))) < 1e-12, form @@ -648,7 +655,7 @@ def test_polar_mass_preconditioner(mapping, bcs, poloidal_solver): from struphy.io.options import DerhamOptions from struphy.topology.grids import TensorProductGrid - num_elements, degree = [8, 12, 4], [3, 3, 2] + num_elements, degree = _pole_safe([8, 12, 4]), [3, 3, 2] domain = getattr(domains, mapping)(num_elements=num_elements[:2], degree=degree[:2]) grid = TensorProductGrid(num_elements=num_elements) derham_opts = DerhamOptions(degree=degree, bcs=bcs, polar_splines=True) @@ -699,7 +706,7 @@ def test_polar_stiffness_preconditioner(mapping, bcs, poloidal_solver): from struphy.io.options import DerhamOptions from struphy.topology.grids import TensorProductGrid - num_elements, degree = [8, 12, 6], [3, 3, 2] + num_elements, degree = _pole_safe([8, 12, 6]), [3, 3, 2] domain = getattr(domains, mapping)(num_elements=num_elements[:2], degree=degree[:2]) grid = TensorProductGrid(num_elements=num_elements) derham_opts = DerhamOptions(degree=degree, bcs=bcs, polar_splines=True) @@ -827,7 +834,7 @@ def test_polar_weighted_average_projection(mapping, stab, directions): from struphy.polar.kron_2dx1d import _interior from struphy.topology.grids import TensorProductGrid - num_elements, degree = [8, 12, 6], [3, 3, 2] + num_elements, degree = _pole_safe([8, 12, 6]), [3, 3, 2] domain = getattr(domains, mapping)(num_elements=num_elements[:2], degree=degree[:2]) grid = TensorProductGrid(num_elements=num_elements) derham_opts = DerhamOptions(degree=degree, bcs=(("free", "dirichlet"), None, None), polar_splines=True) diff --git a/src/struphy/polar/tests/test_polar_cupy.py b/src/struphy/polar/tests/test_polar_cupy.py index 2e7dc2c5f..43eb65c62 100644 --- a/src/struphy/polar/tests/test_polar_cupy.py +++ b/src/struphy/polar/tests/test_polar_cupy.py @@ -1,6 +1,6 @@ """Polar splines on the CuPy backend (struphy-hub/struphy#695, CUDA_STRATEGY.md). -The polar blocks (:class:`~struphy.polar.extraction_operators.PolarExtractionBlocksC1`) are built on the host on +The polar blocks (:class:`~struphy.polar.homotopy.derham_blocks.PolarExtractionBlocksHomotopy`) are built on the host on every backend; on CuPy the polar operators apply device copies of them (``cupyx.scipy.sparse``) to the device coefficients. :func:`check_polar_backends` builds the same polar Derham on the NumPy and the CuPy backend and compares extraction operators, derivatives, ``PolarVector`` arithmetic, mass matrices, projectors and polar @@ -73,14 +73,14 @@ def make_polar_derham(mapping, with_mass=True, poloidal_solver="schur", mpi_dims from struphy.topology.grids import TensorProductGrid size = MPI.COMM_WORLD.Get_size() - num_elements = (max(6, 4 * size), max(6, 3 * size), 4) + # the process at the pole holds the polar rings and two tensor-product rings in eta1 + num_elements = (max(6, 6 * size), max(6, 3 * size), 4) domain = make_domain(mapping, num_elements[:2]) options = DerhamOptions( degree=(domain.degree[0], domain.degree[1], 1), bcs=(("dirichlet", "free"), None, None), nquads_proj=(4, 4, 2), polar_splines=True, - polar_poloidal_solver=poloidal_solver, ) derham = Derham( TensorProductGrid( @@ -300,7 +300,9 @@ def test_device_blocks_numpy_backend(): previous = set_device_sparse_module(DenseDeviceSparse) try: rng = np.random.default_rng(0) - for row in derham.dofs_extraction_ops["1"].blocks_ten_to_ten + E.blocks_ten_to_pol: + # the polar dofs of the homotopy construction are F~^T I^{-1} (a composition without own blocks) + dofs = derham.dofs_extraction_ops["1"] + for row in getattr(dofs, "blocks_ten_to_ten", []) + E.blocks_ten_to_pol: for blk in row: if blk is None: continue diff --git a/src/struphy/polar/tests/test_polar_regularity_one.py b/src/struphy/polar/tests/test_polar_regularity_one.py new file mode 100644 index 000000000..5aa1ec5ed --- /dev/null +++ b/src/struphy/polar/tests/test_polar_regularity_one.py @@ -0,0 +1,183 @@ +"""The polar splines of regularity 1 are the C1 polar splines of Toshniwal et al. (2017) and Holderied & Possanner +(2022): both span the same spaces in every degree. + +The C1 extraction blocks (``PolarExtractionBlocksC1``, removed after commit c0ee1bec) are frozen in +``data/c1_reference.npz`` for spline mappings with the resolution of the grid; the test compares the column spans of +the polar spaces in the tensor-product coefficients of the rings next to the pole (the outer rings are tensor-product +functions in both constructions). +""" + +from dataclasses import dataclass +from pathlib import Path + +import numpy as np +import pytest + +REFERENCE = Path(__file__).resolve().parent / "data" / "c1_reference.npz" + + +@dataclass(frozen=True) +class Case: + domain: str + num_elements: tuple + degree: int + + @property + def key(self): + return f"{self.domain}_{self.num_elements[0]}x{self.num_elements[1]}_p{self.degree}" + + +CASES = [ + Case(domain, num_elements, degree) + for domain in ("cylinder", "torus", "shaped") + for num_elements, degree in (((6, 9), 2), ((8, 12), 3)) +] + + +def shaped_cross_section(eta1, eta2): + """A non-circular, off-centre cross-section (x, y) used for the spline mapping of the "shaped" case.""" + R = eta1 * (1 + 0.2 * eta1 * np.cos(6 * np.pi * eta2)) * (1 + 0.3 * np.sin(2 * np.pi * eta2)) + return R * np.cos(2 * np.pi * eta2) + 0.1 * eta1**2, 0.7 * R * np.sin(2 * np.pi * eta2) + + +def build_domain(case): + from struphy.geometry import domains + from struphy.geometry.base import PoloidalSplineStraight, interp_mapping + + num_elements, degree = case.num_elements, (case.degree, case.degree) + if case.domain == "cylinder": + return domains.IGAPolarCylinder(a=1.0, Lz=1.0, num_elements=num_elements, degree=degree) + if case.domain == "torus": + return domains.IGAPolarTorus(num_elements=num_elements, degree=degree) + cx, cy = interp_mapping( + num_elements, + degree, + (False, True), + lambda e1, e2: shaped_cross_section(e1, e2)[0], + lambda e1, e2: shaped_cross_section(e1, e2)[1], + ) + cx[0], cy[0] = cx[0].mean(), cy[0].mean() + return PoloidalSplineStraight(num_elements=num_elements, degree=degree, cx=cx, cy=cy, Lz=1.0) + + +def build_derham(case, polar=True): + from struphy.feec.psydac_derham import Derham + from struphy.io.options import DerhamOptions + from struphy.topology.grids import TensorProductGrid + + options = DerhamOptions( + degree=(case.degree, case.degree, 1), + bcs=(("free", "free"), None, None), + polar_splines=polar, + polar_regularity=1, + ) + grid = TensorProductGrid(num_elements=(*case.num_elements, 1)) + return Derham(grid, options, domain=build_domain(case)) + + +# coupled components of the polar spaces (the eta1-eta2 components of 1- and 2-forms are coupled at the pole) +GROUPS = {"0": [(0,)], "1": [(0, 1), (2,)], "2": [(0, 1), (2,)], "3": [(0,)]} + + +def _rings(blocks, form, component, n2): + return blocks[form][component][component].shape[1] // n2 + + +def polar_basis(blocks, form, components, rings, n2): + """Rows: the polar basis functions of the coupled ``components`` in the tensor-product coefficients of the first + ``rings[c]`` rings of each component, completed by the tensor-product functions of the rings that the polar + construction leaves untouched.""" + widths = [rings[c] * n2 for c in components] + offsets = np.concatenate([[0], np.cumsum(widths)]) + out = [] + for r in components: + row = [blocks[form][r][c] for c in components] + count = next(b.shape[0] for b in row if b is not None) + M = np.zeros((count, offsets[-1])) + for k, block in enumerate(row): + if block is not None: + M[:, offsets[k] : offsets[k] + block.shape[1]] = block + out.append(M) + for k, c in enumerate(components): + for index in range(_rings(blocks, form, c, n2) * n2, rings[c] * n2): + unit = np.zeros(offsets[-1]) + unit[offsets[k] + index] = 1.0 + out.append(unit[None]) + return np.vstack(out) + + +def frozen_c1_blocks(case): + data = np.load(REFERENCE) + blocks = {} + for form, groups in GROUPS.items(): + components = sorted({c for group in groups for c in group}) + blocks[form] = {r: {c: data.get(f"{case.key}/{form}/{r}{c}") for c in components} for r in components} + return blocks + + +def dense_blocks(ck_blocks): + blocks = {} + for form, groups in GROUPS.items(): + rows = ck_blocks.e_ten_to_pol[form] + components = sorted({c for group in groups for c in group}) + blocks[form] = { + r: { + c: None + if rows[r][c] is None + else (rows[r][c].toarray() if hasattr(rows[r][c], "toarray") else np.asarray(rows[r][c])) + for c in components + } + for r in components + } + return blocks + + +@pytest.mark.parametrize("case", CASES, ids=lambda case: case.key) +def test_regularity_one_spans_the_c1_spaces(case): + derham = build_derham(case) + n2 = case.num_elements[1] + c1, general = frozen_c1_blocks(case), dense_blocks(derham.ck_blocks) + for form, groups in GROUPS.items(): + for components in groups: + rings = {c: max(_rings(b, form, c, n2) for b in (c1, general)) for c in components} + A = polar_basis(c1, form, components, rings, n2) + B = polar_basis(general, form, components, rings, n2) + rank_a, rank_b, rank_ab = (np.linalg.matrix_rank(M, tol=1e-9) for M in (A, B, np.vstack([A, B]))) + assert rank_a == A.shape[0] and rank_b == B.shape[0], (form, components) + assert rank_a == rank_b == rank_ab, f"degree {form}, components {components}: spans differ" + + +@pytest.mark.parametrize("case", CASES, ids=lambda case: case.key) +def test_regularity_one_zero_forms_are_affine_on_the_first_rings(case): + """The characterization of C1 0-forms that does not depend on the mapping: on the rings 0 and 1, every polar + 0-form is the projection of an affine function a + b x + c y.""" + derham = build_derham(case) + n2 = case.num_elements[1] + tensor = build_derham(case, polar=False) + domain = tensor.domain + + def in_plane(e1, e2, e3): + """Cartesian coordinates in the plane of the cross-section ((R, Z) on the torus).""" + x, y, z = domain(e1, e2, e3) + return (np.hypot(x, y), z) if case.domain == "torus" else (x, y) + + affine = [tensor.P0(lambda e1, e2, e3: 1.0 + 0 * e1)] + for k in range(2): + affine.append(tensor.P0(lambda e1, e2, e3, k=k: in_plane(e1, e2, e3)[k])) + Q = np.array([v.toarray().reshape(-1, n2)[:2].ravel() for v in affine]) + E0 = dense_blocks(derham.ck_blocks)["0"][0][0][:, : 2 * n2] + assert np.linalg.matrix_rank(Q) == 3 + assert np.linalg.matrix_rank(np.vstack([Q, E0]), tol=1e-9) == 3 + assert np.linalg.matrix_rank(E0, tol=1e-9) == 3 + + +def test_deprecated_polar_options(): + """``polar_method="c1"`` selects regularity 1 (the C1 spaces) and ``polar_poloidal_solver`` is ignored, both + with a deprecation warning.""" + from struphy.io.options import DerhamOptions + + with pytest.warns(DeprecationWarning, match="polar_method"): + options = DerhamOptions(polar_splines=True, polar_method="c1", polar_regularity=3) + assert options.polar_regularity == 1 + with pytest.warns(DeprecationWarning, match="polar_poloidal_solver"): + DerhamOptions(polar_splines=True, polar_poloidal_solver="ilu") diff --git a/src/struphy/propagators/__init__.py b/src/struphy/propagators/__init__.py index 9a36b4353..71036c9ca 100644 --- a/src/struphy/propagators/__init__.py +++ b/src/struphy/propagators/__init__.py @@ -33,6 +33,7 @@ "PushVinSPHpressure": "struphy.propagators.push_v_in_sph_pressure", "PushVinViscousPotential": "struphy.propagators.push_v_in_viscous_potential", "PushVxB": "struphy.propagators.push_vxb", + "ScalarWaveLeapfrog": "struphy.propagators.scalar_wave_leapfrog", "ShearAlfvenB1": "struphy.propagators.shear_alfven_b1", "ShearAlfvenCurrentCoupling5D": "struphy.propagators.shear_alfven_current_coupling_5d", "ShearAlfvenPropagator": "struphy.propagators.shear_alfven_propagator", @@ -80,6 +81,7 @@ from struphy.propagators.push_v_in_sph_pressure import PushVinSPHpressure from struphy.propagators.push_v_in_viscous_potential import PushVinViscousPotential from struphy.propagators.push_vxb import PushVxB + from struphy.propagators.scalar_wave_leapfrog import ScalarWaveLeapfrog from struphy.propagators.shear_alfven_b1 import ShearAlfvenB1 from struphy.propagators.shear_alfven_current_coupling_5d import ShearAlfvenCurrentCoupling5D from struphy.propagators.shear_alfven_propagator import ShearAlfvenPropagator @@ -126,6 +128,7 @@ "PushVinSPHpressure", "PushVinViscousPotential", "PushVxB", + "ScalarWaveLeapfrog", "ShearAlfvenB1", "ShearAlfvenCurrentCoupling5D", "ShearAlfvenPropagator", diff --git a/src/struphy/propagators/implicit_diffusion.py b/src/struphy/propagators/implicit_diffusion.py index c6e01c065..de18188d3 100644 --- a/src/struphy/propagators/implicit_diffusion.py +++ b/src/struphy/propagators/implicit_diffusion.py @@ -3,7 +3,7 @@ from typing import Callable, Literal import cunumpy as xp -from feectools.linalg.basic import IdentityOperator +from feectools.linalg.basic import IdentityOperator, Vector from feectools.linalg.solvers import inverse from feectools.linalg.stencil import StencilVector from line_profiler import profile @@ -542,7 +542,8 @@ def __call__(self, dt): self._rhs2 *= 0.0 for src, coeff in zip(self.sources, self.coeffs): - if isinstance(src, StencilVector): + # StencilVector, or PolarVector with polar splines + if isinstance(src, Vector): self._rhs2 += sig_3 * coeff * src elif isinstance(src, FEECVariable): v = src.spline.vector diff --git a/src/struphy/propagators/maxwell_weak_ampere.py b/src/struphy/propagators/maxwell_weak_ampere.py index 9c5d45486..16846ba94 100644 --- a/src/struphy/propagators/maxwell_weak_ampere.py +++ b/src/struphy/propagators/maxwell_weak_ampere.py @@ -34,6 +34,17 @@ class MaxwellWeakAmpere(Propagator): - implicit: Crank-Nicolson (implicit mid-point) - explicit: explicit RK methods from ButcherTableau + - leapfrog: symplectic velocity Verlet (one mass solve per half step), + + .. math:: + + \mathbf e^{n+1/2} = \mathbf e^n + \tfrac{\Delta t}{2} \mathbb M^{1,-1} \mathbb C^\top \mathbb M^2 \mathbf b^n\,,\quad + \mathbf b^{n+1} = \mathbf b^n - \Delta t\, \mathbb C \mathbf e^{n+1/2}\,,\quad + \mathbf e^{n+1} = \mathbf e^{n+1/2} + \tfrac{\Delta t}{2} \mathbb M^{1,-1} \mathbb C^\top \mathbb M^2 \mathbf b^{n+1}\,, + + i.e. the staggered leapfrog of the 1-form wave equation :math:`\ddot{\mathbf a} + \mathbb M^{1,-1}\mathbb C^\top + \mathbb M^2 \mathbb C \mathbf a = 0` for :math:`\mathbf b = \mathbb C \mathbf a`, :math:`\mathbf e = -\dot{\mathbf a}` + (stable for :math:`\Delta t < 2/\sqrt{\lambda_\textnormal{max}}`). System size reduction via :class:`~struphy.linear_algebra.schur_solver.SchurSolver`. """ @@ -114,7 +125,7 @@ class Options(OptionsBase): """ # specific literals - OptsAlgo = Literal["implicit", "explicit"] + OptsAlgo = Literal["implicit", "explicit", "leapfrog"] # propagator options algo: OptsAlgo = "implicit" solver: LiteralOptions.OptsSymmSolver = "pcg" @@ -183,6 +194,18 @@ def allocate(self): # pre-allocate arrays self._byn = self._B.codomain.zeros() + elif self.options.algo == "leapfrog": + self._info = False + M1_inv = inverse( + M1, + self.options.solver, + pc=pc, + tol=self.options.solver_params.tol, + maxiter=self.options.solver_params.maxiter, + verbose=self.options.solver_params.verbose, + ) + self._weak_curl = M1_inv @ curl.T @ M2 + self._curl = curl else: self._info = False @@ -240,6 +263,20 @@ def __call__(self, dt): bn1 += bn diffs = self.update_feec_variables(e=en1, b=bn1) + elif self.options.algo == "leapfrog": + # e^{n+1/2} + e_half = self._weak_curl.dot(bn, out=self._e_tmp1) + e_half *= dt / 2 + e_half += en + # b^{n+1} + bn1 = self._curl.dot(e_half, out=self._b_tmp1) + bn1 *= -dt + bn1 += bn + # e^{n+1} + en1 = self._weak_curl.dot(bn1, out=self._e_tmp2) + en1 *= dt / 2 + en1 += e_half + self.update_feec_variables(e=en1, b=bn1) else: self._ode_solver(0.0, dt) diff --git a/src/struphy/propagators/scalar_wave_leapfrog.py b/src/struphy/propagators/scalar_wave_leapfrog.py new file mode 100644 index 000000000..5a7f50eab --- /dev/null +++ b/src/struphy/propagators/scalar_wave_leapfrog.py @@ -0,0 +1,158 @@ +import logging +from dataclasses import dataclass + +from feectools.linalg.solvers import inverse +from line_profiler import profile + +from struphy.feec import preconditioner +from struphy.io.options import LiteralOptions, OptionsBase +from struphy.linear_algebra.solver import SolverParameters +from struphy.models.variables import FEECVariable +from struphy.propagators.base import Propagator +from struphy.utils.utils import check_option + +logger = logging.getLogger("struphy") + + +class ScalarWaveLeapfrog(Propagator): + r"""FEEC discretization of the scalar (0-form) wave equation + :math:`\partial_t^2 \phi - c^2 \Delta \phi = 0` as the first-order system: find + :math:`\phi, \psi \in H^1` such that + + .. math:: + + \frac{\partial \phi}{\partial t} = \psi\,, \qquad + \int_\Omega \frac{\partial \psi}{\partial t}\, \chi \,\textrm d \mathbf x + + c^2 \int_\Omega \nabla \chi \cdot \nabla \phi \,\textrm d \mathbf x = 0 \qquad \forall \,\chi \in H^1\,. + + Time discretization: symplectic velocity Verlet (the staggered leapfrog of the polar wave studies of + R. Remmerswaal), with :math:`\mathbb K = c^2\,\mathbb G^\top \mathbb M^1 \mathbb G`, + + .. math:: + + \boldsymbol \psi^{n+1/2} = \boldsymbol \psi^n - \tfrac{\Delta t}{2} \mathbb M^{0,-1} \mathbb K \boldsymbol \phi^n\,,\quad + \boldsymbol \phi^{n+1} = \boldsymbol \phi^n + \Delta t\, \boldsymbol \psi^{n+1/2}\,,\quad + \boldsymbol \psi^{n+1} = \boldsymbol \psi^{n+1/2} - \tfrac{\Delta t}{2} \mathbb M^{0,-1} \mathbb K \boldsymbol \phi^{n+1}\,, + + which conserves a discrete energy and is stable for :math:`\Delta t < 2/\sqrt{\lambda_\textnormal{max}}`, + :math:`\lambda_\textnormal{max}` the largest eigenvalue of :math:`\mathbb M^{0,-1}\mathbb K`. + """ + + class Variables: + """Container for the wave variables. + + Attributes + ---------- + phi : FEECVariable + The wave field in H1. + psi : FEECVariable + Its time derivative in H1. + """ + + def __init__(self): + self._phi: FEECVariable = None + self._psi: FEECVariable = None + + @property + def phi(self) -> FEECVariable: + return self._phi + + @phi.setter + def phi(self, new): + assert isinstance(new, FEECVariable) + assert new.space == "H1" + self._phi = new + + @property + def psi(self) -> FEECVariable: + return self._psi + + @psi.setter + def psi(self, new): + assert isinstance(new, FEECVariable) + assert new.space == "H1" + self._psi = new + + def __init__(self): + self.variables = self.Variables() + + @dataclass(repr=False) + class Options(OptionsBase): + """Options of :class:`ScalarWaveLeapfrog`. + + Parameters + ---------- + c : float + Wave speed. + solver : LiteralOptions.OptsSymmSolver + Solver of the mass matrix. + precond : LiteralOptions.OptsMassPrecond + Preconditioner of the mass matrix. + solver_params : SolverParameters + Solver controls. + """ + + c: float = 1.0 + solver: LiteralOptions.OptsSymmSolver = "pcg" + precond: LiteralOptions.OptsMassPrecond = "MassMatrixPreconditioner" + solver_params: SolverParameters = None + + def __post_init__(self): + check_option(self.solver, LiteralOptions.OptsSymmSolver) + check_option(self.precond, LiteralOptions.OptsMassPrecond) + if self.solver_params is None: + self.solver_params = SolverParameters() + + @property + def options(self) -> Options: + if not hasattr(self, "_options"): + self._options = self.Options() + return self._options + + @options.setter + def options(self, new): + assert isinstance(new, self.Options) + self._options = new + logger.info(f"\nNew options for propagator '{self.__class__.__name__}':\n{self._options}") + + @profile + def allocate(self): + M0, M1 = self.mass_ops.M0, self.mass_ops.M1 + grad = self.derham.grad + if self.options.precond is None: + pc = None + else: + pc = getattr(preconditioner, self.options.precond)(M0) + M0_inv = inverse( + M0, + self.options.solver, + pc=pc, + tol=self.options.solver_params.tol, + maxiter=self.options.solver_params.maxiter, + verbose=self.options.solver_params.verbose, + ) + # psi' = -M0^{-1} K phi + self._acceleration = -(self.options.c**2) * (M0_inv @ grad.T @ M1 @ grad) + space = self.variables.phi.spline.vector.space + self._psi_half = space.zeros() + self._phi_new = space.zeros() + self._psi_new = space.zeros() + + @profile + def __call__(self, dt): + phi = self.variables.phi.spline.vector + psi = self.variables.psi.spline.vector + + psi_half = self._acceleration.dot(phi, out=self._psi_half) + psi_half *= dt / 2 + psi_half += psi + + phi_new = psi_half.copy(out=self._phi_new) + phi_new *= dt + phi_new += phi + + psi_new = self._acceleration.dot(phi_new, out=self._psi_new) + psi_new *= dt / 2 + psi_new += psi_half + + self.update_feec_variables(phi=phi_new, psi=psi_new) diff --git a/src/struphy/propagators/shear_alfven_current_coupling_5d.py b/src/struphy/propagators/shear_alfven_current_coupling_5d.py index 117cf94f3..e317c537a 100644 --- a/src/struphy/propagators/shear_alfven_current_coupling_5d.py +++ b/src/struphy/propagators/shear_alfven_current_coupling_5d.py @@ -355,8 +355,6 @@ def _initialize_projection_operator_TB(self): [None, None, None], ], transposed=False, - use_cache=True, - polar_shift=True, V_extraction_op=self.derham.extraction_ops[self._u_form], V_boundary_op=self.derham.boundary_ops[self._u_form], P_boundary_op=self.derham.boundary_ops["1"], @@ -370,8 +368,6 @@ def _initialize_projection_operator_TB(self): [None, None, None], ], transposed=True, - use_cache=True, - polar_shift=True, V_extraction_op=self.derham.extraction_ops[self._u_form], V_boundary_op=self.derham.boundary_ops[self._u_form], P_boundary_op=self.derham.boundary_ops["1"], diff --git a/tutorials/dev_tutorial_polar_homotopy.ipynb b/tutorials/dev_tutorial_polar_homotopy.ipynb new file mode 100644 index 000000000..74644d37b --- /dev/null +++ b/tutorials/dev_tutorial_polar_homotopy.ipynb @@ -0,0 +1,630 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "206958e6", + "metadata": {}, + "source": [ + "# Homotopy operators and high-order polar complexes\n", + "\n", + "This developer tutorial walks through the construction of the high-order polar de Rham complex of R. Remmerswaal\n", + "(NMPP seminar 2026; Julia package [PolarFEEC.jl](https://gitlab.mpcdf.mpg.de/rwr/polarfeec.jl)) as implemented in\n", + "`struphy.polar.homotopy`, on the 1D spaces of Struphy `Derham` complexes and on Struphy domains:\n", + "\n", + "0. 1D spaces, tensor-product and mapped forms with commuting projections,\n", + "1. the Poincaré homotopy operators and their matrices,\n", + "2. the polar complex: extraction and dual matrices, cochain/chain property, cohomology,\n", + "3. commuting polar projections: convergence and commutativity,\n", + "4. boundary-compatible homotopies and the 2D curl–curl problem (\"a free two-dimensional lunch\"),\n", + "5. the Maxwell spectrum on the disk,\n", + "6. the 3D total complex and the Struphy `Derham` integration,\n", + "7. the 0- and 1-form wave equations: regularity and convergence,\n", + "8. field lines and the assembly cost in the 3D `Derham` complex,\n", + "9. the 3D wave in a helical tube and the schematics of the talk.\n", + "\n", + "Set `STRUPHY_TUTORIAL_FULL=1` to run the convergence studies at the resolution of the talk and of PolarFEEC.jl." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "aa019215", + "metadata": {}, + "outputs": [], + "source": [ + "import os\n", + "\n", + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", + "import scipy.sparse as sp\n", + "\n", + "from struphy.geometry import domains\n", + "from struphy.polar.homotopy import cumulative_sum_matrix\n", + "from struphy.polar.homotopy.demos import common\n", + "from struphy.polar.homotopy.linalg import as_dense\n", + "\n", + "FULL = os.environ.get(\"STRUPHY_TUTORIAL_FULL\", \"0\") == \"1\"" + ] + }, + { + "cell_type": "markdown", + "id": "f215de09", + "metadata": {}, + "source": [ + "## 0. Spaces, forms and commuting projections (PolarFEEC.jl examples 01–03)\n", + "\n", + "The building blocks are the 1D spaces of a Struphy `Derham` complex (`DerhamSpace1D`: the feectools B-splines and\n", + "their M-splines, clamped or periodic, or the nodal trigonometric functions of `struphy.feec.trig_space`), whose\n", + "derivative is an exact coefficient map, $\\tfrac{d}{dx}B_i = C_{i-1} - C_i$, so that the 1D diagram\n", + "$\\tfrac{d}{dx}\\pi^0 f = \\pi^1 f'$ commutes:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "fcabca06", + "metadata": {}, + "outputs": [], + "source": [ + "from struphy.polar.homotopy.demos import spaces\n", + "\n", + "spaces.spaces_1d()\n", + "spaces.figure_spaces_1d();" + ] + }, + { + "cell_type": "markdown", + "id": "c0462f80", + "metadata": {}, + "source": [ + "Tensor products give 0-, 1- and 2-forms on the unit square with $\\operatorname{curl}\\circ\\operatorname{grad} = 0$\n", + "exactly and projections that commute with the derivatives. The commuting and $L^2$ projections converge at order\n", + "$p + 1$ for 0-forms and $p$ for 1- and 2-forms:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "09088a5b", + "metadata": {}, + "outputs": [], + "source": [ + "spaces.identities()\n", + "square_cells = (8, 16, 32, 64, 128) if FULL else (8, 16, 32)\n", + "square_errors = spaces.convergence(degrees=(1, 2, 3, 4, 5) if FULL else (1, 2, 3), cells=square_cells, verbose=False)\n", + "spaces.figure_convergence(square_errors, square_cells);" + ] + }, + { + "cell_type": "markdown", + "id": "354bf490", + "metadata": {}, + "source": [ + "On a mapped domain (the quarter annulus `HollowCylinder(a1=1, a2=2, poc=4)`) the forms are pushed forward as scalars,\n", + "by the covariant Piola map and as densities; the projections still commute:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5c27603a", + "metadata": {}, + "outputs": [], + "source": [ + "spaces.mapped_identities()\n", + "spaces.figure_mapped_forms();" + ] + }, + { + "cell_type": "markdown", + "id": "a4a9a2f5", + "metadata": {}, + "source": [ + "## 1. Homotopy operators\n", + "\n", + "In 1D the fundamental theorem of calculus reads $P = \\mathfrak P P + \\partial^\\dagger \\partial P$ and\n", + "$A = \\partial \\partial^\\dagger A$ with $\\mathfrak P P = P(0)$ and $\\partial^\\dagger A = \\int_0^x A$. For splines with\n", + "the derivative-scaled 1-form basis ($\\tfrac{d}{dx} B_i = C_{i-1} - C_i$) the derivative is the forward difference\n", + "$\\delta$ and $\\partial^\\dagger$ is the cumulative sum $\\Sigma$, with $\\delta\\Sigma = I$ and\n", + "$\\Sigma\\delta = I - \\mathbb 1 e_1^\\top$:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "c9e76310", + "metadata": {}, + "outputs": [], + "source": [ + "n = 6\n", + "Sigma = cumulative_sum_matrix(n).toarray()\n", + "delta = common.spline_space_1d(3, n - 2).derivative_matrix().toarray()\n", + "print(\"delta Sigma = I:\", np.allclose(delta @ Sigma, np.eye(n)))\n", + "print(\"Sigma delta = I - 1 e1^T:\", np.allclose(Sigma @ delta, np.eye(n + 1) - np.outer(np.ones(n + 1), np.eye(n + 1)[0])))" + ] + }, + { + "cell_type": "markdown", + "id": "0a90d73e", + "metadata": {}, + "source": [ + "In 2D the Poincaré homotopies $(\\operatorname{grad}^\\dagger A)(x) = \\int_0^1 A(tx)\\cdot x\\,dt$ and\n", + "$(\\operatorname{curl}^\\dagger B)(x) = (-y, x)^\\top \\int_0^1 t B(tx)\\,dt$ satisfy\n", + "$P = \\mathfrak P P + \\operatorname{grad}^\\dagger\\operatorname{grad}P$,\n", + "$A = \\operatorname{grad}\\operatorname{grad}^\\dagger A + \\operatorname{curl}^\\dagger\\operatorname{curl}A$,\n", + "$B = \\operatorname{curl}\\operatorname{curl}^\\dagger B$. On a disk-like mapping they are radial integrations in logical\n", + "coordinates: $G^\\dagger = [I \\otimes \\Sigma,\\ 0]$ and $C^\\dagger = [0;\\ I \\otimes \\Sigma]$." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "fdda6d0f", + "metadata": {}, + "outputs": [], + "source": [ + "c = common.polar_complex(3, degree=3, cells=(8, 16)) # PerturbedDiskCylinder, Struphy ordering\n", + "G, C = c.spaces[0].derivative_matrix, c.spaces[1].derivative_matrix\n", + "HG, HC = c.gradient_homotopy_matrix, c.curl_homotopy_matrix\n", + "print(\"C C^dagger = I:\", abs(C @ HC - sp.identity(HC.shape[1])).max())\n", + "print(\"G^dagger C^dagger = 0:\", abs(HG @ HC).max())" + ] + }, + { + "cell_type": "markdown", + "id": "399fa0be", + "metadata": {}, + "source": [ + "## 2. The polar complex\n", + "\n", + "Each degree is represented by an extraction matrix $E$ (its columns are the tensor-product coefficients of the polar\n", + "basis functions, central ones first) and dual matrix $F$ with $F^\\top E = I$. The polar derivatives, homotopies and\n", + "masses are $\\bar D_k = F_{k+1}^\\top D E_k$, $h_k = F_{k-1}^\\top H E_k$ and $\\bar M_k = E_k^\\top M E_k$.\n", + "The 1-form duals are defined by $E F^\\top = \\bar\\Pi^1 = \\operatorname{grad}\\bar\\Pi^0\\operatorname{grad}^\\dagger +\n", + "\\operatorname{curl}^\\dagger\\bar\\Pi^2\\operatorname{curl}$, which makes all projections commute with the derivatives\n", + "and with the homotopies." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "c1409a0d", + "metadata": {}, + "outputs": [], + "source": [ + "print(c)\n", + "print(\"central functions (0-, 1-, 2-forms):\", c.central_counts())\n", + "for k in range(3):\n", + " print(f\"F^T E = I (degree {k}):\", abs(c.duals(k).T @ c.extraction(k) - sp.identity(c.dimension(k))).max())\n", + "print(\"cohomology:\", c.cohomology_dimensions())\n", + "for k in range(3):\n", + " m = c.dimension(k)\n", + " contraction = np.zeros((m, m))\n", + " if k < 2:\n", + " contraction += as_dense(c.homotopy(k + 1) @ c.derivative(k))\n", + " if k > 0:\n", + " contraction += as_dense(c.derivative(k - 1) @ c.homotopy(k))\n", + " error = np.abs(contraction - (np.eye(m) - as_dense(c.cohomology_projection(k)))).max()\n", + " print(f\"d h + h d = I - P on degree {k}: {error:.1e}\")" + ] + }, + { + "cell_type": "markdown", + "id": "59894bb4", + "metadata": {}, + "source": [ + "The enrichment options of the paper: `\"cochain_chain\"` (default), `\"cochain\"` (the reduced spaces of Section 4.2,\n", + "without the compact 2-forms) and `\"none\"` (external functions only):" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a98c77bb", + "metadata": {}, + "outputs": [], + "source": [ + "for enrichment in (\"cochain_chain\", \"cochain\", \"none\"):\n", + " e = common.polar_complex(3, degree=3, cells=(8, 16), enrichment=enrichment)\n", + " print(f\"{enrichment:14s} dimensions {[e.dimension(k) for k in range(3)]}, cohomology {e.cohomology_dimensions()}\")" + ] + }, + { + "cell_type": "markdown", + "id": "174aee0d", + "metadata": {}, + "source": [ + "## 3. Commuting polar projections\n", + "\n", + "The polar projections $\\bar\\pi^k = \\bar\\Pi^{k} \\circ \\pi^k$ (commuting) and the polar $L^2$ projections, for\n", + "$r_x = r_y = \\mathfrak r$ (B-spline degree equal to the regularity) on the perturbed disk:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "fc7ffc6d", + "metadata": {}, + "outputs": [], + "source": [ + "from struphy.polar.homotopy.demos import projections\n", + "\n", + "regularities = (1, 2, 3, 4, 5, 6) if FULL else (1, 2, 3)\n", + "cells = (4, 8, 16, 32) if FULL else (4, 8, 16)\n", + "errors_l2 = projections.convergence(regularities, cells, method=\"l2\")\n", + "errors_commuting = projections.convergence(regularities, cells, method=\"commuting\")\n", + "projections.figure_convergence(errors_l2, cells, title=\"$L^2$ projection\")\n", + "projections.figure_convergence(errors_commuting, cells, title=\"commuting projection\");" + ] + }, + { + "cell_type": "markdown", + "id": "56517c0e", + "metadata": {}, + "source": [ + "Commutativity with $\\operatorname{grad}$, $\\operatorname{curl}$, $\\operatorname{grad}^\\dagger$ and\n", + "$\\operatorname{curl}^\\dagger$: the commuting projections commute to round-off, the $L^2$ projections only up to the\n", + "approximation error." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "c31bdb79", + "metadata": {}, + "outputs": [], + "source": [ + "defects_l2 = projections.commutation(regularities, cells, method=\"l2\")\n", + "defects_commuting = projections.commutation(regularities, cells, method=\"commuting\")\n", + "projections.figure_commutation(defects_l2, cells, title=\"$L^2$ projection\")\n", + "projections.figure_commutation(defects_commuting, cells, title=\"commuting projection\");" + ] + }, + { + "cell_type": "markdown", + "id": "eff76472", + "metadata": {}, + "source": [ + "## 4. Boundary conditions and the 2D curl–curl problem\n", + "\n", + "With `boundary=\"dirichlet\"` the relative spaces ($P = 0$, $A \\times n = 0$ on $\\partial\\Omega$) are built. Their\n", + "homotopies are the boundary-compatible homotopies of the paper, and the degree-2 cohomology projection is the mean\n", + "projection $\\mathfrak B$, so that $B = \\mathfrak B B + \\operatorname{curl}\\operatorname{curl}^\\dagger_\\partial B$.\n", + "With $A = \\operatorname{grad}P + \\operatorname{curl}^\\dagger_\\partial B$ the curl–curl problem\n", + "$(\\operatorname{curl}\\Phi, \\operatorname{curl}A) = (\\Phi, J)$ reduces to **one polar 2-form mass solve**:\n", + "\n", + "$$\n", + "A = h_2\\,(I - \\mathfrak B)\\,\\bar M_2^{-1} h_2^\\top \\ell .\n", + "$$" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "21d9de5d", + "metadata": {}, + "outputs": [], + "source": [ + "from struphy.polar.homotopy.solvers import curl_curl_free_lunch, stiffness_solver\n", + "\n", + "d = common.polar_complex(3, degree=3, cells=(6, 7), angular=\"trig\", boundary=\"dirichlet\")\n", + "hG, hC, B = (as_dense(a) for a in (d.homotopy(1), d.homotopy(2), d.cohomology_projection(2)))\n", + "D0, D1 = d.derivative(0).toarray(), d.derivative(1).toarray()\n", + "print(\"hG D0 = I:\", np.abs(hG @ D0 - np.eye(D0.shape[1])).max())\n", + "print(\"D1 hC = I - B:\", np.abs(D1 @ hC - (np.eye(D1.shape[0]) - B)).max())\n", + "load = d.mass_matrix(1) @ np.cos(np.arange(d.dimension(1)))\n", + "a = curl_curl_free_lunch(d, load)\n", + "b = stiffness_solver(d, 1).range_projection(d.homotopy(2).T @ load)\n", + "print(\"free lunch = saddle-point range projection:\", np.abs(d.derivative(1) @ a - b).max())" + ] + }, + { + "cell_type": "markdown", + "id": "b0f5c663", + "metadata": {}, + "source": [ + "### Poisson and curl–curl on a rounded square (PolarFEEC.jl example 05)\n", + "\n", + "Manufactured solutions on $\\Omega = \\{g > 0\\}$, $g = 1 - x^2 - y^2 - \\varepsilon\\,\\mathrm{Re}((x + iy)^4)$; the\n", + "solvers work in the homotopy gauge (free Poisson: $u_h(0) = 0$; curl–curl: $h A_h = 0$)." + ] + }, + { + "cell_type": "markdown", + "id": "261e5f72", + "metadata": {}, + "source": [ + "The steps of the boundary-compatible homotopy for a zero-mean 2-form: the radial homotopy leaves a tangential trace;\n", + "the periodic primitive of this trace, extended into the outer ring, is subtracted as a gradient:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4e8995aa", + "metadata": {}, + "outputs": [], + "source": [ + "from struphy.polar.homotopy.demos import schematics\n", + "\n", + "schematics.boundary_homotopy();" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "32731ade", + "metadata": {}, + "outputs": [], + "source": [ + "from struphy.polar.homotopy.demos import pde\n", + "\n", + "solution = pde.solve_all(3, 6)\n", + "print(\"L2 errors (u free, u Dirichlet, curl A free, curl A Dirichlet):\", pde.errors(solution))\n", + "print(pde.checks(solution))\n", + "pde.figure_poisson(solution)\n", + "pde.figure_curl_curl(solution);" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "50a2e4d2", + "metadata": {}, + "outputs": [], + "source": [ + "cells05 = pde.CELLS if FULL else (4, 6, 8, 12)\n", + "degrees05 = pde.DEGREES if FULL else (2, 3)\n", + "table = pde.convergence(cells05, degrees05)\n", + "pde.figure_convergence(table, cells05, degrees05);" + ] + }, + { + "cell_type": "markdown", + "id": "45fedc2e", + "metadata": {}, + "source": [ + "## 5. The Maxwell spectrum on the unit disk (PolarFEEC.jl example 06)\n", + "\n", + "The Dirichlet polar 1-forms give a spectrally correct discretization: the kernel of\n", + "$\\bar C^\\top \\bar M_2 \\bar C a = \\lambda \\bar M_1 a$ is exactly $\\operatorname{grad}\\bar V^0$, there are no spurious\n", + "modes, and the eigenvalues converge to the squared zeros of the Bessel derivatives at order $2p$." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "90777ef5", + "metadata": {}, + "outputs": [], + "source": [ + "from struphy.polar.homotopy.demos import maxwell\n", + "\n", + "spectrum = maxwell.study(meshes=maxwell.MESHES if FULL else ((8, 16), (16, 32)), compared=40)\n", + "maxwell.figure_spectrum(spectrum)\n", + "maxwell.figure_modes(spectrum);" + ] + }, + { + "cell_type": "markdown", + "id": "17782964", + "metadata": {}, + "source": [ + "## 6. The 3D total complex\n", + "\n", + "In 3D the polar complex is $\\bar W^m = \\bigoplus_{p+q=m} \\bar V^{p}_{2D} \\otimes V^q_{1D}$ with\n", + "$d(V \\otimes U) = d V \\otimes U + (-1)^p V \\otimes d U$ and the ordered homotopy\n", + "$h(V \\otimes U) = h V \\otimes U + (-1)^p \\mathfrak P V \\otimes h U$. It has the cohomology of the solid torus,\n", + "$(1, 1, 0, 0)$, and $(0, 0, 1, 1)$ with Dirichlet conditions:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "45b2dda3", + "metadata": {}, + "outputs": [], + "source": [ + "from struphy.polar.homotopy.complex3d import PolarComplex3D\n", + "\n", + "section = domains.PerturbedDiskCylinder(R1=0.5, eps=0.1, k=2, Lz=1.0)\n", + "for boundary in (\"free\", \"dirichlet\"):\n", + " c2 = common.polar_complex(2, degree=3, cells=(3, 8), domain=section, boundary=boundary)\n", + " c3 = PolarComplex3D(c2, common.spline_space_1d(3, 8, periodic=True))\n", + " print(boundary, [c3.dimension(m) for m in range(4)], c3.cohomology_dimensions())" + ] + }, + { + "cell_type": "markdown", + "id": "51dad78c", + "metadata": {}, + "source": [ + "### Integration in `Derham`\n", + "\n", + "`PolarExtractionBlocksHomotopy` (in `struphy.polar.homotopy.derham_blocks`) provides the blocks of the polar operators\n", + "of `struphy.polar.linear_operators` for the total complex. Per tensor-product component, the first `n_rings` rings\n", + "hold the polar coefficients (the central polar functions and the external functions of these rings): $\\mathfrak r + 2$\n", + "rings for 0-forms, $(\\mathfrak r + 1, \\mathfrak r + 2)$ for the two components of 2D 1-forms and $\\mathfrak r + 1$ for\n", + "2D 2-forms. With these ring counts the extraction and the duals are the identity outside the polar rings, and the\n", + "derivatives map no tensor-product coefficient into the polar coefficients. The duals of the paper are used with the\n", + "non-conforming ring-0 $y$-component removed, $\\tilde F = Z F$; on the coefficients of physical fields both coincide.\n", + "The commuting projectors are $\\tilde F^\\top \\pi$, so that the polar \"degrees of freedom\" are the polar coefficients and\n", + "the polar inter-/histopolation matrix is the identity." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "67286e89", + "metadata": {}, + "outputs": [], + "source": [ + "from struphy.feec.psydac_derham import Derham\n", + "from struphy.io.options import DerhamOptions\n", + "from struphy.topology.grids import TensorProductGrid\n", + "\n", + "options = DerhamOptions(degree=(3, 3, 1), bcs=((\"free\", \"free\"), None, None), polar_splines=True, polar_regularity=2)\n", + "derham = Derham(TensorProductGrid(num_elements=(8, 12, 2)), options, domain=domains.HollowTorus(a1=0.0, a2=0.3))\n", + "blocks = derham.ck_blocks\n", + "print(\"polar sets (2D):\", {key: len(columns) for key, columns in blocks._sets.items()})\n", + "print(\"grad blocks [Hcurl comp][H1 comp] equal the 2D polar gradient:\",\n", + " np.allclose(blocks.grad_pol_to_pol[0][0].toarray(),\n", + " blocks.complex2d.derivative(0).toarray()[np.ix_(blocks._sets[\"1x\"], blocks._sets[\"0\"])]))" + ] + }, + { + "cell_type": "markdown", + "id": "bd1b1885", + "metadata": {}, + "source": [ + "## 7. Wave equations: regularity and convergence (talk, Sections 1, 3 and 5)\n", + "\n", + "The 0- and 1-form wave equations $\\bar M v'' + \\bar K v = 0$ (Dirichlet; $\\bar K = (DE)^\\top M (DE)$ with $D$ the\n", + "gradient or the curl) are integrated with the staggered leapfrog at $0.9$ times the stability limit. The initial 1-form\n", + "is the homotopy lift $A = h_2 (I - \\mathfrak B) b$ of the projected Gaussian $b$, so that $\\operatorname{curl}A = b$ up\n", + "to its mean. The talk measures the maximum relative $L^\\infty$ error during the final 10% of $[0, 1]$; since the talk\n", + "does not state the reference solution, a finer mesh of the highest regularity is used here. Without\n", + "`STRUPHY_TUTORIAL_FULL=1` the study uses cubic splines, a wider pulse and $T = 0.3$. In this smooth setting all\n", + "regularities converge, and $\\mathfrak r = 0$ (the largest space) has the smallest error; the stagnation of low\n", + "regularities reported in the talk concerns its narrow pulse and quintic splines (`STRUPHY_TUTORIAL_FULL=1`, which\n", + "takes hours and has not been run for this tutorial)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5df12526", + "metadata": {}, + "outputs": [], + "source": [ + "from struphy.polar.homotopy.demos import waves\n", + "\n", + "if FULL:\n", + " wave_args = dict(regularities=(0, 1, 2, 3, 4, 5), radial_cells=(8, 16, 32), reference_cells=64, degree=5)\n", + "else:\n", + " # the pulse of the talk (width 0.04) is not resolved on coarse meshes: use a wider one\n", + " wave_args = dict(regularities=(0, 1, 2), radial_cells=(4, 8, 16, 32), reference_cells=64, degree=3,\n", + " initial=lambda x, y: waves.gaussian(x, y, width=0.15))\n", + "wave_study = waves.convergence(0, final_time=1.0 if FULL else 0.3, verbose=False, **wave_args)\n", + "waves.figure_convergence(wave_study, degree=wave_args[\"degree\"])\n", + "waves.figure_time_traces(wave_study);" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7b15f930", + "metadata": {}, + "outputs": [], + "source": [ + "one_form = waves.run_family(1, regularities=(1, 2) if not FULL else (1, 2, 3, 4), final_time=0.3,\n", + " **(dict(degree=5, radial_cells=8, angular_cells=16) if FULL\n", + " else dict(degree=3, radial_cells=6, angular_cells=12)))\n", + "waves.figure_snapshots(one_form);" + ] + }, + { + "cell_type": "markdown", + "id": "35fc87dd", + "metadata": {}, + "source": [ + "`waves.WaveProblem` also takes `angular=\"trig\"` (the trigonometric angular spaces of the talk), `enrichment=\"none\"`\n", + "or `\"cochain\"` and a `space_mapping` different from the pushforward (the non-matching study).\n", + "\n", + "## 8. Field lines and assembly cost in the 3D `Derham` complex\n", + "\n", + "The field-line study of PolarFEEC.jl example 07 (see the user tutorial) and the assembly cost of the polar mass\n", + "matrices against the degree (the counterpart of example 08; Struphy assembles the tensor-product masses with its own\n", + "kernels and applies $E^\\top M E$):" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7350b0da", + "metadata": {}, + "outputs": [], + "source": [ + "from struphy.polar.homotopy.demos import performance\n", + "\n", + "timings = performance.study(degrees=(1, 2, 3, 4, 5) if FULL else (1, 2, 3), cells=(12, 12, 12) if FULL else (6, 6, 6),\n", + " repeats=3 if FULL else 1)\n", + "performance.figure(timings);" + ] + }, + { + "cell_type": "markdown", + "id": "134f85a7", + "metadata": {}, + "source": [ + "## 9. The 3D wave in a helical tube and the schematics of the talk\n", + "\n", + "The 1-form wave (Maxwell) in the helical tube `HelicalTube(\"stellarator\")` with the 3D polar splines, integrated with\n", + "the leapfrog of Section 7 (`struphy.polar.homotopy.demos.waves3d`; the model version is\n", + "`examples/Maxwell/gaussian_helical_tube`). A cross-section of the initial magnetic field:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "ae947f25", + "metadata": {}, + "outputs": [], + "source": [ + "import warnings\n", + "\n", + "from struphy.polar.homotopy.demos import waves3d\n", + "\n", + "with warnings.catch_warnings():\n", + " warnings.simplefilter(\"ignore\") # the cross-sections rotate along the tube (see derham_blocks)\n", + " wave = waves3d.Wave3D(2, 6 if FULL else 3, degree=3)\n", + "waves3d.figure_cross_section(wave, title=r\"initial $B$, $\\eta_3 = 0$\");" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4afdaa08", + "metadata": {}, + "outputs": [], + "source": [ + "if FULL:\n", + " study3d = waves3d.convergence((1, 2, 3, 4), (5, 8, 12), reference_mesh=16, final_time=0.4)\n", + "else:\n", + " study3d = waves3d.convergence((1, 2), (3, 4), reference_mesh=5, final_time=0.1, degree=3)\n", + "waves3d.figure_convergence(study3d);" + ] + }, + { + "cell_type": "markdown", + "id": "48742d37", + "metadata": {}, + "source": [ + "The schematics of the talk: the coordinate transformation of the helical tube and the de Rham diagrams (drawn with\n", + "`tikzfigure`, which needs LaTeX; `schematics.write_all(folder)` writes all of them):" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "02245664", + "metadata": {}, + "outputs": [], + "source": [ + "schematics.coordinate_transformation_3d();" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/tutorials/tutorial_domains.ipynb b/tutorials/tutorial_domains.ipynb index f8c7c4ab0..794d37cf9 100644 --- a/tutorials/tutorial_domains.ipynb +++ b/tutorials/tutorial_domains.ipynb @@ -273,6 +273,100 @@ "domain.show()" ] }, + { + "cell_type": "markdown", + "id": "e1a20dce", + "metadata": {}, + "source": [ + "## Disk-like domains for high-order polar splines\n", + "\n", + "Three analytic domains with a pole at $\\eta_1 = 0$ come with the high-order polar splines (see the tutorials\n", + "*Polar splines in the Derham complex* and *High-order polar splines*). `PerturbedDiskCylinder` is the cylinder over a\n", + "perturbed disk,\n", + "\n", + "$$\n", + "x = R \\cos(2\\pi\\eta_2), \\quad y = R \\sin(2\\pi\\eta_2), \\quad z = L_z \\eta_3, \\qquad\n", + "R = R_1\\, \\eta_1 \\left(1 + \\varepsilon\\, \\eta_1^{q} \\cos(2\\pi k \\eta_2)\\right),\n", + "$$\n", + "\n", + "with `power` $q = 1$ (default; the perturbation vanishes at the axis) or $q = 0$ (a star-shaped disk)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9459287b", + "metadata": {}, + "outputs": [], + "source": [ + "domain = domains.PerturbedDiskCylinder(R1=1.0, eps=0.2, k=3, Lz=1.0)\n", + "domain.show()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "031296ee", + "metadata": {}, + "outputs": [], + "source": [ + "domain = domains.PerturbedDiskCylinder(R1=1.0, eps=0.2, k=3, Lz=1.0, power=0.0)\n", + "domain.show()" + ] + }, + { + "cell_type": "markdown", + "id": "9392b0b7", + "metadata": {}, + "source": [ + "`RoundedSquareCylinder` is the cylinder over the rounded square $\\{g > 0\\}$,\n", + "$g = 1 - x^2 - y^2 - \\varepsilon\\, \\mathrm{Re}((x + iy)^4)$, whose boundary radius is computed in closed form:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "ca20424e", + "metadata": {}, + "outputs": [], + "source": [ + "domain = domains.RoundedSquareCylinder(eps=0.3)\n", + "domain.show()" + ] + }, + { + "cell_type": "markdown", + "id": "c0d631f4", + "metadata": {}, + "source": [ + "`HelicalTube` is a tube around a helical magnetic axis whose cross-sections rotate along the axis (a spline mapping\n", + "fitted to the analytic geometry). The presets `\"stellarator\"` and `\"islands\"` are the geometries of the 3D wave and the\n", + "field-line studies of the polar splines:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7c9a6ef3", + "metadata": {}, + "outputs": [], + "source": [ + "domain = domains.HelicalTube(\"stellarator\")\n", + "domain.show()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "66a0ff54", + "metadata": {}, + "outputs": [], + "source": [ + "# the parameters of the analytic geometry behind the spline mapping\n", + "for key, val in domain.geometry.items():\n", + " print(key, \"=\", val)" + ] + }, { "cell_type": "markdown", "metadata": {}, diff --git a/tutorials/tutorial_polar_derham.ipynb b/tutorials/tutorial_polar_derham.ipynb new file mode 100644 index 000000000..7fa2b7050 --- /dev/null +++ b/tutorials/tutorial_polar_derham.ipynb @@ -0,0 +1,391 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "ccad13d2", + "metadata": {}, + "source": [ + "# Polar splines in the Struphy de Rham complex\n", + "\n", + "Domains with a magnetic axis (cylinders, tori, stellarators) are parametrized by a mapping $F$ that collapses the\n", + "logical boundary $\\eta_1 = 0$ onto a line, the *pole*. Tensor-product splines pushed forward by $F$ are in general\n", + "discontinuous there. Struphy's **polar splines** are the subspaces of the tensor-product spaces that are smooth at the\n", + "pole and still form a de Rham complex with commuting projectors. They are selected in the `DerhamOptions`:\n", + "\n", + "```python\n", + "DerhamOptions(polar_splines=True, polar_regularity=r)\n", + "```\n", + "\n", + "For regularity $\\mathfrak r$ the 0-forms are $C^{\\mathfrak r}$, the 1-forms $C^{\\mathfrak r - 1}$ and the 2-forms\n", + "$C^{\\mathfrak r - 2}$ at the pole. The construction (R. Remmerswaal, homotopy operators; see the developer tutorial\n", + "*Homotopy operators and high-order polar complexes*) works for every domain with a pole and every regularity\n", + "$\\mathfrak r \\leq p_1$ with $N_2 \\geq 2\\mathfrak r + 1$. **Regularity 1 spans exactly the former C1 polar splines**\n", + "(Toshniwal et al. 2017; Holderied & Possanner 2022).\n", + "\n", + "This tutorial shows\n", + "\n", + "1. domains with a pole,\n", + "2. the polar spaces of different regularities,\n", + "3. why regularity 1 is the C1 construction,\n", + "4. the smoothness of polar splines at the pole,\n", + "5. the commuting projectors,\n", + "6. trigonometric functions in $\\eta_2$ (exact regularity)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "dba3a2bb", + "metadata": {}, + "outputs": [], + "source": [ + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", + "\n", + "from struphy.feec.psydac_derham import Derham\n", + "from struphy.geometry import domains\n", + "from struphy.io.options import DerhamOptions\n", + "from struphy.topology.grids import TensorProductGrid\n", + "\n", + "\n", + "def polar_derham(domain, regularity, num_elements=(8, 24, 1), degree=(5, 5, 1), polar=True):\n", + " # a 2D complex (one cell in eta3) on the cross-section of the domain\n", + " options = DerhamOptions(\n", + " degree=degree,\n", + " bcs=((\"free\", \"free\"), None, None),\n", + " polar_splines=polar,\n", + " polar_regularity=regularity,\n", + " nquads_proj=tuple(p + 3 for p in degree),\n", + " )\n", + " return Derham(TensorProductGrid(num_elements=num_elements), options, domain=domain)" + ] + }, + { + "cell_type": "markdown", + "id": "bf74309b", + "metadata": {}, + "source": [ + "## 1. Domains with a pole\n", + "\n", + "Every Struphy domain whose first logical boundary collapses to a line can be used: analytic mappings such as\n", + "`HollowCylinder(a1=0)` and `HollowTorus(a1=0)`, the star-shaped `PerturbedDiskCylinder`, the `RoundedSquareCylinder`,\n", + "the helical `HelicalTube` and spline mappings such as `IGAPolarCylinder`, `IGAPolarTorus` or `Tokamak`. Unlike the\n", + "former C1 construction, the polar splines do not need a spline mapping with the resolution of the grid, and the number\n", + "of angular cells is arbitrary." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "09a74f67", + "metadata": {}, + "outputs": [], + "source": [ + "for domain in (\n", + " domains.HollowCylinder(a1=0.0, a2=1.0, Lz=1.0),\n", + " domains.PerturbedDiskCylinder(R1=1.0, eps=0.2, k=3, Lz=1.0),\n", + " domains.RoundedSquareCylinder(),\n", + "):\n", + " domain.show(figsize=(9, 4))" + ] + }, + { + "cell_type": "markdown", + "id": "ce07d46d", + "metadata": {}, + "source": [ + "## 2. Polar spaces of regularity $\\mathfrak r$\n", + "\n", + "The polar spline spaces keep the tensor-product splines of the outer rings and replace the inner rings by\n", + "$(\\mathfrak r + 1)(\\mathfrak r + 2)/2$ projected polynomials of degree $\\mathfrak r$ (for 0-forms; similarly for 1- and\n", + "2-forms). In the Struphy data structures (`PolarVector`), the first `n_rings` rings of each component hold the\n", + "`n_polar` polar coefficients and the remaining rings are tensor-product coefficients:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cba7a961", + "metadata": {}, + "outputs": [], + "source": [ + "domain = domains.PerturbedDiskCylinder(R1=1.0, eps=0.2, k=3, Lz=1.0)\n", + "print(\"r space dimension n_polar n_rings\")\n", + "for r in (1, 2, 3):\n", + " derham = polar_derham(domain, r)\n", + " for space_id in (\"H1\", \"Hcurl\", \"Hdiv\", \"L2\"):\n", + " space = derham.polar_coeff_spaces[derham.space_to_form[space_id]]\n", + " print(f\"{r} {space_id:6s} {space.dimension:9d} {str(space.n_polar):13s} {space.n_rings}\")" + ] + }, + { + "cell_type": "markdown", + "id": "c12978b6", + "metadata": {}, + "source": [ + "## 3. Regularity 1 is the C1 construction\n", + "\n", + "A tensor-product 0-form is $C^1$ at the pole exactly when, on the two innermost rings, its coefficients are those of an\n", + "affine function $a + b x + c y$ (Toshniwal et al. 2017). On a spline mapping these coefficients are $a + b\\,c^x_{ij} + c\\,c^y_{ij}$\n", + "with the control points $c_{ij}$ of the mapping, which is how the former C1 basis was built. The polar 0-forms of\n", + "regularity 1 span the same space: on the rings 0 and 1, they lie in the span of the projections of $1$, $x$ and $y$.\n", + "(The test `struphy/polar/tests/test_polar_regularity_one.py` compares all degrees with the former C1 extraction\n", + "matrices.)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "bb99a711", + "metadata": {}, + "outputs": [], + "source": [ + "domain = domains.IGAPolarCylinder(a=1.0, Lz=1.0, num_elements=(8, 12), degree=(3, 3))\n", + "derham = polar_derham(domain, 1, num_elements=(8, 12, 1), degree=(3, 3, 1))\n", + "tensor = polar_derham(domain, 1, num_elements=(8, 12, 1), degree=(3, 3, 1), polar=False)\n", + "n2 = 12\n", + "\n", + "affine = [tensor.P0(lambda e1, e2, e3: 1.0 + 0 * e1)]\n", + "for k in range(2):\n", + " affine.append(tensor.P0(lambda e1, e2, e3, k=k: domain(e1, e2, e3)[k]))\n", + "Q = np.array([v.toarray().reshape(-1, n2)[:2].ravel() for v in affine]) # rings 0 and 1 of 1, x, y\n", + "\n", + "E = derham.extraction_ops[\"0\"].blocks_ten_to_pol[0][0].toarray()[:, : 2 * n2] # polar 0-forms on rings 0 and 1\n", + "print(\"rank of the affine functions on the rings 0, 1: \", np.linalg.matrix_rank(Q))\n", + "print(\"rank of the polar 0-forms of regularity 1 on the rings 0, 1:\", np.linalg.matrix_rank(E, tol=1e-9))\n", + "print(\"joint rank: \", np.linalg.matrix_rank(np.vstack([Q, E]), tol=1e-9))" + ] + }, + { + "cell_type": "markdown", + "id": "4d72973c", + "metadata": {}, + "source": [ + "## 4. Smoothness at the pole\n", + "\n", + "A random element of the polar space of regularity $\\mathfrak r$ is $C^{\\mathfrak r}$ at the pole, while a random\n", + "tensor-product spline is not even continuous. We draw random coefficients and measure the jumps of the derivatives\n", + "$d^k u / dx^k$ across the pole along the line through the rays $\\eta_2 = 0$ and $\\eta_2 = 1/2$ (exact one-sided\n", + "derivatives from the polynomial pieces of the first radial cell, degree-6 splines):" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4e524e29", + "metadata": {}, + "outputs": [], + "source": [ + "from struphy.polar.homotopy.demos.regularity import line_derivative_jumps\n", + "\n", + "domain = domains.PerturbedDiskCylinder(R1=1.0, eps=0.2, k=3, Lz=1.0)\n", + "rng = np.random.default_rng(1)\n", + "\n", + "\n", + "def random_field(derham):\n", + " u = derham.create_spline_function(\"u\", \"H1\")\n", + " v = u.vector\n", + " if hasattr(v, \"pol\"):\n", + " v.pol = [rng.standard_normal(p.shape) for p in v.pol]\n", + " v.tp._data[...] = rng.standard_normal(v.tp._data.shape)\n", + " v.set_tp_coeffs_to_zero()\n", + " else:\n", + " v._data[...] = rng.standard_normal(v._data.shape)\n", + " v.update_ghost_regions()\n", + " u.vector = v\n", + " return u\n", + "\n", + "\n", + "fields = {\"tensor product\": random_field(polar_derham(domain, 1, degree=(6, 6, 1), polar=False))}\n", + "for r in (1, 2, 3):\n", + " fields[f\"r = {r}\"] = random_field(polar_derham(domain, r, degree=(6, 6, 1)))\n", + "\n", + "print(f\"{'':15s}\" + \"\".join(f\" jump d^{k}u\" for k in range(5)))\n", + "for name, u in fields.items():\n", + " print(f\"{name:15s}\" + \"\".join(f\"{j:12.1e}\" for j in line_derivative_jumps(u, domain, num_cells=8, degree=6)))" + ] + }, + { + "cell_type": "markdown", + "id": "7dd61567", + "metadata": {}, + "source": [ + "The kinks are visible along the line $y = 0$ through the pole:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a222e6fb", + "metadata": {}, + "outputs": [], + "source": [ + "e1 = np.linspace(0, 0.3, 301)\n", + "fig, axes = plt.subplots(1, 2, figsize=(11, 4))\n", + "for name, u in fields.items():\n", + " xs, values = [], []\n", + " for e2, sign in ((0.5, -1), (0.0, 1)):\n", + " order = slice(None, None, -1) if sign < 0 else slice(None)\n", + " x = np.asarray(domain(e1, np.array([e2]), np.array([0.0]))[0]).ravel()\n", + " xs.append(x[order])\n", + " values.append(np.asarray(u(e1, np.array([e2]), np.array([0.0]))).ravel()[order])\n", + " x, v = np.concatenate(xs), np.concatenate(values)\n", + " v = v - np.interp(0.0, x, v)\n", + " axes[0].plot(x, v, label=name)\n", + " axes[1].plot(0.5 * (x[1:] + x[:-1]), np.diff(v) / np.diff(x), label=name)\n", + "axes[0].set_title(\"random polar spline along $y = 0$\")\n", + "axes[1].set_title(\"its derivative $du/dx$\")\n", + "for ax in axes:\n", + " ax.axvline(0.0, color=\"gray\", lw=0.5)\n", + " ax.set_xlabel(\"$x$\")\n", + " ax.legend(fontsize=8)\n", + "fig.tight_layout()" + ] + }, + { + "cell_type": "markdown", + "id": "729524da", + "metadata": {}, + "source": [ + "**Exactness.** The polynomials of degree $\\mathfrak r$ are reproduced exactly, and the spaces are exactly\n", + "$C^{\\mathfrak r}$ in every direction, when the pullback of these polynomials lies in the logical spline space. On the\n", + "disk $R = \\eta_1$ this needs radial degree $p_1 \\geq \\mathfrak r$, on the star-shaped disk $p_1 \\geq 2\\mathfrak r$, and\n", + "in both cases the angular harmonics $\\cos(m\\theta)$, $\\sin(m\\theta)$, $m \\leq \\mathfrak r$, which B-splines in $\\eta_2$\n", + "only approximate. Along a single line the angular approximation does not enter; in all directions it does, and the\n", + "remaining non-polynomial angular harmonics $\\delta_k$ of the one-sided derivatives $d^k u(0^+; \\alpha)$ on 48 rays\n", + "(the measure of the talk) are of the size of the angular approximation error:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "b5153533", + "metadata": {}, + "outputs": [], + "source": [ + "from struphy.polar.homotopy.demos.regularity import field_smoothness_defects\n", + "\n", + "print(f\"{'':15s} delta_1 delta_2 delta_3\")\n", + "for name, u in fields.items():\n", + " print(f\"{name:15s}\" + \"\".join(f\"{d:10.1e}\" for d in field_smoothness_defects(u, domain)))" + ] + }, + { + "cell_type": "markdown", + "id": "dfac0eba", + "metadata": {}, + "source": [ + "## 5. Commuting projectors\n", + "\n", + "The commuting projectors of the polar complex are the tensor-product projections followed by the polar dual\n", + "coefficients, $\\bar\\pi = \\tilde F^\\top \\pi$. They commute with the polar derivatives, here for a function given in\n", + "logical coordinates (up to the finite differences used for its gradient):" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d1706320", + "metadata": {}, + "outputs": [], + "source": [ + "derham = polar_derham(domain, 3)\n", + "h = 1e-6\n", + "\n", + "\n", + "def g(a, b, c):\n", + " return np.sin(2 * np.pi * b) * a**3 + a**4 * np.cos(4 * np.pi * b) + 0 * c\n", + "\n", + "\n", + "def d(fun, i):\n", + " def derivative(a, b, c):\n", + " plus, minus = [a, b, c], [a, b, c]\n", + " plus[i] = plus[i] + h\n", + " minus[i] = minus[i] - h\n", + " return (fun(*plus) - fun(*minus)) / (2 * h)\n", + "\n", + " return derivative\n", + "\n", + "\n", + "u = derham.P0(g)\n", + "grad_g = derham.P1([d(g, 0), d(g, 1), d(g, 2)])\n", + "residual = derham.grad.dot(u) - grad_g\n", + "print(\"|| grad P0 g - P1 grad g || / || P1 grad g || =\", np.sqrt(residual.inner(residual) / grad_g.inner(grad_g)))" + ] + }, + { + "cell_type": "markdown", + "id": "460163ca", + "metadata": {}, + "source": [ + "## 6. Trigonometric functions in $\\eta_2$\n", + "\n", + "With `DerhamOptions(angular_basis=\"trig\")` the $\\eta_2$-direction uses `num_elements[1]` (odd) nodal trigonometric\n", + "functions instead of periodic B-splines: the 0-forms are the cardinal functions of the nodes $\\eta_2 = i/n$, the\n", + "1-forms their cell histopolants (so the derivative is the forward difference and the inter- and histopolation matrices\n", + "are the identity). They contain the angular harmonics exactly, so the polar splines are exactly $C^{\\mathfrak r}$ at\n", + "the pole when the radial degree is large enough (here $p_1 = 6 \\geq 2\\mathfrak r$ on the perturbed disk) and\n", + "$n \\geq 2(k + 1)\\mathfrak r + 1$. $\\eta_2$ must not be MPI-decomposed, and particles, local projectors and\n", + "`BasisProjectionOperator` need B-splines in every direction." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a74ebf16", + "metadata": {}, + "outputs": [], + "source": [ + "def trig_derham(domain, regularity, n2=25, degree=(6, 3, 1)):\n", + " options = DerhamOptions(\n", + " degree=degree,\n", + " bcs=((\"free\", \"free\"), None, None),\n", + " polar_splines=True,\n", + " polar_regularity=regularity,\n", + " angular_basis=\"trig\",\n", + " nquads_proj=tuple(p + 3 for p in degree),\n", + " )\n", + " grid = TensorProductGrid(num_elements=(8, n2, 1), mpi_dims_mask=(True, False, False))\n", + " return Derham(grid, options, domain=domain)\n", + "\n", + "\n", + "domain = domains.PerturbedDiskCylinder(R1=1.0, eps=0.2, k=3, Lz=1.0)\n", + "trig_fields = {f\"r = {r}\": random_field(trig_derham(domain, r)) for r in (1, 2, 3)}\n", + "print(f\"{'':15s} delta_1 delta_2 delta_3\")\n", + "for name, u in trig_fields.items():\n", + " print(f\"{name:15s}\" + \"\".join(f\"{d:10.1e}\" for d in field_smoothness_defects(u, domain)))" + ] + }, + { + "cell_type": "markdown", + "id": "aab46dfe", + "metadata": {}, + "source": [ + "The remaining angular harmonics now vanish to round-off for $k \\leq \\mathfrak r$ (compare Section 4). The figures of\n", + "the talk use these spaces; see `struphy.polar.homotopy.demos.talk_figures`.\n", + "\n", + "## Further reading\n", + "\n", + "* Examples with polar splines: `examples/Poisson/polar_rounded_square`, `examples/CurlCurl/polar_rounded_square`,\n", + " `examples/ScalarWave/gaussian_perturbed_disk`, `examples/Maxwell/gaussian_perturbed_disk`,\n", + " `examples/Maxwell/gaussian_helical_tube` and `examples/GuidingCenter/vacuum_near_axis`, and the on-axis cylinders\n", + " and tori of `examples/Poisson`, `examples/CurlCurl` and `examples/GyrokineticPoisson`.\n", + "* *High-order polar splines* (`tutorial_polar_splines`): the polar bases, the wave equations, the spectrum explorer\n", + " and the guiding-centre orbits of the talk.\n", + "* *Homotopy operators and high-order polar complexes* (`dev_tutorial_polar_homotopy`): the construction." + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/tutorials/tutorial_polar_splines.ipynb b/tutorials/tutorial_polar_splines.ipynb new file mode 100644 index 000000000..240003a0d --- /dev/null +++ b/tutorials/tutorial_polar_splines.ipynb @@ -0,0 +1,581 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "710c2a4f", + "metadata": {}, + "source": [ + "# High-order polar splines\n", + "\n", + "Many plasma-physics domains (cylinders, tori, stellarators) are parametrized by a **disk-like coordinate\n", + "transformation**\n", + "\n", + "$$\n", + "\\Phi(\\hat x, \\hat y) = \\begin{bmatrix} R(\\hat x, \\hat y)\\, C(\\hat y) \\\\ R(\\hat x, \\hat y)\\, S(\\hat y)\\end{bmatrix},\n", + "\\qquad R(0, \\hat y) = 0,\n", + "$$\n", + "\n", + "which collapses the logical edge $\\hat x = 0$ onto the centre (the *pole*, or the magnetic axis in 3D). Tensor-product\n", + "splines pushed forward by $\\Phi$ are discontinuous at the pole and their gradients are unbounded. This tutorial shows\n", + "the **high-order polar subspaces** of R. Remmerswaal (talk *Constructing high-order de Rham complexes on 2D and 3D\n", + "domains with disk-like coordinate singularities via homotopy operators*, NMPP seminar 2026, and the Julia package\n", + "[PolarFEEC.jl](https://gitlab.mpcdf.mpg.de/rwr/polarfeec.jl)), which are available in Struphy as\n", + "\n", + "```python\n", + "DerhamOptions(polar_splines=True, polar_regularity=r)\n", + "```\n", + "\n", + "For a regularity $\\mathfrak r$ the polar spaces are conforming ($\\mathfrak r \\geq 1$), the 0-forms are\n", + "$C^{\\mathfrak r}$, the 1-forms $C^{\\mathfrak r - 1}$ and the 2-forms $C^{\\mathfrak r - 2}$ at the pole, they form a\n", + "de Rham complex with the cohomology of the disk and come with commuting projections:\n", + "\n", + "$$\n", + "\\bar V^{0,\\mathfrak r} = \\pi^0 T^{\\mathfrak r} \\oplus V^{0,\\mathfrak r}_{\\rm ext}, \\qquad\n", + "\\bar V^{1,\\mathfrak r} = \\operatorname{grad} \\pi^0 T^{\\mathfrak r} \\oplus \\operatorname{curl}^\\dagger\n", + " \\tilde\\pi^{2,1} T^{\\mathfrak r - 2} \\oplus V^{1,\\mathfrak r}_{\\rm ext}, \\qquad\n", + "\\bar V^{2,\\mathfrak r} = \\tilde\\pi^{2,1} T^{\\mathfrak r - 2} \\oplus V^{2,\\mathfrak r}_{\\rm ext},\n", + "$$\n", + "\n", + "where $T^{\\mathfrak r}$ are the polynomials of total degree $\\mathfrak r$, $\\pi^k$ the commuting tensor-product\n", + "projections and $\\operatorname{curl}^\\dagger$ the Poincaré homotopy operator (a radial integral).\n", + "\n", + "The 2D studies of this tutorial use 2D Struphy `Derham` complexes (one cell in $\\eta_3$) on Struphy domains\n", + "(`struphy.polar.homotopy.demos.common`): the polar complexes are built from the 1D spline spaces of the `Derham`\n", + "complex, or from trigonometric functions in $\\eta_2$ (`DerhamOptions(angular_basis=\"trig\")`). The later sections use\n", + "the same spaces in the 3D `Derham` complex. Set the environment variable `STRUPHY_TUTORIAL_FULL=1` to run the studies\n", + "at the resolution of the talk. All figures of the talk can be written with\n", + "\n", + "```\n", + "python -m struphy.polar.homotopy.demos.talk_figures OUT_FOLDER [--full] [--videos]\n", + "```" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "106c60e2", + "metadata": {}, + "outputs": [], + "source": [ + "import os\n", + "\n", + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", + "\n", + "from struphy.polar.homotopy.demos import basis, common, regularity\n", + "\n", + "FULL = os.environ.get(\"STRUPHY_TUTORIAL_FULL\", \"0\") == \"1\"" + ] + }, + { + "cell_type": "markdown", + "id": "860a2a09", + "metadata": {}, + "source": [ + "## 1. Tensor-product splines on a disk-like mapping\n", + "\n", + "We use the perturbed disk of the talk, $R = \\hat x\\,(1 + 0.2\\,\\hat x \\cos(6\\pi \\hat y))$, $C = \\cos(2\\pi\\hat y)$,\n", + "$S = \\sin(2\\pi\\hat y)$, the Struphy domain `PerturbedDiskCylinder(R1=1, eps=0.2, k=3)`. Its Jacobian\n", + "$J = R\\,\\partial_{\\hat x}R\\,(C S' - S C')$ vanishes only at $\\hat x = 0$." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "f76c059c", + "metadata": {}, + "outputs": [], + "source": [ + "basis.figure_mapping();" + ] + }, + { + "cell_type": "markdown", + "id": "337e05d6", + "metadata": {}, + "source": [ + "The tensor-product 0-forms are products $B_i(\\hat x) U_j(\\hat y)$ of cubic clamped B-splines on $N^x = 5$ cells and\n", + "cubic periodic B-splines on $N^y = 12$ cells (or $13$ nodal trigonometric functions,\n", + "`angular_basis=\"trig\"`). The first ring ($i = 0$) collapses to the pole:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "b3a9e6ee", + "metadata": {}, + "outputs": [], + "source": [ + "basis.figure_tensor_basis(\"bspline\");" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "f2753e1d", + "metadata": {}, + "outputs": [], + "source": [ + "basis.figure_tensor_basis(\"trig\");" + ] + }, + { + "cell_type": "markdown", + "id": "adb61929", + "metadata": {}, + "source": [ + "A generic tensor-product function is discontinuous at the origin, since $P(0) = \\sum_j p_{0j} U_j(\\hat y)$ depends on\n", + "the angle; its gradient is unbounded. The classical remedy is $p_{0j} = P(0)$ ($C^0$) and first-order conditions for\n", + "$C^1$ (Toshniwal et al. 2017; Holderied & Possanner 2022; the former C1 polar splines of Struphy, which are\n", + "exactly the polar splines of regularity 1)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7ba602f5", + "metadata": {}, + "outputs": [], + "source": [ + "basis.figure_discontinuity();" + ] + }, + { + "cell_type": "markdown", + "id": "dca99777", + "metadata": {}, + "source": [ + "## 2. High-order polar 0-forms\n", + "\n", + "The polar subspace of regularity $\\mathfrak r$ keeps the *external* tensor-product functions (rings $i > \\mathfrak r$,\n", + "which vanish to order $\\mathfrak r + 1$ at the pole) and adds back the $(\\mathfrak r + 1)(\\mathfrak r + 2)/2$\n", + "polynomials of degree $\\mathfrak r$, projected onto the tensor-product space and truncated to the inner rings:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1b917b81", + "metadata": {}, + "outputs": [], + "source": [ + "for r in (1, 2, 3):\n", + " basis.figure_polar_zero_forms(r)" + ] + }, + { + "cell_type": "markdown", + "id": "6dbee81f", + "metadata": {}, + "source": [ + "The dimensions and the cohomology $(1, 0, 0)$ of the disk, for every regularity:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8723ba6a", + "metadata": {}, + "outputs": [], + "source": [ + "for r in (0, 1, 2, 3):\n", + " c = common.polar_complex(r, degree=3, cells=(5, 12))\n", + " print(f\"r = {r}: dimensions {[c.dimension(k) for k in range(3)]}, central functions {c.central_counts()}, \"\n", + " f\"cohomology {c.cohomology_dimensions()}\")" + ] + }, + { + "cell_type": "markdown", + "id": "16c83e5d", + "metadata": {}, + "source": [ + "### $\\mathfrak r$-compatibility\n", + "\n", + "The polynomials are reproduced exactly (and the spaces are exactly $C^{\\mathfrak r}$) when the pullback of\n", + "$T^{\\mathfrak r}$ lies in the logical tensor-product space. For $R = R_1\\hat x(1 + \\varepsilon\\hat x\\cos(2\\pi k\\hat y))$\n", + "this needs clamped B-splines of degree $\\geq 2\\mathfrak r$ and the trigonometric modes up to $(k + 1)\\mathfrak r$ in\n", + "$\\hat y$: the trigonometric angular spaces of `angular_basis=\"trig\"`. With B-splines in $\\hat y$ the central\n", + "functions are spline approximations of the polynomials and the regularity holds up to that approximation error.\n", + "\n", + "### Regularity at the centre (PolarFEEC.jl example 04)\n", + "\n", + "A function is $C^k$ at the centre exactly when, for each $j \\leq k$, its $j$-th derivative along the ray at angle\n", + "$\\alpha$ contains only the angular harmonics $j, j-2, \\dots$. We project\n", + "$f(x, y) = \\sin(3x + 1)\\cos(2y - 0.5) + 0.5\\cos(4xy)$ with the commuting projections into the tensor-product space and\n", + "the polar spaces $\\mathfrak r = 0, \\dots, 3$ (degree-5 B-splines, 49 nodal trigonometric functions, the star-shaped\n", + "disk `PerturbedDiskCylinder(power=0)`) and measure the remaining harmonics $\\delta_k$ by polynomial fits along 48\n", + "rays:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9a3a4e14", + "metadata": {}, + "outputs": [], + "source": [ + "results = regularity.study()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cea8b055", + "metadata": {}, + "outputs": [], + "source": [ + "regularity.figure_central_basis(3, results);" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4da7f908", + "metadata": {}, + "outputs": [], + "source": [ + "regularity.figure_regularity(results);" + ] + }, + { + "cell_type": "markdown", + "id": "4a7884ee", + "metadata": {}, + "source": [ + "With a B-spline angular factor the regularity holds only up to the spline approximation of the polynomials:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "bbfd9f33", + "metadata": {}, + "outputs": [], + "source": [ + "print(\"r = 3, B-spline angular factor:\", regularity.bspline_angular_defects())" + ] + }, + { + "cell_type": "markdown", + "id": "82516be9", + "metadata": {}, + "source": [ + "The figures of the talk: the jumps of the derivatives $d^k u/dx^k$ of a random polar 0-form (and 2-form) across the\n", + "centre along the line $y = 0$, for every regularity. The $k$-th derivative is continuous exactly for $k \\leq \\mathfrak r$\n", + "(0-forms) and $k \\leq \\mathfrak r - 2$ (2-forms):" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "77e83c4c", + "metadata": {}, + "outputs": [], + "source": [ + "regularity.figure_regularity_lines(0)\n", + "regularity.figure_regularity_lines(2);" + ] + }, + { + "cell_type": "markdown", + "id": "c013bc9f", + "metadata": {}, + "source": [ + "## 3. Polar 1- and 2-forms\n", + "\n", + "Gradients of polynomials alone do not give a rich enough 1-form space, since the 2-form space contains no\n", + "polynomials. The paper enriches the 2-forms with projected polynomials $\\tilde\\pi^{2,1} T^{\\mathfrak r - 2}$ and adds\n", + "their images under the homotopy operator $\\operatorname{curl}^\\dagger$ to the 1-forms. The resulting spaces form a\n", + "cochain complex (with $\\operatorname{grad}$, $\\operatorname{curl}$) and a chain complex (with the homotopies), with\n", + "$\\bar V^{1,\\mathfrak r} = \\operatorname{grad} \\bar V^{0,\\mathfrak r} \\oplus \\operatorname{curl}^\\dagger\n", + "\\bar V^{2,\\mathfrak r}$. The central basis functions for $\\mathfrak r = 3$:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "18ea019e", + "metadata": {}, + "outputs": [], + "source": [ + "basis.figure_all_forms(3);" + ] + }, + { + "cell_type": "markdown", + "id": "d5fae911", + "metadata": {}, + "source": [ + "## 4. Polar splines in the 3D Struphy complex\n", + "\n", + "In 3D the polar complex is the total complex of the 2D polar complex of the cross-section and the 1D complex in\n", + "$\\eta_3$. In Struphy it is selected with the `DerhamOptions` of the polar splines; the new domains\n", + "`PerturbedDiskCylinder` and `RoundedSquareCylinder` (and all domains with a pole, e.g. `HollowCylinder(a1=0)`,\n", + "`HollowTorus(a1=0)`, `IGAPolarTorus`) can be used:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "dfc4f358", + "metadata": {}, + "outputs": [], + "source": [ + "from struphy.feec.psydac_derham import Derham\n", + "from struphy.geometry import domains\n", + "from struphy.io.options import DerhamOptions\n", + "from struphy.topology.grids import TensorProductGrid\n", + "\n", + "domain = domains.PerturbedDiskCylinder(R1=1.0, eps=0.2, k=3, Lz=1.0)\n", + "options = DerhamOptions(\n", + " degree=(3, 3, 1),\n", + " bcs=((\"free\", \"free\"), None, None),\n", + " polar_splines=True,\n", + " polar_regularity=3,\n", + " nquads_proj=(8, 10, 4),\n", + ")\n", + "derham = Derham(TensorProductGrid(num_elements=(8, 12, 1)), options, domain=domain)\n", + "for space_id in (\"H1\", \"Hcurl\", \"Hdiv\", \"L2\"):\n", + " space = derham.polar_coeff_spaces[derham.space_to_form[space_id]]\n", + " print(f\"{space_id:5s}: polar coefficients per component {space.n_polar}, polar rings {space.n_rings}\")" + ] + }, + { + "cell_type": "markdown", + "id": "4e40f03e", + "metadata": {}, + "source": [ + "The commuting projectors are the paper's $\\bar\\pi = \\bar\\Pi \\circ \\pi$: the tensor-product projection followed by the\n", + "polar dual coefficients. They commute with the polar derivatives (here for logical components, up to the finite\n", + "differences used to differentiate the test functions):" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "15b58cde", + "metadata": {}, + "outputs": [], + "source": [ + "h = 1e-6\n", + "\n", + "\n", + "def f(a, b, c):\n", + " return np.sin(2 * np.pi * b) * a**2 + a**3 * np.cos(4 * np.pi * b) + 0 * c\n", + "\n", + "\n", + "def d(fun, i):\n", + " def g(a, b, c):\n", + " ep, em = [a, b, c], [a, b, c]\n", + " ep[i] = ep[i] + h\n", + " em[i] = em[i] - h\n", + " return (fun(*ep) - fun(*em)) / (2 * h)\n", + "\n", + " return g\n", + "\n", + "\n", + "u = derham.P0(f)\n", + "g = derham.P1([d(f, 0), d(f, 1), d(f, 2)])\n", + "r = derham.grad.dot(u) - g\n", + "print(\"|| grad P0 f - P1 grad f || / || P1 grad f || =\", np.sqrt(r.inner(r) / g.inner(g)))" + ] + }, + { + "cell_type": "markdown", + "id": "ec1bce59", + "metadata": {}, + "source": [ + "### Trigonometric functions in the 3D complex\n", + "\n", + "`angular_basis=\"trig\"` replaces the periodic B-splines in $\\eta_2$ by `num_elements[1]` (odd) nodal trigonometric\n", + "functions in every space of the `Derham` complex (masses, projectors, derivatives, evaluation and the polar\n", + "splines; $\\eta_2$ must not be MPI-decomposed, and particles need B-splines):" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "c995918a", + "metadata": {}, + "outputs": [], + "source": [ + "options = DerhamOptions(\n", + " degree=(6, 3, 1),\n", + " bcs=((\"free\", \"free\"), None, None),\n", + " polar_splines=True,\n", + " polar_regularity=3,\n", + " angular_basis=\"trig\",\n", + " nquads_proj=(9, 6, 4),\n", + ")\n", + "grid = TensorProductGrid(num_elements=(8, 25, 1), mpi_dims_mask=(True, False, False))\n", + "trig_derham = Derham(grid, options, domain=domain)\n", + "u = trig_derham.P0(f)\n", + "g = trig_derham.P1([d(f, 0), d(f, 1), d(f, 2)])\n", + "r = trig_derham.grad.dot(u) - g\n", + "print(\"trig: || grad P0 f - P1 grad f || / || P1 grad f || =\", np.sqrt(r.inner(r) / g.inner(g)))" + ] + }, + { + "cell_type": "markdown", + "id": "9ee50297", + "metadata": {}, + "source": [ + "## 5. Wave equations across the axis\n", + "\n", + "The talk compares the regularities on the 0-form wave equation $\\partial_t^2 P = \\Delta P$ (Dirichlet) with a\n", + "Gaussian pulse centred next to the pole, $P_0 = \\exp(-((x - 0.05)^2 + y^2) / (2\\cdot 0.04^2))$. The semidiscretization\n", + "in polar coefficients is $\\bar M\\,v'' + \\bar K v = 0$ with $\\bar M = E^\\top M E$, $\\bar K = (DE)^\\top M (DE)$; it is\n", + "integrated with the staggered leapfrog at $\\Delta t = 0.9\\cdot 2/\\sqrt{\\lambda_{\\max}}$, which conserves the\n", + "staggered energy exactly. The snapshots for the different regularities:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d3ab7023", + "metadata": {}, + "outputs": [], + "source": [ + "from struphy.polar.homotopy.demos import waves\n", + "\n", + "wave_kwargs = dict(degree=5, radial_cells=8, angular_cells=16) if FULL else dict(degree=3, radial_cells=6, angular_cells=12)\n", + "wave_family = waves.run_family(0, regularities=(0, 1, 2, 3) if FULL else (0, 2), final_time=0.3, **wave_kwargs)\n", + "waves.figure_snapshots(wave_family);" + ] + }, + { + "cell_type": "markdown", + "id": "ef097fd5", + "metadata": {}, + "source": [ + "The same leapfrog is the Struphy model `ScalarWave` (example `examples/ScalarWave/gaussian_perturbed_disk`), and the\n", + "1-form wave is `Maxwell` with `algo=\"leapfrog\"` (examples `examples/Maxwell/gaussian_perturbed_disk` and, in 3D on a\n", + "helical tube, `examples/Maxwell/gaussian_helical_tube`). The convergence studies are in the developer tutorial\n", + "*Homotopy operators and high-order polar complexes*.\n", + "\n", + "## 6. The spectrum explorer of the talk\n", + "\n", + "The interactive slide of the talk shows the eigenpairs of the 2D curl–curl operator\n", + "$(\\operatorname{curl}A, \\operatorname{curl}\\Phi) = \\lambda (A, \\Phi)$ in the Dirichlet polar 1-forms of every\n", + "regularity on the unit disk (`HollowCylinder(a1=0)`, degree 6 on $10 \\times 14$ cells), with the exact eigenvalues\n", + "$(j'_{m,k})^2$ and the *regularity deficit* of every eigenfunction: its relative $L^2$ distance to the polar 1-forms of\n", + "regularity 6. `spectrum.write_explorer` writes the page (open `spectrum_explorer.html` in a browser):" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "f48fdff5", + "metadata": {}, + "outputs": [], + "source": [ + "from struphy.polar.homotopy.demos import spectrum\n", + "\n", + "page = spectrum.write_explorer(\"spectrum_explorer\", count=120 if FULL else 30, pixels=96 if FULL else 64)\n", + "print(page)" + ] + }, + { + "cell_type": "markdown", + "id": "326f7578", + "metadata": {}, + "source": [ + "The spectra contain no spurious modes, and the low eigenvalues agree for every regularity; the deficits show which\n", + "eigenfunctions the spaces of low regularity represent with a kink at the centre:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "0fca28a0", + "metadata": {}, + "outputs": [], + "source": [ + "reference = spectrum.complex_(spectrum.REFERENCE_REGULARITY)\n", + "fig, axes = plt.subplots(1, 2, figsize=(11, 4))\n", + "for r in (0, 1, 2, 3):\n", + " s = spectrum.series(r, reference=reference, count=30)\n", + " exact = np.array(s[\"exact\"], dtype=float)\n", + " n = np.arange(1, len(s[\"lam\"]) + 1)\n", + " axes[0].semilogy(n, np.abs(s[\"lam\"] - exact) / exact, \"o-\", ms=3, label=rf\"$\\mathfrak{{r}} = {r}$\")\n", + " axes[1].semilogy(n, s[\"deficit\"], \"o-\", ms=3, label=rf\"$\\mathfrak{{r}} = {r}$\")\n", + "axes[0].set_title(\"relative eigenvalue error\")\n", + "axes[1].set_title(\"regularity deficit\")\n", + "for ax in axes:\n", + " ax.set_xlabel(\"eigenvalue index\")\n", + " ax.legend(fontsize=8)\n", + "fig.tight_layout()" + ] + }, + { + "cell_type": "markdown", + "id": "f2272dce", + "metadata": {}, + "source": [ + "## 7. Guiding-centre orbits near the magnetic axis\n", + "\n", + "The guiding centres of the talk follow the vacuum field `HelicalVacuumField` projected into the polar splines\n", + "(`struphy.polar.homotopy.demos.orbits`; the particle version is `examples/GuidingCenter/vacuum_near_axis`). Orbits that\n", + "start next to the axis cross it:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "ddcdb270", + "metadata": {}, + "outputs": [], + "source": [ + "from struphy.polar.homotopy.demos import orbits\n", + "\n", + "field = orbits.ProjectedVacuumField(2)\n", + "path, energy_errors = orbits.integrate(field, 0.04, periods=20.0 if FULL else 2.0)\n", + "print(\"max relative energy error:\", energy_errors.max())\n", + "orbits.figure_orbits(path)\n", + "orbits.figure_orbits(path, extent=0.05);" + ] + }, + { + "cell_type": "markdown", + "id": "2fffa11b", + "metadata": {}, + "source": [ + "`orbits.energy_study` measures the energy error against the time step for every regularity (B-spline or\n", + "trigonometric angular spaces). With the field and the markers chosen here (the talk's script was not available) it\n", + "does not reproduce the strong dependence on $\\mathfrak r$ shown in the talk.\n", + "\n", + "## 8. Field lines in a twisted tube (PolarFEEC.jl example 07)\n", + "\n", + "A closed 2-form $B = \\Phi'(r)\\,dr\\wedge d\\theta + d\\zeta\\wedge d\\psi$ on the helical tube `HelicalTube(\"islands\")`\n", + "(rotational transform $\\iota = 0.6 + 0.3 r^2$ and two island chains) is projected with the commuting projector of the\n", + "3D polar complex. Its divergence vanishes in polar coefficients, the magnetic axis included, and the traced field lines\n", + "have the exact rotational transform:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "0c9abef6", + "metadata": {}, + "outputs": [], + "source": [ + "from struphy.polar.homotopy.demos import fieldlines\n", + "\n", + "field = fieldlines.study(turns=150 if FULL else 25, num_elements=(8, 24, 32) if FULL else (6, 18, 16))\n", + "fieldlines.figure_poincare(field);" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +}