diff --git a/atmat/atphysics/TouschekPiwinski/MomAperture_Project2Start.m b/atmat/atphysics/TouschekPiwinski/MomAperture_Project2Start.m new file mode 100644 index 000000000..16c10e8af --- /dev/null +++ b/atmat/atphysics/TouschekPiwinski/MomAperture_Project2Start.m @@ -0,0 +1,247 @@ +function [etn, etp]=MomAperture_Project2Start(THERING, varargin) +% MOMAPERTURE_PROJECT2START calculates the local momentum aperture. +% +% MOMAPERTURE_PROJECT2START is a Bipartition search of the negative and +% positive stability thesholds in the 5th dimension (relative momentum). +% -The 6D closed orbit is taken into account. +% -Particles launched at different REFPTS along the ring are first projected +% to the ring last element so that all particles can be tracked together. +% +% [ETN, ETP] = MOMAPERTURE_PROJECT2START(THERING) +% +% Inputs: +% THERING: ring used for tracking. +% Options: +% REFPTS: REFPTS where to calculate the momentum acceptance. +% Default 1:numel(THERING); +% nturns: Number of turns to track. Default 1000 +% dptol: resolution in momentum acceptance. Default 1e-4 +% dpuguess: unstable momentum threshold guess. Default []. +% If not given it uses the linear momentum acceptance delta_max +% from ringpara. +% troffset: [x y] starting transverse offset for the tracking. +% Default [1e-6 1e-6] +% verbose: boolean indicating verbose mode. Default false. +% epsilon6D: if not passed, all particles are tracked. +% If epsilon6D is given, we track for nturns only +% particles having 6D coordinates different by epsilon6D +% after being projected to the end of the ring. +% Output: +% ETN: stability threshold for positive off momentum particles +% ETP: stability threshold for negative off momentum particles +% +% +% Other functions in the file: +% +% Loste = Multiorigin_ringpass_islost +% Returns a boolean array: tells whether the particle launched at +% the reference point refpts with positive and negative momentum offset is +% lost or not. + +% 2024may30 Z.Marti at ALBA CELLS, original +% 2024jun18 oblanco at ALBA CELLS, rewritten to track positive and negative +% sides in a single call +% See https://github.com/atcollab/at/pull/773 + +% Parse input +p = inputParser; +addOptional(p,'refpts',1:numel(THERING)); +addOptional(p,'nturns',1000); +addOptional(p,'dptol',1e-4); +addOptional(p,'dpuguess',[]); +addOptional(p,'troffset',[1e-6 1e-6]); +addOptional(p,'verbose',false); +addOptional(p,'epsilon6D',0); +parse(p,varargin{:}); +par = p.Results; + +REFPTS=par.refpts; +nturns=par.nturns; +detole=par.dptol; +eu_ini=par.dpuguess; +initcoord=par.troffset; +verbose=par.verbose; +epsilon6D = par.epsilon6D; + +if 0 ~= epsilon6D + fprintf('Particles differing by less than %.3e are considered similar.\n', ... + epsilon6D); +end + +np = numel(REFPTS); +if verbose + fprintf('Using %d reference points.\n',np); + fprintf('Tracking %d turns.\n',nturns); + fprintf('Iteration stops when the momentum step is below %.3f.\n',detole); + fprintf('Using init coords %.3f %.3f um as transverse offsets.\n',1e6*initcoord(1), ... + 1e6*initcoord(2)); +end + +% initial bipartition settings +if isempty(eu_ini) + res = ringpara(THERING); + eu_ini = res.delta_max; + if verbose + fprintf('Using the rf bucket height as unstable momentum limit.\n'); + end +end +es_ini = 0; % lower limit of the stability threshold +et_ini = eu_ini / 2; % starting guess of the stability threshold +if verbose + fprintf('Unstable momentum limit set at start to %.3f%%.\n',100*eu_ini); +end + +% bipatition method for multiple points +orbit=findorbit6(THERING,REFPTS); + +% positive/negative branch +etp = et_ini*ones(np,1); +eup = eu_ini*ones(np,1); +esp = es_ini*ones(np,1); +etn = -et_ini*ones(np,1); +eun = -eu_ini*ones(np,1); +esn = -es_ini*ones(np,1); +de = 1; +iteration = 0; +while de>detole && iteration < 100 + iteration=iteration+1; + if verbose + seconds_initial=datetime('now'); + fprintf('Boundary search, iteration %d ...\n',iteration); + end + % L is true for particles lost on the track + L=Multiorigin_ringpass_islost( THERING, ... + REFPTS, ... + etp, ... + etn, ... + orbit, ... + nturns, ... + initcoord, ... + epsilon6D, ... + verbose ... + ); + % split in positive and negative side of energy offsets + Lp = L(1:2:2*np); + Ln = L(2:2:2*np); + % split in stable (es), unstable (eu) and test (et) energy + esp(Lp==0) = etp(Lp==0); + eup(Lp~=0) = etp(Lp~=0); + etp = (esp+eup)/2; + esn(Ln==0) = etn(Ln==0); + eun(Ln~=0) = etn(Ln~=0); + etn = (esn+eun)/2; + % define new energy step + dep = max(abs(esp-eup)); + den = max(abs(esn-eun)); + de = max(dep,den); + if verbose + elapsed_time=seconds(time(between(seconds_initial,datetime('now')))); + fprintf('%1.3f seconds. Momentum resolution is %1.3e and stops at %1.3e\n',elapsed_time,de,detole); + end +end +end + + +function Loste=Multiorigin_ringpass_islost( THERING, ... + refpts, ... + ep, ... + en, ... + orbit, ... + nturns, ... + initcoord, ... + epsilon6D, ... + verbose ... + ) +% Loste=Multiorigin_ringpass_islost( THERING, ... +% refpts, ... +% ep, ... +% en, ... +% orbit, ... +% nturns, ... +% initcoord, ... +% epsilon6D, ... +% verbose ... +% ) +% Returns a boolean array: tells whether the particle launched at +% the reference point refpts with positive and negative energy offset is +% lost or not. +% Inputs: +% -THERING: cell array Lattice used for the traking. +% -refpts: array [npossx1] with the elements number where to start the traking. +% -ep: array [npossx1] with the positive energy deviation at each refpts. +% -en: array [npossx1] with the negative energy deviation at each refpts. +% -orbit: array [npossx6] with the orbit at each refpts. +% -nturns: number of full turns to track. +% -initcoord: deviation from the 6D closed orbit, same dor each refpts. +% -epsilon6D: minimum 6D distance between particles +% -verbose: print additional info +% Output: +% -Loste : bool array [npossx2], True for lost particles. +% every pair Loste(2*k-1,2*k) for k = 1,... +% corresponds to the positive and negative energy offset. + +nposs=numel(refpts); +Loste=ones(1,2*nposs); +Rin=zeros(6,2*nposs); +Rout=zeros(6,2*nposs); +tinyoffset=epsilon6D; + +% first track the remaining portion of the ring +for ii=1:nposs + Line=THERING(refpts(ii):end); + Rin(:,2*ii-1) = orbit(:,ii) + [initcoord(1) 0 initcoord(2) 0 ep(ii) 0.0]'; + Rin(:,2*ii ) = orbit(:,ii) + [initcoord(1) 0 initcoord(2) 0 en(ii) 0.0]'; + [Rout(:,2*ii-1:2*ii),Loste(2*ii-1:2*ii)] = ringpass( ... + Line, ... + Rin(:,2*ii-1:2*ii) ... + ); +end +nalive1stturn = length(Loste)-sum(Loste); +Ralive1stturn = Rout(:,Loste==0); + +% use particles that have survived to the ring end, +% filter them if necessary, and track them +sizeRalive1turn = size(Ralive1stturn); +trackonly_mask = logical(1:sizeRalive1turn(2)); +similarparticles_index = []; +particles_were_filtered = false; +if (epsilon6D ~= 0) && (nalive1stturn > 1) + particles_were_filtered = true; + % search for non numerically similar particles + DiffR = squeeze(std(repmat(Ralive1stturn,[1 1 nalive1stturn]) ... + - repmat(reshape( ... + Ralive1stturn, ... + [6 1 nalive1stturn] ... + ), ... + [1 nalive1stturn 1] ... + ) ... + ) ... + ); + allposs = (1:nalive1stturn)'*ones(1,nalive1stturn); + similarposs = max(allposs.*(DiffR numpy.ndarray: + """ + :py:func:`momap_project2start` calculates the local momemtum aperture. + + It is a binary search of the negative and positive momentum thresholds + of stability around the closed orbit. + + For a given momentum offset the particles are first tracked from every + reference point to the end of the ring, and then, all or a set of particles + with different 6D coordinates are tracked together. The surviving particles + continue the boundary search with a new energy step until the boundary is + found when the step limit or the momentum limit is reached. + + Usage: + >>> momaperture_project2start(ring) + + Parameters: + ring: list of elements + + Keyword Arguments: + refpts: Selects the location of coordinates output. + See ":ref:`Selecting elements in a lattice `" + nturns: number of turns to be tracked. Default 1000 + dptol: momentum offset resolution. Default 1e-4 + dpuguess: maximum momentum boundary. Default: the rf bucket height. + troffset: (2, N) offsets to be added to the transverse coordinates + on the N reference points. Default 1e-5 m + orbit: (N,6) offsets to be added on the N reference points. + Default, the closed orbit + verbose: print in the standard output additional info. Default False + epsilon6d: float. + If not passed, all particles are tracked. + If epsilon6d is given, we track for nturns only particles + having 6D coordinates different by epsilon6d + + Returns: + dnp: (N,2) array with negative and positive stable momentum boundaries + for the N reference points + + ..note:: + * This function could track in parallel. Set use_mp=True. + Other arguments could be passed, check :py:func:`lattice_track`. + * This function does a quick search, but, it is succeptible to miss + islands of instability due to the varying energy step. + """ + # verboseprint to check flag only once + verbose = kwargs.pop("verbose", False) + verboseprint = print if verbose else lambda *a, **k: None + + rps = kwargs.pop("refpts", ring.uint32_refpts(range(len(ring)))) + rps = ring.get_uint32_index(rps) + nrps = len(rps) + verboseprint(f"Using {nrps} reference points") + + nturns = kwargs.pop("nturns", 1000) + verboseprint(f"Track over {nturns} turns") + + dptol = kwargs.pop("dptol", 1e-4) + verboseprint(f"Momentum resolution {dptol}") + + # set transverse offsets + if "troffset" in kwargs: + add_offset = kwargs.pop("troffset") + verboseprint("Add user offsets") + else: + dxy = 1e-5 + add_offset = numpy.tile(dxy, [2, nrps]) + verboseprint(f"Adding default transverse offsets {dxy} per plane") + + # set the minimum distance btw particles that makes them similar + epsilon6d = kwargs.pop("epsilon6d", 0) + + # first guess + if "dpuguess" in kwargs: + eu_ini = kwargs.pop("dpuguess") + verboseprint(f"Using the users max boundary {eu_ini}") + else: + # use radiation parameters to get the rf bucket + pars = ring.radiation_parameters() + # get energy bucket (deltabucket) max height. + # S.Y.Lee, 4th Edition, Eqs 3.37, 3.38, Sect. II.2 Bucket area + harmonnumber = ring.harmonic_number + etac = pars.etac + theenergy = pars.E0 + phis = pars.phi_s + betar = ring.beta + thevoltage = ring.get_rf_voltage() + yfactor = numpy.sqrt( + numpy.abs(numpy.cos(phis) - 0.5 * (numpy.pi - 2 * phis) * numpy.sin(phis)) + ) + deltabucket = ( + numpy.sqrt( + 2 + * thevoltage + / (numpy.pi * betar**2 * theenergy * harmonnumber * numpy.abs(etac)) + ) + * yfactor + ) + verboseprint(f"Bucket height {deltabucket}") + eu_ini = deltabucket + verboseprint("Using the bucket height as maximum boundary") + es_ini = 0 + et_ini = eu_ini / 2 + + if "orbit" in kwargs: + orbit_s = kwargs.pop("orbit") + verboseprint("Using the users orbit") + else: + _, orbit_s = ring.find_orbit(rps) + verboseprint("Using the closed orbit") + + # start scan of: + # unstable energy + # stable energy + # test energy + etpos = et_ini * numpy.ones(nrps) + eupos = eu_ini * numpy.ones(nrps) + espos = es_ini * numpy.ones(nrps) + etneg = -1 * et_ini * numpy.ones(nrps) + euneg = -1 * eu_ini * numpy.ones(nrps) + esneg = -1 * es_ini * numpy.ones(nrps) + deltae = 1 + iteration = 0 + while iteration < 100 and (deltae > dptol): # safety limit on iterations + iteration = iteration + 1 + t00 = time.time() + # plost is a mask, True for lost particles + plost = multirefpts_track_islost( + ring, + rps, + etpos, + etneg, + orbit_s, + add_offset, + nturns, + epsilon6d, + verbose, + **kwargs, + ) + # split in positive and negative side of energy offsets + plostpos = plost[0::2] + plostneg = plost[1::2] + # split in stable (es), unstable (eu) and test (et) energy + espos[~plostpos] = etpos[~plostpos] + eupos[plostpos] = etpos[plostpos] + etpos = (espos + eupos) / 2 + esneg[~plostneg] = etneg[~plostneg] + euneg[plostneg] = etneg[plostneg] + etneg = (esneg + euneg) / 2 + # define new energy step + depos = max(abs(espos - eupos)) + deneg = max(abs(esneg - euneg)) + deltae = max(depos, deneg) + outmsg = ( + f"Iteration {iteration}", + f" took {format(time.time()-t00):.3} s.", + f" dp_step={deltae}, dptol={dptol}", + ) + verboseprint("".join(outmsg)) + return numpy.vstack([etneg, etpos]).T + + +def projectrefpts( + ring: Lattice, + startrefpts: numpy.ndarray, + particles: numpy.ndarray, + **kwargs: Dict[str, any], +) -> tuple: + """ + :py:fun:`projectrefpts` tracks from multiple reference points. + + Usage: + >>> projectrefpts(ring, startrefpts, particles) + + Parameters: + ring: list of elements + startrefpts: reference point to start tracking from. + particles: (6,N,R,1) array, where N is the number of + particles per reference point, and R is the number of + reference points. + + Keyword arguments: + endrefpt: end reference point. Default: end of last ring element + use_mp: Default True. See :py:fun:`lattice_track` + group: Default False. The starting point info is removed. + All tracked particles are grouped together. + verbose: prints additional info + + Returns: + zout: Tracked particle coordinates at the end ref. point. + Default (6,N,R,1) array, where N is the number of + particles per R references points + If the flag 'group' is used the output becomes a + (6,N*R,1,1) array with all particles. + lostpart: Bool array, True when the particle is lost. + Default (N,R). + If the flag 'group' is used the output becomes a + (N*R) array + """ + lenring = len(ring) + rps = startrefpts + nrps = len(rps) + nparticles = 1 + if len(particles.shape) >= 2: + nparticles = particles.shape[1] + + # verboseprint to check flag only once + verbose = kwargs.pop("verbose", False) + verboseprint = print if verbose else lambda *a, **k: None + + if "endrefpt" in kwargs: + erps = kwargs["endrefpt"] + verboseprint(f"Project particles to start of element index {erps}") + else: + erps = lenring + verboseprint("Project to end point") + + groupparts = kwargs.pop("group", False) + + verboseprint(f"nparticles={nparticles} per reference point") + verboseprint(f"Number of reference points {nrps}") + + # default to parallel + use_mp = kwargs.pop("use_mp", True) + if nparticles == 1: + use_mp = False + + if groupparts: + zout = numpy.zeros((6, nparticles * nrps, 1, 1)) + lostpart = numpy.ones((nparticles * nrps), dtype=bool) + else: + zout = numpy.zeros((6, nparticles, nrps, 1)) + lostpart = numpy.ones((nparticles, nrps), dtype=bool) + + # first, track the remaining portion of the ring + zin = particles.copy() + for i in range(nrps): + ring_downstream = ring.rotate(rps[i]) + zaux = numpy.squeeze(zin[:, :, i, 0]) + verboseprint( + f"Tracking {nparticles} particles on reference point {i+1} of {nrps}" + ) + zoaux, _, dout1 = ring_downstream.track( + zaux, nturns=1, refpts=erps - rps[i], losses=True, use_mp=use_mp + ) + if groupparts: + zout[:, nparticles * i : nparticles * (i + 1), 0, 0] = numpy.reshape( + zoaux, (6, nparticles) + ) + lostpart[nparticles * i : nparticles * (i + 1)] = dout1["loss_map"][ + "islost" + ] + else: + zout[:, :, i, 0] = numpy.reshape(zoaux, (6, nparticles)) + lostpart[:, i] = dout1["loss_map"]["islost"] + return zout, lostpart + + +def multirefpts_track_islost( + ring: Lattice, + refpts: numpy.ndarray, + esetptpos: numpy.ndarray, + esetptneg: numpy.ndarray, + orbit: numpy.ndarray, + initcoord: numpy.ndarray, + nturns: float, + epsilon6d: float, + verbose: bool, + **kwargs: Dict[str, any], +) -> numpy.ndarray: + """ + Tell whether the particle launched is lost. + + Usage: + >>> multirefpts_track_islost(ring, refpts, energysetpt, orbit, initcoord) + + Parameters: + ring: list of elements + + Keyword Arguments: + refpts: Selects the locations. + esetptpos: positive energy set point for tracking. + esetptneg: negative energy set point for tracking. + orbit: (6,N) orbit to be added to the N refpts. + initcoords: (2,N) hor. and ver. transverse offsets in m. + nturns: number of turns to track + epsilon6d: maximum value to consider particles as similar + verbose: print info + + Returns: + Lostpart: (2*N) bool array, for N reference points. + True if the particle is lost. + """ + verboseprint = print if verbose else lambda *a, **k: None + + # track positive and negative side at the same time + nparticles = 2 + + rps = refpts + nrps = len(rps) + lostpart = numpy.ones((nparticles * nrps), dtype=bool) + zin = numpy.zeros((6, nparticles * nrps)) + zout = numpy.zeros((6, nparticles * nrps)) + tinyoffset = epsilon6d + + # project to the end of the ring + erps = len(ring) + + # first, track the remaining portion of the ring + zin[:, 0::2] = orbit.T.copy() + zin[:, 1::2] = orbit.T.copy() + zin[0, 0::2] = zin[0, 0::2] + initcoord[0, :] + zin[0, 1::2] = zin[0, 1::2] + initcoord[0, :] + zin[2, 0::2] = zin[2, 0::2] + initcoord[1, :] + zin[2, 1::2] = zin[2, 1::2] + initcoord[1, :] + zin[4, 0::2] = zin[4, 0::2] + esetptpos + zin[4, 1::2] = zin[4, 1::2] + esetptneg + for i in range(nrps): + ring_downstream = ring.rotate(rps[i]) + zaux = zin[:, (nparticles * i) : (nparticles * (i + 1))] + # verboseprint( + # f"Tracking {nparticles} particles on reference point {i+1} of {nrps}" + # ) + zoaux, _, doutaux = ring_downstream.track( + zaux, nturns=1, refpts=erps - rps[i], losses=True, **kwargs + ) + zout[:, (nparticles * i) : (nparticles * (i + 1))] = numpy.reshape( + zoaux, (6, nparticles) + ) + lostpart[(nparticles * i) : (nparticles * (i + 1))] = doutaux["loss_map"][ + "islost" + ] + + cntalive = len(lostpart) - sum(lostpart) + zinaliveaux = zout[:, ~lostpart] + zinalive_at_ring_end = numpy.asfortranarray(zinaliveaux.copy()) + + # second, use the particles that have survived the ring to the end + # filter them if necessary, and track them + shapealiveatend = numpy.shape(zinalive_at_ring_end) + trackonly_mask = numpy.ones(shapealiveatend[1], dtype=bool) + similarparticles_index = numpy.array([]) + particles_were_filtered = False + if epsilon6d != 0 and cntalive > 1: + particles_were_filtered = True + # search for non numerically similar particles + closenessmatrix = numpy.zeros((cntalive, cntalive), dtype=bool) + for i in range(cntalive): + closenessmatrix[i, :] = [ + numpy.allclose( + zinalive_at_ring_end[:, j], + zinalive_at_ring_end[:, i], + atol=tinyoffset, + ) + for j in range(cntalive) + ] + _, rowidx = numpy.indices((cntalive, cntalive)) + maxidx = numpy.max(rowidx * closenessmatrix, 1) + trackonly_mask, _, similarparticles_index = numpy.unique( + maxidx, return_index=True, return_inverse=True + ) + outmsg = ( + "Speed up when discarding similar particles, ", + f"{100*len(trackonly_mask)/cntalive:.3f}%", + ) + verboseprint("".join(outmsg)) + # track + # dummy track to later reuse the ring + zaux2 = numpy.array([initcoord[0, 0], 0, initcoord[1, 0], 0, 1e-6, 0]).T + ring.track(zaux2, nturns=1) + # track non-numerically similar particles + _, _, dout_multiturn = ring.track( + zinalive_at_ring_end[:, trackonly_mask], + nturns=nturns, + keep_lattice=True, + losses=True, + **kwargs, + ) + if particles_were_filtered: + lostpaux = dout_multiturn["loss_map"]["islost"][similarparticles_index] + else: + lostpaux = dout_multiturn["loss_map"]["islost"] + lostpart[~lostpart] = lostpaux + + return lostpart