From 19f394e3eec5acef8c119b274a7f2407fef1ba77 Mon Sep 17 00:00:00 2001 From: Connor Moreno Date: Fri, 11 Oct 2024 16:44:50 -0500 Subject: [PATCH] Implement custom source mesh grids --- Examples/custom_source_example.py | 10 ++-- Examples/nwl_geom_example.py | 8 +-- Examples/parastell_example.py | 7 +-- Examples/radial_distance_example.py | 7 +-- parastell/parastell.py | 37 ++++++++------ parastell/source_mesh.py | 75 ++++++++++++++++------------- tests/test_parastell.py | 9 ++-- tests/test_source_mesh.py | 29 +++++------ 8 files changed, 104 insertions(+), 78 deletions(-) diff --git a/Examples/custom_source_example.py b/Examples/custom_source_example.py index 53113cfa..7e43d075 100644 --- a/Examples/custom_source_example.py +++ b/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, ) diff --git a/Examples/nwl_geom_example.py b/Examples/nwl_geom_example.py index 0f9b50ba..e9e35da8 100644 --- a/Examples/nwl_geom_example.py +++ b/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) diff --git a/Examples/parastell_example.py b/Examples/parastell_example.py index ba32b6ef..d6ba7aba 100644 --- a/Examples/parastell_example.py +++ b/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) diff --git a/Examples/radial_distance_example.py b/Examples/radial_distance_example.py index 22faf13f..2fe1ad3d 100644 --- a/Examples/radial_distance_example.py +++ b/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) diff --git a/parastell/parastell.py b/parastell/parastell.py index e73628d2..538b3831 100644 --- a/parastell/parastell.py +++ b/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"] = {} diff --git a/parastell/source_mesh.py b/parastell/source_mesh.py index fd28ec82..92daf491 100644 --- a/parastell/source_mesh.py +++ b/parastell/source_mesh.py @@ -142,11 +142,12 @@ 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): @@ -154,14 +155,15 @@ def poloidal_grid(self): @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): @@ -169,14 +171,18 @@ def toroidal_grid(self): @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,11 +569,19 @@ 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() @@ -580,7 +589,7 @@ def generate_source_mesh(): source_mesh.export_mesh( export_dir=args.export_dir, - **(filter_kwargs(source_mesh_dict, ["filename"])) + **(filter_kwargs(source_mesh_dict, ["filename"])), ) diff --git a/tests/test_parastell.py b/tests/test_parastell.py index 70a6f401..62b71642 100644 --- a/tests/test_parastell.py +++ b/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" diff --git a/tests/test_source_mesh.py b/tests/test_source_mesh.py index f0a2fbb1..ee3b8367 100644 --- a/tests/test_source_mesh.py +++ b/tests/test_source_mesh.py @@ -22,30 +22,31 @@ def source_mesh(): vmec_obj = read_vmec.VMECData(vmec_file) - # Set mesh size to minimum that maintains element aspect ratios that do not + # Set mesh grids to minimum that maintains element aspect ratios that do not # result in negative volumes - mesh_size = (6, 41, 51) - toroidal_extent = 90.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) - source_mesh_obj = sm.SourceMesh(vmec_obj, mesh_size, toroidal_extent) + source_mesh_obj = sm.SourceMesh( + vmec_obj, cfs_grid, poloidal_grid, toroidal_grid + ) return source_mesh_obj def test_mesh_basics(source_mesh): - num_s_exp = 6 - num_theta_exp = 41 - num_phi_exp = 51 - tor_ext_exp = 90.0 + cfs_grid_exp = np.linspace(0.0, 1.0, num=6)[1:] + poloidal_grid_exp = np.linspace(0.0, 2 * np.pi, num=41)[:-1] + toroidal_grid_exp = np.linspace(0.0, np.deg2rad(15.0), num=9) 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.toroidal_extent == np.deg2rad(tor_ext_exp) + assert np.allclose(source_mesh.cfs_grid, cfs_grid_exp) + assert np.allclose(source_mesh.poloidal_grid, poloidal_grid_exp) + assert np.allclose(source_mesh.toroidal_grid, toroidal_grid_exp) assert source_mesh.scale == scale_exp remove_files() @@ -55,7 +56,7 @@ def test_vertices(source_mesh): num_s = 6 num_theta = 41 - num_phi = 51 + num_phi = 9 num_verts_exp = num_phi * ((num_s - 1) * (num_theta - 1) + 1) @@ -74,7 +75,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