Skip to content

Commit

Permalink
fixed some PEP warnings in PFS module
Browse files Browse the repository at this point in the history
  • Loading branch information
sergey-tomin committed Nov 6, 2020
1 parent ed68826 commit a6d9d3a
Show file tree
Hide file tree
Showing 10 changed files with 108 additions and 116 deletions.
2 changes: 1 addition & 1 deletion demos/ipython_tutorials/1_introduction.ipynb
Original file line number Diff line number Diff line change
Expand Up @@ -101,7 +101,7 @@
"* [Some useful OCELOT functions](small_useful_features.ipynb). [Web version](http://nbviewer.jupyter.org/github/ocelot-collab/ocelot/blob/dev/demos/ipython_tutorials/small_useful_features.ipynb)\n",
" - Aperture, RK tracking \n",
"* [Example of an accelerator section optimization](accelerator_optim.ipynb). [Web version](https://nbviewer.jupyter.org/github/ocelot-collab/ocelot/blob/master/demos/ipython_tutorials/accelerator_optim.ipynb)\n",
" - A simple demo of accelerator section optimization with a stadard scipy numerical optimization method. "
" - A simple demo of accelerator section optimization with a standard scipy numerical optimization method. "
]
},
{
Expand Down
2 changes: 1 addition & 1 deletion ocelot/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@
general ocelot description
"""

__version__ = '20.11.0'
__version__ = '20.11.2'


__all__ = ['Twiss', "Beam", "Particle", "get_current", "get_envelope", "generate_parray", # beam
Expand Down
2 changes: 1 addition & 1 deletion ocelot/adaptors/genesis.py
Original file line number Diff line number Diff line change
Expand Up @@ -609,7 +609,7 @@ def calc_spec(self, mode='mid', npad=0):
if self.nSlices == 1:
raise AssertionError('Cannot calculate spectrum from steady-state simulation')

if (npad%1 is not 0) or npad < 0:
if (npad%1 != 0) or npad < 0:
raise ValueError('npad should be positive integer')

if mode == 'ff':
Expand Down
2 changes: 1 addition & 1 deletion ocelot/common/math_op.py
Original file line number Diff line number Diff line change
Expand Up @@ -398,7 +398,7 @@ def find_nearest(array, value):

def n_moment(x, counts, c, n):
x = np.squeeze(x)
if x.ndim is not 1:
if x.ndim != 1:
raise ValueError("scale of x should be 1-dimensional")
if x.size not in counts.shape:
raise ValueError("operands could not be broadcast together with shapes %s %s" %(str(x.shape), str(counts.shape)))
Expand Down
64 changes: 29 additions & 35 deletions ocelot/gui/dfl_plot.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@
import time
import matplotlib
import matplotlib.pyplot as plt
import matplotlib.colors as colors # for wigner log scale
import matplotlib.colors as colors # for wigner log scale
import numpy as np
import logging

Expand All @@ -21,8 +21,6 @@
from ocelot.optics.utils import calc_ph_sp_dens
from ocelot.optics.wave import *



from ocelot.gui.colormaps2d.colormap2d import *

# in order to run decorators properly
Expand Down Expand Up @@ -250,7 +248,7 @@ def plot_dfl(dfl, domains=None, z_lim=[], xy_lim=[], figsize=4, cmap=def_cmap, l
suffix += '_log'

if fig_name is None:
if dfl_copy.fileName() is '':
if dfl_copy.fileName() == '':
fig = plt.figure('Radiation distribution' + suffix)
else:
fig = plt.figure('Radiation distribution' + suffix + ' ' + dfl_copy.fileName())
Expand Down Expand Up @@ -558,7 +556,7 @@ def plot_wigner(wig_or_out, z=np.inf, x_units='um', y_units='ev', x_lim=(None, N
raise ValueError('Unknown object for Wigner plot')

if fig_name is None:
if W.fileName() is '':
if W.fileName() == '':
fig_text = 'Wigner distribution'
else:
fig_text = 'Wigner distribution ' + W.fileName()
Expand Down Expand Up @@ -623,24 +621,24 @@ def plot_wigner(wig_or_out, z=np.inf, x_units='um', y_units='ev', x_lim=(None, N

# cmap='RdBu_r'
# axScatter.imshow(wigner, cmap=cmap, vmax=wigner_lim, vmin=-wigner_lim)

if log_scale != 0:
if log_scale==1:
log_scale=0.01
if log_scale == 1:
log_scale = 0.01
wigplot = axScatter.pcolormesh(power_scale[::downsample], spec_scale[::downsample],
wigner[::downsample, ::downsample], cmap=cmap,
norm=colors.SymLogNorm(linthresh=wigner_lim * log_scale, linscale=2,
vmin=-wigner_lim, vmax=wigner_lim),
vmax=wigner_lim, vmin=-wigner_lim)
wigner[::downsample, ::downsample], cmap=cmap,
norm=colors.SymLogNorm(linthresh=wigner_lim * log_scale, linscale=2,
vmin=-wigner_lim, vmax=wigner_lim),
vmax=wigner_lim, vmin=-wigner_lim)
else:
wigplot = axScatter.pcolormesh(power_scale[::downsample], spec_scale[::downsample],
wigner[::downsample, ::downsample], cmap=cmap, vmax=wigner_lim, vmin=-wigner_lim)
wigner[::downsample, ::downsample], cmap=cmap, vmax=wigner_lim, vmin=-wigner_lim)

if plot_cbar:
from mpl_toolkits.axes_grid1.inset_locator import inset_axes
cbaxes = inset_axes(axScatter, width="50%", height="3%", loc=1)
fig.colorbar(wigplot, cax = cbaxes, orientation='horizontal')
cbaxes = inset_axes(axScatter, width="50%", height="3%", loc=1)
fig.colorbar(wigplot, cax=cbaxes, orientation='horizontal')

if plot_text:
if hasattr(wig_or_out, 'is_spectrogram'):
if wig_or_out.is_spectrogram:
Expand Down Expand Up @@ -736,23 +734,20 @@ def plot_wigner(wig_or_out, z=np.inf, x_units='um', y_units='ev', x_lim=(None, N

axHistx.set_xlim(x_lim_appl[0], x_lim_appl[1])
axHisty.set_ylim(y_lim_appl[0], y_lim_appl[1])

if log_scale != 0:
axHistx.set_ylim(np.nanmin(power), np.nanmax(power))
axHisty.set_xlim(np.nanmin(spec), np.nanmax(spec))
axHisty.set_xscale('log')
axHistx.set_yscale('log')

else:
axScatter.axis('tight')
axScatter.set_xlabel(p_label_txt)
axScatter.set_ylabel(f_label_txt)

# axScatter.set_xlim(x_lim[0], x_lim[1])
# axScatter.set_ylim(y_lim[0], y_lim[1])




if savefig != False:
if savefig == True:
Expand Down Expand Up @@ -782,15 +777,15 @@ def plot_dfl_waistscan(sc_res, fig_name=None, figsize=4, showfig=True, savefig=F
_logger.info('plot dfl waist scan')
if showfig == False and savefig == False:
return

if fig_name is None:
if sc_res.fileName() is '':
if sc_res.fileName() == '':
fig = plt.figure('Waist scan')
else:
fig = plt.figure(sc_res.fileName() + ' waist scan')
else:
fig = plt.figure(fig_name)

plt.clf()
fig.set_size_inches((3 * figsize, 2 * figsize), forward=True)
ax_int = fig.add_subplot(1, 1, 1)
Expand Down Expand Up @@ -952,11 +947,11 @@ def plot_stokes_values(S, fig=None, d_pol=0, norm=0, showfig=True, gw=1, directi
if np.size(sc) <= 1:
_logger.warning('plot_stokes_values needs more than a single point to plot (np.size(sc) <= 1)')
return

if d_pol != 0:
gw=0
norm=1
gw = 0
norm = 1

if fig == None:
plt.figure('Stokes S')
plt.clf()
Expand All @@ -970,7 +965,7 @@ def plot_stokes_values(S, fig=None, d_pol=0, norm=0, showfig=True, gw=1, directi
mult = 1e-9
plt.ylabel('$S_0$ [GW]')
elif norm:
mult = 1/np.amax(Scp.s0)
mult = 1 / np.amax(Scp.s0)
else:
mult = 1
plt.ylabel('$S_0$ [W]')
Expand All @@ -986,10 +981,10 @@ def plot_stokes_values(S, fig=None, d_pol=0, norm=0, showfig=True, gw=1, directi
else:
raise ValueError

if d_pol=='lin':
if d_pol == 'lin':
# plt.step(sc, np.sqrt(S.s1**2+S.s2**2), linewidth=2, where='mid',color=[0.5,0.5,0.5], linestyle='--')
plot_function(sc, Scp.deg_pol_l(), linestyle='-', color='#1f77b4', **kwargs)
elif d_pol==1:
elif d_pol == 1:
plot_function(sc, Scp.deg_pol(), linestyle='-', color='#1f77b4', **kwargs)
else:
pass
Expand Down Expand Up @@ -1041,8 +1036,7 @@ def plot_stokes_angles(S, fig=None, showfig=True, direction='z', plot_func='scat
if np.size(sc) <= 1:
_logger.warning('plot_stokes_angles needs more than a single point to plot (np.size(sc) <= 1)')
return



if fig == None:
plt.figure('Stokes angles')
plt.clf()
Expand All @@ -1067,7 +1061,7 @@ def plot_stokes_angles(S, fig=None, showfig=True, direction='z', plot_func='scat

plot_function(sc, Scp.psi(), color='b', **kwargs)
# if plot_func == 'scatter':
# kwargs['s'] = psize
# kwargs['s'] = psize
plot_function(sc, Scp.chi(), color='g', **kwargs)

# if scatter:
Expand Down
8 changes: 4 additions & 4 deletions ocelot/gui/genesis4_plot.py
Original file line number Diff line number Diff line change
Expand Up @@ -280,7 +280,7 @@ def plot_gen4_out_z(out, z=np.inf, params=['rad_power+el_current', 'el_energy+el
# z = out.z[zi]

if fig_name is None:
if out.fileName() is '':
if out.fileName() == '':
fig = plt.figure('Bunch profile at ' + str(z) + 'm')
else:
fig = plt.figure('Bunch profile at ' + str(z) + 'm ' + out.fileName())
Expand Down Expand Up @@ -348,13 +348,13 @@ def plot_gen4_out_z(out, z=np.inf, params=['rad_power+el_current', 'el_energy+el
axf = len(ax) - axt
# ax[0].set_xlim(out.z[0], out.z[-1])
# ax[-1].set_xlabel('z [m]')
if axt is not 0 and axf is not 0:
if axt != 0 and axf != 0:
fig.subplots_adjust(top=0.95, bottom=0.2, right=0.8, left=0.15)
else:
fig.subplots_adjust(top=0.95, bottom=0.1, right=0.8, left=0.15)

for axi in ax[axt:]:
if axt is not 0:
if axt != 0:
pos1 = axi.get_position() # get the original position
pos2 = [pos1.x0 + 0, pos1.y0 - 0.1, pos1.width / 1.0, pos1.height / 1.0]
axi.set_position(pos2)
Expand Down Expand Up @@ -661,7 +661,7 @@ def plot_gen4_out_evo(out, params=['und_quad', 'el_size', 'el_pos', 'el_energy',
out = read_out_file(out, read_level=2)
# add check for output object
if fig_name is None:
if out.fileName() is '':
if out.fileName() == '':
fig = plt.figure(params_str)
_logger.info('plotting ' + params_str)
else:
Expand Down
Loading

0 comments on commit a6d9d3a

Please sign in to comment.