[pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
ECP-WarpX · Jan 25, 2025 · ebc1618 · ebc1618
1 parent cc05d58
commit ebc1618
Show file tree

Hide file tree

Showing 3 changed files with 45 additions and 25 deletions.
diff --git a/examples/scraping_beam/README.rst b/examples/scraping_beam/README.rst
@@ -43,11 +43,11 @@ This example can be run as a Python script (``python3 run_scraping.py``) or with
 Each can also be prefixed with an `MPI executor <https://www.mpi-forum.org>`__, such as ``mpiexec -n 4 ...`` or ``srun -n 4 ...``, depending on the system.
 
 .. tab-set::
-      
+
    .. tab-item:: Python Script
 
        .. literalinclude:: run_scraping.py
-          :language: python3   
+          :language: python3
           :caption: You can copy this file from ``examples/scraping_beam
 /run_scraping.py``.
 
@@ -56,7 +56,7 @@ Each can also be prefixed with an `MPI executor <https://www.mpi-forum.org>`__,
        .. literalinclude:: input_scraping.in
           :language: ini
           :caption: You can copy this file from ``examples/scraping_beam/input_scraping.in``.
-   
+
 
 Analyze
 -------

diff --git a/examples/scraping_beam/analysis_scraping.py b/examples/scraping_beam/analysis_scraping.py
@@ -31,7 +31,17 @@ def get_moments(beam):
     emittance_y = (sigy**2 * sigpy**2 - epstrms["position_y"]["momentum_y"] ** 2) ** 0.5
     emittance_t = (sigt**2 * sigpt**2 - epstrms["position_t"]["momentum_t"] ** 2) ** 0.5
 
-    return (sigx, sigy, sigt, sigpx, sigpy, sigpt, emittance_x, emittance_y, emittance_t)
+    return (
+        sigx,
+        sigy,
+        sigt,
+        sigpx,
+        sigpy,
+        sigpt,
+        emittance_x,
+        emittance_y,
+        emittance_t,
+    )
 
 
 # initial/final beam
@@ -50,10 +60,12 @@ def get_moments(beam):
 beam_radius = 2.0e-3
 aperture_radius = 3.5e-3
 correlation_k = 0.5
-drift_distance = 6.0   
+drift_distance = 6.0
 
 print("Initial Beam:")
-sigx, sigy, sigt, sigpx, sigpy, sigpt, emittance_x, emittance_y, emittance_t = get_moments(initial)
+sigx, sigy, sigt, sigpx, sigpy, sigpt, emittance_x, emittance_y, emittance_t = (
+    get_moments(initial)
+)
 print(f"  sigx={sigx:e} sigy={sigy:e} sigt={sigt:e}")
 print(f"  sigpx={sigpx:e} sigpy={sigpy:e} sigpt={sigpt:e}")
 print(
@@ -67,16 +79,16 @@ def get_moments(beam):
 assert np.allclose(
     [sigx, sigy, sigt],
     [
-        beam_radius/2.0,
-        beam_radius/2.0,
-        beam_radius/2.0,
+        beam_radius / 2.0,
+        beam_radius / 2.0,
+        beam_radius / 2.0,
     ],
     rtol=rtol,
     atol=atol,
 )
 
 atol = 1.0e-11  # ignored
-print(f"  atol={atol}")  
+print(f"  atol={atol}")
 
 assert np.allclose(
     [sigpx, sigpy, sigpt, emittance_x, emittance_y, emittance_t],
@@ -93,23 +105,25 @@ def get_moments(beam):
 
 
 # calculation of predicted final beam parameters
-beam_radius_no_aperture = beam_radius*(1.0 + correlation_k * drift_distance)
+beam_radius_no_aperture = beam_radius * (1.0 + correlation_k * drift_distance)
 beam_radius_with_aperture = min(beam_radius_no_aperture, aperture_radius)
 
-fractional_loss = 1.0 - min(1.0, (aperture_radius/beam_radius_no_aperture)**2)
-sigma_x_final = beam_radius_with_aperture/2.0
-sigma_px_final = correlation_k/(1.0 + correlation_k * drift_distance) * sigma_x_final
+fractional_loss = 1.0 - min(1.0, (aperture_radius / beam_radius_no_aperture) ** 2)
+sigma_x_final = beam_radius_with_aperture / 2.0
+sigma_px_final = correlation_k / (1.0 + correlation_k * drift_distance) * sigma_x_final
 
 print("")
 print("Predicted Final Beam:")
-print(f"  sigx={sigma_x_final:e} sigy={sigma_x_final:e} sigt={beam_radius/2.0:e}")
+print(f"  sigx={sigma_x_final:e} sigy={sigma_x_final:e} sigt={beam_radius / 2.0:e}")
 print(f"  sigpx={sigma_px_final:e} sigpy={sigma_px_final:e} sigpt=0.0")
 print(f"  fractional_loss={fractional_loss:e}")
 
 
 print("")
 print("Final Beam:")
-sigx, sigy, sigt, sigpx, sigpy, sigpt, emittance_x, emittance_y, emittance_t = get_moments(final)
+sigx, sigy, sigt, sigpx, sigpy, sigpt, emittance_x, emittance_y, emittance_t = (
+    get_moments(final)
+)
 print(f"  sigx={sigx:e} sigy={sigy:e} sigt={sigt:e}")
 print(f"  sigpx={sigpx:e} sigpy={sigpy:e} sigpt={sigpt:e}")
 print(
@@ -125,7 +139,7 @@ def get_moments(beam):
     [
         sigma_x_final,
         sigma_x_final,
-        beam_radius/2.0,
+        beam_radius / 2.0,
         sigma_px_final,
         sigma_px_final,
     ],
@@ -141,7 +155,7 @@ def get_moments(beam):
         0.0,
         0.0,
         0.0,
-    ],      
+    ],
     atol=atol,
 )
 
@@ -156,6 +170,6 @@ def get_moments(beam):
 print(f"  atol={atol}")
 assert np.allclose(
     [loss_pct],
-    [100*fractional_loss],
+    [100 * fractional_loss],
     atol=atol,
 )
diff --git a/examples/scraping_beam/run_scraping.py b/examples/scraping_beam/run_scraping.py
@@ -48,9 +48,9 @@
 
 #   particle bunch
 distr = distribution.Kurth4D(
-    lambdaX=beam_radius/2.0,
-    lambdaY=beam_radius/2.0,
-    lambdaT=beam_radius/2.0,
+    lambdaX=beam_radius / 2.0,
+    lambdaY=beam_radius / 2.0,
+    lambdaT=beam_radius / 2.0,
     lambdaPx=0.0,
     lambdaPy=0.0,
     lambdaPt=0.0,
@@ -63,11 +63,17 @@
 # initialize the linear map
 Iden = elements.LinearMap.Map6x6.identity()
 Rmat = Iden
-Rmat[2,1] = correlation_k
-Rmat[4,3] = correlation_k
+Rmat[2, 1] = correlation_k
+Rmat[4, 3] = correlation_k
 
 # elements
-drift1 = elements.Drift(name="d1", ds=drift_distance, aperture_x = aperture_radius, aperture_y = aperture_radius, nslice=40)
+drift1 = elements.Drift(
+    name="d1",
+    ds=drift_distance,
+    aperture_x=aperture_radius,
+    aperture_y=aperture_radius,
+    nslice=40,
+)
 map1 = elements.LinearMap(R=Rmat)
 
 # design the accelerator lattice