Implement custom source mesh grids

Examples/custom_source_example.py

@@ -25,13 +25,15 @@ def my_custom_reaction_rate(n_i, T_i):
 
 vmec_obj = read_vmec.VMECData(vmec_file)
 
-mesh_size = (11, 81, 61)
-toroidal_extent = 90.0
+cfs_grid = np.linspace(0.0, 1.0, num=11)
+poloidal_grid = np.linspace(0.0, 360.0, num=81)
+toroidal_grid = np.linspace(0.0, 90.0, num=61)
 
 source_mesh_obj = sm.SourceMesh(
     vmec_obj,
-    mesh_size,
-    toroidal_extent,
+    cfs_grid,
+    poloidal_grid,
+    toroidal_grid,
     plasma_conditions=my_custom_plasma_conditions,
     reaction_rate=my_custom_reaction_rate,
 )

Examples/nwl_geom_example.py

@@ -1,6 +1,7 @@
 from pathlib import Path
 import parastell.parastell as ps
 from parastell.cubit_io import tag_surface_legacy
+import numpy as np
 
 # Define directory to export all output files to
 export_dir = ""
@@ -30,10 +31,11 @@
 )
 
 # Define source mesh parameters
-mesh_size = (11, 81, 61)
-toroidal_extent = 90.0
+cfs_grid = np.linspace(0.0, 1.0, num=11)
+poloidal_grid = np.linspace(0.0, 360.0, num=81)
+toroidal_grid = np.linspace(0.0, 90.0, num=61)
 # Construct source
-stellarator.construct_source_mesh(mesh_size, toroidal_extent)
+stellarator.construct_source_mesh(cfs_grid, poloidal_grid, toroidal_grid)
 # Export source file
 stellarator.export_source_mesh(filename="source_mesh", export_dir=export_dir)
 

Examples/parastell_example.py

@@ -70,10 +70,11 @@
 )
 
 # Define source mesh parameters
-mesh_size = (11, 81, 61)
-toroidal_extent = 90.0
+cfs_grid = np.linspace(0.0, 1.0, num=11)
+poloidal_grid = np.linspace(0.0, 360.0, num=81)
+toroidal_grid = np.linspace(0.0, 90.0, num=61)
 # Construct source
-stellarator.construct_source_mesh(mesh_size, toroidal_extent)
+stellarator.construct_source_mesh(cfs_grid, poloidal_grid, toroidal_grid)
 # Export source file
 stellarator.export_source_mesh(filename="source_mesh", export_dir=export_dir)
 

Examples/radial_distance_example.py

@@ -100,10 +100,11 @@
 stellarator.export_magnets()
 
 # Define source mesh parameters
-mesh_size = (11, 81, 61)
-toroidal_extent = 90.0
+cfs_grid = np.linspace(0.0, 1.0, num=11)
+poloidal_grid = np.linspace(0.0, 360.0, num=81)
+toroidal_grid = np.linspace(0.0, 90.0, num=61)
 # Construct source
-stellarator.construct_source_mesh(mesh_size, toroidal_extent)
+stellarator.construct_source_mesh(cfs_grid, poloidal_grid, toroidal_grid)
 # Export source file
 stellarator.export_source_mesh(filename="source_mesh", export_dir=export_dir)
 

parastell/parastell.py

@@ -265,19 +265,18 @@ def export_magnets(
                 mesh_filename=mesh_filename, export_dir=export_dir
             )
 
-    def construct_source_mesh(self, mesh_size, toroidal_extent, **kwargs):
+    def construct_source_mesh(
+        self, cfs_grid, poloidal_grid, toroidal_grid, **kwargs
+    ):
         """Constructs SourceMesh class object.
 
         Arguments:
-            mesh_size (tuple of int): number of grid points along each axis of
-                flux-coordinate space, in the order (num_s, num_theta, num_phi).
-                'num_s' is the number of closed flux surfaces for vertex
-                locations in each toroidal plane. 'num_theta' is the number of
-                poloidal angles for vertex locations in each toroidal plane.
-                'num_phi' is the number of toroidal angles for planes of
-                vertices.
-            toroidal_extent (float) : extent of source mesh in toroidal
-                direction [deg].
+            cfs_grid (iterable of float): mesh grid points along closed flux
+                surface (CFS) dimension of flux space.
+            poloidal_grid (iterable of float): mesh grid points along poloidal
+                angle dimension of flux space [deg].
+            toroidal_grid (iterable of float): mesh grid points along toroidal
+                angle dimension of flux space [deg].
 
         Optional attributes:
             scale (float): a scaling factor between the units of VMEC and [cm]
@@ -292,8 +291,9 @@ def construct_source_mesh(self, mesh_size, toroidal_extent, **kwargs):
         """
         self.source_mesh = sm.SourceMesh(
             self._vmec_obj,
-            mesh_size,
-            toroidal_extent,
+            cfs_grid,
+            poloidal_grid,
+            toroidal_grid,
             logger=self._logger,
             **kwargs,
         )
@@ -646,7 +646,16 @@ def parastell():
 
     if args.source:
         source_mesh = all_data["source_mesh"]
-        stellarator.construct_source_mesh(**source_mesh)
+
+        mesh_size = source_mesh["mesh_size"]
+        toroidal_extent = source_mesh["toroidal_extent"]
+        cfs_grid = np.linspace(0.0, 1.0, mesh_size[0])
+        poloidal_grid = np.linspace(0.0, 360.0, num=mesh_size[1])
+        toroidal_grid = np.linspace(0.0, toroidal_extent, num=mesh_size[2])
+
+        stellarator.construct_source_mesh(
+            cfs_grid, poloidal_grid, toroidal_grid, **source_mesh
+        )
         stellarator.export_source_mesh(
             export_dir=args.export_dir,
             **(filter_kwargs(source_mesh, sm.export_allowed_kwargs)),
@@ -679,7 +688,7 @@ def parastell():
         nwl_required_keys = ["toroidal_angles", "poloidal_angles", "wall_s"]
 
         nwl_build = {}
-        for key in nwl_keys:
+        for key in nwl_required_keys:
             nwl_build[key] = invessel_build[key]
         nwl_build["radial_build"] = {}
 

parastell/source_mesh.py

@@ -142,41 +142,47 @@ def cfs_grid(self):
 
     @cfs_grid.setter
     def cfs_grid(self, array):
-        self._cfs_grid = array
-        if self._cfs_grid[0] != 0 or self._cfs_grid[-1] != 1:
+        if array[0] != 0 or array[-1] != 1:
             e = ValueError("CFS grid values must span the range [0, 1].")
             self._logger.error(e.args[0])
             raise e
+        # Don't include magnetic axis in list of s values
+        self._cfs_grid = array[1:]
 
     @property
     def poloidal_grid(self):
         return self._poloidal_grid
 
     @poloidal_grid.setter
     def poloidal_grid(self, array):
-        self._poloidal_grid = np.deg2rad(array)
-        if self._poloidal_grid[-1] - self._poloidal_grid[0] != 360.0:
+        if array[-1] - array[0] != 360.0:
             e = ValueError(
                 "Poloidal extent spanned by poloidal_grid must be exactly 360 "
                 "degrees."
             )
             self._logger.error(e.args[0])
             raise e
+        # Exclude repeated entry at end of closed loop
+        self._poloidal_grid = np.deg2rad(array[:-1])
 
     @property
     def toroidal_grid(self):
         return self._toroidal_grid
 
     @toroidal_grid.setter
     def toroidal_grid(self, array):
-        self._toroidal_grid = np.deg2rad(array)
-        if self._toroidal_grid[-1] - self._toroidal_grid[0] > 360.0:
+        if array[-1] - array[0] > 360.0:
             e = ValueError(
                 "Toroidal extent spanned by toroidal_grid cannot exceed 360 "
                 "degrees."
             )
             self._logger.error(e.args[0])
             raise e
+        # Exclude repeated entry at end if closed loop
+        if array[-1] - array[0] == 2 * np.pi:
+            self._toroidal_grid = np.deg2rad(array[:-1])
+        else:
+            self._toroidal_grid = np.deg2rad(array)
 
     @property
     def logger(self):
@@ -225,28 +231,20 @@ def create_vertices(self):
         """
         self._logger.info("Computing source mesh point cloud...")
 
-        phi_list = np.linspace(0, self._toroidal_extent, num=self.num_phi)
-        # don't include magnetic axis in list of s values
-        s_list = np.linspace(0.0, 1.0, num=self.num_s)[1:]
-        # don't include repeated entry at 0 == 2*pi
-        theta_list = np.linspace(0, 2 * np.pi, num=self.num_theta)[:-1]
-
-        # don't include repeated entry at 0 == 2*pi
-        if self._toroidal_extent == 2 * np.pi:
-            phi_list = phi_list[:-1]
-
-        self.verts_per_ring = theta_list.shape[0]
-        # add one vertex per plane for magenetic axis
-        self.verts_per_plane = s_list.shape[0] * self.verts_per_ring + 1
+        self.verts_per_ring = self._poloidal_grid.shape[0]
+        # add one vertex per plane for magnetic axis
+        self.verts_per_plane = (
+            self._cfs_grid.shape[0] * self.verts_per_ring + 1
+        )
+        num_verts = self._toroidal_grid.shape[0] * self.verts_per_plane
 
-        num_verts = phi_list.shape[0] * self.verts_per_plane
         self.coords = np.zeros((num_verts, 3))
         self.coords_s = np.zeros(num_verts)
 
         # Initialize vertex index
         vert_idx = 0
 
-        for phi in phi_list:
+        for phi in self._toroidal_grid:
             # vertex coordinates on magnetic axis
             self.coords[vert_idx, :] = (
                 np.array(self.vmec_obj.vmec2xyz(0, 0, phi)) * self.scale
@@ -256,8 +254,8 @@ def create_vertices(self):
             vert_idx += 1
 
             # vertex coordinate away from magnetic axis
-            for s in s_list:
-                for theta in theta_list:
+            for s in self._cfs_grid:
+                for theta in self._poloidal_grid:
                     self.coords[vert_idx, :] = (
                         np.array(self.vmec_obj.vmec2xyz(s, theta, phi))
                         * self.scale
@@ -354,7 +352,10 @@ def _get_vertex_id(self, vertex_idx):
         ma_offset = phi_idx * self.verts_per_plane
 
         # Wrap around if final plane and it is 2*pi
-        if self._toroidal_extent == 2 * np.pi and phi_idx == self.num_phi - 1:
+        if (
+            self._toroidal_grid[-1] == 2 * np.pi
+            and phi_idx == len(self._toroidal_grid) - 1
+        ):
             ma_offset = 0
 
         # Compute index offset from closed flux surface
@@ -363,7 +364,7 @@ def _get_vertex_id(self, vertex_idx):
         theta_offset = theta_idx
 
         # Wrap around if theta is 2*pi
-        if theta_idx == self.num_theta:
+        if theta_idx == len(self._poloidal_grid) + 1:
             theta_offset = 1
 
         id = ma_offset + s_offset + theta_offset
@@ -490,17 +491,17 @@ def create_mesh(self):
         self.mesh_set = self.mbc.create_meshset()
         self.mbc.add_entity(self.mesh_set, self.verts)
 
-        for phi_idx in range(self.num_phi - 1):
+        for phi_idx, _ in enumerate(self._toroidal_grid[:-1]):
             # Set alternation flag to true at beginning of each toroidal block
             self.alt_flag = True
             # Create tetrahedra for wedges at center of plasma
-            for theta_idx in range(1, self.num_theta):
+            for theta_idx, _ in enumerate(self._poloidal_grid, 1):
                 self._create_tets_from_wedge(theta_idx, phi_idx)
                 self.alt_flag = not self.alt_flag
 
             # Create tetrahedra for hexahedra beyond center of plasma
-            for s_idx in range(self.num_s - 2):
-                for theta_idx in range(1, self.num_theta):
+            for s_idx, _ in enumerate(self._cfs_grid[:-1]):
+                for theta_idx, _ in enumerate(self._poloidal_grid, 1):
                     self._create_tets_from_hex(s_idx, theta_idx, phi_idx)
                     self.alt_flag = not self.alt_flag
 
@@ -568,19 +569,27 @@ def generate_source_mesh():
 
     source_mesh_dict = all_data["source_mesh"]
 
+    mesh_size = source_mesh_dict["mesh_size"]
+    toroidal_extent = source_mesh_dict["toroidal_extent"]
+    cfs_grid = np.linspace(0.0, 1.0, mesh_size[0])
+    poloidal_grid = np.linspace(0.0, 360.0, num=mesh_size[1])
+    toroidal_grid = np.linspace(0.0, toroidal_extent, num=mesh_size[2])
+
     source_mesh = SourceMesh(
         vmec_obj,
-        source_mesh_dict["mesh_size"],
-        source_mesh_dict["toroidal_extent"],
-        logger=logger**source_mesh_dict,
+        cfs_grid,
+        poloidal_grid,
+        toroidal_grid,
+        logger=logger,
+        **source_mesh_dict,
     )
 
     source_mesh.create_vertices()
     source_mesh.create_mesh()
 
     source_mesh.export_mesh(
         export_dir=args.export_dir,
-        **(filter_kwargs(source_mesh_dict, ["filename"]))
+        **(filter_kwargs(source_mesh_dict, ["filename"])),
     )
 
 

tests/test_parastell.py

@@ -83,7 +83,7 @@ def test_parastell(stellarator):
     coils_file = Path("files_for_tests") / "coils.example"
     width = 40.0
     thickness = 50.0
-    toroidal_extent = 90.0
+    toroidal_extent = 15.0
     sample_mod = 6
 
     stellarator.construct_magnets(
@@ -103,10 +103,11 @@ def test_parastell(stellarator):
     assert Path(step_filename_exp).with_suffix(".step").exists()
     assert Path(mesh_filename_exp).with_suffix(".h5m").exists()
 
-    mesh_size = (4, 8, 4)
-    toroidal_extent = 15.0
+    cfs_grid = np.linspace(0.0, 1.0, num=6)
+    poloidal_grid = np.linspace(0.0, 360.0, num=41)
+    toroidal_grid = np.linspace(0.0, 15.0, num=9)
 
-    stellarator.construct_source_mesh(mesh_size, toroidal_extent)
+    stellarator.construct_source_mesh(cfs_grid, poloidal_grid, toroidal_grid)
 
     filename_exp = "source_mesh"