Add ellaume kickmap class

[Gabrielrezende-asc]

This PR adds a new class for calculating kickmaps based on the Elleaume formalism. It also includes some linter-driven formatting changes.

[fernandohds564]

Very nice, @Gabrielrezende-asc . I have a few suggestions in the code.

[fernandohds564]

This method coulde be a staticmethod

[fernandohds564]

you can use the function polyvander2d from numpy, then you won't even need the method above to calculate the matrix for the fitting.
Besides, I think there was an error in the calculation of the residue, right? the square should be inside the sum.
Additionally, I think lstsq already return the residue, so we can get directly from it

edit: I noticed you were only calculating the fitting matrix up to a certain total degree of the polynomial, so I adjusted the suggestion to do the same. I'm also keeping the coefficients as a 2D array, to make its meaning more clearer and to ease the calculation of the derivatives bellow.

Suggested change

	def fit_2d_polyfit(self, degree):
	if self.posx is None or self.posy is None:
	raise ValueError('posx and posy must be set before calculating potential.')
	if self.potential is None:
	raise ValueError('Potential has not been calculated yet. Call calc_full_potential first.')
	potential = self.potential
	posx = self.posx
	posy = self.posy
	matrix = self.calc_2d_polyfit_matrix(degree)
	invmat = _np.linalg.pinv(matrix)
	pot_vec = _np.reshape(potential, len(posx)*len(posy), order='C')
	coefs = _np.dot(invmat, pot_vec)
	potential_fit = _np.reshape(_np.dot(matrix, coefs), (len(posx), len(posy)), order='C')
	residue = _np.sqrt(_np.sum(potential_fit-potential)**2)
	self.matrix_poly = matrix
	self.fit_coefs = coefs
	return coefs, residue
	def fit_2d_polyfit(self, degree):
	if self.posx is None or self.posy is None:
	raise ValueError('posx and posy must be set before calculating potential.')
	if self.potential is None:
	raise ValueError('Potential has not been calculated yet. Call calc_full_potential first.')
	potential = self.potential
	posx = self.posx
	posy = self.posy
	X, Y = _np.meshgrid(posx, posy)
	matrix = _np.polynomial.polynomial.polyvander2d(
	X.ravel(), Y.ravel(), [degree, degree]
	)
	coefs = np.zeros((degree + 1, degree + 1))
	pows = _np.arange(degree + 1)
	idcs = (pows[:, None] + pows[None, :]).ravel() <= degree
	matrix = matrix[:, idcs]
	coefs.ravel()[idcs], residue, *_ = _np.linalg.lstsq(
	matrix, potential.ravel(), rcond=None
	)
	self.matrix_poly = matrix
	self.fit_coefs = coefs
	return coefs, residue[0]

[fernandohds564]

Using numpy features:

Suggested change

	def calc_potential_fit(self, degree):
	if self.posx_fit is None or self.posy_fit is None:
	raise ValueError('posx_fit and posy_fit must be set before calculating potential.')
	posx = 1e3*self.posx_fit # convert [m] to [mm]
	posy = 1e3*self.posy_fit # convert [m] to [mm]
	coefs = self.fit_coefs
	matrix = self.calc_2d_polyfit_matrix(degree, posx, posy)
	potential_fit = _np.reshape(_np.dot(matrix, coefs), (len(posx), len(posy)), order='C')
	self.matrix_poly = matrix
	self.potential_fit = potential_fit
	return potential_fit
	def calc_potential_fit(self, degree):
	if self.posx_fit is None or self.posy_fit is None:
	raise ValueError('posx_fit and posy_fit must be set before calculating potential.')
	posx = 1e3*self.posx_fit # convert [m] to [mm]
	posy = 1e3*self.posy_fit # convert [m] to [mm]
	X, Y = _np.meshgrid(posx, posy)
	coefs = self.fit_coefs
	potential_fit = _np.polynomial.polynomial.polyval2d(X.ravel(), Y.ravel(), coefs)
	potential_fit = potential_fit.reshape(posx.size, posy.size)
	self.potential_fit = potential_fit
	return potential_fit

[fernandohds564]

@Gabrielrezende-asc, could you explain to me later in person what this method do? does it creates a simple difference matrix with nearest neighboors normalized by the distance of the local y coordinate? If so, you can calculate derivatives without this matrix, by only doing something similar to np.diff(matrix, axis=0) / np.diff(y)[:, None].

[fernandohds564]

I think I understood. This is a polynomial derivative, right?

I think you can also use numpy features to do this by reshaping the coefficients array and using numpy.polynomial.polynomial.polyder(self.coefs, axis=0) for y derivatives or with axis=1 for x derivatives, considering you accepted my suggestion above of saving the coefficients as a 2D array.

[fernandohds564]

Suggested change

	def calc_dely(self):
	dy = self.calc_dy_operator()
	potential_fit = self.potential_fit
	coefs = self.fit_coefs
	matrix = self.matrix_poly
	coefs_dy = _np.dot(dy, coefs)
	dp_dy = _np.dot(matrix, coefs_dy)
	dp_dy = _np.reshape(dp_dy, potential_fit.shape, order='C')
	return dp_dy
	def calc_dely(self):
	der = _np.polynomial.polynomial.polyder(self.fit_coefs, axis=0)
	X, Y = _np.meshgrid(self.posx, self.posy)
	return _np.polynomial.polynomial.polyval2d(X.ravel(), Y.ravel(), der)