Lint and light reformat of code

Author: szymonlopaciuk

tests/test_beamstrahlung_pic.py

@@ -23,15 +23,15 @@ def test_beambeam3d_beamstrahlung_pic(test_context):
 
     print(repr(test_context))
 
-    # fcc ttbar 4 IP
+    # fcc ttbar 4 IP
     # https://indico.cern.ch/event/1202105/contributions/5408583/attachments/2659051/4608141/FCCWeek_Optics_Oide_230606.pdf
-    bunch_intensity     =  1.55e11  # [e]
+    bunch_intensity     =  1.55e11  # [e]
     energy              =  182.5  # [GeV]
     p0c                 =  182.5e9  # [eV]
     mass0               =  511000  # [eV]
     phi                 =  15e-3  # [rad] half xing
     physemit_x          =  1.59e-09  # [m]
-    physemit_y          =  9e-13  # [m]
+    physemit_y          =  9e-13  # [m]
     beta_x              =  1  # [m]
     beta_y              =  1.6e-3  # [m]
     sigma_x             =  np.sqrt(physemit_x*beta_x)  # [m]
@@ -46,7 +46,7 @@ def test_beambeam3d_beamstrahlung_pic(test_context):
     n_macroparticles_b2 = int(1e6)
 
     ################
-    # create beams #
+    # create beams #
     ################
 
     #e-
@@ -124,7 +124,7 @@ def test_beambeam3d_beamstrahlung_pic(test_context):
     bbpic_ip1_b2.pipeline_manager = pipeline_manager
 
     ######################
-    # set up xtrack line #
+    # set up xtrack line #
     ######################
 
     line_b1 = xt.Line(elements=[bbpic_ip1_b1])
@@ -139,7 +139,7 @@ def test_beambeam3d_beamstrahlung_pic(test_context):
         enable_debug_log=True, verbose=True)
 
     ################
-    # configure BS #
+    # configure BS #
     ################
 
     assert line_b1._needs_rng == False
@@ -156,7 +156,7 @@ def test_beambeam3d_beamstrahlung_pic(test_context):
         xf.BeamBeamPIC3D, capacity={"beamstrahlungtable": int(3e5)})
 
     #####################
-    # track 1 collision #
+    # track 1 collision #
     #####################
 
     multitracker.track(num_turns=1)
@@ -218,7 +218,7 @@ def test_beambeam3d_beamstrahlung_pic(test_context):
     photons_pic_b1 = set(record_b1_pic.beamstrahlungtable.photon_energy)
     photons_pic_b2 = set(record_b2_pic.beamstrahlungtable.photon_energy)
 
-    # compare the upper part of the energy spectrum bc it has better statistics than low end tail
+    # compare the upper part of the energy spectrum bc it has better statistics than low end tail
     hist_pic_b1 = np.histogram(np.array([*photons_pic_b1]), bins=bins)[0][int(n_bins/2):]
     hist_pic_b2 = np.histogram(np.array([*photons_pic_b2]), bins=bins)[0][int(n_bins/2):]
 

xfields/beam_elements/beambeam3dpic.py

@@ -5,15 +5,13 @@
 
 import numpy as np
 from scipy.constants import e as qe
-from scipy.constants import c as clight
 
 import xobjects as xo
 import xtrack as xt
-import pickle as pkl
-
-from ..general import _pkg_root
-from .beambeam3d import _init_alpha_phi
 from xfields import TriLinearInterpolatedFieldMap
+from .beambeam3d import _init_alpha_phi
+from ..general import _pkg_root
+
 
 class BeamstrahlungTable(xo.HybridClass):
     """
@@ -30,7 +28,7 @@ class BeamstrahlungTable(xo.HybridClass):
      particle_id: [1] array index in the xpart.Particles object of the primary macroparticle emitting the photon
      primary_energy: [eV] total energy of primary macroparticle before emission of this photon
      photon_id: [1] counter for photons emitted from the same primary in the same collision with a single slice
-     photon_energy: [eV] total energy of a beamstrahlung photon
+     photon_energy: [eV] total energy of a beamstrahlung photon
      photon_critical_energy (with quantum BS only): [eV] critical energy of a beamstrahlung photon
      rho_inv (with quantum BS only): [m^-1] (Fr/dz) inverse bending radius of the primary macroparticle
     """
@@ -44,7 +42,7 @@ class BeamstrahlungTable(xo.HybridClass):
       'photon_energy': xo.Float64[:],
       'photon_critical_energy': xo.Float64[:],
       'rho_inv': xo.Float64[:],
-        }
+    }
 
 class LumiTable(xo.HybridClass):
     """
@@ -61,24 +59,23 @@ class LumiTable(xo.HybridClass):
      luminosity: [m^-2] integrated luminosity per bunch crossing for one turn, obtained from the charge density grid. Turn by turn lumi. This is overwritten at every call.
     """
     _xofields = {
-      '_index': xt.RecordIndex,
-      'at_turn': xo.Int64[:],
-      'at_element': xo.Int64[:],  # size: n_turns
-      'luminosity': xo.Float64[:],  # size: n_turns
-        }
+        '_index': xt.RecordIndex,
+        'at_turn': xo.Int64[:],
+        'at_element': xo.Int64[:],  # size: n_turns
+        'luminosity': xo.Float64[:],  # size: n_turns
+    }
 
 # currently not possible to have a record table with elements of different size so lumigrid is an attribute
 
 class BeamBeamPIC3DRecord(xo.HybridClass):
     _xofields = {
         'beamstrahlungtable': BeamstrahlungTable,
         'lumitable': LumiTable,
-       }
+    }
 
 class BeamBeamPIC3D(xt.BeamElement):
 
     _xofields = {
-
         '_sin_phi': xo.Float64,
         '_cos_phi': xo.Float64,
         '_tan_phi': xo.Float64,
@@ -112,7 +109,7 @@ class BeamBeamPIC3D(xt.BeamElement):
         # beamstrahlung
         'flag_beamstrahlung': xo.Int64,
 
-        # luminosity
+        # luminosity
         'flag_luminosity': xo.Int64,
         'flag_lumigrid': xo.Int64,
     }
@@ -325,16 +322,16 @@ def track(self, particles):
             pp.x = -pp.x
             self.fieldmap_self.update_from_particles(particles=pp, update_phi=False)
 
-            # for visualizing luminous region, size (2, nt=nz, nx, ny, nz), 0: self, 1: other
+            # for visualizing luminous region, size (2, nt=nz, nx, ny, nz), 0: self, 1: other
             at_turn = pp._xobject.at_turn[0] # On CPU there is always an active particle in position 0
             nx, ny, nz, dx, dy, dz = (self.fieldmap_self.nx, self.fieldmap_self.ny, self.fieldmap_self.nz,
                                       self.fieldmap_self.dx, self.fieldmap_self.dy, self.fieldmap_self.dz)
 
-            # scaling from [C] to [macroparts]
+            # scaling from [C] to [macroparts]
             weight = pp._xobject.weight[0]  # number of elementary charges per macroparticle
-            pwei = (dx*dy*dz) / (qe * weight)  # qe=1.6e-19 [C] from scipy.constants
+            pwei = (dx*dy*dz) / (qe * weight)  # qe=1.6e-19 [C] from scipy.constants
 
-            # other beam: fixed in center of grid, self beam: moves in and out of grid
+            # other beam: fixed in center of grid, self beam: moves in and out of grid
             if self.flag_lumigrid:
 
                 # unit: [1] (macroparticle dist.)
@@ -343,13 +340,13 @@ def track(self, particles):
                              [self.fieldmap_self.rho.flatten(), self.fieldmap_other.rho.flatten()]
                              ) * pwei)
 
-            # 2: kinematic factor, dt(=dz): integral over the time, unit: [m^-2]
+            # 2: kinematic factor, dt(=dz): integral over the time, unit: [m^-2]
             num_macroparts_in_grid_self  = np.abs(np.sum(self.fieldmap_self.rho )*pwei)
             num_macroparts_in_grid_other = np.abs(np.sum(self.fieldmap_other.rho)*pwei)
             self.record.lumitable.luminosity[at_turn] += (dz*2*weight**2 * num_macroparts_in_grid_self * num_macroparts_in_grid_other *
                     self.compute_lumi_integral_3d(self.fieldmap_self.rho, self.fieldmap_other.rho, dx, dy, dz))
 
-            # move self beam back to center of its own grid
+            # move self beam back to center of its own grid
             pp.zeta = -pp.zeta + z_step_other + z_step_self
             pp.x = -pp.x
             self.fieldmap_self.update_from_particles(particles=pp, update_phi=False)
@@ -442,24 +439,27 @@ def compute_lumi_integral_3d(self, lumigrid_my_beam, lumigrid_other_beam, dx, dy
                 h_multiplied[0,    0, nz-1] + h_multiplied[nx-1,    0, nz-1] +
                 h_multiplied[0, ny-1, nz-1] + h_multiplied[nx-1, ny-1, nz-1])
 
-        # interior points
+        # interior points
         secondPart = self._buffer.context.nplike_lib.sum(
                 h_multiplied[1:nx-1, 1:ny-1, 1:nz-1])
 
-        # x boundaries
+        # x boundaries
         thirdPart = self._buffer.context.nplike_lib.sum(
                 h_multiplied[1:nx-1, 0, 1:nz-1] + h_multiplied[1:nx-1, ny-1, 1:nz-1])
 
-        # y boundaries
+        # y boundaries
         fourthPart = self._buffer.context.nplike_lib.sum(
                 h_multiplied[0, 1:ny-1, 1:nz-1] + h_multiplied[nx-1, 1:ny-1, 1:nz-1])
 
-        # z boundaries
+        # z boundaries
         fifthPart = self._buffer.context.nplike_lib.sum(
                 h_multiplied[1:nx-1, 1:ny-1, 0] + h_multiplied[1:nx-1, 1:ny-1, nz-1])
 
-        # 3D trapezoid integral
+        # 3D trapezoid integral
         integralf = integral + 0.125 * dx * dy * dz * (
                 8 * secondPart + 4 * thirdPart + 4 * fourthPart + 4 * fifthPart)
 
-        return integralf if not self._buffer.context.nplike_lib.isnan(integralf) else 0
+        if self._buffer.context.nplike_lib.isnan(integralf):
+            return 0
+
+        return integralf