diff --git a/src/sourcemesh.py b/src/sourcemesh.py new file mode 100644 index 00000000..c78cce75 --- /dev/null +++ b/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() \ No newline at end of file diff --git a/tests/test_sourcemesh.py b/tests/test_sourcemesh.py new file mode 100644 index 00000000..26d05678 --- /dev/null +++ b/tests/test_sourcemesh.py @@ -0,0 +1,43 @@ +import src.sourcemesh as sm +import read_vmec as rv +import numpy as np + +vmec_file = 'files_for_tests/wout_vmec.nc' + +def test_meshbasics(): + + vmec = rv.vmec_data(vmec_file) + + num_s_exp = 4 + num_theta_exp = 8 + num_phi_exp = 4 + tor_ext_exp = 90.0 + + source_mesh = sm.SourceMesh(vmec, num_s_exp, num_theta_exp, num_phi_exp, + tor_ext_exp) + + 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.tor_ext == np.deg2rad(tor_ext_exp) + +def test_vertices(): + + vmec = rv.vmec_data(vmec_file) + + num_s_exp = 4 + num_theta_exp = 8 + num_phi_exp = 4 + tor_ext_exp = 90.0 + + num_verts_exp = num_phi_exp * ( (num_s_exp - 1) * (num_theta_exp - 1) + 1) + + source_mesh = sm.SourceMesh(vmec, num_s_exp, num_theta_exp, num_phi_exp, + tor_ext_exp) + + source_mesh.create_vertices() + + assert source_mesh.coord.shape == (num_verts_exp, 3) + assert source_mesh.coord_s.shape == (num_verts_exp,) + + assert source_mesh.verts.shape == (num_verts_exp,)