Merge pull request #59 from gonuke/ivc_oo_pphw

Some final input on the IVC OO update

src/cubit_io.py

@@ -1,22 +1,18 @@
-import cubit
 from pathlib import Path
 import os
 import inspect
 import subprocess
 
+import cubit
 
-def init_cubit(cubit_initialized=False):
-    """Initializes Coreform Cubit with the DAGMC plugin.
+initialized = False
 
-    Arguments:
-        cubit_initialized (bool): flag to indicate whether Coreform Cubit has
-            been initialized.
-
-    Returns:
-        cubit_initialized (bool): flag to indicate whether Coreform Cubit has
-            been initialized.
+def init_cubit():
+    """Initializes Coreform Cubit with the DAGMC plugin.
     """
-    if not cubit_initialized:
+    global initialized
+
+    if not initialized:
         cubit_plugin_dir = (
             Path(os.path.dirname(inspect.getfile(cubit))) / Path('plugins')
         )
@@ -29,9 +25,7 @@ def init_cubit(cubit_initialized=False):
             '-commandplugindir',
             str(cubit_plugin_dir)
         ])
-        cubit_initialized = True
-
-    return cubit_initialized
+        initialized = True
 
 
 def import_step_cubit(filename, import_dir):
@@ -61,6 +55,8 @@ def export_step_cubit(filename, export_dir=''):
         export_dir (str): directory to which to export the STEP output file
             (defaults to empty string).
     """
+    init_cubit()
+
     export_path = Path(export_dir) / Path(filename).with_suffix('.step')
     cubit.cmd(f'export step "{export_path}" overwrite')
 
@@ -74,6 +70,8 @@ def export_mesh_cubit(filename, export_dir=''):
         export_dir (str): directory to which to export the H5M output file
             (defaults to empty string).
     """
+    init_cubit()
+
     exo_path = Path(export_dir) / Path(filename).with_suffix('.exo')
     h5m_path = Path(export_dir) / Path(filename).with_suffix('.h5m')
 
@@ -101,6 +99,8 @@ def export_dagmc_cubit_legacy(
         export_dir (str): directory to which to export the DAGMC output file
             (defaults to empty string).
     """
+    init_cubit()
+
     tol_str = ''
 
     if faceting_tolerance is not None:
@@ -134,6 +134,8 @@ def export_dagmc_cubit_native(
         export_dir (str): directory to which to export the DAGMC output file
             (defaults to empty string).
     """
+    init_cubit()
+
     cubit.cmd(
         f'set trimesher coarse on ratio {anisotropic_ratio} '
         f'angle {deviation_angle}'

src/invessel_build.py

@@ -1,17 +1,17 @@
-import cubit
-import cadquery as cq
-import cad_to_dagmc
-import read_vmec
-from src.utils import expand_ang_list, normalize, def_default_params
 import log
-import numpy as np
-from scipy.interpolate import RegularGridInterpolator
-from pathlib import Path
 import argparse
 import yaml
+from pathlib import Path
 
-m2cm, _, invessel_build_def, _, _, _ = def_default_params()
+import numpy as np
+from scipy.interpolate import RegularGridInterpolator
 
+import cubit
+import cadquery as cq
+import cad_to_dagmc
+import read_vmec
+
+from src.utils import expand_ang_list, normalize, m2cm, invessel_build_def
 
 def orient_spline_surfaces(volume_id):
     """Extracts the inner and outer surface IDs for a given ParaStell in-vessel
@@ -24,6 +24,7 @@ def orient_spline_surfaces(volume_id):
         inner_surface_id (int): Cubit ID of in-vessel component inner surface.
         outer_surface_id (int): Cubit ID of in-vessel component outer surface.
     """
+
     surfaces = cubit.get_relatives('volume', volume_id, 'surface')
 
     spline_surfaces = []
@@ -48,47 +49,6 @@ def orient_spline_surfaces(volume_id):
 
     return inner_surface_id, outer_surface_id
 
-
-def merge_layer_surfaces(components_dict):
-    """Merges ParaStell in-vessel component surfaces in Coreform Cubit based on
-    surface IDs rather than imprinting and merging all. Assumes that the
-    components dictionary is ordered radially outward. Note that overlaps
-    between magnet volumes and in-vessel components will not be merged in this
-    workflow.
-
-    Arguments:
-        components_dict (dict): dictionary of ParaStell components. This
-            dictionary must have the form
-            {
-                'component_name': {
-                    'vol_id': Coreform Cubit volume ID(s) for component (int or
-                        iterable of int)
-                    (additional keys are allowed)
-                }
-            }
-    """
-    # Tracks the surface id of the outer surface of the previous layer
-    prev_outer_surface_id = None
-
-    for data in components_dict.values():
-        # Skip merging for magnets
-        if isinstance(data['vol_id'], list):
-            continue
-
-        inner_surface_id, outer_surface_id = (
-            orient_spline_surfaces(data['vol_id'])
-        )
-
-        # Conditionally skip merging (first iteration only)
-        if prev_outer_surface_id is None:
-            prev_outer_surface_id = outer_surface_id
-        else:
-            cubit.cmd(
-                f'merge surface {inner_surface_id} {prev_outer_surface_id}'
-            )
-            prev_outer_surface_id = outer_surface_id
-
-
 class InVesselBuild(object):
     """Parametrically models fusion stellarator in-vessel components using
     plasma equilibrium data and user-defined parameters. In-vessel component
@@ -175,16 +135,15 @@ def __init__(
         if self.logger == None or not self.logger.hasHandlers():
             self.logger = log.init()
 
-        self.Surfaces = []
-        self.Components = []
+        self.Surfaces = {}
+        self.Components = {}
 
-        try:
-            assert (self.repeat + 1) * self.toroidal_angles[-1] <= 360.0, (
+        if (self.repeat + 1) * self.toroidal_angles[-1] > 360.0:
+            e = AssertionError(
                 'Total toroidal extent requested with repeated geometry '
                 'exceeds 360 degrees. Please examine phi_list and the repeat '
                 'parameter.'
             )
-        except AssertionError as e:
             self.logger.error(e.args[0])
             raise e
 
@@ -255,12 +214,11 @@ def populate_surfaces(self):
         ))
 
         for name, layer_data in self.radial_build.items():
-            try:
-                assert np.all(np.array(layer_data['thickness_matrix']) >= 0), (
+            if np.any(np.array(layer_data['thickness_matrix']) < 0):
+                e = AssertionError(
                     'Component thicknesses must be greater than or equal to 0. '
                     'Check thickness inputs for negative values.'
                 )
-            except AssertionError as e:
                 self.logger.error(e.args[0])
                 raise e
 
@@ -277,22 +235,20 @@ def populate_surfaces(self):
                 offset_mat
             )
 
-            self.Surfaces.append(
-                Surface(
+            self.Surfaces[name] = Surface(
                     self.vmec, s, self.theta_list, self.phi_list,
                     interpolated_offset_mat, self.scale
                 )
-            )
-
-        [surface.populate_ribs() for surface in self.Surfaces]
+
+        [surface.populate_ribs() for surface in self.Surfaces.values()]
 
     def calculate_loci(self):
         """Calls calculate_loci method in Surface class for each component
         specified in the radial build.
         """
         self.logger.info('Computing point cloud for in-vessel components...')
 
-        [surface.calculate_loci() for surface in self.Surfaces]
+        [surface.calculate_loci() for surface in self.Surfaces.values()]
 
     def generate_components(self):
         """Constructs a CAD solid for each component specified in the radial
@@ -311,7 +267,7 @@ def generate_components(self):
             num=self.repeat
         )
 
-        for surface in self.Surfaces:
+        for name, surface in self.Surfaces.items():
             outer_surface = surface.generate_surface()
 
             if interior_surface is not None:
@@ -325,14 +281,40 @@ def generate_components(self):
                 rot_segment = segment.rotate((0, 0, 0), (0, 0, 1), angle)
                 component = component.union(rot_segment)
 
-            self.Components.append(component)
+            self.Components[name] = component
             interior_surface = outer_surface
 
     def get_loci(self):
         """Returns the set of point-loci defining the outer surfaces of the
         components specified in the radial build.
         """
-        return np.array([surface.get_loci() for surface in self.Surfaces])
+        return np.array([surface.get_loci() for surface in self.Surfaces.values()])
+
+    def merge_layer_surfaces(self):
+        """Merges ParaStell in-vessel component surfaces in Coreform Cubit based on
+        surface IDs rather than imprinting and merging all. Assumes that the
+        radial_build dictionary is ordered radially outward. Note that overlaps
+        between magnet volumes and in-vessel components will not be merged in this
+        workflow.
+        """
+        # Tracks the surface id of the outer surface of the previous layer
+        prev_outer_surface_id = None
+
+        for data in self.invessel_build.radial_build.values():
+
+            inner_surface_id, outer_surface_id = (
+                orient_spline_surfaces(data['vol_id'])
+            )
+
+            # Conditionally skip merging (first iteration only)
+            if prev_outer_surface_id is None:
+                prev_outer_surface_id = outer_surface_id
+            else:
+                cubit.cmd(
+                    f'merge surface {inner_surface_id} {prev_outer_surface_id}'
+                )
+                prev_outer_surface_id = outer_surface_id
+
 
     def export_step(self, export_dir=''):
         """Export CAD solids as STEP files via CadQuery.
@@ -345,9 +327,7 @@ def export_step(self, export_dir=''):
 
         self.export_dir = export_dir
 
-        for component, name in zip(
-            self.Components, self.radial_build.keys()
-        ):
+        for name, component in self.Components.items():
             export_path = (
                 Path(self.export_dir) / Path(name).with_suffix('.step')
             )
@@ -372,11 +352,9 @@ def export_cad_to_dagmc(self, filename='dagmc', export_dir=''):
 
         model = cad_to_dagmc.CadToDagmc()
 
-        for component, (_, layer_data) in zip(
-            self.Components, self.radial_build.items()
-        ):
+        for name, component in self.Components.items():
             model.add_cadquery_object(
-                component, material_tags=[layer_data['mat_tag']]
+                component, material_tags=[self.radial_build[name]['mat_tag]']]
             )
 
         export_path = Path(export_dir) / Path(filename).with_suffix('.h5m')

src/magnet_coils.py

@@ -1,13 +1,12 @@
-import cubit
-from src.cubit_io import init_cubit, export_step_cubit, export_mesh_cubit
 import log
-import numpy as np
 import yaml
 import argparse
-from src.utils import normalize, def_default_params
 
-m2cm, _, _, magnets_def, _, _ = def_default_params()
+import numpy as np
 
+import cubit
+import src.cubit_io as cubit_io 
+from src.utils import normalize, m2cm, magnets_def
 
 class MagnetSet(object):
     """An object representing a set of modular stellarator magnet coils.
@@ -29,8 +28,6 @@ class MagnetSet(object):
             neutronics model (defaults to 'magnets').
         logger (object): logger object (defaults to None). If no logger is
             supplied, a default logger will be instantiated.
-        cubit_initialized (bool): flag to indicate whether Coreform Cubit has
-            been initialized (defaults to False).
     """
 
     def __init__(
@@ -43,7 +40,6 @@ def __init__(
         scale=m2cm,
         mat_tag=None,
         logger=None,
-        cubit_initialized=False
     ):
         self.coils_file_path = coils_file_path
         self.cross_section = cross_section
@@ -58,7 +54,7 @@ def __init__(
         else:
             self.logger = logger
 
-        self.cubit_initialized = init_cubit(cubit_initialized)
+        cubit_io.init_cubit()
 
         self._extract_filaments()
         self._extract_cross_section()
@@ -340,7 +336,7 @@ def export_step(self, filename='magnets', export_dir=''):
         self.logger.info('Exporting STEP file for magnet coils...')
 
         self.step_filename = filename
-        export_step_cubit(filename=filename, export_dir=export_dir)
+        cubit_io.export_step_cubit(filename=filename, export_dir=export_dir)
 
     def mesh_magnets(self):
         """Creates tetrahedral mesh of magnet volumes via Coreform Cubit.
@@ -362,7 +358,7 @@ def export_mesh(self, filename='magnet_mesh', export_dir=''):
         """
         self.logger.info('Exporting mesh H5M file for magnet coils...')
 
-        export_mesh_cubit(filename=filename, export_dir=export_dir)
+        cubit_io.export_mesh_cubit(filename=filename, export_dir=export_dir)
 
 
 class MagnetCoil(object):