def fit_plane(x, y, z):
xx, yy = m.meshgrid(x, y)
pts = m.isfinite(z)
xx_, yy_ = xx[pts].flatten(), yy[pts].flatten()
flat = m.ones(xx_.shape)
coefs = m.lstsq(m.stack([xx_, yy_, flat]).T, z[pts].flatten(),
rcond=None)[0]
plane_fit = coefs[0] * xx + coefs[1] * yy + coefs[2]
return plane_fit
def fit_sphere(z):
x, y = m.linspace(-1, 1, z.shape[1]), m.linspace(-1, 1, z.shape[0])
xx, yy = m.meshgrid(x, y)
pts = m.isfinite(z)
xx_, yy_ = xx[pts].flatten(), yy[pts].flatten()
rho, phi = cart_to_polar(xx_, yy_)
focus = defocus(rho, phi)
coefs = m.lstsq(m.stack([focus, m.ones(focus.shape)]).T, z[pts].flatten(), rcond=None)[0]
rho, phi = cart_to_polar(xx, yy)
sphere = defocus(rho, phi) * coefs[0]
return sphere
def fit_plane(x, y, z):
pts = m.isfinite(z)
if len(z.shape) > 1:
x, y = m.meshgrid(x, y)
xx, yy = x[pts].flatten(), y[pts].flatten()
else:
xx, yy = x, y
flat = m.ones(xx.shape)
coefs = m.lstsq(m.stack([xx, yy, flat]).T, z[pts].flatten(), rcond=None)[0]
plane_fit = coefs[0] * x + coefs[1] * y + coefs[2]
return plane_fit
def fit(data, num_terms=16, rms_norm=False, round_at=6):
''' Fits a number of Zernike coefficients to provided data by minimizing
the root sum square between each coefficient and the given data. The
data should be uniformly sampled in an x,y grid.
Args:
data (`numpy.ndarray`): data to fit to.
num_terms (`int`): number of terms to fit, fits terms 0~num_terms.
rms_norm (`bool`): if true, normalize coefficients to unit RMS value.
round_at (`int`): decimal place to round values at.
Returns:
numpy.ndarray: an array of coefficients matching the input data.
'''
if num_terms > len(zernmap):
raise ValueError(f'number of terms must be less than {len(zernmap)}')
# precompute the valid indexes in the original data
pts = m.isfinite(data)
# set up an x/y rho/phi grid to evaluate Zernikes on
x, y = m.linspace(-1, 1, data.shape[1]), m.linspace(-1, 1, data.shape[0])
xx, yy = m.meshgrid(x, y)
rho = m.sqrt(xx**2 + yy**2)[pts].flatten()
phi = m.arctan2(xx, yy)[pts].flatten()
# compute each Zernike term
zernikes = []
for i in range(num_terms):
func = z.zernikes[zernmap[i]]
base_zern = func(rho, phi)
if rms_norm:
base_zern *= func.norm
zernikes.append(base_zern)
zerns = m.asarray(zernikes).T
# use least squares to compute the coefficients
coefs = m.lstsq(zerns, data[pts].flatten(), rcond=None)[0]
return coefs.round(round_at)
def zernikefit(data, x=None, y=None, rho=None, phi=None, terms=16, norm=False, residual=False, round_at=6, map_='fringe'):
"""Fits a number of Zernike coefficients to provided data.
Works by minimizing the mean square error between each coefficient and the
given data. The data should be uniformly sampled in an x,y grid.
Parameters
----------
data : `numpy.ndarray`
data to fit to.
x : `numpy.ndarray`, optional
x coordinates, same shape as data
y : `numpy.ndarray`, optional
y coordinates, same shape as data
rho : `numpy.ndarray`, optional
radial coordinates, same shape as data
phi : `numpy.ndarray`, optional
azimuthal coordinates, same shape as data
terms : `int`, optional
number of terms to fit, fits terms 0~terms
norm : `bool`, optional
if True, normalize coefficients to unit RMS value
residual : `bool`, optional
if True, return a tuple of (coefficients, residual)
round_at : `int`, optional
decimal place to round values at.
map_ : `str`, optional, {'fringe', 'noll'}
which ordering of Zernikes to use
Returns
-------
coefficients : `numpy.ndarray`
an array of coefficients matching the input data.
residual : `float`
RMS error between the input data and the fit.
Raises
------
ValueError
too many terms requested.
"""
map_ = maps[map_]
if terms > len(fringemap):
raise ValueError(f'number of terms must be less than {len(fringemap)}')
data = data.T # transpose to mimic transpose of zernikes
# precompute the valid indexes in the original data
pts = m.isfinite(data)
if x is None and rho is None:
# set up an x/y rho/phi grid to evaluate Zernikes on
rho, phi = make_rho_phi_grid(*reversed(data.shape))
rho = rho[pts].flatten()
phi = phi[pts].flatten()
elif rho is None:
rho, phi = cart_to_polar(x, y)
rho, phi = rho[pts].flatten(), phi[pts].flatten()
# compute each Zernike term
zerns_raw = []
for i in range(terms):
func = zernikes[map_[i]]
base_zern = func(rho, phi)
if norm:
base_zern *= func.norm
zerns_raw.append(base_zern)
zerns = m.asarray(zerns_raw).T
# use least squares to compute the coefficients
meas_pts = data[pts].flatten()
coefs = m.lstsq(zerns, meas_pts, rcond=None)[0]
if round_at is not None:
coefs = coefs.round(round_at)
if residual is True:
components = []
for zern, coef in zip(zerns_raw, coefs):
components.append(coef * zern)
_fit = m.asarray(components)
_fit = _fit.sum(axis=0)
rmserr = rms(data[pts].flatten() - _fit)
return coefs, rmserr
else:
return coefs