Merge pull request #4 from Edgar-21/integrate_model_merging

Integrate model merging

Examples/parastell_example.py

@@ -1,5 +1,5 @@
 import numpy as np
-import openmc
+
 import parastell.parastell as ps
 
 
@@ -20,10 +20,7 @@
 uniform_unit_thickness = np.ones((len(toroidal_angles), len(poloidal_angles)))
 
 radial_build_dict = {
-    "first_wall": {
-        "thickness_matrix": uniform_unit_thickness * 5,
-        "mat_tag": "iron",
-    },
+    "first_wall": {"thickness_matrix": uniform_unit_thickness * 5},
     "breeder": {
         "thickness_matrix": (
             [
@@ -37,40 +34,51 @@
                 [75.0, 75.0, 75.0, 75.0, 25.0, 25.0, 75.0, 75.0, 75.0],
                 [75.0, 75.0, 75.0, 25.0, 25.0, 25.0, 75.0, 75.0, 75.0],
             ]
-        ),
-        "mat_tag": "iron",
-    },
-    "back_wall": {
-        "thickness_matrix": uniform_unit_thickness * 5,
-        "mat_tag": "iron",
-    },
-    "shield": {
-        "thickness_matrix": uniform_unit_thickness * 50,
-        "mat_tag": "iron",
+        )
     },
+    "back_wall": {"thickness_matrix": uniform_unit_thickness * 5},
+    "shield": {"thickness_matrix": uniform_unit_thickness * 50},
     "vacuum_vessel": {
         "thickness_matrix": uniform_unit_thickness * 10,
-        "mat_tag": "tungsten",
+        "mat_tag": "vac_vessel",
     },
 }
 # Construct in-vessel components
 stellarator.construct_invessel_build(
-    toroidal_angles,
-    poloidal_angles,
-    wall_s,
-    radial_build_dict,
-    use_pydagmc=True,
+    toroidal_angles, poloidal_angles, wall_s, radial_build_dict
+)
+# Export in-vessel component files
+stellarator.export_invessel_build(
+    export_cad_to_dagmc=False, export_dir=export_dir
 )
 
-for surf in stellarator.invessel_build.dag_model.surfaces:
-    print(surf)
-    print(surf.surf_sense)
+# Define build parameters for magnet coils
+coils_file = "coils.example"
+width = 40.0
+thickness = 50.0
+toroidal_extent = 90.0
+# Construct magnets
+stellarator.construct_magnets(
+    coils_file, width, thickness, toroidal_extent, sample_mod=6
+)
+# Export magnet files
+stellarator.export_magnets(
+    step_filename="magnets",
+    export_mesh=True,
+    mesh_filename="magnet_mesh",
+    export_dir=export_dir,
+)
 
-print(stellarator.invessel_build.dag_model.volumes)
+# Define source mesh parameters
+mesh_size = (11, 81, 61)
+toroidal_extent = 90.0
+# Construct source
+stellarator.construct_source_mesh(mesh_size, toroidal_extent)
+# Export source file
+stellarator.export_source_mesh(filename="source_mesh", export_dir=export_dir)
 
-print(stellarator.invessel_build.dag_model.groups)
+# Build Cubit model of Parastell Components
+stellarator.build_cubit_model(skip_imprint=False, legacy_faceting=True)
 
-stellarator.construct_magnets(
-    "../tests/files_for_tests/coils.example", 10, 10, 90
-)
-stellarator.export_magnets()
+# Export DAGMC neutronics H5M file
+stellarator.export_dagmc(filename="dagmc", export_dir=export_dir)

Examples/pydagmc_example.py

@@ -0,0 +1,87 @@
+# Parastell can also generate invessel components with PyMOAB and PyDAGMC.
+# For complex geometry this may be more reliable than the CAD workflow, but
+# results in faceted solids, rather than the smooth spline surfaces in the
+# CAD workflow.
+import numpy as np
+import parastell.parastell as ps
+from parastell.utils import merge_dagmc_files
+
+# Define directory to export all output files to
+export_dir = ""
+# Define plasma equilibrium VMEC file
+vmec_file = "wout_vmec.nc"
+
+# Instantiate ParaStell build
+stellarator = ps.Stellarator(vmec_file)
+
+# Define build parameters for in-vessel components
+toroidal_angles = [0.0, 11.25, 22.5, 33.75, 45.0, 56.25, 67.5, 78.75, 90.0]
+poloidal_angles = [0.0, 45.0, 90.0, 135.0, 180.0, 225.0, 270.0, 315.0, 360.0]
+wall_s = 1.08
+# Use more points in the point cloud to smooth the ivb components
+num_ribs = 150
+num_rib_pts = 160
+
+# Define a matrix of uniform unit thickness
+uniform_unit_thickness = np.ones((len(toroidal_angles), len(poloidal_angles)))
+
+radial_build_dict = {
+    "first_wall": {"thickness_matrix": uniform_unit_thickness * 5},
+    "breeder": {
+        "thickness_matrix": (
+            [
+                [75.0, 75.0, 75.0, 25.0, 25.0, 25.0, 75.0, 75.0, 75.0],
+                [75.0, 75.0, 75.0, 25.0, 25.0, 75.0, 75.0, 75.0, 75.0],
+                [75.0, 75.0, 25.0, 25.0, 75.0, 75.0, 75.0, 75.0, 75.0],
+                [65.0, 25.0, 25.0, 65.0, 75.0, 75.0, 75.0, 75.0, 65.0],
+                [45.0, 45.0, 75.0, 75.0, 75.0, 75.0, 75.0, 45.0, 45.0],
+                [65.0, 75.0, 75.0, 75.0, 75.0, 65.0, 25.0, 25.0, 65.0],
+                [75.0, 75.0, 75.0, 75.0, 75.0, 25.0, 25.0, 75.0, 75.0],
+                [75.0, 75.0, 75.0, 75.0, 25.0, 25.0, 75.0, 75.0, 75.0],
+                [75.0, 75.0, 75.0, 25.0, 25.0, 25.0, 75.0, 75.0, 75.0],
+            ]
+        )
+    },
+    "back_wall": {"thickness_matrix": uniform_unit_thickness * 5},
+    "shield": {"thickness_matrix": uniform_unit_thickness * 40},
+    "vacuum_vessel": {
+        "thickness_matrix": uniform_unit_thickness * 10,
+        "mat_tag": "vac_vessel",
+    },
+}
+# Construct in-vessel components
+stellarator.construct_invessel_build(
+    toroidal_angles,
+    poloidal_angles,
+    wall_s,
+    radial_build_dict,
+    use_pydagmc=True,
+    num_ribs=num_ribs,
+    num_rib_pts=num_rib_pts,
+)
+# this file contains only the in vessel components
+stellarator.invessel_build.dag_model.write_file("dagmc.h5m")
+
+# Define build parameters for magnet coils
+coils_file = "coils.example"
+width = 40.0
+thickness = 50.0
+toroidal_extent = 90.0
+
+# Now, generate the magnet CAD files, and create a separate DAGMC model
+stellarator.construct_magnets(
+    coils_file, width, thickness, toroidal_extent, sample_mod=6
+)
+stellarator.export_magnets(
+    step_filename="magnets",
+    export_mesh=False,
+    export_dir=export_dir,
+)
+
+stellarator.build_cubit_model(skip_imprint=False, legacy_faceting=False)
+stellarator.export_dagmc("magnets.h5m")
+merge_dagmc_files(["dagmc.h5m", "magnets.h5m"], "merged_dagmc.h5m")
+
+# it is recommended to use the DAGMC command line tool overlap_check
+# to verify that the geometry is valid, and that the magnets and ivb
+# components do not intersect.

parastell/invessel_build.py

@@ -105,6 +105,8 @@ class InVesselBuild(object):
             defined.
         logger (object): logger object (optional, defaults to None). If no
             logger is supplied, a default logger will be instantiated.
+        use_pydagmc (bool): If True, generate components with pydagmc, rather
+            than CADQuery
 
     Optional attributes:
         repeat (int): number of times to repeat build segment for full model
@@ -121,13 +123,16 @@ class InVesselBuild(object):
             (defaults to m2cm = 100).
     """
 
-    def __init__(self, vmec_obj, radial_build, logger=None, **kwargs):
+    def __init__(
+        self, vmec_obj, radial_build, use_pydagmc, logger=None, **kwargs
+    ):
 
         self.mbc = core.Core()
         self.dag_model = dagmc.DAGModel(self.mbc)
         self.logger = logger
         self.vmec_obj = vmec_obj
         self.radial_build = radial_build
+        self.use_pydagmc = use_pydagmc
 
         self.repeat = 0
         self.num_ribs = 61
@@ -193,7 +198,7 @@ def _interpolate_offset_matrix(self, offset_mat):
                 self.radial_build.poloidal_angles,
             ),
             offset_mat,
-            method="pchip",
+            method="linear" if self.use_pydagmc else "pchip",
         )
 
         interpolated_offset_mat = np.array(
@@ -378,15 +383,11 @@ def generate_components_pydagmc(self):
             self.dag_model.volumes,
             list(self.radial_build.radial_build.items())[1:],
         ):
-            print(vol)
-            print(layer_data)
+
             mat = layer_data.get("mat_tag", layer_name)
             group = dagmc.Group.create(self.dag_model, name="mat:" + mat)
             group.add_set(vol)
 
-        self.dag_model.write_file("all_surfaces.vtk")
-        self.dag_model.write_file("dagmc.h5m")
-
     def connect_ribs_with_tris_moab(self, rib1, rib2, reverse=False):
         mb_tris = []
         for rib_loci_index, _ in enumerate(rib1.rib_loci[0:-1]):

parastell/parastell.py

@@ -11,7 +11,7 @@
 from . import magnet_coils as mc
 from . import source_mesh as sm
 from . import cubit_io
-from .utils import read_yaml_config, filter_kwargs, m2cm
+from .utils import read_yaml_config, filter_kwargs, m2cm, merge_dagmc_files
 
 build_cubit_model_allowed_kwargs = ["skip_imprint", "legacy_faceting"]
 export_dagmc_allowed_kwargs = [
@@ -171,6 +171,7 @@ def construct_invessel_build(
 
         self.invessel_build.populate_surfaces()
         self.invessel_build.calculate_loci()
+        self.use_pydagmc = use_pydagmc
         if use_pydagmc:
             self.invessel_build.generate_components_pydagmc()
         else:
@@ -401,7 +402,7 @@ def build_cubit_model(self, skip_imprint=False, legacy_faceting=True):
         else:
             cubit_io.init_cubit()
 
-        if self.invessel_build:
+        if self.invessel_build and not self.use_pydagmc:
             self.invessel_build.import_step_cubit()
 
         if self.magnet_set:

parastell/utils.py

@@ -1,8 +1,12 @@
 import yaml
+import tempfile
+from pathlib import Path
 
 import numpy as np
 import math
 from scipy.ndimage import gaussian_filter
+from pymoab import core, types
+
 
 m2cm = 100
 m3tocm3 = m2cm * m2cm * m2cm
@@ -202,3 +206,97 @@ def smooth_matrix(matrix, steps, sigma):
         previous_matrix = smoothed_matrix
 
     return smoothed_matrix
+
+
+class DAGMCMerger:
+    """Class to facilitate renumbering of entities to merge dagmc models"""
+
+    def __init__(self):
+        self.mb = core.Core()
+        self.geom_dim_tag = self.mb.tag_get_handle(
+            "GEOM_DIMENSION", 1, types.MB_TYPE_INTEGER, types.MB_TAG_DENSE
+        )
+        self.global_id_tag = self.mb.tag_get_handle(
+            "GLOBAL_ID", 1, types.MB_TYPE_INTEGER, types.MB_TAG_DENSE
+        )
+
+    def load_file(self, filename):
+        self.mb.load_file(filename)
+
+    def get_max_ids(self):
+        """Get the maximum id in the dagmc model for each entity type
+
+        returns:
+            max_ids (dict): {entity dimension (int): max_id (int)}
+        """
+        # dim 0: vertex
+        # dim 1: edge
+        # dim 2: surface
+        # dim 3: volume
+        # dim 4: group
+        max_ids = {dim: 0 for dim in range(5)}
+        root_set = self.mb.get_root_set()
+        geom_sets = self.mb.get_entities_by_type_and_tag(
+            root_set, types.MBENTITYSET, self.geom_dim_tag, [None]
+        )
+
+        for entity_set in geom_sets:
+            dim = self.mb.tag_get_data(self.geom_dim_tag, entity_set)[0][0]
+            gid = self.mb.tag_get_data(self.global_id_tag, entity_set)[0][0]
+            max_ids[dim] = max(max_ids[dim], gid)
+
+        return max_ids
+
+    def reassign_global_ids(self, offsets):
+        """Renumber entities, starting from the offset value for that dimension"""
+        root_set = self.mb.get_root_set()
+        geom_sets = self.mb.get_entities_by_type_and_tag(
+            root_set, types.MBENTITYSET, self.geom_dim_tag, [None]
+        )
+
+        for entity_set in geom_sets:
+            dim = self.mb.tag_get_data(self.geom_dim_tag, entity_set)[0][0]
+            current_gid = self.mb.tag_get_data(self.global_id_tag, entity_set)[
+                0
+            ][0]
+            self.mb.tag_set_data(
+                self.global_id_tag,
+                entity_set,
+                current_gid + offsets.get(dim, 0),
+            )
+
+    def write_file(self, filename):
+        self.mb.write_file(filename)
+
+
+def merge_dagmc_files(files_to_merge, output_file):
+    """Takes a list of dagmc models, and renumbers entity ids such that they
+    will no longer clash, allowing the models to be loaded into the same
+    PyMOAB core instance, and saved to a single model.
+
+    Arguments:
+        files_to_merge (list of str): List of paths to dagmc models to be
+            merged.
+        output_file (str): path to which the merged model should be saved.
+    """
+    temp_files = []
+    max_ids = {dim: 0 for dim in range(5)}
+    for file in files_to_merge:
+        merger = DAGMCMerger()
+        merger.load_file(file)
+        merger.reassign_global_ids(max_ids)
+        max_ids = merger.get_max_ids()
+        with tempfile.NamedTemporaryFile(
+            delete=False, suffix=".h5m"
+        ) as temp_file:
+            temp_filename = temp_file.name
+        merger.write_file(temp_filename)
+        temp_files.append(temp_filename)
+
+    merged_mbc = core.Core()
+    for file in temp_files:
+        merged_mbc.load_file(file)
+    merged_mbc.write_file(output_file)
+
+    Path.unlink(temp_filename)
+    return True