#
# xoppy_calc_tcpy: pure-python reimplementation of the TC code
# (on-axis brilliance tuning curves of an ideal undulator)
#
# Reproduces the method and the results of TC v1.97 by R.J. Dejus (APS, ANL),
# the binary distributed with xoppylib,
# as driven by xoppylib.xoppy_run_binaries.xoppy_calc_xtc, without calling
# the external binary. Implemented method: "infinite-N with convolution"
# (Dejus' method, METHOD=1 of XTC), for planar (HELICAL=0) and helical
# (HELICAL=1) devices.
#
# Physics summary
# ---------------
# 1) The angular spectral flux of harmonic i of an ideal undulator is computed
# in the Bessel function approximation: at observation angle (theta, phi),
# with alpha = gamma*theta, a = 1 + (Kx^2+Ky^2)/2 + alpha^2, the (complex)
# field amplitudes are expansions in products of Bessel functions
# Jn(y)*J_{i+2n+m}(x) (m = -1, 0, +1) -- see _s0_harmonic() below.
# 2) In the infinite-N limit, photons of energy E at harmonic i are emitted
# on the cone alpha_i^2 = i*ER/E - (1 + K^2/2), ER = 2*gamma^2*(hc/lambda0).
# The convolution of the angular flux density with the Gaussian angular
# distribution of the electron beam, evaluated on axis, reduces to a 1D
# integral over the azimuth phi along that cone.
# 3) The finite number of periods is restored by convolving the spectrum
# with the sinc^2 line shape of width DEW = E1/N.
# 4) The result is divided by the effective source area
# 2*pi*sqrt((sigx^2+sigr^2)*(sigy^2+sigr^2)), where sigr is the
# detuning-dependent effective source size (Walker's parametrization),
# giving the on-axis brilliance in ph/s/mrad^2/mm^2/0.1%bw.
# 5) The beam energy spread is applied as a Gaussian convolution (rms width
# 2*sige*E) and, for each K of the scan, the peak of the resulting curve
# near harmonic i is stored: this builds the tuning curve.
#
import numpy
import scipy
# ----------------------------------------------------------------------
# physical constants (updated from values as in tc.f/usb.f, Physics Today Aug. 1990)
# ----------------------------------------------------------------------
_C_LIGHT = scipy.constants.c # 2.99792458e8 # speed of light [m/s]
_ME = scipy.constants.m_e # 9.1093897e-31 # electron rest mass [kg]
_EC = scipy.constants.e # 1.60217733e-19 # elementary charge [C]
_MEE = _ME * _C_LIGHT ** 2 / _EC * 1e-6 # 0.51099906 # electron rest mass [MeV]
_HBAR = scipy.constants.hbar # 1.05457266e-34 # hbar [Js]
_H = scipy.constants.h # 6.6260755e-34 # Planck constant [Js]
_EPSZ = scipy.constants.epsilon_0 # 8.854187817e-12 # permittivity of vacuum [F/m]
_PI = numpy.pi
_BW = 1.0e-3 # 0.1% bandwidth
_FSC = _EC * _EC / (4.0 * _PI * _EPSZ * _HBAR * _C_LIGHT) # fine structure
_C_EVANG = _H * _C_LIGHT / _EC * 1e10 # 12398.42
_PTOT_FAC = _PI / 3.0 * _EC / _EPSZ / _MEE**2 * 1e6 # 0.07257
_PD_FAC = 21.0 / (16.0 * _PI * _MEE**2) * _PTOT_FAC # 0.11611
_CV = 1.0 / (8.0 * _PI * _PI) * 1e-10 * 1e6 # 1.2665e-6
# numerical parameters (as in usb.f/tc.f)
_NPHI = 20 # azimuthal points in (0, pi/2)
_NSIGMA = 4 # angular integration extent of the e-beam Gaussian
_NOMEGA = 64 # base number of points per 2*EW for the sinc^2 convolution
_COMEGA = 8.0 # half-extent of the sinc^2 line shape in units of DEW
_EPSH = 1.0e-2 # relative cutoff for the sum over harmonics
_EPSE = 1.0e-8 # bracketing distance at harmonic edges [eV]
_SPECLIM = 1.0e10 # minimum brilliance to retain
_IB = 3 # peak validity margin (points)
_CUTOFF = 1.0e-3 # harmonic intensity cutoff rel. to previous harmonic
# (TC v1.97; TC v1.95 used 5.0e-2)
_NSMALL = 2 # consecutive small harmonics before truncating the
# harmonic sum (TC v1.97; v1.95 stopped after one)
def _bessjn_all(mmax, x):
"""
J_m(x) for all integer orders m = 0..mmax, vectorized over x (x >= 0).
Miller's downward recurrence with renormalization (the same algorithm as
bright_bessjn in brighte.f / Numerical Recipes bessjn), evaluated for the
whole x array at once. Returns an array of shape (mmax+1,) + x.shape.
"""
x = numpy.asarray(x, dtype=float)
xf = x.ravel()
out = numpy.zeros((mmax + 1, xf.size))
tiny = xf < 1e-12
xs = numpy.where(tiny, 1.0, xf)
m0 = int(max(mmax, numpy.ceil(xs.max()))) + 2
m0 = m0 + int(numpy.sqrt(40.0 * m0))
if m0 % 2:
m0 += 1 # even starting order
bjp = numpy.zeros_like(xs)
bj = numpy.full_like(xs, 1e-30)
bsum = numpy.zeros_like(xs)
for m in range(m0, 0, -1):
bjm = (2.0 * m / xs) * bj - bjp # J_{m-1} (unnormalized)
bjp, bj = bj, bjm
big = numpy.abs(bj) > 1e10 # rescale to avoid overflow
if big.any():
bj[big] *= 1e-10
bjp[big] *= 1e-10
bsum[big] *= 1e-10
out[:, big] *= 1e-10
if m - 1 <= mmax:
out[m - 1] = bj
if (m - 1) % 2 == 0 and (m - 1) > 0: # accumulate even orders
bsum += bj
out /= (2.0 * bsum + out[0]) # J0 + 2*sum_k J_{2k} = 1
out[:, tiny] = 0.0
out[0, tiny] = 1.0
return out.reshape((mmax + 1,) + x.shape)
def _jn_signed(table, m):
"""J_m from a table over orders 0..mmax (first axis), for integer m of
any sign, using J_{-m} = (-1)^m J_m. Returns orders on the last axis."""
m = numpy.asarray(m)
am = numpy.abs(m)
sgn = numpy.where((m < 0) & (am % 2 == 1), -1.0, 1.0)
return sgn * numpy.moveaxis(table[am], 0, -1)
def _s0_harmonic(i, kx, ky, alpha, cosphi, sinphi):
"""
|A|^2 angular function of harmonic i of an ideal planar/helical/elliptical
undulator (Bessel function approximation). Equivalent to s0 of BRIGHTE
(bright1/bright3 in brighte.f), in fully vectorized complex form:
S_m = sum_n Jn(y) * J_{i+2n+m}(x) * exp(-i(i+2n+m)*phi'), m=-1,0,+1
Ax = (i/a) * (2*alpha*cos(phi)*S_0 - ky*(S_+1 + S_-1))
Ay = (i/a) * (2*alpha*sin(phi)*S_0 - 1j*kx*(S_+1 - S_-1))
s0 = |Ax|^2 + |Ay|^2
with a = 1+(kx^2+ky^2)/2+alpha^2, x = (2i/a)*alpha*sqrt((kx sinphi)^2 +
(ky cosphi)^2), y = (i/4a)*(ky^2-kx^2), phi' = atan2(kx sinphi, ky cosphi).
alpha: array (...,), cosphi/sinphi: arrays broadcastable to alpha.
"""
alpha2 = alpha * alpha
a = 1.0 + 0.5 * (kx * kx + ky * ky) + alpha2
x = (2.0 * i / a) * alpha * numpy.sqrt((kx * sinphi)**2 + (ky * cosphi)**2)
y = (0.25 * i / a) * (ky * ky - kx * kx)
nlim = max(int(6.2 + 1.41 * float(numpy.max(numpy.abs(y)))) + 2, 3)
n = numpy.arange(-nlim, nlim + 1) # (nn,)
ordr = i + 2 * n
Jy = _jn_signed(_bessjn_all(nlim, numpy.abs(y)), n) # (..., nn)
if ky * ky - kx * kx < 0.0: # J_n(-|y|)
Jy = Jy * numpy.where(numpy.abs(n) % 2 == 1, -1.0, 1.0)
xtab = _bessjn_all(i + 2 * nlim + 1, x)
Jx0 = _jn_signed(xtab, ordr)
Jxp = _jn_signed(xtab, ordr + 1)
Jxm = _jn_signed(xtab, ordr - 1)
if kx == 0.0: # planar: all real
S0 = numpy.sum(Jy * Jx0, axis=-1)
Ss = numpy.sum(Jy * (Jxp + Jxm), axis=-1)
Ax = (i / a) * (2.0 * alpha * cosphi * S0 - ky * Ss)
Ay = (i / a) * (2.0 * alpha * sinphi * S0)
return Ax * Ax + Ay * Ay
phi = numpy.arctan2(kx * sinphi, ky * cosphi)
ph = numpy.exp(-1j * ordr * phi[..., None])
S0 = numpy.sum(Jy * Jx0 * ph, axis=-1)
Sp = numpy.sum(Jy * Jxp * ph * numpy.exp(-1j * phi[..., None]), axis=-1)
Sm = numpy.sum(Jy * Jxm * ph * numpy.exp(+1j * phi[..., None]), axis=-1)
Ax = (i / a) * (2.0 * alpha * cosphi * S0 - ky * (Sp + Sm))
Ay = (i / a) * (2.0 * alpha * sinphi * S0 - 1j * kx * (Sp - Sm))
return (Ax * Ax.conj() + Ay * Ay.conj()).real
def _usb(energy, cur, sigx, sigy, sigx1, sigy1, period, n, kx, ky,
eminu, emaxu, nek):
"""
On-axis brilliance spectrum of an ideal undulator including the electron
beam emittance: equivalent to subroutine USB (usb.f) with METHOD=4
(Dejus' infinite-N + sinc^2 convolution).
Returns (e, spec): nek+1 energies [eV] from eminu to emaxu and the
brilliance [ph/s/mrad^2/mm^2/0.1%bw].
"""
gamma = energy / _MEE * 1e3
lamdar = period * 1e8 / (2.0 * gamma**2) # reduced wavelength [A]
er = _C_EVANG / lamdar # reduced energy [eV]
k3 = 1.0 + 0.5 * (kx * kx + ky * ky)
e1z = er / k3 # first harmonic on axis [eV]
length = n * period * 1e-2 # device length [m]
sigu = sigx1 * 1e-3 # divergences [rad]
sigv = sigy1 * 1e-3
fu = 0.5 / sigu**2 if sigu > 0 else 0.0
fv = 0.5 / sigv**2 if sigv > 0 else 0.0
ap2max = gamma**2 * ((_NSIGMA * sigu)**2 + (_NSIGMA * sigv)**2)
argmax = 0.5 * _NSIGMA**2
sigx2, sigy2 = sigx**2, sigy**2 # sizes [mm^2]
# prefactor: angular convolution normalization, per mrad^2 (via 1e6)
c1 = n * n * _FSC * _BW * cur * 1e-3 / _EC / (2.0 * _PI * sigu * sigv * 1e6)
# azimuthal grid (first quadrant midpoints), 4-fold symmetry
dphi = (_PI / 2.0) / _NPHI
phic = (numpy.arange(_NPHI) + 0.5) * dphi
cphi, sphi = numpy.cos(phic), numpy.sin(phic)
# sinc^2 line shape parameters
dew = e1z / n # natural line width [eV]
ew = _COMEGA * dew # convolution half-extent
pe = _PI / dew
# ------------------------------------------------------------------
# master energy grid: harmonic-reachable range, refined at the edges
# ------------------------------------------------------------------
e1min = e1z * k3 / (k3 + ap2max) # lowest reachable E1
e1max = e1z
emin, emax = eminu - ew, emaxu + ew
ie1 = max(1, int(numpy.ceil(emin / e1max - 1e-12)))
emin = max(ie1 * e1min, emin)
ie2 = int(numpy.ceil(emax / e1min - 1e-12)) - 1
if ie2 < ie1:
raise RuntimeError("usb: no harmonics reachable in the energy range")
emax = min(ie2 * e1max, emax)
if emax <= emin:
raise RuntimeError("usb: no harmonics reachable in the energy range")
dw = 2.0 * ew / _NOMEGA # base step = dew/4
# Dejus' variable-step master grid: walk from emin to emax with the base
# step dw, halving the step (down to dw/128) when approaching each
# harmonic upper edge ep = i*e1max (the spectrum is discontinuous there);
# bracket each edge with points at ep -/+ EPSE and jump to the next
# harmonic lower edge i*e1min when the harmonic bands do not overlap.
pts = [emin]
ih = ie1
ep = ih * e1max
gap = -1.0
while pts[-1] < emax:
idw = 1
while idw <= 128 and pts[-1] > ep - dew / idw:
idw *= 2
while pts[-1] < ep - _EPSE and pts[-1] < emax - _EPSE:
if idw <= 128 and pts[-1] > ep - dew / idw:
idw *= 2
pts.append(pts[-1] + dw / min(idw, 128))
sl = min(ep, emax)
pts[-1] = sl - _EPSE # place point just below edge
pts.append(sl + _EPSE) # and just above
ih += 1
gap = ih * e1min - pts[-1]
if gap > 0.0:
pts.append(ih * e1min) # jump across the gap
ep = ih * e1max
if gap > 0.0:
pts = pts[:-1] # drop point beyond emax
e_m = numpy.asarray(pts)
nm = e_m.size
# ------------------------------------------------------------------
# infinite-N spectrum on the master grid (sum over contributing harmonics)
# ------------------------------------------------------------------
spec_m = numpy.zeros(nm)
r = er / e_m # (nm,)
imin = numpy.floor(k3 / r).astype(int) + 1
imax = numpy.floor((ap2max + k3) / r).astype(int)
first = numpy.ones(nm, dtype=bool) # first harmonic flag
active = numpy.ones(nm, dtype=bool)
strikes = numpy.zeros(nm, dtype=int) # consecutive small harmonics
for ih in range(int(imin.min()), int(imax.max()) + 1):
sel = active & (ih >= imin) & (ih <= imax)
if not sel.any():
continue
rs = r[sel]
alpha2i = rs * ih - k3
alphai = numpy.sqrt(numpy.maximum(alpha2i, 0.0))
thetai = alphai / gamma
# effective source size at this detuning (Walker)
delta = alpha2i * n / rs
sigr2 = numpy.where(delta < 2.15,
(1.29 + 1.229 * (delta - 0.8)**2)**2,
5.81 * delta)
sigr2 = sigr2 * _CV * _C_EVANG / e_m[sel] * length # [mm^2]
const = c1 / (2.0 * _PI * numpy.sqrt((sigx2 + sigr2) * (sigy2 + sigr2)))
# ring integral over phi with the Gaussian angular weight
u = thetai[:, None] * cphi[None, :]
v = thetai[:, None] * sphi[None, :]
arg = u * u * fu + v * v * fv
w = numpy.where(arg < argmax, numpy.exp(-arg), 0.0)
s0 = _s0_harmonic(ih, kx, ky, alpha=alphai[:, None],
cosphi=cphi[None, :], sinphi=sphi[None, :])
contr = 4.0 * const * (rs / (2.0 * n)) * dphi * numpy.sum(w * s0, axis=1)
idx = numpy.where(sel)[0]
spec_m[idx] += contr
# stop summing after _NSMALL consecutive (non-first) harmonics each
# adding less than EPS of the running total (TC v1.97)
small = (contr <= _EPSH * spec_m[idx]) & (~first[idx])
strikes[idx] = numpy.where(small, strikes[idx] + 1, 0)
active[idx[strikes[idx] >= _NSMALL]] = False
first[idx] = False
# ------------------------------------------------------------------
# convolution with the sinc^2 line shape, onto the user energy grid.
# The window endpoints EU -/+ EW carry sinc^2 = 0 exactly, so only the
# trapezoid weights of the first/last grid points inside each window
# need the partial-interval correction.
# ------------------------------------------------------------------
de = (emaxu - eminu) / nek
eu = eminu + numpy.arange(nek + 1) * de
he = numpy.diff(e_m) # (nm-1,)
ha = numpy.empty(nm) # trapezoid weights
ha[1:-1] = 0.5 * (he[:-1] + he[1:])
ha[0], ha[-1] = 0.5 * he[0], 0.5 * he[-1]
arg = pe * (eu[:, None] - e_m[None, :]) # (neu, nm)
inwin = numpy.abs(arg) < (_COMEGA * _PI) * (1.0 - 1e-12)
snc = numpy.ones_like(arg)
nz = numpy.abs(arg) > 1e-6
snc[nz] = (numpy.sin(arg[nz]) / arg[nz]) ** 2
W = numpy.where(inwin, ha[None, :] * snc, 0.0)
# partial-interval weights at the window boundaries
a_w, b_w = eu - ew, eu + ew
j1 = numpy.searchsorted(e_m, a_w, side='right') - 1 # E[j1] <= a < E[j1+1]
j2 = numpy.searchsorted(e_m, b_w, side='right') - 1 # E[j2] <= b < E[j2+1]
rows = numpy.arange(nek + 1)
s1 = (j1 >= 0) & (j1 < nm - 1) # left edge inside grid
jl = numpy.clip(j1 + 1, 0, nm - 1)
W[rows[s1], jl[s1]] = snc[rows[s1], jl[s1]] * 0.5 * (
(e_m[jl[s1]] - a_w[s1]) +
he[numpy.clip(jl[s1], 0, nm - 2)] * (jl[s1] < nm - 1))
s2 = (j2 > 0) & (j2 < nm - 1) # right edge inside grid
W[rows[s2], j2[s2]] = snc[rows[s2], j2[s2]] * 0.5 * (
(b_w[s2] - e_m[j2[s2]]) + he[j2[s2] - 1])
empty = (j2 <= 0) | (j1 >= nm - 1) | (j2 - j1 < 1)
W[empty, :] = 0.0
spec = (W @ spec_m) / dew
return eu, spec
def _gauss_convolve(e, spec, sige, eiz):
"""Gaussian beam-energy-spread convolution (rms width 2*sige*eiz),
equivalent to gauss_convolve in tc.f; trims the array edges."""
de = e[1] - e[0]
sigp = 2.0 * sige * eiz / de # sigma in grid units
if sigp < 5.0: # NPPSIGMA-1
raise RuntimeError("gauss_convolve: too few points (sige too small)")
np_g = int(6.0 * sigp + 1.0) # 2*NSIGMA*sigp
if np_g % 2 == 0:
np_g += 1
xk = numpy.arange(np_g) - (np_g - 1) / 2.0
gs = numpy.exp(-xk**2 / (2.0 * sigp**2))
spec2 = numpy.convolve(spec, gs / gs.sum(), mode='valid')
nh = np_g // 2
return e[nh:len(e) - nh], spec2
def _peak(e, spec):
"""Abscissa and ordinate of the maximum; invalid (as in tc.f) if the
maximum falls within IB points of either end of the array."""
ip = int(numpy.argmax(spec))
if ip >= len(spec) - 1 - _IB or ip <= _IB - 1:
return None, None
return e[ip], spec[ip]
[docs]def xoppy_calc_tcpy(
ENERGY = 7.0, # ring energy [GeV]
CURRENT = 100.0, # ring current [mA]
ENERGY_SPREAD = 0.00096, # sigma(E)/E
SIGX = 0.274, # rms horizontal beam size [mm]
SIGY = 0.011, # rms vertical beam size [mm]
SIGX1 = 0.0113, # rms horizontal divergence [mrad]
SIGY1 = 0.0036, # rms vertical divergence [mrad]
PERIOD = 3.23, # undulator period [cm]
NP = 70, # number of periods
EMIN = 2950.0, # min energy of the FIRST harmonic [eV]
EMAX = 13500.0, # max energy of the FIRST harmonic [eV]
N = 40, # number of points of the tuning curve
HARMONIC_FROM = 1,
HARMONIC_TO = 15,
HARMONIC_STEP = 2,
HELICAL = 0, # 0: planar, 1: helical
METHOD = 1, # kept for interface compatibility (Dejus')
NEKS = 100, # energy points for the peak search
output_file = "tcpy.out",
verbose = True,
):
"""
Pure python version of xoppylib.xoppy_run_binaries.xoppy_calc_xtc
(TC: undulator on-axis brilliance tuning curves, infinite-N method
with convolution). Returns (data, harmonics_data) with the same
structure: data is the stack of all rows
[E_no_emittance, E_peak, Brilliance, Ky, Ptot, Pd], and harmonics_data
is a list of [harmonic_number, rows-of-that-harmonic].
"""
if METHOD not in (0, 1):
raise ValueError("only METHOD=0,1 (Dejus, infinite-N + convolution) implemented")
sige = ENERGY_SPREAD
neks = NEKS if NEKS != 0 else 100
ne = N if N != 0 else 20
ihmin = HARMONIC_FROM if HARMONIC_FROM != 0 else 1
ihmax = HARMONIC_TO if HARMONIC_TO != 0 else 5
ihstep= HARMONIC_STEP if HARMONIC_STEP != 0 else 2
gamma = ENERGY / _MEE * 1e3
lamdar = PERIOD * 1e8 / (2.0 * gamma**2)
er = _C_EVANG / lamdar
emin, emax = EMIN, EMAX
if emax >= er: # reset to K=0.2 value (as in tc.f)
emax = _C_EVANG / (lamdar * (1.0 + 0.20**2 / 2.0))
if verbose:
print("Warning: emax reset to %.1f eV" % emax)
e1zmin, e1zmax = emin, emax
de = (e1zmax - e1zmin) / (ne - 1) if ne > 1 else 0.0
def k_of(ek): # deflection parameter(s) at E1=ek
kyv = numpy.sqrt(2.0 * (er / ek - 1.0))
if HELICAL == 1:
kyv /= numpy.sqrt(2.0)
return kyv, kyv # kx, ky
return 0.0, kyv
harmonics = list(range(ihmin, ihmax + 1, ihstep))
nharm = len(harmonics)
gk_planar = lambda k: k * (k**6 + 24.0 * k**4 / 7.0 + 4.0 * k * k + 16.0 / 7.0) / (1.0 + k * k)**3.5
gk_helical = lambda k: 32.0 / 7.0 * k / (1.0 + k * k)**3
# ------------------------------------------------------------------
# peak shifts dep1, dep2 of harmonics 1 and 2 at K = Kmin
# ------------------------------------------------------------------
ek = e1zmax
kx, ky = k_of(ek)
dep = []
for i in (1, 2):
eiz = i * ek
ekmin, ekmax = (0.95 * eiz, 1.01 * eiz) if i == 1 else (0.820 * eiz, 1.002 * eiz)
e_, s_ = _usb(ENERGY, CURRENT, SIGX, SIGY, SIGX1, SIGY1, PERIOD, NP,
kx, ky, ekmin, ekmax, 500)
ep, _sp = _peak(e_, s_)
if ep is None:
raise RuntimeError("tc: no peak found for the initial shift search (harmonic %d)" % i)
ep = min(ep, 0.9995 * eiz)
dep.append(eiz - ep)
dep1, dep2 = dep
# ------------------------------------------------------------------
# main loop over harmonics and K values
# ------------------------------------------------------------------
ei = numpy.zeros((ne, nharm))
eb = numpy.zeros((ne, nharm))
sb = numpy.zeros((ne, nharm))
kyb = numpy.zeros(ne)
kxb = numpy.zeros(ne)
ptot = numpy.zeros(ne)
pd = numpy.zeros(ne)
shmax = 0.0
aborted = False
for ih, i in enumerate(harmonics):
if aborted:
break
lodd = (i % 2 == 1)
specmin = _SPECLIM if ih == 0 else _CUTOFF * shmax
ein = 1.10 * (i + ihstep) * e1zmin
eiz = 1.01 * ein
je = ne + 1
smax = 0.0
e_, s_ = None, None
def store_k(j, kxv, kyv): # K, Ptot, Pd columns (first harmonic)
if ih == 0:
kk2 = kxv * kxv + kyv * kyv
kyb[j] = kyv
kxb[j] = kxv
ptot[j] = _PTOT_FAC * NP * kk2 * ENERGY**2 * CURRENT * 1e-3 / (PERIOD * 1e-2)
gk = gk_helical(kyv) if HELICAL == 1 else gk_planar(kyv)
pd[j] = _PD_FAC * NP * kyv * gk * ENERGY**4 * CURRENT * 1e-3 / (PERIOD * 1e-2)
# scan down from Kmin to find the highest K-point with intensity
while smax < specmin and je > 1 and eiz > ein:
je -= 1
ek = e1zmin + (je - 1) * de
kx, ky = k_of(ek)
store_k(je - 1, kx, ky)
e1z = ek
eiz = i * e1z
if lodd:
if HELICAL == 1:
ekmin = eiz - i**2 * 1.5 * dep1
else:
ekmin = eiz - i * dep1
if i == 1:
ekmin -= dep1
ekmax = eiz + i * dep1 / 2.0
if i > int(e1z / dep1):
if verbose:
print("Warning: overlapping range for initial peak search, harmonic %d" % i)
aborted = True
break
else:
ekmin = (eiz - i**2 * 0.5 * dep2) if HELICAL == 1 else (eiz - 4.0 * dep2)
ekmax = eiz
e_, s_ = _usb(ENERGY, CURRENT, SIGX, SIGY, SIGX1, SIGY1, PERIOD, NP,
kx, ky, ekmin, ekmax, neks)
smax = s_.max()
ei[je - 1, ih] = eiz
eb[je - 1, ih] = eiz
sb[je - 1, ih] = 0.0
if aborted:
break
if smax < _SPECLIM:
if verbose:
print("Warning: harmonic intensity too small for harmonic %d" % i)
break
ep, _sp = _peak(e_, s_)
if ep is None:
break
fc = 0.990 * ep / eiz
if i > int(1.0 / (1.0 - fc)):
if verbose:
print("Warning: overlapping range for peak search, harmonic %d" % i)
break
shmax = 0.0
for j in range(1, je + 1):
ek = e1zmin + (j - 1) * de
kx, ky = k_of(ek)
store_k(j - 1, kx, ky)
fc2 = 1.002 if lodd else 1.000
eiz = i * ek
ekmin, ekmax = fc * eiz, fc2 * eiz
nek = neks
if sige > 0.0:
dek = (ekmax - ekmin) / nek
sigee = 2.0 * sige * eiz
ekmin -= sigee * 3.0 # NSIGMA=3 for the espread convolution
ekmax += sigee * 3.0
if sigee / 6.0 < dek: # NPPSIGMA=6
dek = sigee / 6.0
nek = int((ekmax - ekmin) / dek + 1.0)
e_, s_ = _usb(ENERGY, CURRENT, SIGX, SIGY, SIGX1, SIGY1, PERIOD, NP,
kx, ky, ekmin, ekmax, nek)
if sige > 0.0:
e_, s_ = _gauss_convolve(e_, s_, sige, eiz)
ep, sp = _peak(e_, s_)
if ep is None:
aborted = True
break
shmax = max(shmax, sp)
ei[j - 1, ih] = eiz
eb[j - 1, ih] = ep
sb[j - 1, ih] = sp
if aborted:
break
if verbose:
print("Harmonic %d completed" % i)
# ------------------------------------------------------------------
# assemble output (same structure as xoppy_calc_xtc) and write file
# ------------------------------------------------------------------
harmonics_data = []
rows_all = []
f = open(output_file, "w") if output_file is not None else None
if f is not None:
f.write("# TCPY: python tuning curves (infinite-N + convolution)\n")
f.write("# On-axis Brilliance (ph/s/mrad^2/mm^2/0.1%bw)\n")
f.write("# Energy(eV no emittance) Energy(eV) Brilliance Ky Ptot(W) Pd(W/mr^2)\n")
for ih, i in enumerate(harmonics):
if HELICAL == 1:
rows = numpy.column_stack((ei[:, ih], eb[:, ih], sb[:, ih],
kxb, kyb, ptot, pd))
fmt = "%12.3f %12.3f %14.5e %8.3f %8.3f %9.1f %12.1f\n"
else:
rows = numpy.column_stack((ei[:, ih], eb[:, ih], sb[:, ih],
kyb, ptot, pd))
fmt = "%12.3f %12.3f %14.5e %8.3f %9.1f %12.1f\n"
harmonics_data.append([i, rows])
rows_all.append(rows)
if f is not None:
f.write("# Harmonic %d\n" % i)
for r in rows:
f.write(fmt % tuple(r))
if f is not None:
f.close()
if verbose:
print("File written to disk: %s" % output_file)
data = numpy.vstack(rows_all)
return data, harmonics_data
if __name__ == "__main__":
if 0: # simple comparison
args = {
'ENERGY' : 6.0,
'CURRENT' : 200.0,
'ENERGY_SPREAD' : 0.001,
'SIGX' : 0.0334281,
'SIGY' : 0.0072813899999999996,
'SIGX1' : 0.00451097,
'SIGY1' : 0.00194034,
'PERIOD' : 1.7999999999999998,
'NP' : 111.111,
'EMIN' : 8008.164712466435,
'EMAX' : 40040.823562332174,
'N' : 40,
'HARMONIC_FROM' : 1,
'HARMONIC_TO' : 15,
'HARMONIC_STEP' : 2,
'HELICAL' : 0,
'METHOD' : 1,
'NEKS' : 100,
}
# NOTE on comparing with the Fortran via xoppylib.xoppy_calc_xtc:
# the wrapper writes tc.inp with rounded values ("%d" for NP, "%10.4f"
# for the sigmas, "%10.1f" for EMIN/EMAX), so the binary never sees the
# full-precision args above (e.g. NP=111.111 -> 111: +0.2% in B since
# B ~ N^2; SIGY1=0.00194034 -> 0.0019: -2.1%). For an apples-to-apples
# comparison, feed the python code the same rounded values:
args_as_fortran_sees_them = dict(args)
args_as_fortran_sees_them.update(
NP = int(args['NP']),
SIGX = float("%10.4f" % args['SIGX']),
SIGY = float("%10.4f" % args['SIGY']),
SIGX1 = float("%10.4f" % args['SIGX1']),
SIGY1 = float("%10.4f" % args['SIGY1']),
EMIN = float("%10.1f" % args['EMIN']),
EMAX = float("%10.1f" % args['EMAX']),
)
# python (use plain **args instead for full-precision results)
data, harmonics_data = xoppy_calc_tcpy( **args_as_fortran_sees_them )
print("PY: Number of harmonics calculated: ", len(harmonics_data))
print(harmonics_data[0][1][:5])
# fortran
from xoppylib.xoppy_run_binaries import xoppy_calc_xtc
data_f, harmonics_data_f = xoppy_calc_xtc( **args )
print("F: Number of harmonics calculated: ", len(harmonics_data_f))
print(harmonics_data_f[0][1][:5])
#
# example plot
#
if True:
from srxraylib.plot.gol import plot
#
#
print("Number of harmonics calculated: ", len(harmonics_data))
plot((harmonics_data[0][1])[:, 0],
(harmonics_data[0][1])[:, 2],
(harmonics_data_f[0][1])[:, 0],
(harmonics_data_f[0][1])[:, 2],
title="harmonic number = %d" % (harmonics_data[0][0]), xtitle="Energy[eV]", ytitle="Brilliance",
ylog=0, legend=['python', 'fortran'])
#
# end script
#`
if 1: #claude run
# =============================================================================
# make_tc_comparison.py
#
# Generates tc_comparison.png: tuning curves of the Fortran TC v1.97 binary
# (solid) overlaid with xoppy_calc_tcpy (dashed red), plus relative-difference
# panels, for two machine configurations.
#
# Requirements:
# - xoppy_calc_tcpy.py importable (same directory is fine)
# - one Fortran tc.out per case. Either take it from a previous
# xoppy_calc_xtc run, or write a tc.inp and run the xoppylib binary:
# <python-site-packages>/xoppylib/bin/linux/tc tc.inp
# Now: Run the previous example (case 1)
# Run the XOPPY/tc default and rename it (now /home/srio/Oasys2/tc_default.out
# IMPORTANT: the python call below must receive the *same values the binary
# read from tc.inp*. xoppylib writes tc.inp with "%d" for NP, "%10.4f" for
# the four sigmas and "%10.1f" for EMIN/EMAX, so pass the rounded values
# (e.g. NP=111, SIGY1=0.0019), not the full-precision ones.
# =============================================================================
import numpy as np
import matplotlib
# matplotlib.use("Agg")
import matplotlib.pyplot as plt
# from xoppy_calc_tcpy import xoppy_calc_tcpy
def parse_tc_out(path, ncol=6):
"""Parse tc.out into a list of (ne, ncol) arrays, one per 'Harmonic' block.
Columns: E_zero_emittance, E_peak, Brilliance, Ky, Ptot, Pdensity
(7 columns for helical runs, which insert Kx). Non-numeric 6-field
header lines are skipped by the float() try/except.
"""
blocks, cur = [], None
for line in open(path):
if line.strip().startswith("Harmonic"):
if cur:
blocks.append(np.array(cur))
cur = []
else:
t = line.split()
if cur is not None and len(t) == ncol:
try:
cur.append([float(v) for v in t])
except ValueError:
pass
if cur:
blocks.append(np.array(cur))
return blocks
# ----------------------------------------------------------------------------
# The two cases: (title, fortran tc.out path, args exactly as in tc.inp)
# ----------------------------------------------------------------------------
cases = [
("APS-like: 7 GeV, $\\lambda_0$=3.23 cm, N=70",
"/home/srio/Oasys2/tc_default.out", # <- adapt path
dict(ENERGY=7.0, CURRENT=100.0, ENERGY_SPREAD=0.00096,
SIGX=0.274, SIGY=0.011, SIGX1=0.0113, SIGY1=0.0036,
PERIOD=3.23, NP=70, EMIN=2950.0, EMAX=13500.0)),
("MBA-like: 6 GeV, $\\lambda_0$=1.8 cm, N=111",
"tc.out", # <- adapt path (your run)
dict(ENERGY=6.0, CURRENT=200.0, ENERGY_SPREAD=0.001,
SIGX=0.0334, SIGY=0.0073, SIGX1=0.0045, SIGY1=0.0019,
PERIOD=1.8, NP=111, EMIN=8008.2, EMAX=40040.8)),
]
fig, axes = plt.subplots(2, 2, figsize=(11.5, 8))
for col, (title, ftout, kw) in enumerate(cases):
b_fortran = parse_tc_out(ftout)
_, hd = xoppy_calc_tcpy(N=40, HARMONIC_FROM=1, HARMONIC_TO=15,
HARMONIC_STEP=2, HELICAL=0, METHOD=1, NEKS=100,
verbose=False, output_file=None, **kw)
ax0, ax1 = axes[0, col], axes[1, col]
for ihx, blk in enumerate(b_fortran):
py = hd[ihx][1]
m = blk[:, 2] > 0 # skip cutoff (zeroed) rows
# top: tuning curves, peak energy vs peak brilliance
ax0.plot(blk[m, 1], blk[m, 2], "-", lw=2.2, alpha=0.85, color="C0",
label="Fortran" if ihx == 0 else None)
ax0.plot(py[m, 1], py[m, 2], "--", lw=1.1, color="red",
label="python" if ihx == 0 else None)
# bottom: relative difference in units of 1e-5
ax1.plot(blk[m, 1], 1e5 * (py[m, 2] / blk[m, 2] - 1), ".-",
ms=3, lw=0.7, label="h%d" % hd[ihx][0])
ax0.set_xscale("log");
ax0.set_yscale("log")
ax0.set_title(title, fontsize=10)
ax0.set_ylabel("Brilliance (ph/s/mrad$^2$/mm$^2$/0.1%bw)", fontsize=8)
ax0.legend(fontsize=8);
ax0.tick_params(labelsize=8)
ax1.set_xscale("log")
ax1.set_xlabel("Photon energy (eV)", fontsize=9)
ax1.set_ylabel("rel. difference ($\\times 10^{-5}$)", fontsize=8)
ax1.legend(fontsize=6, ncol=4);
ax1.tick_params(labelsize=8)
fig.suptitle("xoppy_calc_tcpy vs TC v1.97 binary (xoppylib): odd harmonics 1-15",
fontsize=11)
fig.tight_layout(rect=(0, 0, 1, 0.97))
# fig.savefig("tc_comparison.png", dpi=140)
# print("tc_comparison.png written")
plt.show()