Analysis: pandas.concat all files

ECP-WarpX · Apr 11, 2022 · 268a6aa · 268a6aa
1 parent 88e43d6
commit 268a6aa
Show file tree

Hide file tree

Showing 4 changed files with 80 additions and 8 deletions.
diff --git a/examples/cfchannel/analysis_cfchannel.py b/examples/cfchannel/analysis_cfchannel.py
@@ -35,9 +35,27 @@ def get_moments(beam):
         emittance_x, emittance_y, emittance_t)
 
 
+def read_all_files(file_pattern):
+    """Read in all CSV files from each MPI rank (and potentially OpenMP
+    thread. Concatinate into one Pandas dataframe.
+
+    Returns
+    -------
+    pandas.DataFrame
+    """
+    return pd.concat(
+        (
+            pd.read_csv(filename, delimiter=r"\s+")
+            for filename in glob.glob(file_pattern)
+        ),
+        axis=0,
+        ignore_index=True,
+    )
+
+
 # initial/final beam on rank zero
-initial = pd.read_csv("diags/initial_beam.txt.0.0", delimiter=r"\s+")
-final = pd.read_csv("diags/output_beam.txt.0.0", delimiter=r"\s+")
+initial = read_all_files("diags/initial_beam.txt.*")
+final = read_all_files("diags/output_beam.txt.*")
 
 # compare number of particles
 num_particles = 10000

diff --git a/examples/chicane/analysis_chicane.py b/examples/chicane/analysis_chicane.py
@@ -35,9 +35,27 @@ def get_moments(beam):
         emittance_x, emittance_y, emittance_t)
 
 
+def read_all_files(file_pattern):
+    """Read in all CSV files from each MPI rank (and potentially OpenMP
+    thread. Concatinate into one Pandas dataframe.
+
+    Returns
+    -------
+    pandas.DataFrame
+    """
+    return pd.concat(
+        (
+            pd.read_csv(filename, delimiter=r"\s+")
+            for filename in glob.glob(file_pattern)
+        ),
+        axis=0,
+        ignore_index=True,
+    )
+
+
 # initial/final beam on rank zero
-initial = pd.read_csv("diags/initial_beam.txt.0.0", delimiter=r"\s+")
-final = pd.read_csv("diags/output_beam.txt.0.0", delimiter=r"\s+")
+initial = read_all_files("diags/initial_beam.txt.*")
+final = read_all_files("diags/output_beam.txt.*")
 
 # compare number of particles
 num_particles = 10000

diff --git a/examples/fodo/analysis_fodo.py b/examples/fodo/analysis_fodo.py
@@ -35,9 +35,27 @@ def get_moments(beam):
         emittance_x, emittance_y, emittance_t)
 
 
+def read_all_files(file_pattern):
+    """Read in all CSV files from each MPI rank (and potentially OpenMP
+    thread. Concatinate into one Pandas dataframe.
+
+    Returns
+    -------
+    pandas.DataFrame
+    """
+    return pd.concat(
+        (
+            pd.read_csv(filename, delimiter=r"\s+")
+            for filename in glob.glob(file_pattern)
+        ),
+        axis=0,
+        ignore_index=True,
+    )
+
+
 # initial/final beam on rank zero
-initial = pd.read_csv("diags/initial_beam.txt.0.0", delimiter=r"\s+")
-final = pd.read_csv("diags/output_beam.txt.0.0", delimiter=r"\s+")
+initial = read_all_files("diags/initial_beam.txt.*")
+final = read_all_files("diags/output_beam.txt.*")
 
 # compare number of particles
 num_particles = 10000

diff --git a/examples/kurth/analysis_kurth.py b/examples/kurth/analysis_kurth.py
@@ -35,9 +35,27 @@ def get_moments(beam):
         emittance_x, emittance_y, emittance_t)
 
 
+def read_all_files(file_pattern):
+    """Read in all CSV files from each MPI rank (and potentially OpenMP
+    thread. Concatinate into one Pandas dataframe.
+
+    Returns
+    -------
+    pandas.DataFrame
+    """
+    return pd.concat(
+        (
+            pd.read_csv(filename, delimiter=r"\s+")
+            for filename in glob.glob(file_pattern)
+        ),
+        axis=0,
+        ignore_index=True,
+    )
+
+
 # initial/final beam on rank zero
-initial = pd.read_csv("diags/initial_beam.txt.0.0", delimiter=r"\s+")
-final = pd.read_csv("diags/output_beam.txt.0.0", delimiter=r"\s+")
+initial = read_all_files("diags/initial_beam.txt.*")
+final = read_all_files("diags/output_beam.txt.*")
 
 # compare number of particles
 num_particles = 10000