Create object-oriented version of source mesh generation

src/sourcemesh.py

@@ -0,0 +1,381 @@
+import argparse
+import yaml
+import read_vmec
+import numpy as np
+from pymoab import core, types
+from pathlib import Path
+
+m2cm = 100
+
+def rxn_rate(s):
+    """Calculates fusion reaction rate in plasma.
+
+    Arguments:
+        s (float): closed magnetic flux surface index in range of 0 (magnetic
+            axis) to 1 (plasma edge).
+
+    Returns:
+        rr (float): fusion reaction rate (1/cm^3/s). Equates to neutron source
+            density.
+    """
+    # Define m^3 to cm^3 constant
+    m3tocm3 = 1e6
+
+    if s == 1:
+        rr = 0
+    else:
+        # Temperature
+        T = 11.5 * (1 - s)
+        # Ion density
+        n = 4.8e20 * (1 - s**5)
+        # Reaction rate in 1/m^3/s
+        rr = 3.68e-18 * (n**2) / 4 * T**(-2/3) * np.exp(-19.94 * T**(-1/3))
+
+    return rr/m3tocm3
+
+class SourceMesh(object):
+    """
+    Generates a source mesh that describes the relative source intensity of
+    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.
+
+    Parameters:
+        vmec (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 plasma parameter, s, poloidal angle, theta,
+                and toroidal angle, phi.
+        num_s (int) : number of rings of vertices in each toroidal plane
+        num_theta (int) : number of poloidal angles for vertex locations
+        num_phi (int) : number of toroidal angles for planes of vertices
+        tor_ext (float) : extend of source mesh in toroidal direction in degrees
+                  assumed in degrees
+    """
+    def __init__(
+            self,
+            vmec, 
+            num_s,
+            num_theta,
+            num_phi,
+            tor_ext
+    ):
+
+        self.vmec = vmec
+        self.num_s = num_s
+        self.num_theta = num_theta
+        self.num_phi = num_phi
+        self.tor_ext = np.deg2rad(tor_ext)
+        self.strengths = []
+
+        self._create_mbc()
+
+
+    def _create_mbc(self):
+        """
+        Creates PyMOAB core instance with source strength tag.
+        (Internal function not intended to be called externally.)
+
+        Returns:
+            mbc (object): PyMOAB core instance.
+            tag_handle (TagHandle): PyMOAB source strength tag.
+        """
+        self.mbc = core.Core()
+
+        tag_type = types.MB_TYPE_DOUBLE
+        tag_size = 1
+        storage_type = types.MB_TAG_DENSE
+        tag_name = "SourceStrength"
+        self.tag_handle = self.mbc.tag_get_handle(
+            tag_name, tag_size,tag_type, storage_type,
+            create_if_missing = True
+        )
+
+    def create_vertices(self):
+        """
+        Creates mesh vertices and adds them to PyMOAB core.
+
+        The grid of mesh vertices is generated from the user input
+        defining the number of meshes in each of the plasma
+        coordinate directions. Care is taken to manage the
+        mesh at the 0 == 2 * pi wrap so that everything
+        is closed and consistent.
+        """
+        phi_list = np.linspace(0, self.tor_ext, 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.tor_ext == 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  
+
+        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:
+            # vertex coordinates on magnetic axis
+            self.coords[vert_idx,:] = np.array(self.vmec.vmec2xyz(0, 0, phi)) * m2cm
+            self.coords_s[vert_idx] = 0
+
+            vert_idx += 1
+
+            # vertex coordinate away from magnetic axis
+            for s in s_list:
+                for theta in theta_list:
+                    self.coords[vert_idx,:] = np.array(self.vmec.vmec2xyz(s, theta, phi)) * m2cm
+                    self.coords_s[vert_idx] = s
+
+                    vert_idx += 1
+
+        self.verts = self.mbc.create_vertices(self.coords)
+
+
+    def _source_strength(self, tet_ids):
+        """
+        Computes neutron source strength for a tetrahedron using five-node
+        Gaussian quadrature.
+        (internal function not intended to be called externally.)
+
+        Arguments:
+            ids (list of int): tetrahedron vertex indices.
+
+        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 = [ rxn_rate(self.coords_s[id]) for id in tet_ids ]
+
+        # Define barycentric coordinates for integration points
+        bary_coords = np.array([
+            [0.25, 0.25, 0.25, 0.25],
+            [0.5, 1/6, 1/6, 1/6],
+            [1/6, 0.5, 1/6, 1/6],
+            [1/6, 1/6, 0.5, 1/6],
+            [1/6, 1/6, 1/6, 0.5]
+        ])
+
+        # Define weights for integration points
+        int_w = np.array([-0.8, 0.45, 0.45, 0.45, 0.45])
+
+        # Interpolate source strength at integration points
+        ss_int_pts = np.dot(bary_coords, vertex_strengths)
+
+        # Compute edge vectors between tetrahedron vertices
+        edge_vectors = np.subtract(tet_coords[:3], tet_coords[3]).T
+
+        tet_vol = np.abs(np.linalg.det(edge_vectors))/6
+
+        ss = tet_vol * np.dot(int_w, ss_int_pts)
+
+        return ss
+
+    def _create_tet(self, tet_ids):
+        """
+        Creates tetrahedron and adds to pyMOAB core.
+        (internal function not intended to be called externally.)
+
+        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)
+
+        # Compute source strength for tetrahedron
+        ss = self._source_strength(tet_ids)
+        self.strengths.append(ss)
+
+        # Tag tetrahedra with source strength data
+        self.mbc.tag_set_data(self.tag_handle, tet, [ss])
+
+    def _get_vertex_id(self, vertex_idx):
+        """Computes vertex index in row-major order as stored by MOAB from
+        three-dimensional n x 3 matrix indices.
+        (internal function not intended to be called externally.)
+
+        Arguments:
+            vert_idx (list of int): list of vertex 
+                [toroidal angle index, flux surface index, poloidal angle index]
+
+        Returns:
+            id (int): vertex index in row-major order as stored by MOAB
+        """
+
+        s_idx, theta_idx, phi_idx = vertex_idx
+
+        ma_offset = phi_idx * self.verts_per_plane
+
+        # Wrap around if final plane and it is 2*pi
+        if self.tor_ext == 2*np.pi and phi_idx == self.num_phi - 1:
+            ma_offset = 0
+
+        # Compute index offset from closed flux surface
+        s_offset = s_idx * self.verts_per_ring
+
+        theta_offset = theta_idx
+
+        # Wrap around if theta is 2*pi
+        if theta_idx == self.num_theta:
+            theta_offset = 1
+
+        id = ma_offset + s_offset + theta_offset
+
+        return id
+
+    def _create_tets_from_hex(self, s_idx, theta_idx, phi_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.
+        """
+
+        # relative offsets of vertices in a 3-D index space
+        hex_vertex_stencil = np.array([
+            [0, 1, 0],
+            [0, 0, 0],
+            [0, 1, 1],
+            [0, 0, 1],
+            [1, 1, 0],
+            [1, 0, 0],
+            [1, 1, 1],
+            [1, 0, 1] ])
+
+        # Ids of hex vertices applying offset stencil to current point
+        hex_idx_data = np.array([s_idx, theta_idx, phi_idx]) + hex_vertex_stencil
+
+        idx_list = [ self._get_vertex_id(vertex_idx) for vertex_idx in hex_idx_data ]
+
+        # Define MOAB canonical ordering of hexahedron vertex indices
+        hex_canon_ids = [
+            [idx_list[0], idx_list[3], idx_list[1], idx_list[4]],
+            [idx_list[7], idx_list[4], idx_list[6], idx_list[3]],
+            [idx_list[2], idx_list[1], idx_list[3], idx_list[6]],
+            [idx_list[5], idx_list[6], idx_list[4], idx_list[1]],
+            [idx_list[3], idx_list[1], idx_list[4], idx_list[6]]
+        ]
+
+        for vertex_ids in hex_canon_ids:
+            self._create_tet(vertex_ids)
+
+    def _create_tets_from_wedge(self, theta_idx, phi_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.
+        """
+
+        # 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],
+            [0, 0            ,  1],
+            [0, theta_idx    ,  1],
+            [0, theta_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
+
+        idx_list = [ self._get_vertex_id(vertex_idx) for vertex_idx in wedge_idx_data ]
+
+        # Define MOAB canonical ordering of wedge vertex indices
+        wedge_canon_ids = [
+            [idx_list[1], idx_list[2], idx_list[4], idx_list[0]],
+            [idx_list[5], idx_list[4], idx_list[2], idx_list[3]],
+            [idx_list[0], idx_list[2], idx_list[4], idx_list[3]]
+        ]
+
+        for vertex_ids in wedge_canon_ids:
+            self._create_tet(vertex_ids)
+
+    def create_mesh(self):
+        """Creates volumetric source mesh in real space.
+        """
+
+        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):
+            # 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)
+
+            # 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)
+
+    def write(self, filename):
+        """Use pyMOAB interface to write source mesh with 
+        source strengths tagged.
+        """
+
+        self.mbc.write_file(str(filename))
+
+def parse_args():
+    """
+    Parser for running as a script
+    """
+    parser = argparse.ArgumentParser(prog='sourcemesh')
+
+    parser.add_argument('filename', help='YAML file defining this case')
+
+    return parser.parse_args()
+
+def read_yaml_src(filename):
+    """
+    Read YAML file describing this case and extract data relevant for source
+    definition.
+    """
+    with open(filename) as yaml_file:
+        all_data =  yaml.safe_load(yaml_file)
+
+    # extract data to define source mesh
+    return all_data['plasma_eq'], all_data['source']            
+
+def generate_source_mesh():
+    """
+    Main method when run as a command line script.
+    """
+    args = parse_args()
+
+    vmec_file, src_data = read_yaml_src(args.filename)
+
+    vmec = read_vmec.vmec_data(vmec_file)
+
+    source_mesh = SourceMesh(vmec, src_data['num_s'], src_data['num_theta'], src_data['num_phi'], src_data['tor_ext'])
+    source_mesh.create_vertices()
+    source_mesh.create_mesh()
+
+    source_mesh.write(Path(src_data['export_dir']) / Path(src_data['mesh_file']))
+
+if __name__ == "__main__":
+    generate_source_mesh()