fix #12 use lume standard method name, write_initial_particles() (#13)

rslume/elegant.py

@@ -37,9 +37,10 @@ def load_output(self):
         # TODO(pjm): load other sdds output files
         # TODO(pjm): load warnings and errors from log
 
-    # -- RS addition
-
-    def set_particle_input(self, particle_group, filename="in.sdds"):
+    def write_initial_particles(self, filename="in.sdds"):
+        particle_group = self.initial_particles
+        if not particle_group:
+            return
         filepath = os.path.join(self.workdir, filename)
         pmd_beamphysics.interfaces.elegant.write_elegant(
             particle_group,

rslume/genesis2.py

@@ -13,7 +13,7 @@
 
 class Genesis2(genesis.genesis2.Genesis2):
 
-    # override lume-genesis final_particles()
+    # override lume-genesis final_particles() and write_input()
 
     def final_particles(self):
         # compute required wavelengths to hold the whole beam and center at 0
@@ -33,9 +33,10 @@ def final_particles(self):
         v.t -= numpy.mean(v.t)
         return v
 
-    # -- RS addition
-
-    def set_particle_input(self, particle_group, filename="inpart.dat"):
+    def write_initial_particles(self, filename="inpart.dat"):
+        particle_group = self.initial_particles
+        if not particle_group:
+            return
         filepath = os.path.join(self.workdir, filename)
 
         # gamma
@@ -71,3 +72,7 @@ def set_particle_input(self, particle_group, filename="inpart.dat"):
 
         with open(filepath, "wb") as f:
             d.tofile(f)
+
+    def write_input(self):
+        self.write_initial_particles()
+        super().write_input()

rslume/opal.py

@@ -43,9 +43,10 @@ def load_output(self):
                 ),
             )
 
-    # -- RS addition
-
-    def set_particle_input(self, particle_group, filename="in.dat"):
+    def write_initial_particles(self, filename="in.dat"):
+        particle_group = self.initial_particles
+        if not particle_group:
+            return
         filepath = os.path.join(self.workdir, filename)
         pmd_beamphysics.interfaces.opal.write_opal(
             particle_group,

rslume/wrapper.py

@@ -63,6 +63,7 @@ def plot(self):
         raise NotImplementedError("plot() not yet implemented.")
 
     def write_input(self):
+        self.write_initial_particles()
         self._input.write_files(self.path)
 
     # Sirepo lattice and command accessors