diff --git a/pyat/at/physics/energy_loss.py b/pyat/at/physics/energy_loss.py index 929484afc..d68733b3d 100644 --- a/pyat/at/physics/energy_loss.py +++ b/pyat/at/physics/energy_loss.py @@ -70,9 +70,16 @@ def tracking(ring): sextupole_pass=None, octupole_pass=None, copy=True) - o0 = numpy.zeros(6) - o6 = numpy.squeeze(lattice_pass(ringtmp, o0, refpts=len(ringtmp))) - return -o6[4] * ring.energy + + o6 = numpy.squeeze(lattice_pass(ringtmp, numpy.zeros(6), refpts=len(ringtmp))) + if numpy.isnan(o6[0]): + dp = 0 + for e in ringtmp: + ot = numpy.squeeze(lattice_pass([e], numpy.zeros(6))) + dp += -ot[4] * ring.energy + return dp + else: + return -o6[4] * ring.energy if isinstance(method, str): method = ELossMethod[method.upper()] @@ -150,7 +157,7 @@ def deq(x): return timelag, ts -def set_cavity_phase(ring, method=ELossMethod.INTEGRAL, +def set_cavity_phase(ring, method=ELossMethod.TRACKING, refpts=None, cavpts=None, copy=False): """ Adjust the TimeLag attribute of RF cavities based on frequency, diff --git a/pyat/at/physics/radiation.py b/pyat/at/physics/radiation.py index f6fe5222c..103392dc7 100644 --- a/pyat/at/physics/radiation.py +++ b/pyat/at/physics/radiation.py @@ -12,6 +12,7 @@ from at.tracking import lattice_pass from at.physics import find_orbit6, find_m66, find_elem_m66 from at.physics import find_mpole_raddiff_matrix, get_tunes_damp +from at.physics import ELossMethod __all__ = ['ohmi_envelope', 'get_radiation_integrals', 'quantdiffmat', 'gen_quantdiff_elem', 'tapering'] @@ -397,6 +398,8 @@ def tapering(ring, multipoles=True, niter=1, **kwargs): KEYWORDS multipoles=True scale all multipoles + method Method for energy loss computation + (see get_energy_loss) niter=1 number of iteration XYStep=1.0e-8 transverse step for numerical computation DPStep=1.0E-6 momentum deviation used for computation of orbit6 @@ -404,6 +407,7 @@ def tapering(ring, multipoles=True, niter=1, **kwargs): xy_step = kwargs.pop('XYStep', DConstant.XYStep) dp_step = kwargs.pop('DPStep', DConstant.DPStep) + method = kwargs.pop('method', ELossMethod.TRACKING) dipoles = get_refpts(ring, Dipole) b0 = get_value_refpts(ring, dipoles, 'BendingAngle') k0 = get_value_refpts(ring, dipoles, 'PolynomB', index=0) @@ -411,7 +415,7 @@ def tapering(ring, multipoles=True, niter=1, **kwargs): for i in range(niter): _, o6 = find_orbit6(ring, refpts=range(len(ring)+1), - XYStep=xy_step, DPStep=dp_step) + XYStep=xy_step, DPStep=dp_step, method=method) dpps = (o6[dipoles, 4] + o6[dipoles+1, 4]) / 2 set_value_refpts(ring, dipoles, 'PolynomB', b0/ld*dpps+k0*(1+dpps), index=0) @@ -420,7 +424,7 @@ def tapering(ring, multipoles=True, niter=1, **kwargs): mults = get_refpts(ring, Multipole) k0 = get_value_refpts(ring, dipoles, 'PolynomB', index=0) _, o6 = find_orbit6(ring, refpts=range(len(ring)+1), - XYStep=xy_step, DPStep=dp_step) + XYStep=xy_step, DPStep=dp_step, method=method) dpps = (o6[mults, 4] + o6[mults+1, 4]) / 2 for dpp, el in zip(dpps, ring[mults]): el.PolynomB *= 1+dpp