From 29883a8ff568290f1ec6ff8421968a77836c8905 Mon Sep 17 00:00:00 2001
From: Connor Moreno <camoreno@Denali.ep.wisc.edu>
Date: Mon, 29 Apr 2024 12:26:00 -0500
Subject: [PATCH] Make build_full_model and export_full_model global functions
 and fix use of dictionaries as inputs

---
 parastell/parastell.py | 138 +++++++++++++++++++----------------------
 1 file changed, 63 insertions(+), 75 deletions(-)

diff --git a/parastell/parastell.py b/parastell/parastell.py
index 10bc2137..40cab17c 100644
--- a/parastell/parastell.py
+++ b/parastell/parastell.py
@@ -412,77 +412,6 @@ def export_dagmc(
                 **kwargs
             )
 
-    def build_full_model(self, invessel_build, magnet_coils, source_mesh):
-        """Use all of the input to make a complete stellarator model with
-        in-vessel components, magnets, and a source mesh.
-
-        Arguments:
-            invessel_build (dict): dictionary of RadialBuild and InVesselBuild
-                parameters.
-            magnet_coils (dict): dictionary of MagnetSet parameters.
-            source_mesh (dict): dictionary of SourceMesh parameters.
-        """
-        self.construct_invessel_build(
-            invessel_build['toroidal_angles'],
-            invessel_build['poloidal_angles'],
-            invessel_build['wall_s'],
-            invessel_build['radial_build'],
-            **invessel_build
-        )
-
-        self.construct_magnets(
-            magnet_coils['coils_file'],
-            magnet_coils['cross_section'],
-            magnet_coils['toroidal_extent'],
-            **magnet_coils
-        )
-
-        self.construct_source_mesh(
-            source_mesh['mesh_size'],
-            source_mesh['toroidal_extent'],
-            **source_mesh
-        )
-
-    def export_full_model(
-        self, invessel_build, magnet_coils, source_mesh, dagmc_export,
-        export_dir=''
-    ):
-        """Exports a complete stellarator model when in-vessel components,
-        magnets, and source mesh have been constructed.
-
-        Arguments:
-            invessel_build (dict): dictionary of RadialBuild and InVesselBuild
-                parameters.
-            magnet_coils (dict): dictionary of MagnetSet parameters.
-            source_mesh (dict): dictionary of SourceMesh parameters.
-            dagmc_export (dict): dictionary of DAGMC export parameters.
-            export_dir (str): directory to which output files are exported
-                (optional, defaults to empty string).
-        """
-        export_cad_to_dagmc = invessel_build.get('export_cad_to_dagmc', False)
-        dagmc_filename = invessel_build.get('dagmc_filename', 'dagmc')
-
-        self.export_invessel_build(
-            export_dir=export_dir,
-            export_cad_to_dagmc=export_cad_to_dagmc,
-            dagmc_filename=dagmc_filename
-        )
-        
-        self.export_magnets(
-            export_dir=export_dir,
-            **(filter_kwargs(magnet_coils, mc.export_allowed_kwargs))
-        )
-
-        self.export_source_mesh(
-            export_dir=export_dir,
-            **(filter_kwargs(source_mesh, sm.export_allowed_kwargs))
-        )
-        
-        self.export_dagmc(
-            export_dir=export_dir,
-            **dagmc_export
-        )
-
 
 def parse_args():
     """Parser for running as a script.
@@ -600,6 +529,63 @@ def check_inputs(
             )
             logger.error(e.args[0])
             raise e
+        
+
+def build_full_model(stellarator, invessel_build, magnet_coils, source_mesh):
+    """Use all of the input to make a complete stellarator model with in-vessel
+    components, magnets, and a source mesh.
+
+    Arguments:
+        stellarator (object): Stellarator class object.
+        invessel_build (dict): dictionary of RadialBuild and InVesselBuild
+            parameters.
+        magnet_coils (dict): dictionary of MagnetSet parameters.
+        source_mesh (dict): dictionary of SourceMesh parameters.
+    """
+    stellarator.construct_invessel_build(**invessel_build)
+    stellarator.construct_magnets(**magnet_coils)
+    stellarator.construct_source_mesh(**source_mesh)
+
+
+def export_full_model(
+    stellarator, export_dir, invessel_build, magnet_coils, source_mesh,
+    dagmc_export
+):
+    """Exports a complete stellarator model when in-vessel components, magnets,
+    and source mesh have been constructed.
+
+    Arguments:
+        stellarator (object): Stellarator class object.
+        export_dir (str): directory to which output files are exported.
+        invessel_build (dict): dictionary of RadialBuild and InVesselBuild
+            parameters.
+        magnet_coils (dict): dictionary of MagnetSet parameters.
+        source_mesh (dict): dictionary of SourceMesh parameters.
+        dagmc_export (dict): dictionary of DAGMC export parameters.
+    """
+    export_cad_to_dagmc = invessel_build.get('export_cad_to_dagmc', False)
+    dagmc_filename = invessel_build.get('dagmc_filename', 'dagmc')
+
+    stellarator.export_invessel_build(
+        export_dir=export_dir,
+        export_cad_to_dagmc=export_cad_to_dagmc,
+        dagmc_filename=dagmc_filename
+    )
+    
+    stellarator.export_magnets(
+        export_dir=export_dir,
+        **(filter_kwargs(magnet_coils, mc.export_allowed_kwargs))
+    )
+
+    stellarator.export_source_mesh(
+        export_dir=export_dir,
+        **(filter_kwargs(source_mesh, sm.export_allowed_kwargs))
+    )
+    
+    stellarator.export_dagmc(
+        export_dir=export_dir,
+        **dagmc_export
+    )
 
 
 def parastell():
@@ -631,18 +617,20 @@ def parastell():
         logger=logger
     )
 
-    stellarator.build_full_model(
+    build_full_model(
+        stellarator,
         all_data['invessel_build'],
         all_data['magnet_coils'],
         all_data['source_mesh']
     )
 
-    stellarator.export_full_model(
+    export_full_model(
+        stellarator,
+        args.export_dir,
         all_data['invesel_build'],
         all_data['magnet_coils'],
         all_data['source_mesh'],
-        all_data['dagmc_export'],
-        export_dir=args.export_dir,
+        all_data['dagmc_export']
     )