Dask: 2D Phase Space Example (#958)

Add a phase-space example for Dask

Author: ax3l

examples/11_particle_dataframe.py

@@ -36,7 +36,7 @@
 print(df)
 
 
-# Dask example: process large datasets
+# Particles
 if found_dask:
     df = electrons.to_dask()
     print(df)
@@ -54,15 +54,28 @@
                            weights=df["weighting"].to_dask_array())
     print(h.compute())
 
-    #   example3: save all data data to parquet files
+    #   example3: longitudinal phase space (dask 2021.04.0+)
+    z_min = df["position_z"].min().compute()
+    z_max = df["position_z"].max().compute()
+
+    z_pz, z_pz_bins = da.histogramdd(
+        df[['position_z', 'momentum_z']].to_dask_array(),
+        bins=[80, 80],
+        range=[[z_min, z_max], [-8.0e-23, 8.0e-23]],
+        weights=df["weighting"].to_dask_array()
+    )
+    print(z_pz.compute())
+
+    #   example4: save all data data to parquet files
     delayed_save = delayed(df.to_parquet("electrons.parquet"))
     delayed_save.compute()
 
-# Dask example 2: meshes
+# Meshes
 if found_dask:
-    E_x = s.iterations[400].meshes["E"]["x"]
-    E_y = s.iterations[400].meshes["E"]["y"]
-    E_z = s.iterations[400].meshes["E"]["z"]
+    E = s.iterations[400].meshes["E"]
+    E_x = E["x"]
+    E_y = E["y"]
+    E_z = E["z"]
     darr_x = E_x.to_dask_array()
     darr_y = E_y.to_dask_array()
     darr_z = E_z.to_dask_array()
@@ -71,4 +84,7 @@
     Intensity = darr_x * darr_x + darr_y * darr_y + darr_z * darr_z
     Intensity_max = Intensity.max().compute()
     idx_max = da.argwhere(Intensity == Intensity_max).compute()[0]
-    print("maximum intensity I={} at index={}".format(Intensity_max, idx_max))
+    pos_max = E.grid_unit_SI * 1.0e6 * (
+        idx_max * E.grid_spacing + E.grid_global_offset)
+    print("maximum intensity I={} at index={} z={}mu".format(
+        Intensity_max, idx_max, pos_max[2]))