diff --git a/parastell/source_mesh.py b/parastell/source_mesh.py index ae10f06..0d7d932 100644 --- a/parastell/source_mesh.py +++ b/parastell/source_mesh.py @@ -66,30 +66,31 @@ class SourceMesh(object): neutrons in a magnetically confined plasma described by a VMEC plasma equilibrium. - The mesh will be defined on a regular grid in the plasma coordinates of s, - theta, phi. Mesh vertices will be defined on circular grid at each toroidal - plane, and connected between toroidal planes. This results in wedge elements - along the magnetic axis and hexagonal elements throughout the remainder of - the mesh. Each of these elements will be subdivided into tetrahedra (4 for - the wedges and 5 for the hexahedra) to result in a mesh that is simpler to - use. - - Each tetrahedron will be tagged with the volumetric neutron source intensity - in n/cm3/s, using on a finite-element based quadrature of the source - intensity evaluated at each vertex. + The mesh will be defined on a regular grid in the flux coordinates of + (CFS value, poloidal angle, toroidal angle). Mesh vertices will be defined + on circular grid at each toroidal plane, and connected between toroidal + planes. This results in wedge elements along the magnetic axis and + hexagonal elements throughout the remainder of the mesh. Each of these + elements will be subdivided into tetrahedra (3 for the wedges and 5 for the + hexahedra) to result in a mesh that is simpler to use. + + Each tetrahedron will be tagged with the volumetric neutron source + intensity in n/cm3/s, using on a finite-element based quadrature of the + source intensity evaluated at each vertex. Arguments: vmec_obj (object): plasma equilibrium VMEC object as defined by the PyStell-UW VMEC reader. Must have a method - 'vmec2xyz(s, theta, phi)' that returns an (x,y,z) coordinate for - any closed flux surface label, s, poloidal angle, theta, and - toroidal angle, phi. + 'vmec2xyz(cfs, poloidal_ang, toroidal_ang)' that returns + a (x,y,z) coordinate for any closed flux surface value, cfs, + poloidal angle, poloidal_ang, and toroidal angle, toroidal_ang. 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. + flux-coordinate space, in the order (num_cfs_pts, num_poloidal_pts, + num_toroidal_pts). 'num_cfs_pts' is the number of closed flux + surfaces for vertex locations in each toroidal plane. + 'num_poloidal_pts' is the number of poloidal angles for vertex + locations in each toroidal plane. 'num_toroidal_pts' is the number + of toroidal angles for planes of vertices. toroidal_extent (float): extent of source mesh in toroidal direction [deg]. logger (object): logger object (optional, defaults to None). If no @@ -112,9 +113,9 @@ def __init__( self.logger = logger self.vmec_obj = vmec_obj - self.num_s = mesh_size[0] - self.num_theta = mesh_size[1] - self.num_phi = mesh_size[2] + self.num_cfs_pts = mesh_size[0] + self.num_poloidal_pts = mesh_size[1] + self.num_toroidal_pts = mesh_size[2] self.toroidal_extent = toroidal_extent self.scale = m2cm @@ -133,18 +134,56 @@ def __init__( self._create_mbc() + @property + def num_poloidal_pts(self): + return self._num_poloidal_pts + + @num_poloidal_pts.setter + def num_poloidal_pts(self, value): + if value % 2 != 1: + e = AttributeError( + "To ensure that tetrahedral faces are coincident at the end of " + "the closed poloidal loop, the number of poloidal intervals " + "must be even. To ensure this, the number of poloidal grid " + "points must be odd." + ) + self._logger.error(e.args[0]) + raise e + self._num_poloidal_pts = value + + @property + def num_toroidal_pts(self): + return self._num_toroidal_pts + + @num_toroidal_pts.setter + def num_toroidal_pts(self, value): + self._num_toroidal_pts = value + @property def toroidal_extent(self): return self._toroidal_extent @toroidal_extent.setter def toroidal_extent(self, angle): - self._toroidal_extent = np.deg2rad(angle) - if self._toroidal_extent > 360.0: - e = ValueError("Toroidal extent cannot exceed 360.0 degrees.") + angle = np.deg2rad(angle) + + if angle > 2 * np.pi: + e = AttributeError("Toroidal extent cannot exceed 360.0 degrees.") self._logger.error(e.args[0]) raise e + if angle == 2 * np.pi and self._num_toroidal_pts % 2 != 1: + e = AttributeError( + "To ensure that tetrahedral faces are coincident at the end of " + "the closed toroidal loop, the number of toroidal intervals " + "must be even. To ensure this, the number of toroidal grid " + "points must be odd." + ) + self._logger.error(e.args[0]) + raise e + + self._toroidal_extent = angle + @property def logger(self): return self._logger @@ -192,44 +231,54 @@ 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] + # Exclude entry at magnetic axis + cfs_grid_pts = np.linspace(0.0, 1.0, num=self.num_cfs_pts)[1:] + + # Exclude repeated entry at 0 == 2*pi + poloidal_grid_pts = np.linspace( + 0, 2 * np.pi, num=self._num_poloidal_pts + )[:-1] - # don't include repeated entry at 0 == 2*pi + toroidal_grid_pts = np.linspace( + 0, self._toroidal_extent, num=self._num_toroidal_pts + ) + # Conditionally exclude repeated entry at 0 == 2*pi if self._toroidal_extent == 2 * np.pi: - phi_list = phi_list[:-1] + toroidal_grid_pts = toroidal_grid_pts[:-1] - self.verts_per_ring = theta_list.shape[0] + self.verts_per_ring = poloidal_grid_pts.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_plane = cfs_grid_pts.shape[0] * self.verts_per_ring + 1 - num_verts = phi_list.shape[0] * self.verts_per_plane + num_verts = toroidal_grid_pts.shape[0] * self.verts_per_plane self.coords = np.zeros((num_verts, 3)) - self.coords_s = np.zeros(num_verts) + self.coords_cfs = np.zeros(num_verts) # Initialize vertex index vert_idx = 0 - for phi in phi_list: + for toroidal_ang in toroidal_grid_pts: # vertex coordinates on magnetic axis self.coords[vert_idx, :] = ( - np.array(self.vmec_obj.vmec2xyz(0, 0, phi)) * self.scale + np.array(self.vmec_obj.vmec2xyz(0, 0, toroidal_ang)) + * self.scale ) - self.coords_s[vert_idx] = 0 + self.coords_cfs[vert_idx] = 0 vert_idx += 1 # vertex coordinate away from magnetic axis - for s in s_list: - for theta in theta_list: + for cfs in cfs_grid_pts: + for poloidal_ang in poloidal_grid_pts: self.coords[vert_idx, :] = ( - np.array(self.vmec_obj.vmec2xyz(s, theta, phi)) + np.array( + self.vmec_obj.vmec2xyz( + cfs, poloidal_ang, toroidal_ang + ) + ) * self.scale ) - self.coords_s[vert_idx] = s + self.coords_cfs[vert_idx] = cfs vert_idx += 1 @@ -246,13 +295,12 @@ def _source_strength(self, tet_ids): Returns: ss (float): integrated source strength for tetrahedron. """ - # Initialize list of vertex coordinates for each tetrahedron vertex tet_coords = [self.coords[id] for id in tet_ids] # Initialize list of source strengths for each tetrahedron vertex vertex_strengths = [ - self.reaction_rate(*self.plasma_conditions(self.coords_s[id])) + self.reaction_rate(*self.plasma_conditions(self.coords_cfs[id])) for id in tet_ids ] @@ -289,7 +337,6 @@ def _create_tet(self, tet_ids): Arguments: tet_ids (list of int): tetrahedron vertex indices. """ - tet_verts = [self.verts[int(id)] for id in tet_ids] tet = self.mbc.create_element(types.MBTET, tet_verts) self.mbc.add_entity(self.mesh_set, tet) @@ -315,36 +362,42 @@ def _get_vertex_id(self, vertex_idx): Returns: id (int): vertex index in row-major order as stored by MOAB """ + cfs_idx, poloidal_idx, toroidal_idx = vertex_idx - s_idx, theta_idx, phi_idx = vertex_idx - - ma_offset = phi_idx * self.verts_per_plane + ma_offset = toroidal_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_extent == 2 * np.pi + and toroidal_idx == self._num_toroidal_pts - 1 + ): ma_offset = 0 - # Compute index offset from closed flux surface - s_offset = s_idx * self.verts_per_ring - - theta_offset = theta_idx + # Compute index offset from closed flux surface, taking single vertex + # at magnetic axis into account + if cfs_idx == 0: + cfs_offset = cfs_idx + else: + cfs_offset = (cfs_idx - 1) * self.verts_per_ring + 1 - # Wrap around if theta is 2*pi - if theta_idx == self.num_theta: - theta_offset = 1 + poloidal_offset = poloidal_idx + # Wrap around if poloidal angle is 2*pi + if poloidal_idx == self._num_poloidal_pts - 1: + poloidal_offset = 0 - id = ma_offset + s_offset + theta_offset + id = ma_offset + cfs_offset + poloidal_offset return id - def _create_tets_from_hex(self, s_idx, theta_idx, phi_idx): + def _create_tets_from_hex(self, cfs_idx, poloidal_idx, toroidal_idx): """Creates five tetrahedra from defined hexahedron. (Internal function not intended to be called externally) Arguments: - idx_list (list of int): list of hexahedron vertex indices. + cfs_idx (int): index defining location along CFS axis. + poloidal_idx (int): index defining location along poloidal angle axis. + toroidal_idx (int): index defining location along toroidal angle axis. """ - # relative offsets of vertices in a 3-D index space hex_vertex_stencil = np.array( [ @@ -361,7 +414,8 @@ def _create_tets_from_hex(self, s_idx, theta_idx, phi_idx): # Ids of hex vertices applying offset stencil to current point hex_idx_data = ( - np.array([s_idx, theta_idx, phi_idx]) + hex_vertex_stencil + np.array([cfs_idx, poloidal_idx, toroidal_idx]) + + hex_vertex_stencil ) idx_list = [ @@ -375,39 +429,52 @@ def _create_tets_from_hex(self, s_idx, theta_idx, phi_idx): # first, ordered clockwise relative to the thumb, followed by the # remaining vertex at the end of the thumb. # See Moreno, Bader, Wilson 2024 for hexahedron splitting - hex_canon_ids = [ - [idx_list[0], idx_list[2], idx_list[1], idx_list[5]], - [idx_list[0], idx_list[3], idx_list[2], idx_list[7]], - [idx_list[0], idx_list[7], idx_list[5], idx_list[4]], - [idx_list[7], idx_list[2], idx_list[5], idx_list[6]], - [idx_list[0], idx_list[2], idx_list[5], idx_list[7]], + canonical_ordering_schemes = [ + [ + [idx_list[0], idx_list[3], idx_list[1], idx_list[4]], + [idx_list[1], idx_list[3], idx_list[2], idx_list[6]], + [idx_list[1], idx_list[4], idx_list[6], idx_list[5]], + [idx_list[3], idx_list[6], idx_list[4], idx_list[7]], + [idx_list[1], idx_list[3], idx_list[6], idx_list[4]], + ], + [ + [idx_list[0], idx_list[2], idx_list[1], idx_list[5]], + [idx_list[0], idx_list[3], idx_list[2], idx_list[7]], + [idx_list[0], idx_list[7], idx_list[5], idx_list[4]], + [idx_list[7], idx_list[2], idx_list[5], idx_list[6]], + [idx_list[0], idx_list[2], idx_list[5], idx_list[7]], + ], ] - for vertex_ids in hex_canon_ids: + # Alternate canonical ordering schemes defining hexahedron splitting to + # avoid gaps and overlaps between non-planar hexahedron faces + scheme_idx = (cfs_idx + poloidal_idx + toroidal_idx) % 2 + + for vertex_ids in canonical_ordering_schemes[scheme_idx]: self._create_tet(vertex_ids) - def _create_tets_from_wedge(self, theta_idx, phi_idx): + def _create_tets_from_wedge(self, poloidal_idx, toroidal_idx): """Creates three tetrahedra from defined wedge. (Internal function not intended to be called externally) Arguments: - idx_list (list of int): list of wedge vertex indices. + poloidal_idx (int): index defining location along poloidal angle axis. + toroidal_idx (int): index defining location along toroidal angle axis. """ - # relative offsets of wedge vertices in a 3-D index space wedge_vertex_stencil = np.array( [ [0, 0, 0], - [0, theta_idx, 0], - [0, theta_idx + 1, 0], + [1, poloidal_idx, 0], + [1, poloidal_idx + 1, 0], [0, 0, 1], - [0, theta_idx, 1], - [0, theta_idx + 1, 1], + [1, poloidal_idx, 1], + [1, poloidal_idx + 1, 1], ] ) # Ids of wedge vertices applying offset stencil to current point - wedge_idx_data = np.array([0, 0, phi_idx]) + wedge_vertex_stencil + wedge_idx_data = np.array([0, 0, toroidal_idx]) + wedge_vertex_stencil idx_list = [ self._get_vertex_id(vertex_idx) for vertex_idx in wedge_idx_data @@ -420,13 +487,24 @@ def _create_tets_from_wedge(self, theta_idx, phi_idx): # first, ordered clockwise relative to the thumb, followed by the # remaining vertex at the end of the thumb. # See Moreno, Bader, Wilson 2024 for wedge splitting - wedge_canon_ids = [ - [idx_list[0], idx_list[2], idx_list[1], idx_list[3]], - [idx_list[3], idx_list[2], idx_list[4], idx_list[5]], - [idx_list[3], idx_list[2], idx_list[1], idx_list[4]], + canonical_ordering_schemes = [ + [ + [idx_list[0], idx_list[2], idx_list[1], idx_list[3]], + [idx_list[1], idx_list[3], idx_list[5], idx_list[4]], + [idx_list[1], idx_list[3], idx_list[2], idx_list[5]], + ], + [ + [idx_list[0], idx_list[2], idx_list[1], idx_list[3]], + [idx_list[3], idx_list[2], idx_list[4], idx_list[5]], + [idx_list[3], idx_list[2], idx_list[1], idx_list[4]], + ], ] - for vertex_ids in wedge_canon_ids: + # Alternate canonical ordering schemes defining wedge splitting to + # avoid gaps and overlaps between non-planar wedge faces + scheme_idx = (poloidal_idx + toroidal_idx) % 2 + + for vertex_ids in canonical_ordering_schemes[scheme_idx]: self._create_tet(vertex_ids) def create_mesh(self): @@ -436,15 +514,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 toroidal_idx in range(self._num_toroidal_pts - 1): # Create tetrahedra for wedges at center of plasma - for theta_idx in range(1, self.num_theta): - self._create_tets_from_wedge(theta_idx, phi_idx) + for poloidal_idx in range(self._num_poloidal_pts - 1): + self._create_tets_from_wedge(poloidal_idx, toroidal_idx) # 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): - self._create_tets_from_hex(s_idx, theta_idx, phi_idx) + for cfs_idx in range(1, self.num_cfs_pts - 1): + for poloidal_idx in range(self._num_poloidal_pts - 1): + self._create_tets_from_hex( + cfs_idx, poloidal_idx, toroidal_idx + ) def export_mesh(self, filename="source_mesh", export_dir=""): """Use PyMOAB interface to write source mesh with source strengths diff --git a/tests/test_parastell.py b/tests/test_parastell.py index 70a6f40..b05095d 100644 --- a/tests/test_parastell.py +++ b/tests/test_parastell.py @@ -103,7 +103,7 @@ 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) + mesh_size = (6, 41, 9) toroidal_extent = 15.0 stellarator.construct_source_mesh(mesh_size, toroidal_extent) diff --git a/tests/test_source_mesh.py b/tests/test_source_mesh.py index f0a2fbb..2a64b7f 100644 --- a/tests/test_source_mesh.py +++ b/tests/test_source_mesh.py @@ -24,8 +24,8 @@ def source_mesh(): # Set mesh size to minimum that maintains element aspect ratios that do not # result in negative volumes - mesh_size = (6, 41, 51) - toroidal_extent = 90.0 + mesh_size = (6, 41, 9) + toroidal_extent = 15.0 source_mesh_obj = sm.SourceMesh(vmec_obj, mesh_size, toroidal_extent) @@ -34,17 +34,17 @@ def source_mesh(): def test_mesh_basics(source_mesh): - num_s_exp = 6 - num_theta_exp = 41 - num_phi_exp = 51 - tor_ext_exp = 90.0 + num_cfs_exp = 6 + num_poloidal_pts_exp = 41 + num_toroidal_pts_exp = 9 + tor_ext_exp = 15.0 scale_exp = 100 remove_files() - assert source_mesh.num_s == num_s_exp - assert source_mesh.num_theta == num_theta_exp - assert source_mesh.num_phi == num_phi_exp + assert source_mesh.num_cfs_pts == num_cfs_exp + assert source_mesh.num_poloidal_pts == num_poloidal_pts_exp + assert source_mesh.num_toroidal_pts == num_toroidal_pts_exp assert source_mesh.toroidal_extent == np.deg2rad(tor_ext_exp) assert source_mesh.scale == scale_exp @@ -53,18 +53,20 @@ def test_mesh_basics(source_mesh): def test_vertices(source_mesh): - num_s = 6 - num_theta = 41 - num_phi = 51 + num_cfs_exp = 6 + num_poloidal_pts_exp = 41 + num_toroidal_pts_exp = 9 - num_verts_exp = num_phi * ((num_s - 1) * (num_theta - 1) + 1) + num_verts_exp = num_toroidal_pts_exp * ( + (num_cfs_exp - 1) * (num_poloidal_pts_exp - 1) + 1 + ) remove_files() source_mesh.create_vertices() assert source_mesh.coords.shape == (num_verts_exp, 3) - assert source_mesh.coords_s.shape == (num_verts_exp,) + assert source_mesh.coords_cfs.shape == (num_verts_exp,) assert len(source_mesh.verts) == num_verts_exp remove_files() @@ -74,7 +76,7 @@ def test_mesh_generation(source_mesh): num_s = 6 num_theta = 41 - num_phi = 51 + num_phi = 9 tets_per_wedge = 3 tets_per_hex = 5