Corrections for using lume-base, cathode start status with distgen, e…

…xamples.

astra/astra.py

@@ -57,11 +57,13 @@ class Astra(CommandWrapper):
     COMMAND = "$ASTRA_BIN"
     INPUT_PARSER = parsers.parse_astra_input_file
 
-    def __init__(self, *,
+    def __init__(self, 
+                 input_file=None,
+                 *,                 
                  group=None,
                  **kwargs
                  ):
-        super().__init__(**kwargs)
+        super().__init__(input_file=input_file, **kwargs)
         # Save init
         self.original_input_file = self._input_file
 

astra/astra_distgen.py

@@ -10,6 +10,8 @@
 from distgen.writers import write_astra
 from distgen.tools import update_nested_dict
 
+from lume import tools as lumetools
+
 from pmd_beamphysics import ParticleGroup
 
 from h5py import File
@@ -65,7 +67,7 @@ def run_astra_with_distgen(settings=None,
         return run_astra(settings=settings, 
                          astra_input_file=astra_input_file, 
                          workdir=workdir,
-                         astra_bin=astra_bin, 
+                         command=astra_bin, 
                          timeout=timeout, 
                          verbose=verbose)
 
@@ -74,7 +76,7 @@ def run_astra_with_distgen(settings=None,
         print('run_astra_with_generator') 
 
     # Distgen generator
-    G = Generator(input = distgen_input_file, verbose=verbose)  
+    G = Generator(distgen_input_file, verbose=verbose)  
 
     # Make astra objects
     if astra_input_file.endswith('.yaml'):
@@ -86,7 +88,7 @@ def run_astra_with_distgen(settings=None,
             A.configure() # again to make sure things are set properly  
 
     else:
-        A = Astra(astra_bin=astra_bin, input_file=astra_input_file, workdir=workdir)
+        A = Astra(command=astra_bin, input_file=astra_input_file, workdir=workdir)
 
 
 
@@ -130,6 +132,9 @@ def run_astra_with_distgen(settings=None,
     # Run distgen
     G.run()
     P = G.particles
+    # Special flag for cathode start
+    if G['start:type'] == 'cathode':
+        P.status[:] -1
 
     # Attach to Astra object
     A.initial_particles = P
@@ -192,7 +197,7 @@ def evaluate_astra_with_distgen(settings,
     output['fingerprint'] = fingerprint
 
     if archive_path:
-        path = tools.full_path(archive_path)
+        path = lumetools.full_path(archive_path)
         assert os.path.exists(path), f'archive path does not exist: {path}'
         archive_file = os.path.join(path, fingerprint+'.h5')
         output['archive'] = archive_file
@@ -211,7 +216,7 @@ def fingerprint_astra_with_distgen(astra_object, distgen_object):
     f1 = astra_object.fingerprint()
     f2 = distgen_object.fingerprint()
     d = {'f1':f1, 'f2':2}
-    return tools.fingerprint(d)
+    return lumetools.fingerprint(d)
 
 
 

astra/evaluate.py

@@ -191,8 +191,8 @@ def evaluate_astra_with_generator(settings,
                                     astra_input_file=astra_input_file,
                                     generator_input_file=generator_input_file,
                                     workdir=workdir, 
-                                    astra_bin=astra_bin,
-                                    generator_bin=generator_bin,
+                                    command=astra_bin,
+                                    command_generator=generator_bin,
                                     timeout=timeout,
                                     auto_set_spacecharge_mesh=auto_set_spacecharge_mesh,
                                     verbose=verbose)

astra/generator.py

@@ -23,8 +23,8 @@ class AstraGenerator(CommandWrapper):
     """
     COMMAND = "$GENERATOR_BIN"
 
-    def __init__(self, **kwargs):
-        super().__init__(**kwargs)
+    def __init__(self, input_file, **kwargs):
+        super().__init__(input_file=input_file, **kwargs)
         # Save init
         self.original_input_file = self.input_file
 
@@ -41,9 +41,11 @@ def load_input(self, input_filepath, absolute_paths=True, **kwargs):
         super().load_input(input_filepath, **kwargs)
         if absolute_paths:
             parsers.fix_input_paths(self.input, root=self.original_path)
-
+        self.input = self.input['input']
+
+
     def input_parser(self, path):
-        return parsers.parse_astra_input_file(path)['input']
+        return parsers.parse_astra_input_file(path)
 
     def configure(self):
         # Check that binary exists

astra/parsers.py

@@ -375,7 +375,6 @@ def fix_input_paths(input_dict, root='', prefixes=['file_', 'distribution', 'q_t
     
     """
     for nl in input_dict:
-        #print(nl)
         for key in input_dict[nl]:
             if any([key.startswith(prefix) for prefix in prefixes]):
                 val = input_dict[nl][key]