Adding middle stage extraction for HiPACE.

abel/classes/stage/impl/stage_hipace.py

@@ -283,7 +283,9 @@ def track(self, beam_incoming, savedepth=0, runnable=None, verbose=False):
         dz = beta_matched/20
 
         # convert to number of steps (and re-adjust timestep to be divisible)
-        self.num_steps = np.ceil(self.length_flattop/dz)
+        self.num_steps = int(np.ceil(self.length_flattop / dz))
+        # Adds 1 to make it even if odd
+        self.num_steps += self.num_steps % 2  # make it even (add 1 if odd)
 
         if self.output is not None:
             remainder = self.num_steps % self.output
@@ -351,7 +353,7 @@ def track(self, beam_incoming, savedepth=0, runnable=None, verbose=False):
 
         # extract insitu diagnostics and wakefield data
         self.__extract_evolution(tmpfolder, beam0, runnable)
-        self.__extract_initial_and_final_step(tmpfolder, beam0, runnable)
+        self.__extract_initial_middle_and_final_step(tmpfolder, beam0, runnable)
 
         # delete temp folder
         shutil.rmtree(tmpfolder)
@@ -495,7 +497,7 @@ def __extract_evolution(self, tmpfolder, beam0, runnable):
             shutil.move(insitu_path, destination_path)
 
 
-    def __extract_initial_and_final_step(self, tmpfolder, beam0, runnable):
+    def __extract_initial_middle_and_final_step(self, tmpfolder, beam0, runnable):
 
         from openpmd_viewer import OpenPMDTimeSeries
 
@@ -512,18 +514,41 @@ def __extract_initial_and_final_step(self, tmpfolder, beam0, runnable):
         Ez, metadata = ts.get_field(field='Ez', slice_across=['x'], iteration=max(ts.iterations))
         self.final.plasma.wakefield.onaxis.zs = metadata.z
         self.final.plasma.wakefield.onaxis.Ezs = Ez
-
+        print(ts.iterations)
+        # extract middle field information if there are more than two steps
+        if len(ts.iterations) > 2:
+            # Get step closest to middle
+            middle_step = ts.iterations[np.argmin(abs(ts.iterations - int(max(ts.iterations)/2)))]
+
+            # extract middle on-axis wakefield
+            Ez, metadata = ts.get_field(field='Ez', slice_across=['x'], iteration=middle_step)
+            self.middle.plasma.wakefield.onaxis.zs = metadata.z
+            self.middle.plasma.wakefield.onaxis.Ezs = Ez
+
+            # extract middle beam density
+            jz0_beam, metadata0_beam = ts.get_field(field='jz_beam', iteration=middle_step)
+            self.middle.beam.density.extent = metadata0_beam.imshow_extent[[2,3,0,1]]
+            self.middle.beam.density.rho = -jz0_beam.T/(SI.c*SI.e)
+
+            # extract middle plasma density
+            rho_plasma, metadata_plasma = ts.get_field(field='rho', iteration=middle_step)
+            self.middle.plasma.density.extent = metadata_plasma.imshow_extent[[2,3,0,1]]
+            self.middle.plasma.density.rho = -(rho_plasma.T/SI.e-self.plasma_density)
+
+            # Extract step 
+            self.middle.n_step = middle_step
+            print(self.middle.n_step)
         # extract initial plasma density
         rho0_plasma, metadata0_plasma = ts.get_field(field='rho', iteration=min(ts.iterations))
         self.initial.plasma.density.extent = metadata0_plasma.imshow_extent[[2,3,0,1]]
         self.initial.plasma.density.rho = -(rho0_plasma.T/SI.e-self.plasma_density)
 
-        # extract final beam density
+        # extract initial beam density
         jz0_beam, metadata0_beam = ts.get_field(field='jz_beam', iteration=min(ts.iterations))
         self.initial.beam.density.extent = metadata0_beam.imshow_extent[[2,3,0,1]]
         self.initial.beam.density.rho = -jz0_beam.T/(SI.c*SI.e)
 
-        # extract initial plasma density
+        # extract final plasma density
         rho_plasma, metadata_plasma = ts.get_field(field='rho', iteration=max(ts.iterations))
         self.final.plasma.density.extent = metadata_plasma.imshow_extent[[2,3,0,1]]
         self.final.plasma.density.rho = -(rho_plasma.T/SI.e-self.plasma_density)
@@ -542,8 +567,20 @@ def __extract_initial_and_final_step(self, tmpfolder, beam0, runnable):
     def _make_ramp_profile(self, tmpfolder):
         """Prepare the ramps (local to HiPACE)."""
 
-        # check that there is not already a plasma density profile set
-        assert self.plasma_density_from_file is None
+        # If there is already a density file open and make the plasma profile
+        if self.plasma_density_from_file:
+            ss, ns = [], []
+            with open(self.plasma_density_from_file, 'r') as f:
+                for line in f:
+                    s, n = line.strip().split()
+                    ss.append(float(s))
+                    ns.append(float(n))
+
+                ss = np.array(ss)
+                ns = np.array(ns)
+                self.plasma_profile.ss = ss
+                self.plasma_profile.ns = ns
+            return
 
         # make the plasma ramp profile
         if self.has_ramp():