Skip to content

Commit

Permalink
Calculation of full Jacobian and removed transpositions on "user side"
Browse files Browse the repository at this point in the history 
    included BPM and corrector scalings
    code cleanup
    Some transposition which may not be necesary remain inside the code
  • Loading branch information
@lmalina
lmalina committed Jan 15, 2024
1 parent 94aaa5f commit 5209961
Show file tree
Hide file tree
Showing 2 changed files with 28 additions and 145 deletions.
131 changes: 21 additions & 110 deletions pySC/correction/loco.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@
LOGGER = logging_tools.get_logger(__name__)


def calculate_jacobian(SC, C_model, dkick, used_cor_ind, bpm_indexes, cor_ind, quads_ind, dk, trackMode=TRACK_ORB,
def calculate_jacobian(SC, C_model, dkick, used_cor_ind, bpm_indexes, quads_ind, dk, trackMode=TRACK_ORB,
useIdealRing=True, skewness=False, order=1, method=SETTING_ADD, includeDispersion=False, rf_step=1E3,
cav_ords=None, full_jacobian=True):
pool = multiprocessing.Pool()
Expand All @@ -25,37 +25,18 @@ def calculate_jacobian(SC, C_model, dkick, used_cor_ind, bpm_indexes, cor_ind, q
pool.join()

results = [result / dk for result in results]
if full_jacobian: # # Construct the complete Jacobian matrix for the LOCO
# TODO modify for calibration errors of given size
if full_jacobian: # Construct the complete Jacobian matrix for the LOCO
# assuming only linear scaling errors of BPMs and corrector magnets
n_correctors = len(np.concatenate(used_cor_ind))
n_bpms = len(bpm_indexes) * 2 # in both planes

return np.concatenate((results, np.tile(C_model, (n_correctors + n_bpms, 1, 1))))

'''
pool = multiprocessing.Pool()
cor_args = [(cor_ind, SC, C_model, dkick, used_cor_ind, bpm_indexes, dk, trackMode, useIdealRing,
skewness, order, method, includeDispersion, rf_step, cav_ords) for quad_index in quads_ind]
results = pool.map(generating_cor_response_matrices, cor_args)
pool.close()
pool.join()
Jc = [result / dkick for result in results]
pool = multiprocessing.Pool()
bpm_args = [(bpm_indexes, SC, C_model, dkick, used_cor_ind, bpm_indexes, dk, trackMode, useIdealRing,
skewness, order, method, includeDispersion, rf_step, cav_ords) for quad_index in quads_ind]
results = pool.map(generating_bpm_response_matrices, bpm_args)
pool.close()
pool.join()
Jb = [result / dkick for result in results]
return np.concatenate((results,Jc, Jb), axis=0)))
'''

j_cor = np.zeros((n_correctors,) + C_model.shape)
for i in range(n_correctors):
j_cor[i, :, i] = C_model[:, i] # a single column of response matrix corresponding to a corrector
j_bpm = np.zeros((n_bpms,) + C_model.shape)
for i in range(n_bpms):
j_bpm[i, i, :] = C_model[i, :] # a single row of response matrix corresponding to a given plane of BPM
# return np.concatenate((results, np.tile(C_model, (n_correctors + n_bpms, 1, 1))))
return np.concatenate((results, j_cor, j_bpm), axis=0)
return results


Expand All @@ -79,58 +60,6 @@ def generating_quads_response_matrices(args):
SC.set_magnet_setpoints(quad_index, -dk, skewness, order, method)
return np.hstack((C_measured - C_model, (dispersion_meas - dispersion_model).reshape(-1, 1)))

def generating_cor_response_matrices(args):
(cor_index, SC, C_model, correctors_kick, used_cor_indexes, used_bpm_indexes, correctors_kick, trackMode, useIdealRing,
skewness, order, method, includeDispersion, rf_step, cav_ords) = args
LOGGER.debug('generating response to cor of index', cor_index)
if not includeDispersion:
#SC.set_magnet_setpoints(cor_index, dk, skewness, order, method)
err1 = SC.RING[cor_index].CalErrorA[0]
err2 = SC.RING[cor_index].CalErrorB[0]

SC.register_magnets(cor_index, CalErrorA=np.array([correctors_kick, 0]), CalErrorB=np.array([correctors_kick, 0]))
C_measured = SCgetModelRM(SC, used_bpm_indexes, used_cor_indexes, dkick=correctors_kick,
useIdealRing=useIdealRing,
trackMode=trackMode)
SC.register_magnets(cor_index, CalErrorA=np.array([err1, 0]), CalErrorB=np.array([err2, 0]))
#SC.set_magnet_setpoints(cor_index, -dk, skewness, order, method)
return C_measured - C_model

dispersion_model = SCgetModelDispersion(SC, used_bpm_indexes, CAVords=cav_ords)
SC.register_magnets(cor_index, CalErrorA=np.array([correctors_kick, 0]), CalErrorB=np.array([correctors_kick, 0]))
C_measured = SCgetModelRM(SC, used_bpm_indexes, used_cor_indexes, dkick=correctors_kick, useIdealRing=useIdealRing,
trackMode=trackMode)
dispersion_meas = SCgetModelDispersion(SC, used_bpm_indexes, CAVords=cav_ords, rfStep=rf_step)
SC.register_magnets(cor_index, CalErrorA=np.array([err1, 0]), CalErrorB=np.array([err2, 0]))
return np.hstack((C_measured - C_model, (dispersion_meas - dispersion_model).reshape(-1, 1)))


def generating_bpm_response_matrices(args):
(bpm_indexes, SC, C_model, correctors_kick, used_cor_indexes, used_bpm_indexes, correctors_kick, trackMode, useIdealRing,
skewness, order, method, includeDispersion, rf_step, cav_ords) = args
LOGGER.debug('generating response to bpm of index', bpm_indexes)
if not includeDispersion:
# SC.set_magnet_setpoints(cor_index, dk, skewness, order, method)
err1 = SC.RING[bpm_indexes].CalError[0]
err2 = SC.RING[bpm_indexes].CalError[1]

SC.register_magnets(bpm_indexes, CalError=correctors_kick)
C_measured = SCgetModelRM(SC, used_bpm_indexes, used_cor_indexes, dkick=correctors_kick,
useIdealRing=useIdealRing,
trackMode=trackMode)
SC.register_magnets(bpm_indexes, CalError=correctors_kick)
# SC.set_magnet_setpoints(cor_index, -dk, skewness, order, method)
return C_measured - C_model

dispersion_model = SCgetModelDispersion(SC, used_bpm_indexes, CAVords=cav_ords)
SC.register_magnets(bpm_indexes, CalError=correctors_kick)
C_measured = SCgetModelRM(SC, used_bpm_indexes, used_cor_indexes, dkick=correctors_kick, useIdealRing=useIdealRing,
trackMode=trackMode)
dispersion_meas = SCgetModelDispersion(SC, used_bpm_indexes, CAVords=cav_ords, rfStep=rf_step)
SC.register_magnets(bpm_indexes, CalError=correctors_kick)
return np.hstack((C_measured - C_model, (dispersion_meas - dispersion_model).reshape(-1, 1)))



def measure_closed_orbit_response_matrix(SC, bpm_ords, cm_ords, dkick=1e-5):
LOGGER.info('Calculating Measure response matrix')
Expand Down Expand Up @@ -166,18 +95,15 @@ def loco_correction_ng(initial_guess0, orm_model, orm_measured, J, lengths, incl
initial_guess = initial_guess0.copy()
mask = _get_parameters_mask(including_fit_parameters, lengths)
residuals = objective(initial_guess[mask], orm_measured - orm_model, J[mask, :, :], weights)
r = residuals.reshape(orm_model.shape)
r = residuals.reshape(np.transpose(orm_model).shape)
t2 = np.zeros([len(initial_guess), 1])
for i in range(len(initial_guess)):
t2[i] = np.sum(np.dot(np.dot(J[i], weights), r.T))
results = get_inverse(J, t2, s_cut, weights)

return results

t2[i] = np.sum(np.dot(np.dot(J[i].T, weights), r.T))
return get_inverse(J, t2, s_cut, weights)


def objective(masked_params, orm_residuals, masked_jacobian, weights):
return np.dot(orm_residuals - np.einsum("ijk,i->jk", masked_jacobian, masked_params),
return np.dot(np.transpose(orm_residuals - np.einsum("ijk,i->jk", masked_jacobian, masked_params)),
np.sqrt(weights)).ravel()


Expand Down Expand Up @@ -235,24 +161,9 @@ def select_equally_spaced_elements(total_elements, num_elements):

def get_inverse(jacobian, B, s_cut, weights, plot=False):
n_resp_mats = len(jacobian)
#matrix = np.zeros([n_resp_mats, n_resp_mats])
#for i in range(n_resp_mats):
# for j in range(n_resp_mats):
# matrix[i, j] = np.sum(np.dot(np.dot(jacobian[i], weights), jacobian[j].T))
sum_ = np.sum(jacobian, axis=1) # Sum over i and j for all planes
matrix = sum_ @ weights @ sum_.T
matrixi = SCgetPinv(matrix, num_removed=n_resp_mats - min(n_resp_mats, s_cut), plot=plot)

'''
u, s, v = np.linalg.svd(matrix, full_matrices=True)
smat = 0.0 * matrix
si = s ** -1
n_sv = s_cut
si[n_sv:] *= 0.0
smat[:n_resp_mats, :n_resp_mats] = np.diag(si)
matrixi = np.dot(v.transpose(), np.dot(smat.transpose(), u.transpose()))
'''
r = (np.dot(matrixi, B)).reshape(-1)
e = np.dot(matrix, r).reshape(-1) - B.reshape(-1)

return r
sum_corr = np.sum(jacobian, axis=2) # Sum over i and j for all planes
matrix = np.dot(np.dot(sum_corr, weights), sum_corr.T)
inv_matrix = SCgetPinv(matrix, num_removed=n_resp_mats - min(n_resp_mats, s_cut), plot=plot)
results = np.ravel(np.dot(inv_matrix, B))
e = np.ravel(np.dot(matrix, results)) - np.ravel(B)
return results
42 changes: 7 additions & 35 deletions pySC/example_hmba.py
View file
Original file line number Diff line number Diff line change
@@ -1,9 +1,3 @@
#!/usr/bin/env python
# coding: utf-8

# In[2]:


import at
import numpy as np
from at import Lattice
Expand Down Expand Up @@ -65,27 +59,12 @@ def create_sc_class():
orbit_response_matrix_model = SCgetModelRM(SC, used_bpms_ords, used_cor_ords, trackMode='ORB', useIdealRing=True, dkick=CMstep)
model_dispersion = SCgetModelDispersion(SC, used_bpms_ords, cav_ords, trackMode='ORB', Z0=np.zeros(6), nTurns=1,
rfStep=RFstep, useIdealRing=True)
Jn = loco.calculate_jacobian(SC, orbit_response_matrix_model, CMstep, used_cor_ords, used_bpms_ords, cor_ords, np.concatenate(quads_ords), dk,
Jn = loco.calculate_jacobian(SC, orbit_response_matrix_model, CMstep, used_cor_ords, used_bpms_ords, np.concatenate(quads_ords), dk,
trackMode='ORB', useIdealRing=False, skewness=False, order=1, method='add',
includeDispersion=False, rf_step=RFstep, cav_ords=cav_ords)
Jn = np.transpose(Jn, (0, 2, 1))
#weights = 1
weights = np.eye(len(used_bpms_ords) * 2)
tmp = np.sum(Jn, axis=1)
A = tmp @ weights @ tmp.T
u, s, v = np.linalg.svd(A, full_matrices=True)
import matplotlib.pyplot as plt

plt.plot(np.log(s), 'd--')
plt.title('singular value')
plt.xlabel('singular values')
plt.ylabel('$\log(\sigma_i)$')
plt.show()

n_singular_values = 20

#Jt = loco.get_inverse(Jn, n_singular_values, weights)

_, _, twiss_err = at.get_optics(SC.RING, SC.ORD.BPM)
bx_rms_err, by_rms_err = loco.model_beta_beat(SC.RING, twiss, SC.ORD.BPM, plot=False)
info_tab = 14 * " "
Expand All @@ -108,14 +87,14 @@ def create_sc_class():
initial_guess[lengths[0] + lengths[1]:] = 1e-6

# method lm (least squares)
#fit_parameters = loco.loco_correction_lm(initial_guess, np.transpose(orbit_response_matrix_model),
# np.transpose(orbit_response_matrix_measured), Jn, lengths,
# including_fit_parameters, bounds=(-0.03, 0.03), weights=weights, verbose=2)
# fit_parameters = loco.loco_correction_lm(initial_guess, orbit_response_matrix_model,
# orbit_response_matrix_measured, Jn, lengths,
# including_fit_parameters, bounds=(-0.03, 0.03), weights=weights, verbose=2)

# method ng
fit_parameters = loco.loco_correction_ng(initial_guess, np.transpose(orbit_response_matrix_model),
np.transpose(orbit_response_matrix_measured), Jn, lengths,
including_fit_parameters, n_singular_values, weights=weights)
fit_parameters = loco.loco_correction_ng(initial_guess, orbit_response_matrix_model,
orbit_response_matrix_measured, Jn, lengths,
including_fit_parameters, n_singular_values, weights=weights)

dg = fit_parameters[:lengths[0]] if len(fit_parameters) > n_quads else fit_parameters
SC = loco.set_correction(SC, dg, np.concatenate(quads_ords))
Expand All @@ -125,10 +104,3 @@ def create_sc_class():
LOGGER.info(f"Correction reduction: \n"
f" beta_x: {(1 - bx_rms_cor / bx_rms_err) * 100:.2f}%\n"
f" beta_y: {(1 - by_rms_cor / by_rms_err) * 100:.2f}%\n")


# In[ ]:




0 comments on commit 5209961

Please sign in to comment.