Fix Transient Beam Loading Computation + Energy Loss Bug

[lcarver]

There was a problem with the computation of the beam induced voltage when gaps were present in the ring. This bug has been present since January #1025 . Basically, the cavity phasor was adjusted by one bucket whether a bunch was present or not, but then when the next bunch arrived, the full previous gap was applied again (incorrectly).

This is now fixed.

This PR also uncovers a problem with get_energy_loss since the recent fast ring upgrade. get_timelag_fromU0 does not return the same number for ts for each cavity. There is a temporary fix in place until @swhite2401 can provide a cleaner fix.

Below is an example script for BessyII (my TBL benchmark case).

import numpy as np
import matplotlib.pyplot as plt
from at.constants import clight
import at
from at.collective import Wake, WakeElement, WakeType, WakeComponent
from scipy.interpolate import interp1d

def gen_fill_pattern(h):
    gap = 100
    fill_pattern = np.zeros(h)
    fill_pattern[0:int((h-gap)/2)] = 1
    fill_pattern[int((h+gap)/2):h] = 1

    fill_pattern[3] = 3
    fill_pattern[9] = 3
    fill_pattern[15] = 3

    #fill_pattern[200] = 4
    fill_pattern[201] = 4
    fill_pattern[243] = 3
    fill_pattern[379] = 3
    fill_pattern[385] = 3
    fill_pattern[391] = 3
    fill_pattern[397] = 3

    return fill_pattern/np.sum(fill_pattern)
    
ring_dict = {
            'circumference': 240,
            'harmonic_number': 400,
            'ac': 7.3e-4,
            'energy': 1.7e9,
            'sigma_e': 0.000695451,
            'U0': 2.3e5,
            # from lattice file
            #U0 = 1.749192e+05 

            # from python script
            #U0 = 2.3e5
            }


beta_main = 2.2
Q0_main = 29600
QL_main = Q0_main / (1+beta_main)
Nmain = 4
Rsl_main = 114*QL_main*Nmain # R/Q * QL * Nmain

# 3rd harmonic passive harmonic cavities (there is one cavity that is different)
n = 3
RoQ_3h = 63

Q0_3h_1 = 13750
N_3h_1 = 3
Rs_3h_1 = RoQ_3h * Q0_3h_1 * N_3h_1
detune_3h_1 = 0.38e6

Q0_3h_2 = 9250
N_3h_2 = 1
Rs_3h_2 = RoQ_3h * Q0_3h_2 * N_3h_2
detune_3h_2 = 0.46e6



Vmain = 1.4e6

fillp = gen_fill_pattern(ring_dict['harmonic_number'])


t0 = ring_dict['circumference'] / clight
tauz = 0.00775692 / t0

simple_ring = at.simple_ring(ring_dict['energy'], ring_dict['circumference'], ring_dict['harmonic_number'], 0.1, 0.1, Vmain, ring_dict['ac'], U0=ring_dict['U0'], tauz=tauz/100, espread=ring_dict['sigma_e'])
simple_ring.pop(-1) #remove pointless nonlinear element
simple_ring.pop(-2) #remove quantdiff for single particle model

simple_ring.enable_6d()
simple_ring.set_cavity_phase()

# set fill pattern, and beam current
I0 = 0.3
simple_ring.set_fillpattern(fillp)
simple_ring.set_beam_current(I0)
Nbunches = simple_ring.bunch_list.shape[0]

# some simulation parameters
Nslice = 1 #only 1 slice as only one macroparticle per bunch!
Nturns = 30000

Nparts = Nbunches


# for pyat you must first add all rf cavities as elements, then you add beamloading later.
# here I just make two RF cavities with 0 voltage and 0 length, but with frequency=3*rf_frequency
harm_cavity_1 = at.RFCavity('HC1', 0.0, 0.0, n*simple_ring.rf_frequency, n*simple_ring.harmonic_number, simple_ring.energy)
harm_cavity_2 = at.RFCavity('HC2', 0.0, 0.0, n*simple_ring.rf_frequency, n*simple_ring.harmonic_number, simple_ring.energy)

simple_ring.insert(1, harm_cavity_2) #puts them in order, main ->HC1 -> HC2
simple_ring.insert(1, harm_cavity_1)

# main beamloading
at.add_beamloading(simple_ring, QL_main, Rsl_main, cavpts=[0], Nslice=Nslice, VoltGain=1e-3, PhaseGain=1e-3, cavitymode=at.CavityMode.ACTIVE, fbmode=at.FeedbackMode.ONETURN, buffersize=1)#, buffersize=1000, windowlength=200)

# passive cavity for HC type 1
at.add_beamloading(simple_ring, Q0_3h_1, Rs_3h_1, detune=detune_3h_1, cavpts=[1], Nslice=Nslice, cavitymode=at.CavityMode.PASSIVE, system_harmonic=n)
   
# passive cavity for HC type 2
at.add_beamloading(simple_ring, Q0_3h_2, Rs_3h_2, detune=detune_3h_2, cavpts=[2], Nslice=Nslice, cavitymode=at.CavityMode.PASSIVE, system_harmonic=n)




# add wake element with Nslice=1 to model the losses to the machine wake
kq = 1.25e13 #in V/C
sr = np.arange(-0.1,0.1,1e-3)
wake_z = np.heaviside(sr,0.5)*kq

wa = Wake(sr)
wa.add(WakeType.TABLE, WakeComponent.Z, sr, wake_z)

welem = WakeElement('waa', simple_ring, wa, Nslice=1)
simple_ring.insert(3,welem)


parts = np.zeros((6, Nparts))

all_z = np.zeros((Nturns, Nparts))
all_Vg_main = np.zeros((Nturns,3))
all_Vb_harm_1 = np.zeros((Nturns,2))
all_Vb_harm_2 = np.zeros((Nturns,2))

msk = fillp!=0
for i in np.arange(Nturns):
        
    simple_ring.track(parts, nturns=1, in_place=True, refpts=None, keep_counter=True)
    
    
    all_z[i,:] = parts[5,:]
    all_Vg_main[i,:] = simple_ring[0].Vgen[:3]    
    all_Vb_harm_1[i,:] = simple_ring[1].Vbeam
    all_Vb_harm_2[i,:] = simple_ring[2].Vbeam

last_turn_z = all_z[-1,:]
main_phase_transient = last_turn_z*2*np.pi*simple_ring[0].ResFrequency/clight

ifun = interp1d(simple_ring.bunch_list, main_phase_transient, fill_value=0, bounds_error=False)
# now some plotting

fig, (ax1,ax2) = plt.subplots(2,1,figsize=(12,6))
ax1.plot(simple_ring.bunch_list, 1e12*(all_z[-1,:] - np.mean(all_z[-1,:]))/clight, color='r', marker='.', linestyle='None')
ax2.plot(simple_ring.bunch_list, 1e3*simple_ring.bunch_currents, color='r', marker='.', linestyle='None')
ax1.grid(color='k', linestyle='dashed', alpha=0.4)

ax1.set_ylabel('Arrival Time (centered) [ps]')
ax1.xaxis.set_visible(False)
ax2.set_ylabel(r'$I_{b}$ [mA]')
ax2.set_xlabel('Bunch #')
ax2.set_ylim(0, np.amax(1e3*simple_ring.bunch_currents)*1.1)

plt.show()

[lcarver]

ready for review

[swhite2401]

@lcarver I have just pushed a possible simplification for the get_energy_loss function

[swhite2401]

Please make sure you run ruff checks on the modified files, thanks!

[lcarver]

all fine for me. My tests work, I ran ruff and black. I also cleaned up the C a little bit (but more work to be done on this soon!) and my test still works fine.

[lcarver]

ready for merge