style: apply RUFFs style suggestions

Author: elenafuengar

wakis/routines.py

@@ -6,15 +6,13 @@
 import numpy as np
 import h5py
 from tqdm import tqdm
-from scipy.constants import c as c_light
 from wakis.sources import Beam
-from wakis.field_monitors import FieldMonitor
 
 class RoutinesMixin():
 
     def emsolve(self, Nt, source=None, callback=None,
-                save=False, fields=['E'], components=['Abs'], save_every=1, subdomain=None, 
-                plot=False, plot_every=1, use_etd=False, 
+                save=False, fields=['E'], components=['Abs'], save_every=1, subdomain=None,
+                plot=False, plot_every=1, use_etd=False,
                 plot3d=False, **kwargs):
         '''
         Run the simulation and save the selected field components in HDF5 files
@@ -26,13 +24,13 @@ def emsolve(self, Nt, source=None, callback=None,
         Nt: int
             Number of timesteps to run
         source: source object
-            source object from `sources.py` defining the time-dependednt source. 
+            source object from `sources.py` defining the time-dependednt source.
             It should have an update function `source.update(solver, t)`
         save: bool
             Flag to enable saving the field in HDF5 format
         fields: list, default ['E']
             3D field magnitude ('E', 'H', or 'J') to save
-            'Ex', 'Hy', etc., is also accepted and will override 
+            'Ex', 'Hy', etc., is also accepted and will override
             the `components` parameter.
         components: list, default ['z']
             Field compoonent ('x', 'y', 'z', 'Abs') to save. It will be overriden
@@ -49,10 +47,10 @@ def emsolve(self, Nt, source=None, callback=None,
             Number of timesteps between consecutive plots
         **kwargs:
             Keyword arguments to be passed to the Plot2D function.
-            * Default kwargs used for 2D plotting: 
+            * Default kwargs used for 2D plotting:
                 {'field':'E', 'component':'z',
-                'plane':'ZY', 'pos':0.5, 'title':'Ez', 
-                'cmap':'rainbow', 'patch_reverse':True, 
+                'plane':'ZY', 'pos':0.5, 'title':'Ez',
+                'cmap':'rainbow', 'patch_reverse':True,
                 'off_screen': True, 'interpolation':'spline36'}
             * Default kwargs used for 3D plotting:
                 {'field':'E', 'component':'z',
@@ -71,8 +69,9 @@ def emsolve(self, Nt, source=None, callback=None,
         h5py
         '''
         self.Nt = Nt
-        if source is not None: self.source = source
-        
+        if source is not None:
+            self.source = source
+
         if save:
 
             hfs = {}
@@ -94,8 +93,8 @@ def emsolve(self, Nt, source=None, callback=None,
 
         if plot:
             plotkw = {'field':'E', 'component':'z',
-                    'plane':'ZY', 'pos':0.5, 'cmap':'rainbow', 
-                    'patch_reverse':True, 'title':'Ez', 
+                    'plane':'ZY', 'pos':0.5, 'cmap':'rainbow',
+                    'patch_reverse':True, 'title':'Ez',
                     'off_screen': True, 'interpolation':'spline36'}
             plotkw.update(kwargs)
 
@@ -105,28 +104,28 @@ def emsolve(self, Nt, source=None, callback=None,
                     'title':'Ez', 'cmap':'jet', 'clip_volume':False, 'clip_normal':'-y',
                     'field_on_stl':True, 'field_opacity':1.0,
                     'off_screen':True, 'zoom':1.0, 'nan_opacity':1.0}
-            
+
             plotkw.update(kwargs)
 
         # get ABC values
         if self.activate_abc:
             E_abc_2, H_abc_2 = self.get_abc()
             E_abc_1, H_abc_1 = self.get_abc()
 
-        # Time loop 
+        # Time loop
         for n in tqdm(range(Nt)):
 
-            if source is not None: 
+            if source is not None:
                 source.update(self, n*self.dt)
 
             if save and n%save_every == 0:
                 for field in hfs.keys():
                     try:
                         d = getattr(self, field[0])[xx,yy,zz,field[1:]]
-                    except:
+                    except AttributeError:
                         raise(f'Component {field} not valid. Input must have a \
                               field ["E", "H", "J"] and a component ["x", "y", "z", "Abs"]')
-                    
+
                     # Save timestep in HDF5
                     hfs[field]['#'+str(n).zfill(5)] = d
 
@@ -155,38 +154,38 @@ def emsolve(self, Nt, source=None, callback=None,
             for hf in hfs:
                 hf.close()
 
-    def wakesolve(self, wakelength, 
-                  wake=None, 
+    def wakesolve(self, wakelength,
+                  wake=None,
                   callback=None,
-                  compute_plane='both', 
+                  compute_plane='both',
                   plot=False, plot_from=None, plot_every=1, plot_until=None,
-                  save_J=False, 
+                  save_J=False,
                   add_space=None, #for legacy
                   use_field_monitor=False,
                   field_monitor=None,
-                  use_edt=None, #deprecated 
+                  use_edt=None, #deprecated
                   **kwargs):
         '''
         Run the EM simulation and compute the longitudinal (z) and transverse (x,y)
-        wake potential WP(s) and impedance Z(s). 
-        
-        The `Ez` field is saved every timestep in a subdomain (xtest, ytest, z) around 
+        wake potential WP(s) and impedance Z(s).
+
+        The `Ez` field is saved every timestep in a subdomain (xtest, ytest, z) around
         the beam trajectory in HDF5 format file `Ez.h5`.
 
-        The computed results are available as Solver class attributes: 
+        The computed results are available as Solver class attributes:
             - wake potential: WP (longitudinal), WPx, WPy (transverse) [V/pC]
             - impedance: Z (longitudinal), Zx, Zy (transverse) [Ohm]
             - beam charge distribution: lambdas (distance) [C/m] lambdaf (spectrum) [C]
 
         Parameters:
         -----------
         wakelength: float
-            Desired length of the wake in [m] to be computed 
-            
+            Desired length of the wake in [m] to be computed
+
             Maximum simulation time in [s] can be computed from the wakelength parameter as:
-            .. math::    t_{max} = t_{inj} + (wakelength + (z_{max}-z_{min}))/c 
+            .. math::    t_{max} = t_{inj} + (wakelength + (z_{max}-z_{min}))/c
         wake: Wake obj, default None
-            `Wake()` object containing the information needed to run 
+            `Wake()` object containing the information needed to run
             the wake solver calculation. See Wake() docstring for more information.
             Can be passed at `Solver()` instantiation as parameter too.
         save_J: bool, default False
@@ -201,11 +200,11 @@ def wakesolve(self, wakelength,
             Number of timesteps between consecutive plots
         **kwargs:
             Keyword arguments to be passed to the Plot2D function.
-            Default kwargs used: 
-                {'plane':'ZY', 'pos':0.5, 'title':'Ez', 
-                'cmap':'rainbow', 'patch_reverse':True, 
+            Default kwargs used:
+                {'plane':'ZY', 'pos':0.5, 'title':'Ez',
+                'cmap':'rainbow', 'patch_reverse':True,
                 'off_screen': True, 'interpolation':'spline36'}
-        
+
         Raises:
         -------
         AttributeError:
@@ -218,7 +217,8 @@ def wakesolve(self, wakelength,
         h5py
         '''
 
-        if wake is not None: self.wake = wake
+        if wake is not None:
+            self.wake = wake
         if self.wake is None:
             raise('Wake solver information not passed to the solver instantiation')
 
@@ -228,26 +228,26 @@ def wakesolve(self, wakelength,
         # plot params defaults
         if plot:
             plotkw = {'plane':'ZY', 'pos':0.5, 'title':'Ez',
-                    'cmap':'rainbow', 'patch_reverse':True,  
+                    'cmap':'rainbow', 'patch_reverse':True,
                     'off_screen': True, 'interpolation':'spline36'}
             plotkw.update(kwargs)
-        
+
         # integration path (test position)
         self.xtest, self.ytest = self.wake.xtest, self.wake.ytest
         self.ixt, self.iyt = np.abs(self.x-self.xtest).argmin(), np.abs(self.y-self.ytest).argmin()
         if compute_plane.lower() == 'longitudinal':
             xx, yy = self.ixt, self.iyt
         else:
-            xx, yy = slice(self.ixt-1, self.ixt+2), slice(self.iyt-1, self.iyt+2)            
+            xx, yy = slice(self.ixt-1, self.ixt+2), slice(self.iyt-1, self.iyt+2)
 
-        # Compute simulation time 
+        # Compute simulation time
         self.wake.wakelength = wakelength
         self.ti = self.wake.ti
         self.v = self.wake.v
         if self.use_mpi:  #E- should it be zmin, zmax instead?
             z = self.Z # use global coords
             zz = slice(0, self.NZ)
-        else: 
+        else:
             z = self.z
             zz = slice(0, self.Nz)
 
@@ -256,13 +256,13 @@ def wakesolve(self, wakelength,
         self.tmax, self.Nt = tmax, Nt
 
         # Add beam source
-        beam = Beam(q=self.wake.q, 
-                    sigmaz=self.wake.sigmaz, 
+        beam = Beam(q=self.wake.q,
+                    sigmaz=self.wake.sigmaz,
                     beta=self.wake.beta,
-                    xsource=self.wake.xsource, ysource=self.wake.ysource, 
+                    xsource=self.wake.xsource, ysource=self.wake.ysource,
                     )
-        
-        # hdf5 
+
+        # hdf5
         self.Ez_file = self.wake.Ez_file
         hf = None  # needed for MPI
         if self.use_mpi:
@@ -291,10 +291,12 @@ def save_to_h5(self, hf, field, x, y, z, comp, n):
                 if self.rank == 0:
                     hf['#'+str(n).zfill(5)] = _field
             else:
-                hf['#'+str(n).zfill(5)] = getattr(self, field)[x, y, z, comp] 
-        
-        if plot_until is None: plot_until = Nt
-        if plot_from is None: plot_from = int(self.ti/self.dt)
+                hf['#'+str(n).zfill(5)] = getattr(self, field)[x, y, z, comp]
+
+        if plot_until is None:
+            plot_until = Nt
+        if plot_from is None:
+            plot_from = int(self.ti/self.dt)
 
         print('Running electromagnetic time-domain simulation...')
         for n in tqdm(range(Nt)):
@@ -303,10 +305,10 @@ def save_to_h5(self, hf, field, x, y, z, comp, n):
             beam.update(self, n*self.dt)
 
             # Save
-            save_to_h5(self, hf, 'E', xx, yy, zz, 'z', n) 
+            save_to_h5(self, hf, 'E', xx, yy, zz, 'z', n)
             if save_J:
-                save_to_h5(self, hfJ, 'J', xx, yy, zz, 'z', n) 
-            
+                save_to_h5(self, hfJ, 'J', xx, yy, zz, 'z', n)
+
             # Advance
             self.one_step()
 
@@ -324,7 +326,7 @@ def save_to_h5(self, hf, field, x, y, z, comp, n):
             if callback is not None:
                 callback(self, n*self.dt)
 
-        # End of time loop       
+        # End of time loop
         if self.use_mpi:
             if self.rank==0:
                 hf.close()
@@ -340,4 +342,4 @@ def save_to_h5(self, hf, field, x, y, z, comp, n):
 
             # Compute wakefield magnitudes is done inside WakeSolver
             self.wake.solve(compute_plane=compute_plane)
-            
+