fix format

docs/source/run/parameters.rst

@@ -889,19 +889,16 @@ Option: ``gaussian``
     Spatial chirp at focus defined by `S. Akturk et al., Optics Express 12, 4399 (2004) <https://doi.org/10.1364/OPEX.12.004399>`__.
 
 * ``<laser name>.phi2`` (`float`) optional (default `pi/2`)
-    The amount of temporal chirp :math:`\phi^{(2)}` at focus (in the lab frame). Namely, a wave packet
-    centered on the frequency :math:`(\omega_0 + \delta \omega)` will reach its peak intensity
-    at :math:`z(\delta \omega) = z_0 - c \phi^{(2)} \, \delta \omega`. Thus, a positive
-    :math:`\phi^{(2)}` corresponds to positive chirp, i.e. red part of the spectrum in the
-    front of the pulse and blue part of the spectrum in the back. More specifically, the electric
-    field in the focal plane is of the form:
+    The amount of temporal chirp :math:`\phi^{(2)}` at focus (in the lab frame).
+    Namely, a wave packetcentered on the frequency :math:`(\omega_0 + \delta \omega)` will reach its peak intensity at :math:`z(\delta \omega) = z_0 - c \phi^{(2)} \, \delta \omega`.
+    Thus, a positive :math:`\phi^{(2)}` corresponds to positive chirp, i.e. red part of the spectrum in the front of the pulse and blue part of the spectrum in the back.
+    More specifically, the electric field in the focal plane is of the form:
 
     .. math::
 
         E(\boldsymbol{x},t) \propto Re\left[ \exp\left(  -\frac{(t-t_{peak})^2}{\tau^2 + 2i\phi^{(2)}} + i\omega_0 (t-t_{peak}) + i\phi_0 \right) \right]
 
-    where :math:`\tau` is given by ``<laser_name>.tau`` and represents the
-    Fourier-limited duration of the laser pulse. Thus, the actual duration of the chirped laser pulse is:
+    where :math:`\tau` is given by ``<laser_name>.tau`` and represents the Fourier-limited duration of the laser pulse. Thus, the actual duration of the chirped laser pulse is:
 
     .. math::
 

examples/laser/analysis_laser_init_chirp.py

@@ -1,76 +1,77 @@
-#! /usr/bin/env python3
+#!/usr/bin/env python3
 
-# Copyright 2024
-#
-# This file is part of HiPACE++.
-#
-# Authors: Xingjian Hui
-# License: BSD-3-Clause-LBNL
-
-
-import matplotlib.pyplot as plt
-import matplotlib
-import argparse
 import numpy as np
 import scipy.constants as scc
-import scipy
 from openpmd_viewer.addons import LpaDiagnostics
+import argparse
+
+# Constants
+C = scc.c  # Speed of light (precomputing constant)
+
 
-def get_zeta(Ar,m, w0,L):
-    nu = 0
-    sum=0
-    laser_module=np.abs(Ar)
-    phi_envelop=np.array(np.arctan2(Ar.imag, Ar.real))
-    #unwrap phi_envelop
-    phi_envelop = np.unwrap(np.unwrap(phi_envelop, axis=0),axis=1)
-    #calculate pphi_pz/
+def get_zeta(Ar, m, w0, L):
+    laser_module = np.abs(Ar)
+    phi_envelope = np.unwrap(np.angle(Ar), axis=1)
+
+    # Vectorized operations to calculate pphi_pz and pphi_pzpy
     z_diff = np.diff(m.z)
     x_diff = np.diff(m.x)
-    pphi_pz = (np.diff(phi_envelop, axis=0)).T/ (z_diff/scc.c)
-    pphi_pzpy = ((np.diff(pphi_pz, axis=0)).T/(x_diff))
-    for i in range(len(m.z)-2):
-        for j in range(len(m.x)-2):
-            nu=nu+pphi_pzpy[i,j]*laser_module[i,j]
-            sum=sum+laser_module[i,j]
-    nu= nu / scc.c / sum
+
+    pphi_pz = np.diff(phi_envelope, axis=0) / (z_diff[:, None] / C)
+    pphi_pzpy = np.diff(pphi_pz, axis=1) / x_diff
+
+    # Vectorized summation
+    weighted_sum = np.sum(pphi_pzpy * laser_module[:-2, :-2])
+    total_sum = np.sum(laser_module[:-2, :-2])
+
+    nu = weighted_sum / C / total_sum
     a = 4 * nu * w0**2 * L**4
-    b = -4 * scc.c
+    b = -4 * C
     c = nu * w0**2 * L**2
+
     zeta_solutions = np.roots([a, b, c])
     return np.min(zeta_solutions)
 
-def get_phi2 (Ar,m,tau):
-    #get temporal chirp phi2
-    temp_chirp = 0
-    sum=0
-    laser_module1=np.abs(Ar)
-    phi_envelop=np.unwrap(np.array(np.arctan2(Ar.imag, Ar.real)),axis=0)
-    #calculate pphi_pz
+
+def get_phi2(Ar, m, tau):
+    laser_module = np.abs(Ar)
+    phi_envelope = np.unwrap(np.angle(Ar), axis=0)
+
     z_diff = np.diff(m.z)
-    pphi_pz = (np.diff(phi_envelop, axis=0)).T/ (z_diff/scc.c)
-    pphi_pz2 = ((np.diff(pphi_pz, axis=1))/(z_diff[:len(z_diff)-1])/scc.c).T
-    for i in range(len(m.z)-2):
-        for j in range(len(m.x)-2):
-            temp_chirp=temp_chirp+pphi_pz2[i,j]*laser_module1[i,j]
-            sum=sum+laser_module1[i,j]
-    x=temp_chirp*scc.c**2/sum
-    a=4*x
-    b=-4
-    c=tau**4*x
+
+    # Vectorized operations to calculate pphi_pz and pphi_pz2
+    pphi_pz = np.diff(phi_envelope, axis=0) / (z_diff[:, None] / C)
+    pphi_pz2 = np.diff(pphi_pz, axis=1) / (z_diff[:-1][:, None] / C)
+
+    # Vectorized summation
+    temp_chirp = np.sum(pphi_pz2 * laser_module[:-2, :-2])
+    total_sum = np.sum(laser_module[:-2, :-2])
+
+    x = temp_chirp * C**2 / total_sum
+    a = 4 * x
+    b = -4
+    c = tau**4 * x
+
     zeta_solutions = np.roots([a, b, c])
     return np.max(zeta_solutions)
 
-def get_centroids (F, x, z):
-    index_array = np.mgrid[0:F.shape[0],0:F.shape[1]][1]
-    centroids = np.sum(index_array * np.abs(F**2), axis=1)/np.sum(np.abs(F**2),axis=1)
+
+def get_centroids(F, x, z):
+    # Vectorized computation of centroids
+    index_array = np.arange(F.shape[1])[None, :]  # Faster with broadcasting
+    centroids = np.sum(index_array * np.abs(F)**2, axis=1) / np.sum(np.abs(F)**2, axis=1)
     return z[centroids.astype(int)]
 
-def get_beta (F,m):
-    k0 = 2 *scc.pi / 800e-9
+
+def get_beta(F, m, k0):
     z_centroids = get_centroids(F.T, m.x, m.z)
-    weight = np.mean(np.abs(F.T)**2,axis=np.ndim(F)-1)
-    derivative = np.gradient(z_centroids) / ( m.x[1] - m.x[0] )
-    return (np.sum(derivative * weight) / np.sum(weight))/k0/scc.c
+    weight = np.mean(np.abs(F.T)**2, axis=np.ndim(F) - 1)
+
+    # Vectorized derivative calculation
+    derivative = np.gradient(z_centroids) / (m.x[1] - m.x[0])
+
+    return np.sum(derivative * weight) / np.sum(weight) / k0 / C
+
 
 parser = argparse.ArgumentParser(description='Verify the chirp initialization')
 parser.add_argument('--output-dir',
@@ -80,22 +81,23 @@ def get_beta (F,m):
 parser.add_argument('--chirp_type',
                     dest='chirp_type',
                     default='phi2',
-                    help='the type of the initialised chirp')
+                    help='The type of the initialized chirp')
 args = parser.parse_args()
 
 ts = LpaDiagnostics(args.output_dir)
 
 Ar, m = ts.get_field(field='laserEnvelope', iteration=0)
-lambda0=.8e-6            # Laser wavelength
-w0 = 30.e-6              # Laser waist
+lambda0 = 0.8e-6  # Laser wavelength
+w0 = 30.e-6       # Laser waist
 L0 = 5e-6
-tau = L0 / scc.c         # Laser duration
-if args.chirp_type == 'phi2' :
+tau = L0 / C  # Laser duration
+k0 = 2 * np.pi / 800e-9
+if args.chirp_type == 'phi2':
     phi2 = get_phi2(Ar, m, tau)
-    assert(np.abs(phi2-2.4e-26)/2.4e-26 < 1e-2)
-elif args.chirp_type == 'zeta' :
+    assert np.abs(phi2 - 2.4e-26) / 2.4e-26 < 1e-2
+elif args.chirp_type == 'zeta':
     zeta = get_zeta(Ar, m, w0, L0)
-    assert(np.abs(zeta-2.4e-19)/2.4e-19 < 1e-2)
-elif args.chirp_type == 'beta' :
-    beta = get_beta(Ar, m)
-    assert(np.abs(beta-2e-17)/2e-17 < 1e-2)
+    assert np.abs(zeta - 2.4e-19) / 2.4e-19 < 1e-2
+elif args.chirp_type == 'beta':
+    beta = get_beta(Ar, m, k0)
+    assert np.abs(beta - 2e-17) / 2e-17 < 1e-2