"""Building a tokamak by tensorisation, and checking it is one. The mesh, on its own -- what it is made of, what it costs, and the two controls that say it is right. The Gaussian projected on it lives in ``tokamak_projection_dg.py``, which imports from here; keeping the two apart means the geometry can be re-run and looked at without paying for a solve, and that anything else built on this tokamak starts from one definition of it. A solid torus is a poloidal section times the toroidal angle, curved by a map on the PRODUCT space: ``(x, y, phi) -> ((R + x) cos phi, (R + x) sin phi, y)``. The angle couples with ``R``, so this geometry is *not* a product of geometries -- but it still lives on a mesh that IS one, which is the whole point of the ``mapping`` argument of :func:`tensor_mesh`. Two sections go through the same factory, and that they do is the claim being made: a circular disk, cheap to iterate on, and the **JET wall** read from its EQDSK. The wall's coordinates are the machine's own ``(R, Z)``, so it is the same map with ``major_radius = 0`` and nothing else changes. **The two controls.** The volume against **Pappus**, ``2 pi`` times the first moment of the section, read off the SECTION's own quadrature -- better than ``2 pi^2 R`` because it also exists for a wall that has no closed form, and because it separates the two errors instead of adding them: what is left over is the polynomial approximation of the toroidal arc, not of the wall. Then the divergence theorem, ``div (x, y, z) = 3``, which is what says the faces are listed once each with their normals pointing outwards. The faces come for free: ``_build_faces`` matches faces in any dimension, so this mesh does fluxes as well as volumes. Both are closed toroidally (``periodic_directions=(2,)``), so their only boundary is the outer surface: the cells at ``phi = 2pi`` list the nodes of those at ``phi = 0``. ⚠ The same product left UNMAPPED -- a cylinder -- cannot be, and pairs the very same nodes: only the geometry knows that its two ends are ``2 pi`` apart in space rather than on top of each other. Gluing it folds the last layer of cells back on itself, measured as a volume of 41.69 instead of 19.72 with nothing raised. Hence the check, and hence: what makes a geometry periodic is the map, not the declaration. """ import pathlib import jax import jax.numpy as jnp import matplotlib.pyplot as plt import numpy as np from mpl_toolkits.mplot3d.art3d import Line3DCollection from scimba_jax.linear_approximation.meshes.mesh import Mesh from scimba_jax.linear_approximation.meshes.tensor_mesh import tensor_mesh from scimba_jax.linear_approximation.meshes.unstructured_mesh import UnstructuredMesh from scimba_jax.linear_approximation.quad.gauss_quad import UnitSquareTensorized from scimba_jax.mapping.macro_mesh import macro_mesh_ogrid_disk from scimba_jax.mapping.mapping import InvertibleFunction, Mapping GEOMETRY_ORDER = 2 QUAD_ORDER = 4 SECTION_CELLS = 4 N_TOROIDAL = 32 MAJOR_RADIUS = 3.0 # The JET wall, meshed once by GMSH and cached by # `fem/solve/classical_approach/solve_laplacian_unstructured_2d.py` -- reused # rather than re-meshed, because GMSH does not repeat itself (see that file). JET_MESH = ( pathlib.Path(__file__).resolve().parents[1] / "fem" / "solve" / "classical_approach" / "mesh_jet_p3_h0.13.npz" ) JET_TOROIDAL = 16 # ── The two factors, and the map that curves their product ────────────────── def poloidal_section(cells=SECTION_CELLS): """The circular section: a disk, as an unstructured mesh of curved quads.""" macro = macro_mesh_ogrid_disk(radius=1.0, n=cells, order=GEOMETRY_ORDER) return UnstructuredMesh.from_macro_mesh( macro, UnitSquareTensorized(dim=2, order=QUAD_ORDER) ) def jet_section(): """The real thing: the JET wall, as read from its EQDSK and meshed. ⚠ Its coordinates are the machine's own ``(R, Z)`` in metres, with ``R`` from 1.84 to 3.89 -- an ABSOLUTE major radius, not an offset from one. So the same product map applies with ``major_radius = 0``, and nothing else changes: same factory, same faces, same periodicity. That is the point of keeping the geometry in the map rather than in the mesh. """ if not JET_MESH.exists(): raise FileNotFoundError( f"{JET_MESH} is missing; run " "fem/solve/classical_approach/solve_laplacian_unstructured_2d.py " "once to mesh the wall and cache it." ) stored = np.load(JET_MESH) return UnstructuredMesh.from_macro_mesh( _NodesAndCells(stored["nodes"], stored["cells"]), UnitSquareTensorized(dim=2, order=QUAD_ORDER), ) class _NodesAndCells: """The three fields ``from_macro_mesh`` reads, from a cached array pair.""" def __init__(self, nodes, cells): self.nodes = np.asarray(nodes) self.cells = np.asarray(cells) self.order = int(round(cells.shape[1] ** 0.5)) - 1 def toroidal_angle(cells=N_TOROIDAL): """``[0, 2pi]``, as a 1-D structured mesh -- the map carries the ``2pi``.""" return Mesh( dim=1, n_cells=(cells,), ref_quad=UnitSquareTensorized(dim=1, order=QUAD_ORDER), mapping=Mapping( mappings=[ InvertibleFunction( lambda x: 2.0 * jnp.pi * x, lambda y: y / (2.0 * jnp.pi) ) ] ), ) def torus_map(major_radius=MAJOR_RADIUS): """``(x, y, phi)`` in the product space -> a point of the torus. Invertible away from the axis, and the inverse is in closed form: the mesh needs it to locate a physical point in a cell. """ def forward(z): x, y, phi = z[..., 0], z[..., 1], z[..., 2] return jnp.stack( [ (major_radius + x) * jnp.cos(phi), (major_radius + x) * jnp.sin(phi), y, ], axis=-1, ) def inverse(p): big_x, big_y, big_z = p[..., 0], p[..., 1], p[..., 2] return jnp.stack( [ jnp.sqrt(big_x**2 + big_y**2) - major_radius, big_z, jnp.arctan2(big_y, big_x) % (2.0 * jnp.pi), ], axis=-1, ) return InvertibleFunction(forward, inverse) def revolve(section, toroidal_cells, major_radius, closed=True): """Turn any poloidal section into a solid torus. Args: section: The 2-D mesh to revolve. toroidal_cells: Cells around the torus. major_radius: Added to the section's first coordinate. ``3`` for a disk centred on the origin, ``0`` for a wall already given in machine ``(R, Z)``. closed: Glue ``phi = 2pi`` onto ``phi = 0``, leaving only the outer surface as boundary. Returns: The 3-D :class:`UnstructuredMesh`, with its faces and their quadrature stored. """ mesh = tensor_mesh( section, toroidal_angle(toroidal_cells), mapping=torus_map(major_radius), storage=True, periodic_directions=(2,) if closed else (), ) # ⚠ `storage=True` does not reach an unstructured product, and a mesh now # stores only what someone declares reading (`request_storage`): the # divergence control reads the boundary faces' quadrature, so it asks. mesh.request_storage("faces") return mesh def tokamak(section_cells=SECTION_CELLS, toroidal_cells=N_TOROIDAL, closed=True): """The circular-section tokamak, the cheap one to iterate on.""" return revolve( poloidal_section(section_cells), toroidal_cells, MAJOR_RADIUS, closed ) def jet(toroidal_cells=JET_TOROIDAL, closed=True): """The JET tokamak: the real wall, revolved. ⚠ Much bigger -- 483 poloidal cells of degree 3 rather than a few dozen of degree 2 -- so both the volume and the divergence check are the same controls run on a mesh that is not a toy. """ return revolve(jet_section(), toroidal_cells, 0.0, closed) # ── What the mesh believes about itself ───────────────────────────────────── def all_cells(mesh): """``jnp.arange`` of the cells -- the argument every ``vmap`` maps on. ⚠ Not a Python loop. A cell index arriving as a Python ``int`` is a compile-time constant, so a loop over the cells traces the cell map once per cell; the same code under ``vmap`` traces once. Measured on this example as the difference between two minutes and two seconds. """ return jnp.arange(mesh.n_cells_total) def volume(mesh): """Quadrature of ``1``: what the mesh believes its own volume to be.""" weights = jax.vmap(lambda c: mesh._local_weights_points(c)[0])(all_cells(mesh)) return float(jnp.sum(weights)) def divergence_defect(mesh): """``|flux of (x,y,z) - 3 x volume|``, the faces' own control. Compared against the mesh's OWN volume, so what is being tested is the face listing rather than the degree-2 disk: a face counted twice, one left off, a normal pointing inwards or a Nanson weight missing its ``|J|`` each move this by a finite amount. """ weights = mesh._store_ext_face_wts points = mesh._store_ext_face_pts normals = mesh._store_ext_face_nrm flux = float(jnp.sum(weights * jnp.sum(points * normals, -1))) return abs(flux - 3.0 * volume(mesh)) def pappus_volume(section, major_radius): """``2 pi`` times the first moment of the section -- Pappus, exactly. ⚠ The right reference for a revolved mesh, and better than ``2 pi^2 R`` because it also exists for a wall that has no closed form. Read off the SECTION's own quadrature, so comparing it to the 3-D mesh's volume tests the product and the map while saying nothing about how well the section approximates the machine -- the two errors are separated instead of added. """ weights, points = jax.vmap(section._local_weights_points)(all_cells(section)) return float(2.0 * jnp.pi * jnp.sum(weights * (major_radius + points[..., 0]))) def report(name, mesh, reference): """One line: size, volume against Pappus, divergence defect.""" got = volume(mesh) print( f"{name:9s} {mesh.n_cells_total:6d} mailles " f"{len(mesh.external_faces_idx):5d} faces de bord " f"volume {got:11.6f} Pappus {reference:11.6f} " f"ecart relatif {abs(got - reference) / abs(reference):.2e} " f"|flux - 3 vol| {divergence_defect(mesh):.2e}" ) # ── Drawing the mesh ──────────────────────────────────────────────────────── def face_wireframe(mesh, faces, n_samples=7): """The four curved edges of every listed face, ``(n, 4, n_samples, 3)``. Drawn from the MAP rather than from the corner nodes, so a curved face looks curved -- on the tokamak the toroidal edges are arcs, and joining corners by segments would hide exactly the coupling being illustrated. ⚠ Vectorised over the faces, not looped. On JET's 7728 cells the loop version traced the cell map once per face and did not finish in two minutes. """ line = np.linspace(0.0, 1.0, n_samples) edges = [] for axis in range(2): for fixed in (0.0, 1.0): tangential = np.empty((n_samples, 2)) tangential[:, axis] = line tangential[:, 1 - axis] = fixed edges.append(tangential) edges = jnp.asarray(np.stack(edges)) # (4, n_samples, 2) def per_face(face_id, cell): reference = jax.vmap(mesh._face_reference_points, in_axes=(None, 0))( face_id, edges ) return jax.vmap(lambda points: mesh._unit_hypercube_to_cell(cell, points))( reference ) return np.asarray( jax.vmap(per_face)( jnp.asarray(mesh.faces_left_edge)[faces], jnp.asarray(mesh.faces_left)[faces], ) ) def surface_of(mesh, keep=None): """Faces bounding the kept part: the outer ones, plus those the cut exposes. A wireframe of every cell is unreadable in 3-D -- the far side shows through the near one. The boundary of the region is what one actually wants to see, and after the face listing it is free: an outer face if its cell is kept, an interior face if exactly one of its two cells is. """ kept = np.ones(mesh.n_cells_total, dtype=bool) if keep is not None: centres = np.asarray( jax.vmap( lambda c: mesh._unit_hypercube_to_cell(c, jnp.full((1, 3), 0.5))[0] )(all_cells(mesh)) ) kept = np.array([keep(centre) for centre in centres]) left, right = mesh.faces_left, mesh.faces_right on_boundary = (right < 0) & kept[left] exposed = (right >= 0) & (kept[left] != kept[np.maximum(right, 0)]) return np.where(on_boundary | exposed)[0] def quarter_cut(point): """Three quarters of the solid, so the poloidal section stays visible.""" return not (point[0] > 0.0 and point[1] < 0.0) def draw(axes, mesh, title, keep=quarter_cut): """The surface of the mesh, cut away so the poloidal section shows. One :class:`Line3DCollection` rather than one ``plot`` per edge: JET's surface has thousands of them, and matplotlib charges per artist. """ curves = face_wireframe(mesh, surface_of(mesh, keep)) axes.add_collection3d( Line3DCollection( curves.reshape(-1, curves.shape[-2], 3), colors="0.25", linewidths=0.4 ) ) everything = curves.reshape(-1, 3) for setter, column in zip((axes.set_xlim, axes.set_ylim, axes.set_zlim), range(3)): setter(everything[:, column].min(), everything[:, column].max()) tidy_3d(axes, title) def tidy_3d(axes, title): axes.set_title(title) axes.set_box_aspect((1, 1, 1)) axes.set_xticks([]) axes.set_yticks([]) axes.set_zticks([]) def main(): circular_section = poloidal_section() circular = revolve(circular_section, N_TOROIDAL, MAJOR_RADIUS) report("circulaire", circular, pappus_volume(circular_section, MAJOR_RADIUS)) wall = jet_section() machine = revolve(wall, JET_TOROIDAL, 0.0) report("JET", machine, pappus_volume(wall, 0.0)) figure = plt.figure(figsize=(11, 5.5)) draw( figure.add_subplot(121, projection="3d"), circular, f"section circulaire ({circular.n_cells_total} mailles)", ) draw( figure.add_subplot(122, projection="3d"), machine, f"paroi JET ({machine.n_cells_total} mailles)", ) figure.suptitle( "tokamak par tensorisation : section poloidale x segment [0, 2pi], " "referme en phi" ) plt.tight_layout() plt.show() if __name__ == "__main__": main()