def test_sliced_input(self):
# cython code chokes on non C contiguous arrays
xx = np.linspace(-1, 1, 100)
x = xx[::3]
y = xx[::3]
t = _augknt(x, 1)
make_lsq_spline(x, y, t, k=1)
def robust_spline_fit(y):
bspl = interp.make_lsq_spline(x=x, y=y, t=knots, k=order)
resid = np.abs(bspl(x) - y)
keep_idx = resid <= np.percentile(resid, (1 - q) * 100)
bspl = interp.make_lsq_spline(
x=x[keep_idx], y=y[keep_idx], t=knots, k=order)
return bspl(x)
def test_sliced_input(self):
# cython code chokes on non C contiguous arrays
xx = np.linspace(-1, 1, 100)
x = xx[::3]
y = xx[::3]
t = _augknt(x, 1)
make_lsq_spline(x, y, t, k=1)
def test_complex(self):
# cmplx-valued `y`
x, t, k = self.x, self.t, self.k
yc = self.y * (1. + 2.j)
b = make_lsq_spline(x, yc, t, k)
b_re = make_lsq_spline(x, yc.real, t, k)
b_im = make_lsq_spline(x, yc.imag, t, k)
assert_allclose(b(x), b_re(x) + 1.j * b_im(x), atol=1e-15, rtol=1e-15)
def test_complex(self):
# cmplx-valued `y`
x, t, k = self.x, self.t, self.k
yc = self.y * (1. + 2.j)
b = make_lsq_spline(x, yc, t, k)
b_re = make_lsq_spline(x, yc.real, t, k)
b_im = make_lsq_spline(x, yc.imag, t, k)
assert_allclose(b(x), b_re(x) + 1.j*b_im(x), atol=1e-15, rtol=1e-15)
def test_weights(self):
# weights = 1 is same as None
x, y, t, k = self.x, self.y, self.t, self.k
w = np.ones_like(x)
b = make_lsq_spline(x, y, t, k)
b_w = make_lsq_spline(x, y, t, k, w=w)
assert_allclose(b.t, b_w.t, atol=1e-14)
assert_allclose(b.c, b_w.c, atol=1e-14)
assert_equal(b.k, b_w.k)
def test_weights(self):
# weights = 1 is same as None
x, y, t, k = self.x, self.y, self.t, self.k
w = np.ones_like(x)
b = make_lsq_spline(x, y, t, k)
b_w = make_lsq_spline(x, y, t, k, w=w)
assert_allclose(b.t, b_w.t, atol=1e-14)
assert_allclose(b.c, b_w.c, atol=1e-14)
assert_equal(b.k, b_w.k)
def approximate_spline(time_from_starts, path, k=3, approx=INF):
from scipy.interpolate import make_interp_spline, make_lsq_spline
x = time_from_starts
if approx == INF:
positions = make_interp_spline(time_from_starts,
path,
k=k,
t=None,
bc_type='clamped')
positions.x = positions.t[positions.k:-positions.k]
else:
# TODO: approximation near the endpoints
# approx = min(approx, len(x) - 2*k)
assert approx <= len(x) - 2 * k
t = np.r_[
(x[0], ) * (k + 1),
# np.linspace(x[0]+1e-3, x[-1]-1e-3, num=approx, endpoint=True),
np.linspace(x[0], x[-1], num=2 + approx, endpoint=True)[1:-1],
(x[-1], ) * (k + 1)]
# t = positions.t # Need to slice
# w = np.zeros(...)
w = None
positions = make_lsq_spline(x, path, t, k=k, w=w)
positions.x = positions.t[positions.k:-positions.k]
return positions
def lsq_interp(cls, x, y, degree=4, fraction_n=0.1):
"""Caller fot the scipy least squares method interpolator
To call scipy.interpolate_make_lsq_spline():
- x: abcissas
- y: ordinates
- t: knots, array-like with shape(n+k+1)
- k: b-spline degree
- w: weights for spline fitting
- axis: interpolation axis, default is 0
- check_finite: whether to check if the inputs contains only finite
elements
Inputs
- x,y: one dimensional arrays
- degree: integer, defines the polynomial fitting
- fraction_n: number of points to be used, a fraction of the total
number of points
Returns
- object, the interpolator
NOTES:
(*) number of data points must be larger than the spline degree
(*) knots must be a subset of data points of x[j] such that
t[j] < x[j] < t[j+k+1], for j=0,1..n-k-2
(*) degree 4 works slightly better than lower values
"""
# Number of points to be used
naux = np.int(np.ceil(x.shape[0] * fraction_n))
# By hand I set the use of all but 2 points at each end of the sample
p1, p2 = 2, x.shape[0] - 3
t = x[p1 : p2 : naux]
t = np.r_[(x[0],) * (degree + 1), t, (x[-1],) * (degree + 1)]
# Interpolate with the subset of points
lsq_spline = interpolate.make_lsq_spline(x, y, t, degree)
return lsq_spline
def smooth(y, num_interactions):
x = np.linspace(0, num_interactions, num_interactions)
k = 4
knots = np.linspace(x[0], x[-1], 1000)
t = np.r_[(x[0], ) * k, knots, (x[-1], ) * k]
spl = make_lsq_spline(x, y, t, k)
return spl(x)
def find_n_best_knots(self, init_spline, ranks):
test_x = self.test_x
knots = self.knots
last_knot = len(knots) - 1
indices_max = ranks.argsort()
new_knots = []
lsq_spline = ()
lsq_err = 0
for num in range(self.min_n, max(last_knot, len(self.x))):
indices = np.sort(indices_max[-num:][::-1])
new_knots = np.array([knots[i] for i in indices if i != last_knot and i != 0])
new_knots = np.r_[(self.x[0],) * (self.degree + 1)
, new_knots
, (self.x[-1],) * (self.degree + 1)]
lsq_spline = make_lsq_spline(self.x, self.y, new_knots)
lsq_err = max(abs(lsq_spline(test_x) - init_spline(test_x)))
if lsq_err <= self.cubic_spl_err:
return new_knots[self.degree + 1: -(self.degree + 1)], lsq_spline, lsq_err
return new_knots[self.degree + 1: -(self.degree + 1)], lsq_spline, lsq_err
def test_minimize():
"""
Compare least square minimization procedure to scipy's
"""
return
npts = 60
px = np.linspace(-3 + 1e-2, 3 - 1e-2, npts)
py = np.exp(-px**2) + 0.01 * np.random.randn(npts)
from scipy.interpolate import make_lsq_spline, BSpline
n = 9
p = 3
t = np.linspace(-2, 2, 20)
U = np.r_[(px[0], ) * (p + 1), t, (px[-1], ) * (p + 1)]
P = sf.bspline.bspline_lsq(px, py, U, p)
u = np.linspace(-3, 3, 100)
z = []
for ui in u:
z.append(sf.bspline.curvepoint(p, U, P, ui))
spl = make_lsq_spline(px, py, U, p)
plt.clf()
plt.plot(px, py, 'bo')
plt.plot(u, spl(u), 'g-', lw=3, label='LSQ spline')
plt.plot(u, z, 'k')
def calligraphic_fit(points, loopiness):
"""Adds `loopiness` numer of loops along the curve where the curvature is highest."""
CURVATURE_FLAT = 1
LOOP_SIZE = 8
## normalize the points so that the paramters and constatns make more sense across different datasets
points_means = points.mean(axis=0, keepdims=True)
points_stds = points.std(axis=0, keepdims=True)
points = (points - points_means) / points_stds
## fit a spline
num_points = points.shape[0]
u = np.linspace(0, 1, num_points)
k = 3
# the number of control points is about one third of the original number of points
# determines the extend of the smoothing
t = np.linspace(0, 1, num_points // 3)
t = np.concatenate([[0] * k, t, [1] * k])
spl = interp.make_lsq_spline(u, points, t, k=k)
## calc curvature and find extrema
u2 = np.linspace(0, 1, num_points * 10)
x = interp.splev(u2, spl)
xdot = interp.splev(u2, spl, 1)
xddot = interp.splev(u2, spl, 2)
# see: https://www.math.tugraz.at/~wagner/Dreibein
curvature = cross(xdot, xddot) / norm(xdot)**3
mins, maxes = peakdet(curvature, delta=0)
## add `loopiness` number of loops where the curvature is highest
peaks = np.concatenate([mins, maxes])
peaks = sorted(peaks, key=lambda p: abs(curvature[p]), reverse=True)
slices = []
Slice = namedtuple("Slice", ["i", "j", "x"])
for peak in peaks[:loopiness]:
# find the first points left and right of the peak that have a low enough curvature
for i in reversed(range(0, peak)):
if np.sign(curvature[peak]) * curvature[i] <= CURVATURE_FLAT:
break
for j in range(peak + 1, len(curvature)):
if np.sign(curvature[peak]) * curvature[j] <= CURVATURE_FLAT:
break
# extend the loop out from these low-curvature points
C = line_line_intersect(x[i - 1], x[i], x[j], x[j + 1])
# todo: maybe handle the case when there is no intersection
loop = loop_the_loop(x[i], x[j], C, LOOP_SIZE)
slices.append(Slice(i, j + 1, loop))
x_new = replace_slices(x, slices)
## undo the normalization
x_new = x_new * points_stds + points_means
return x_new
def simple_lsq_spline(arr: xr.DataArray, order=3):
xs = arr.coords[arr.dims[0]].values
ys = arr.values
knots = np.linspace(xs[0], xs[-1], order)
t = np.r_[(xs[0],) * (order - 1), knots, (xs[-1],) * (order - 1)]
spl = make_lsq_spline(xs, ys, t, order - 1)
return spl
def trajectory_to_coef(y, basis, basis_features, basis_dimension):
"""
Given a trajectory, compute its associated coefficients for each
state with respect to a functional basis.
Inputs:
- y: DataFrame
Trajectory - Index has to start at 0
- basis: string
Name of the functional basis
- basis_features: dict
Contain information on the basis for each state
- basis_dimension: dict
Give the dimension of the basis for each state
Output:
- coef: list of pd.Series
Each element of the list contains the coefficients of a
state
"""
# Define data on [0, 1] because each trajectory is considered as being
# defined on [0,1]
evaluation_points_nb = y.shape[0] - 1
x = y.index / evaluation_points_nb
coef = []
if basis == 'legendre':
# Compute coefficients for each state
for state in basis_dimension:
# NB: Use Legendre class to fix the domain of the basis
least_square_fit = Legendre.fit(x,
y[state],
deg=basis_dimension[state]-1,
domain=[0, 1])
s = pd.Series(least_square_fit.coef, name=state)
coef.append(s)
elif basis == 'bspline':
# Get internal knots
t = basis_features['knots']
# Compute coefficients for each state
for state in basis_dimension:
# Get degree
k_state = basis_features[state]
# Add external knots depending on the degree
t_state = np.r_[(0,)*(k_state+1), t, (1,)*(k_state+1)]
# Interpolate
spl = make_lsq_spline(x, y[state], t_state, k_state)
s = pd.Series(spl.c, name=state)
coef.append(s)
coef = np.array([c for series in coef for c in series.values])
return coef
def test_lstsq(self):
# check LSQ construction vs a full matrix version
x, y, t, k = self.x, self.y, self.t, self.k
c0, AY = make_lsq_full_matrix(x, y, t, k)
b = make_lsq_spline(x, y, t, k)
assert_allclose(b.c, c0)
assert_equal(b.c.shape, (t.size - k - 1,))
# also check against numpy.lstsq
aa, yy = AY
c1, _, _, _ = np.linalg.lstsq(aa, y)
assert_allclose(b.c, c1)
def test_lstsq(self):
# check LSQ construction vs a full matrix version
x, y, t, k = self.x, self.y, self.t, self.k
c0, AY = make_lsq_full_matrix(x, y, t, k)
b = make_lsq_spline(x, y, t, k)
assert_allclose(b.c, c0)
assert_equal(b.c.shape, (t.size - k - 1, ))
# also check against numpy.lstsq
aa, yy = AY
c1, _, _, _ = np.linalg.lstsq(aa, y, rcond=-1)
assert_allclose(b.c, c1)
def test_1d_make_lsq(ndspline):
N = 100
stddev = 1E-3
sample_x = np.sort(get_query_points(ndspline, n=N).squeeze())
sample_y = ndspline(sample_x)
signal_rms = (sample_y**2).sum(axis=0) / N
snr_ratio = 10
sample_y = sample_y + (signal_rms / snr_ratio) * np.random.random(
sample_y.shape)
# it was non-trivial to figure out the proper parameters for
# scipy.interpolate. It needed specific knot sequence (possibly other
# solutions) and sorted sample data. ndspline did not need either.
for k in range(0, 4):
knots = np.r_[(0.0, ) * (k + 1), 0.25, 0.5, 0.75, (1.0, ) * (k + 1)]
# unweighted
nspl = ndsplines.make_lsq_spline(sample_x, sample_y.copy(), [knots],
[k])
try:
ispl = interpolate.make_lsq_spline(sample_x, sample_y.copy(),
knots, k)
except np.linalg.linalg.LinAlgError as e:
if "leading minor" in e.__repr__():
print(e)
else:
assert_allclose(nspl.coefficients.reshape(-1), ispl.c.reshape(-1))
# random weights
w = np.random.random(N)
nspl = ndsplines.make_lsq_spline(sample_x, sample_y, [knots], [k], w)
try:
ispl = interpolate.make_lsq_spline(sample_x, sample_y, knots, k, w)
except np.linalg.linalg.LinAlgError as e:
if "leading minor" in e.__repr__():
print(e)
else:
assert_allclose(nspl.coefficients.reshape(-1), ispl.c.reshape(-1))
def bspline(X,Y,knots=8,k=3,lowclamp=False, highclamp=False):
'''
Returns a BSpline interpolation function
of a provided 1D curve.
With fewer knots, this will provide a
smooth curve that ignores local wiggles.
Parameters:
-----------
X,Y : np.ndarray
1D arrays for the curve being interpolated.
knots : int
Number of INTERNAL knots, i.e. the number
of breakpoints that are being considered
when generating the BSpline.
k : int
Degree of the BSpline. Recommended to leave
at 3.
lowclamp : bool
Enables or disables clamping at the lowest
X-value.
highclamp : bool
Enables or disables clamping at the highest
X-value.
Returns:
--------
spl : scipy.interpolate._bsplines.BSpline
Interpolation function that works over X's
domain.
'''
# Creating the knots
t_int = np.linspace(X.min(),X.max(),knots) # Internal knots, incl. beginning and end points of domain.
t_begin = np.linspace(X.min(),X.min(),k)
t_end = np.linspace(X.max(),X.max(),k)
t = np.r_[t_begin,t_int,t_end] # The entire knot vector.
# Generating the spline
w = np.zeros(X.shape)+1 # Weights.
if lowclamp==True:
w[0]=X.max()*1000000 # Setting a high weight for the X.min() term.
if highclamp==True:
w[-1]=X.max()*1000000 # Setting a high weight for the X.max() term.
spl = make_lsq_spline(X, Y, t, k,w)
return spl
def compute_spline(data, time, idx1, idx2, k):
'''
:param data: data
:param time: time
:param idx1: current index start
:param idx2: current index end
:param k: order of spline
:return: t: knots
'''
idx2 = idx2-1
variable = data[idx1:idx2]
time = time[idx1:idx2]
t = np.r_[(time[0],)*(k+1), (time[-1],)*(k+1)]
spl_lsq = make_lsq_spline(time, variable, t, k=k)
residuals = np.power((variable - spl_lsq(time)),2)
avg_sum_residuals = np.sum(residuals) #/ variable.shape[0] if uncomment do average per point
return spl_lsq, avg_sum_residuals
def test_int_xy(self):
x = np.arange(10).astype(np.int_)
y = np.arange(10).astype(np.int_)
t = _augknt(x, k=1)
# cython chokes on "buffer type mismatch"
make_lsq_spline(x, y, t, k=1)
def __getitem__(self, input_index):
index = int(input_index % self.num_images)
label_key = self.db_keys[index]
label_key_idx = int(label_key.split("_")[1])
images_start = label_key_idx - self.context_length + 1
images_end = label_key_idx + 1
packetrange = range(images_start, images_end)
keys = ["image_%d" % (i, ) for i in packetrange]
assert (keys[-1] == label_key)
label = self.label_db_wrapper.getMultiAgentLabel(keys[-1])
assert (keys[-1] + ".jpg" == label.image_tag.image_file)
rtn_session_times = np.array(
[p.session_time for p in label.ego_agent_trajectory.poses],
dtype=np.float64)
egopose = np.eye(4, dtype=np.float64)
egopose[0:3, 3] = np.array([
label.ego_agent_pose.translation.x,
label.ego_agent_pose.translation.y,
label.ego_agent_pose.translation.z
],
dtype=np.float64)
egopose[0:3, 0:3] = Rot.from_quat(
np.array([
label.ego_agent_pose.rotation.x,
label.ego_agent_pose.rotation.y,
label.ego_agent_pose.rotation.z,
label.ego_agent_pose.rotation.w
],
dtype=np.float64)).as_matrix()
egoposeinv = np.linalg.inv(egopose)
raceline_local = np.matmul(egoposeinv,
self.raceline_global)[0:3].transpose()
# print("raceline_local.shape: %s " % (str(raceline_local.shape),))
raceline_distances = np.linalg.norm(raceline_local, ord=2, axis=1)
# print("raceline_distances.shape: %s " % (str(raceline_distances.shape),))
closestidx = np.argmin(raceline_distances)
idxsamp = np.arange(closestidx - int(round(self.raceline_buffer / 3)),
closestidx + self.raceline_buffer + 1,
step=1,
dtype=np.int64) % raceline_local.shape[0]
# print("idxsamp.shape: %s " % (str(idxsamp.shape),))
raceline_close = raceline_local[idxsamp]
# print("raceline_close.shape: %s " % (str(raceline_close.shape),))
raceline_close = raceline_close[
raceline_close[:, self.position_indices[0]] >= 0.0]
raceline_close_dists = np.hstack([
np.zeros(1, dtype=np.float64),
np.cumsum(
np.linalg.norm(raceline_close[1:] - raceline_close[:-1],
ord=2,
axis=1))
])
k = 13
spl: BSpline = make_lsq_spline(raceline_close_dists,
raceline_close,
sensibleKnots(raceline_close_dists, k),
k=k)
dsamp = np.linspace(0,
raceline_close_dists[-1],
num=rtn_session_times.shape[0])
raceline_label = spl(dsamp)
transform = self.transforms[int(input_index / self.num_images)]
images_pil = [
PILImage.fromarray(self.image_db_wrapper.getImage(key))
for key in keys
]
images_pil = [
F.resize(transform(
impil.crop([
0,
int(round(self.row_crop_ratio * impil.height)),
impil.width - 1, impil.height - 1
])),
self.image_size,
interpolation=PIL.Image.LANCZOS) for impil in images_pil
]
images_torch = torch.stack([self.totensor(img) for img in images_pil])
return images_torch, raceline_label[:, self.
position_indices], dsamp, 0, 0, packetrange[
-1]
def test_multiple_rhs(self):
x, t, k, n = self.x, self.t, self.k, self.n
y = np.random.random(size=(n, 5, 6, 7))
b = make_lsq_spline(x, y, t, k)
assert_equal(b.c.shape, (t.size-k-1, 5, 6, 7))
def extract_oned_spec(path,
imgtype,
slit_along,
method='aper',
aper_param=None,
seeing=None,
order=None,
show=True,
save=True):
'''
'''
filelist = sorted(glob(F'{path}/{imgtype}*.fits'))
if len(filelist) == 0:
raise FileExistsError(
F'File {F"{path}/{imgtype}*.fits"} do not exist.')
cnts_sc = list()
errs_sc = list()
cnts_bg = list()
errs_bg = list()
for i, f in enumerate(filelist):
# Load image
(img, err), _ = load_image(f, 0)
# Transpose
if slit_along == 'row':
img = img.T
err = err.T
x = np.arange(img.shape[1])
y = np.arange(img.shape[0])
if method == 'median':
print(
F'[Extract 1-D Spectrum] Extract {i+1} of {len(filelist)} {imgtype} image [{method}]'
)
cnt_sc = np.median(img, axis=0)
err_sc = np.sqrt((err**2).sum(axis=0)) / err.shape[0]
elif method == 'aper':
print(
F'[Extract 1-D Spectrum] Estimate background for {i+1} of {len(filelist)} {imgtype} image [{method}]'
)
naper, aper_sc, aper_bg = aper_param
guessidx = np.median(img, axis=1).argmax()
maxidx = np.zeros(img.shape[1])
for j in range(img.shape[1]):
idxmin, idxmax = (guessidx - 1.5 * seeing).astype(int), (
guessidx + 1.5 * seeing).astype(int)
maxidx[j] = idxmin + gaussian_filter(img[idxmin:idxmax, j],
sigma=seeing).argmax()
# p = np.poly1d( np.polyfit( x, maxidx, 3 ) )
knots = np.r_[(x[0], ) * (order + 1), (x[-1], ) * (order + 1)]
spl = make_lsq_spline(x, maxidx, t=knots, k=order)
maxloc = spl(x)
locbgmin = (maxloc - aper_bg / 2)
locbgmax = (maxloc + aper_bg / 2)
cnt_tt = np.zeros([naper, img.shape[1]])
err_tt = np.zeros([naper, img.shape[1]])
cnt_bg = np.zeros([naper, img.shape[1]])
err_bg = np.zeros([naper, img.shape[1]])
img_bgsb = np.zeros(img.shape)
apers = np.zeros([naper + 1, img.shape[1]])
for j in range(img.shape[1]):
# Background fitting
idxbg = (y < locbgmin[j]) | (locbgmax[j] < y)
p = np.poly1d(np.polyfit(y[idxbg], img[idxbg, j], 1))
bg = p(y)
# Extract 1-D spectra
apers[:, j] = np.linspace(-aper_sc / 2, aper_sc / 2,
naper + 1) + maxloc[j]
locscmins, locscmaxs = apers[:-1, j], apers[1:, j]
for k, (locscmin,
locscmax) in enumerate(zip(locscmins, locscmaxs)):
idxsc = (locscmin < y) & (y < locscmax)
# Background estimation
cnt_bg[k, j] = bg[idxsc].sum() + \
bg[y[idxsc][ 0]-1] * ( y[idxsc][ 0] - locscmin ) + \
bg[y[idxsc][-1]+1] * ( locscmax - y[idxsc][-1] )
# Background uncertainty
err_bg[k, j] = np.median(
err[idxbg, j]) * np.sqrt(locscmax - locscmin + 1)
# Total flux estimation
cnt_tt[k, j] = img[idxsc, j].sum() + \
img[y[idxsc][ 0]-1, j] * ( y[idxsc][ 0] - locscmin ) +\
img[y[idxsc][-1]+1, j] * ( locscmax - y[idxsc][-1] )
# Total flux uncertainty
err_tt[k, j] = np.sqrt( ( err[idxsc, j]**2 ).sum() + \
( err[y[idxsc][ 0]-1, j] * ( y[idxsc][ 0] - locscmin ) )**2 + \
( err[y[idxsc][-1]+1, j] * ( locscmax - y[idxsc][-1] ) )**2 )
# Background subtraction
img_bgsb[:, j] = img[:, j] - bg
print(
F'[Extract 1-D Spectrum] Estimate target flux and uncertainty for {i+1} of {len(filelist)} {imgtype} image [{method}]'
)
# Source flux estimation
cnt_sc = cnt_tt - cnt_bg
# Source flux uncertainty
err_sc = np.sqrt(err_tt**2 + err_bg**2)
cnts_bg.append(cnt_bg)
errs_bg.append(err_bg)
# 2-D plot
title = F'2-D Background Subtracted {imgtype} image {str(i+1).zfill(len(str(len(filelist))))}'
print(
F'[Extract 1-D Spectrum] Plot 2-D background subtracted {imgtype} image from {i+1} of {len(filelist)} {imgtype} image',
end='')
fig, ax = plt.subplots(1, 1, figsize=(10, 8))
fig.subplots_adjust(right=0.8)
# Image
vmin = -3 * np.std(img_bgsb, ddof=1)
im = ax.imshow(img_bgsb,
vmin=vmin,
cmap='Greys_r',
origin='lower',
extent=(0.5, img_bgsb.shape[1] + 0.5, 0.5,
img_bgsb.shape[0] + 0.5))
# Background aperture
ax.plot(x + 1, locbgmin + 1, 'lightskyblue', ls='--')
ax.plot(x + 1, locbgmax + 1, 'lightskyblue', ls='--')
# Target aperture
for aper in apers:
ax.plot(x + 1, aper + 1, 'yellow', ls='-')
# Colorbar
cax = fig.add_axes([
ax.get_position().x1 + 0.02,
ax.get_position().y0, 0.04,
ax.get_position().height
])
cb = fig.colorbar(im, cax=cax)
# Settings
ax.tick_params(which='major',
direction='in',
color='w',
top=True,
right=True,
length=7,
width=1.5,
labelsize=18)
if slit_along == 'col':
ax.set_xlabel('Dispersion', fontsize=22)
ax.set_ylabel('Slit', fontsize=22)
if slit_along == 'row':
ax.set_xlabel('Slit', fontsize=22)
ax.set_ylabel('Dispersion', fontsize=22)
ax.set_title(title, fontsize=24)
cb.ax.tick_params(which='major',
direction='in',
color='w',
right=True,
length=7,
width=1.5,
labelsize=18)
cb.ax.set_ylabel('ADU', fontsize=22)
if save:
print(F' to `{ F"figs/{title}.png".replace( " ", "_" ) }`')
if not os.path.exists('figs'): os.makedirs('figs')
plt.savefig(F'figs/{title}.png'.replace(' ', '_'), dpi=144)
else:
print('')
if show:
print(F'[Extract 1-D Spectrum] Show plot')
plt.show()
plt.close()
else:
raise ValueError(
'Unknown method. Only methods in [`median`, `aper`] are available.'
)
cnts_sc.append(cnt_sc)
errs_sc.append(err_sc)
# 1-D plot
title = F'1-D Extracted {imgtype} Spectrum {str(i+1).zfill(len(str(len(filelist))))}'
print(
F'[Extract 1-D Spectrum] Plot 1-D spectrum extracted from {i+1} of {len(filelist)} {imgtype} image',
end='')
fig, ax = plt.subplots(1, 1, figsize=(12, 6))
if method == 'median':
# Target
ax.step(x + 1, cnt_sc, where='mid', label='Target', zorder=3)
ax.fill_between(x + 1,
cnt_sc - err_sc,
cnt_sc + err_sc,
color='C0',
alpha=0.2,
zorder=2)
else:
for m in range(cnt_sc.shape[0]):
# Target
if m == 0:
ax.step(x + 1,
cnt_sc[m],
where='mid',
color='C0',
zorder=3,
label='Target')
else:
ax.step(x + 1,
cnt_sc[m],
where='mid',
color='C0',
zorder=3)
ax.fill_between(x + 1,
cnt_sc[m] - err_sc[m],
cnt_sc[m] + err_sc[m],
color='C0',
alpha=0.2,
zorder=2)
# Background
ax.step(x + 1,
np.median(cnt_bg, axis=0),
where='mid',
color='C1',
label='Background',
zorder=1)
# Settings
ax.grid(axis='both', color='0.95', zorder=0)
ax.set_xlim(x.min() + 1, x.max() + 1)
ax.tick_params(which='major',
direction='in',
top=True,
right=True,
length=7,
width=1.5,
labelsize=18)
ax.set_xlabel('Dispersion Direction [px]', fontsize=22)
ax.set_ylabel('Counts', fontsize=22)
ax.legend(fontsize=22)
ax.set_title(title, fontsize=24)
if save:
print(F' to `{ F"figs/{title}.png".replace( " ", "_" ) }`')
if not os.path.exists('figs'): os.makedirs('figs')
plt.savefig(F'figs/{title}.png'.replace(' ', '_'), dpi=144)
else:
print('')
if show:
print('[Extract 1-D Spectrum] Show plot')
plt.show()
plt.close()
# Write to file
if not os.path.exists('bak'): os.makedirs('bak')
file_path = F'bak/1-D_Extracted_{imgtype}_Spectrum_{str(i+1).zfill(len(str(len(filelist))))}.dat'
print(
F'[Extract 1-D Spectrum] Write 1-D spectrum extracted from {i+1} of {len(filelist)} {imgtype} image to `{file_path}`'
)
if method == 'median':
np.savetxt(file_path, np.vstack([cnt_sc, err_sc]).T, fmt='%15.8e')
else:
np.savetxt(file_path,
np.vstack([cnt_sc, err_sc, cnt_bg, err_bg]).T,
fmt='%15.8e')
if method == 'median':
return (cnts_sc, errs_sc)
else:
return (cnts_sc, errs_sc), (cnts_bg, errs_bg)
def compute_statistics(self, tf_dir=None):
"""
Compute statistics of the transfer functions in a given directory.
Statistics are:
* one-lag autocorrelation coefficient, estimator for smoothness
* average of errors on components
* fit to a least-squres smooth curve
* normalized standard deviation of the first derivative, another
smoothness estimator
:param tf_dir: path to directory of transfer functions
:type tf_dir: string
:returns: data frame of all the statistics estimated
:rtype: pandas.DataFrame
.. note:: Writes a file to the tf_dir named tf_quality_statistics.csv
"""
if tf_dir is not None:
self.tf_dir = tf_dir
edi_list = glob.glob('{0}\*.edi'.format(self.tf_dir))
stat_array = np.zeros(len(edi_list),
dtype=[(key, np.float)
for key in sorted(self.types)])
station_list = []
for kk, edi in enumerate(edi_list):
mt_obj = mt.MT(edi)
station_list.append(mt_obj.station)
for ii in range(2):
for jj in range(2):
flip = False
comp = self.z_dict[(ii, jj)]
### locate bad points
bad_points_res = self.locate_bad_res_points(
mt_obj.Z.resistivity[:, ii, jj])
stat_array[kk]['bad_points_res_{0}'.format(comp)] = max(
[1, len(bad_points_res)])
bad_points_phase = self.locate_bad_phase_points(
mt_obj.Z.phase[:, ii, jj])
stat_array[kk]['bad_points_phase_{0}'.format(comp)] = max(
[1, len(bad_points_res)])
### need to get the data points that are within the reasonable range
### and not 0
nz_index = np.nonzero(mt_obj.Z.resistivity[:, ii, jj])
nz_index = np.delete(nz_index, bad_points_res)
nz_index = np.delete(nz_index, bad_points_phase)
f = mt_obj.Z.freq[nz_index]
res = mt_obj.Z.resistivity[nz_index, ii, jj]
res_err = mt_obj.Z.resistivity_err[nz_index, ii, jj]
phase = mt_obj.Z.phase[nz_index, ii, jj]
phase_err = mt_obj.Z.phase_err[nz_index, ii, jj]
if len(f) < 2:
print(mt_obj.station, comp, nz_index)
continue
# need to sort the array to be ordered with assending
# frequency. Check to see if f is ascending, if not flip
if f[0] > f[1]:
flip = True
f = f[::-1]
res = res[::-1]
res_err = res_err[::-1]
phase = phase[::-1]
phase_err = phase_err[::-1]
### make parameter for least squares fit
k = 7 # order of the fit
# knots, has to be at least to the bounds of f
t = np.r_[(f[0], ) * (k + 1), [min(1, f.mean())],
(f[-1], ) * (k + 1)]
### estimate a least squares fit
try:
ls_res = interpolate.make_lsq_spline(f, res, t, k)
ls_phase = interpolate.make_lsq_spline(f, phase, t, k)
### compute a standard deviation between the ls fit and data
stat_array[kk]['res_{0}_fit'.format(comp)] = (
res - ls_res(f)).std()
stat_array[kk]['phase_{0}_fit'.format(comp)] = (
phase - ls_phase(f)).std()
except (ValueError, np.linalg.LinAlgError) as error:
stat_array[kk]['res_{0}_fit'.format(comp)] = np.NaN
stat_array[kk]['phase_{0}_fit'.format(comp)] = np.NaN
print('{0} {1} {2}'.format(mt_obj.station, comp,
error))
### taking median of the error is more robust
stat_array[kk]['res_{0}_std'.format(comp)] = np.median(
res_err)
stat_array[kk]['phase_{0}_std'.format(comp)] = np.median(
phase_err)
### estimate smoothness
stat_array[kk]['res_{0}_corr'.format(comp)] = np.corrcoef(
res[0:-1], res[1:])[0, 1]
stat_array[kk]['phase_{0}_corr'.format(
comp)] = np.corrcoef(phase[0:-1], phase[1:])[0, 1]
### estimate smoothness with difference
stat_array[kk]['res_{0}_diff'.format(comp)] = np.abs(
np.median(np.diff(res)))
stat_array[kk]['phase_{0}_diff'.format(comp)] = np.abs(
np.median(np.diff(phase)))
### compute tipper
if ii == 0:
tcomp = self.t_dict[(0, jj)]
t_index = np.nonzero(mt_obj.Tipper.amplitude[:, 0, jj])
bad_points_t = self.locate_bad_tipper_points(
mt_obj.Tipper.amplitude[:, 0, jj])
stat_array[kk]['bad_points_tipper_{0}'.format(
tcomp)] = max([1, len(bad_points_t)])
t_index = np.delete(t_index, bad_points_t)
if t_index.size == 0:
continue
else:
tmag = mt_obj.Tipper.amplitude[t_index, 0, jj]
tmag_err = mt_obj.Tipper.amplitude_err[t_index, 0,
jj]
tip_f = mt_obj.Tipper.freq[t_index]
if flip:
tmag = tmag[::-1]
tmag_err = tmag_err[::-1]
tip_f = tip_f[::-1]
tip_t = np.r_[(tip_f[0], ) * (k + 1),
[min(1, tip_f.mean())],
(tip_f[-1], ) * (k + 1)]
try:
ls_tmag = interpolate.make_lsq_spline(
tip_f, tmag, tip_t, k)
stat_array[kk]['tipper_{0}_fit'.format(
tcomp)] = np.std(tmag - ls_tmag(tip_f))
except (ValueError,
np.linalg.LinAlgError) as error:
stat_array[kk]['tipper_{0}_fit'.format(
tcomp)] = np.NaN
print('{0} {1} {2}'.format(
mt_obj.station, tcomp, error))
stat_array[kk]['tipper_{0}_std'.format(
tcomp)] = tmag_err.mean()
stat_array[kk]['tipper_{0}_corr'.format(
tcomp)] = np.corrcoef(tmag[0:-1], tmag[1:])[0,
1]
stat_array[kk]['tipper_{0}_diff'.format(
tcomp)] = np.std(np.diff(tmag)) / abs(
np.mean(np.diff(tmag)))
### write file
df = pd.DataFrame(stat_array, index=station_list)
df = df.replace(0, np.NAN)
df.to_csv(os.path.join(self.tf_dir, 'tf_quality_statistics.csv'),
index=True,
na_rep='NaN')
return df
def sens_func(wav,
cnt,
exp,
seeing,
sw,
airmass,
extfile,
stdfile,
wave_range,
order,
show=1,
save=0):
'''
'''
# Speed of light
c = 2.99792458e+10 # [cm/s]
inverse = False
if wav[1] < wav[0]:
wav = wav[::-1]
cnt = cnt[::-1]
inverse = True
# Slit loss correction
# --------------------
print(
F'[Sensitivity function] Slit loss correction: FWHM = {seeing:.2f} [px], Slit Width = {sw:.2f} [px]'
)
sl = get_slit_loss(fwhm=seeing, sw=sw)
cnt = cnt / (1 - sl)
# Atmospheric extinction correction
# ---------------------------------
print(
'[Sensitivity function] Extinction correction: Load extinction coefficients'
)
wav_ext, ext = np.loadtxt(extfile).T
print('[Sensitivity function] Extinction correction: Dereddening')
ext = interp1d(wav_ext,
ext,
kind='quadratic',
bounds_error=False,
fill_value='extrapolate')(wav)
ext_corr_factor = 10**(0.4 * airmass * ext)
cnt = cnt * ext_corr_factor
# Convert to counts/A/s
# ---------------------
# 1. Bandpass
dwav = np.abs(np.diff(wav))
dwav = np.hstack([dwav[0], dwav, dwav[-1]])
dwav = (dwav[:-1] + dwav[1:]) / 2
# 2. Convert
cnt = cnt / dwav / exp
# Sensitivity function
# --------------------
if not os.path.exists(os.path.join(os.getcwd(), stdfile)):
if not os.path.exists(
os.path.join(
os.path.split(os.path.realpath(__file__))[0], 'onedstds/',
stdfile)):
raise FileNotFoundError('No standard file found.')
sys.exit()
else:
stdfile = os.path.join(
os.path.split(os.path.realpath(__file__))[0], 'onedstds/',
stdfile)
else:
stdfile = os.path.join(os.getcwd(), stdfile)
# 1. Load standard spectrum
print('[Sensitivity function] Load archived standard spectrum')
if os.path.split(os.path.split(stdfile)[0])[1] == 'calspec':
tab = Table.read(stdfile)
wav_mag = tab['WAVELENGTH'].data
flx_mod = tab['FLUX'].data
bp = tab['FWHM'].data
else:
stdspec = np.loadtxt(stdfile, skiprows=1).T
if stdspec.shape[0] == 3:
wav_mag, mag, bp = stdspec
elif stdspec.shape[0] == 2:
wav_mag, mag = stdspec
dwav_mag = np.abs(np.diff(wav_mag))
dwav_mag = np.hstack([dwav_mag[0], dwav_mag, dwav_mag[-1]])
bp = (dwav_mag[:-1] + dwav_mag[1:]) / 2
flx_mod = 10**(-0.4 *
mag) * 3631e-23 * c / wav_mag**2 * 1e8 # [erg/cm2/s/A]
# 2. Comparision
print('[Sensitivity function] Comparision')
flx_obs = np.zeros(flx_mod.shape[0])
bins = np.vstack([wav_mag - bp / 2, wav_mag + bp / 2]).T
for i, bin in enumerate(bins):
idx = (bin[0] < wav) & (wav < bin[1])
edges = interp1d(wav,
cnt,
'linear',
bounds_error=False,
fill_value=np.nan)(bin)
flx_obs[i] = np.trapz(np.hstack([edges[0], cnt[idx], edges[1]]),
x=np.hstack([bin[0], wav[idx], bin[1]
])) / (bin[1] - bin[0])
sen = flx_obs / flx_mod
# 3. Filter NaN
idx = np.isnan(sen)
wav_sen = wav_mag[~idx]
sen = 2.5 * np.log10(sen[~idx])
print('[Sensitivity function] Fitting')
mask = ~np.isnan(sen) & (wave_range[0] < wav_sen) & (wav_sen <
wave_range[1])
for i in range(5):
knots = np.r_[(wav_sen[mask][0], ) * (order + 1),
(wav_sen[mask][-1], ) * (order + 1)]
spl = make_lsq_spline(wav_sen[mask], sen[mask], t=knots, k=order)
mask = mask & ~sigma_clip(
sen - spl(wav_sen), sigma=0.5, maxiters=1, masked=True).mask
sen_fit = spl(wav)
# Plot
print('[Sensitivity function] Plot fitted sensitivity function', end='')
fig, ax = plt.subplots(2, 1, figsize=(12, 12))
# Sensitivity function
ax[0].plot(wav_sen[mask], sen[mask], '+', color='black', ms=10)
ax[0].plot(wav_sen[~mask], sen[~mask], '+', color='lightgrey', ms=10)
ax[0].plot(wav, sen_fit, '-', c='red', lw=2)
# Settings
ax[0].grid(axis='both', color='0.95', zorder=0)
ax[0].set_xlim(wav.min(), wav.max())
ax[0].tick_params(which='major',
direction='in',
top=True,
right=True,
length=7,
width=1.5,
labelsize=18)
ax[0].set_ylabel('Sensitivity Function', fontsize=22)
ax[0].set_title('Sensitivity Function', fontsize=24)
ax[1].plot(wav_sen[mask], (sen - spl(wav_sen))[mask],
'+',
color='black',
ms=10,
zorder=1)
ax[1].plot(wav_sen[~mask], (sen - spl(wav_sen))[~mask],
'+',
color='lightgrey',
ms=10,
zorder=0)
ax[1].axhline(y=0, c='red', ls='--', zorder=2)
# Settings
ax[1].grid(axis='both', color='0.95')
ax[1].set_xlim(wav.min(), wav.max())
ax[1].set_ylim((sen - spl(wav_sen))[mask].min() * 0.8,
(sen - spl(wav_sen))[mask].max() * 1.2)
ax[1].tick_params(which='major',
direction='in',
top=True,
right=True,
length=7,
width=1.5,
labelsize=18)
ax[1].set_xlabel('Wavelength [$\\mathrm{\\AA}$]', fontsize=22)
ax[1].set_ylabel('Residuals', fontsize=22)
fig.align_ylabels()
if save:
print(' to `figs/Sensitivity_Function.png`')
plt.savefig('figs/Sensitivity_Function.png', dpi=144)
else:
print('')
if show:
print('[Sensitivity function] Show plot')
plt.show()
plt.close()
if inverse:
sen_fit = sen_fit[::-1]
wav = wav[::-1]
cnt = cnt[::-1]
# Write to file
if not os.path.exists('cal'): os.makedirs('cal')
file_path = F'cal/SensFunc.dat'
print(
F'[Sensitivity function] Write sensitivity function to `{file_path}`')
np.savetxt(file_path, np.vstack([wav, sen_fit]).T, fmt='%15.8e')
return sen_fit
def compute_statistics(self, tf_dir=None):
"""
compute some statistics
"""
if tf_dir is not None:
self.tf_dir = tf_dir
edi_list = glob.glob("{0}\*.edi".format(self.tf_dir))
stat_array = np.zeros(len(edi_list),
dtype=[(key, np.float)
for key in sorted(self.types)])
station_list = []
for kk, edi in enumerate(edi_list):
mt_obj = mt.MT(edi)
station_list.append(mt_obj.station)
for ii in range(2):
for jj in range(2):
comp = self.z_dict[(ii, jj)]
res = mt_obj.Z.resistivity[:, ii, jj]
### need to get the data points that are within the reasonable range
### and not 0
nz_index = np.where((res > 10e-5) & (res < 10e9)
& (res != 0.0))
res = res[nz_index][::-1]
res_err = mt_obj.Z.resistivity_err[nz_index, ii,
jj][0][::-1]
phase = mt_obj.Z.phase[nz_index, ii, jj][0][::-1]
phase_err = mt_obj.Z.phase_err[nz_index, ii, jj][0][::-1]
f = mt_obj.Z.freq[nz_index][::-1]
### make parameter for least squares fit
k = 7 # order of the fit
# knots, has to be at least to the bounds of f
t = np.r_[(f[0], ) * (k + 1), [1], (f[-1], ) * (k + 1)]
### estimate a least squares fit
ls_res = interpolate.make_lsq_spline(f, res, t, k)
ls_phase = interpolate.make_lsq_spline(f, phase, t, k)
### compute a standard deviation between the ls fit and data
stat_array[kk]["res_{0}_fit".format(comp)] = (
res - ls_res(f)).std()
stat_array[kk]["phase_{0}_fit".format(comp)] = (
phase - ls_phase(f)).std()
stat_array[kk]["res_{0}_std".format(comp)] = res_err.mean()
stat_array[kk]["phase_{0}_std".format(
comp)] = phase_err.mean()
### estimate smoothness
stat_array[kk]["res_{0}_corr".format(comp)] = np.corrcoef(
res[0:-1], res[1:])[0, 1]
stat_array[kk]["phase_{0}_corr".format(
comp)] = np.corrcoef(phase[0:-1], phase[1:])[0, 1]
### estimate smoothness with difference
stat_array[kk]["res_{0}_diff".format(comp)] = np.std(
np.diff(res)) / abs(np.mean(np.diff(res)))
stat_array[kk]["phase_{0}_diff".format(comp)] = np.std(
np.diff(phase)) / abs(np.mean(np.diff(phase)))
### compute tipper
if ii == 0:
tcomp = self.t_dict[(0, jj)]
t_index = np.nonzero(mt_obj.Tipper.amplitude[:, 0, jj])
if t_index[0].size == 0:
continue
else:
tmag = mt_obj.Tipper.amplitude[t_index, 0,
jj][0][::-1]
tmag_err = mt_obj.Tipper.amplitude_err[t_index, 0,
jj][0][::-1]
tf = mt_obj.Tipper.freq[t_index][::-1]
ls_tmag = interpolate.make_lsq_spline(
tf, tmag, t, k)
stat_array[kk]["tipper_{0}_fit".format(
tcomp)] = np.std(tmag - ls_tmag(tf))
stat_array[kk]["tipper_{0}_std".format(
tcomp)] = tmag_err.mean()
stat_array[kk]["tipper_{0}_corr".format(
tcomp)] = np.corrcoef(tmag[0:-1], tmag[1:])[0,
1]
stat_array[kk]["tipper_{0}_diff".format(
tcomp)] = np.std(np.diff(tmag)) / abs(
np.mean(np.diff(tmag)))
### write file
df = pd.DataFrame(stat_array, index=station_list)
df.to_csv(os.path.join(edi_dir, "data_quality_statistics.csv"),
index=True)
return df
def get_zp_shift( self, order, show, save ):
'''
A method to get zeropoint shift along dispersion axis.
Parameters
----------
order : int
Order used in polyfit.
show : bool
If `True`, the plot will be shown.
save : bool
If `True`, the plot will be written to file.
Returns
-------
self.shift : array_like
Fitted zeropoint shift.
'''
# Derive zeropoint shift
# ----------------------
shift = np.array([])
for k, img in enumerate( self.imgs ):
if self.slit_along == 'col': data = img * 1
if self.slit_along == 'row': data = img.T
# Find peaks
# ----------
print( F'[Zeropoint correction] Seeking for peaks in the {k+1}/{self.num} images' )
peaks, properties = find_peaks( data.mean( axis = 0 ) / data.mean( axis = 0 ).max(),
height = ( 0.3, 0.8 ),
distance = data.shape[1] // 10,
width = 0 )
# Print peak info.
tab = PrettyTable( hrules = HEADER, vrules = NONE )
tab.field_names = [ 'CENTER', 'HEIGHT', 'WIDTH' ]
for i, peak in enumerate( peaks ):
tab.add_row([ peak, round( properties['peak_heights'][i], 2 ), int( properties['widths'][i] ) ])
print( '\n' + tab.get_string() + '\n' )
# Gaussian fitting
x = np.arange( data.shape[1] )
mu = np.zeros([ data.shape[0], len( peaks ) ])
for i in range( data.shape[0] ):
print( F'\r[Zeropoint correction] Peak center fitting in the ({i+1}/{data.shape[0]})-th row of {k+1}/{self.imgs.shape[0]} images', end = '', flush = True )
for j, peak in enumerate( peaks ):
idxmin, idxmax = peak - int( properties['widths'][j] ), peak + int( properties['widths'][j] )
popt, pcov = curve_fit( Gaussian,
x[idxmin:idxmax],
data[i, idxmin:idxmax],
bounds = ( [data[i, idxmin:idxmax].max()*0.5, x[idxmin:idxmax].min(), 0 ],
[data[i, idxmin:idxmax].max()*1.5, x[idxmin:idxmax].max(), x[idxmin:idxmax].shape[0] ] ) )
mu[i, j] = popt[1]
# Convert to relative shift
shift = np.hstack([ shift.reshape( mu.shape[0], -1 ), ( mu - mu[mu.shape[0]//2] ) ])
print( '' )
# Mean
shift_mean = shift.mean( axis = 1 )
# Polyfit to shift curve
# ----------------------
print( F'[Zeropoint correction] Fitting zeropoint shift iteratively (order = {order}, 5 iterations, nsigma = 1)' )
y = np.arange( shift_mean.shape[0] ) + 1
mask = np.ones( shift_mean.shape[0], dtype = bool )
for k in range( 5 ):
knots = np.r_[ ( y[mask][0], ) * ( order + 1 ), ( y[mask][-1], ) * ( order + 1 ) ]
spl = make_lsq_spline( y[mask], shift_mean[mask], t = knots, k = order )
mask = mask & ~sigma_clip( shift_mean - spl( y ), sigma = 1, maxiters = 1, masked = True ).mask
# p = np.poly1d( np.polyfit( y[mask], shift_mean[mask], order ) )
# shift_fit = p( y )
# mask = mask & ( ~sigma_clip( shift_mean - shift_fit, sigma = 1, maxiters = 1, masked = True ).mask )
self.shift = spl( y )
# Write to file
# -------------
if not os.path.exists( 'bak' ): os.makedirs( 'bak' )
np.savetxt( 'bak/zp_shift.dat', self.shift, fmt = '%15.8e' )
# Plot
# ----
fig, ax = plt.subplots( 2, 1, figsize = ( 10, 8 ) )
fig.subplots_adjust( hspace = 0 )
# Shifts
for i in range( shift.shape[1] ):
ax[0].plot( y, shift[:, i], 'r+' )
ax[0].plot( y[~mask], shift_mean[~mask], '+', c = 'grey' )
ax[0].plot( y[mask], shift_mean[mask], 'k+' )
# Fitted shifts
ax[0].plot( self.shift, '-', c = 'yellow', lw = 2 )
# Settings
ax[0].set_xlim( y.min(), y.max() )
ax[0].tick_params( which = 'major', direction = 'in', top = True, right = True, length = 7, width = 1.5, labelsize = 18 )
ax[0].set_xticklabels([])
ax[0].set_ylabel( 'Displacement [px]', fontsize = 22 )
ax[0].set_title( 'Zeropoint Shift Curve Fitting', fontsize = 24 )
# Residuals
ax[1].plot( y[mask], shift_mean[mask] - self.shift[mask], 'k+' )
ax[1].plot( y[~mask], shift_mean[~mask] - self.shift[~mask], '+', c = 'grey' )
# Settings
ax[1].axhline( y = 0, ls = '--', c = 'yellow', lw = 2 )
ax[1].set_xlim( y.min(), y.max() )
ax[1].tick_params( which = 'major', direction = 'in', top = True, right = True, length = 7, width = 1.5, labelsize = 18 )
ax[1].set_xlabel( 'Slit', fontsize = 22 )
ax[1].set_ylabel( 'Residuals [px]', fontsize = 22 )
fig.align_ylabels()
if save:
fig_path = 'figs'
print( F'[Zeropoint correction] Plot zeropoint shift curve fitting to `{ os.path.join( fig_path, "Zeropoint_shift_fitting.png" ) }`' )
plt.savefig( os.path.join( fig_path, 'Zeropoint_shift_curve_fitting.png' ), dpi = 144 )
if show:
print( F'[Zeropoint correction] Show plot' )
plt.show()
plt.close()
return deepcopy( self.shift )
def test_multiple_rhs(self):
x, t, k, n = self.x, self.t, self.k, self.n
y = np.random.random(size=(n, 5, 6, 7))
b = make_lsq_spline(x, y, t, k)
assert_equal(b.c.shape, (t.size - k - 1, 5, 6, 7))
def compute_bspline_dot_product_derivatives(basis_features, basis_dimension):
"""
Compute dot products of B-splines and their derivatives.
Input:
- basis_features: dict
Contain information on the basis for each state
- basis_dimension: dict
Give the number of basis functions for each state
Outputs:
- dot_product_12: ndarray
Array containing the dot products of Legendre polynomials
with their derivatives
- dot_product_22: ndarray
Array containing the dot products of Legendre polynomials
derivatives
"""
# Compute the dimension of the problem
dimension = np.sum([basis_dimension[elt] for elt in basis_dimension])
# Get the knots
t = basis_features['knots']
# FIXME: Consider small parameter to avoid vanishing of the last B-spline
# at 1
eps = 1e-16
dot_product_12 = np.zeros([dimension, dimension])
dot_product_22 = np.zeros([dimension, dimension])
i, j = 0, 0
# Loop over states
for state1 in basis_dimension:
# Get degree of the B-splines of state1
k1 = basis_features[state1]
# Add external knots depending on the degree
t1 = np.r_[(0, ) * (k1 + 1), t, (1, ) * (k1 + 1)]
for state2 in basis_dimension:
# Get degree of the B-splines of state2
k2 = basis_features[state2]
# Add external knots depending on the degree
t2 = np.r_[(0, ) * (k2 + 1), t, (1, ) * (k2 + 1)]
for m in range(basis_dimension[state1]):
# Define m-th B-spline of the state1 basis
spl_m = BSpline.basis_element(t1[m:m + k1 + 2])
# Reproduce the same spline for differenciation because of
# differenciation problems with BSpline.basis_element()
# FIXME: simplify if possible
# Construct knots by first finding the internal knots and then
# by adding the right numbers of external knots
t1m = t1[m:m + k1 + 2]
ind_min1 = np.max(np.argwhere(t1m == t1[m]))
ind_max1 = np.min(np.argwhere(t1m == t1[m + k1 + 1]))
t_m = np.r_[(t1m[ind_min1], ) * k1, t1m[ind_min1:ind_max1 + 1],
(t1m[ind_max1], ) * k1]
x_m = np.linspace(t1m[0], t1m[-1] - eps, 50)
spl_m = make_lsq_spline(x_m, spl_m(x_m), t_m, k1)
# Compute derivative
spl_m_deriv = spl_m.derivative(nu=1)
for n in range(basis_dimension[state2]):
# Define n-th B-spline of the state2 basis
spl_n = BSpline.basis_element(t2[n:n + k2 + 2])
# FIXME: simplify if possible
# Construct knots by first finding the internal knots and
# then by adding the right numbers of external knots
t2n = t2[n:n + k2 + 2]
ind_min2 = np.max(np.argwhere(t2n == t2[n]))
ind_max2 = np.min(np.argwhere(t2n == t2[n + k2 + 1]))
t_n = np.r_[(t2n[ind_min2], ) * k2,
t2n[ind_min2:ind_max2 + 1],
(t2n[ind_max2], ) * k2]
x_n = np.linspace(t2n[0], t2n[-1] - eps, 50)
spl_n = make_lsq_spline(x_n, spl_n(x_n), t_n, k2)
# Compute derivative
spl_n_deriv = spl_n.derivative(nu=1)
max_t = max(t1[m], t2[n])
min_t = min(t1[m + k1 + 1], t2[n + k2 + 1])
# If intersection of supports then do computations
if max_t < min_t:
# Numerical integration
quad_int_12 = quad(lambda x: spl_m(x) * spl_n_deriv(x),
max_t, min_t)
quad_int_22 = quad(
lambda x: spl_m_deriv(x) * spl_n_deriv(x), max_t,
min_t)
dot_product_12[i + m, j + n] += quad_int_12[0]
dot_product_22[i + m, j + n] += quad_int_22[0]
j += basis_dimension[state2]
j = 0
i += basis_dimension[state1]
return dot_product_12, dot_product_22
def spline_fit(y):
bspl = interp.make_lsq_spline(x=x, y=y, t=knots, k=order)
return bspl(x)
import numpy as np x = np.linspace(-15, 15, 200) y = np.exp(-x**2) + 0.1 * np.random.randn(200) from scipy.interpolate import make_lsq_spline, BSpline t = [-1, 0, 1] k = 4 t = np.r_[(x[0], ) * (k + 1), t, (x[-1], ) * (k + 1)] print(t) spl = make_lsq_spline(x, y, t, k) from scipy.interpolate import make_interp_spline spl_i = make_interp_spline(x, y) import matplotlib.pyplot as plt # xs = np.linspace(-3, 3, 100) xs = x plt.plot(x, y, 'ro', ms=5) plt.plot(xs, spl(xs), 'g-', lw=3, label='LSQ spline') plt.plot(xs, spl_i(xs), 'b-', lw=3, alpha=0.7, label='interp spline') plt.legend(loc='best') plt.show()
# Generate some noisy data: x = np.linspace(-3, 3, 50) y = np.exp(-x**2) + 0.1 * np.random.randn(50) # Now fit a smoothing cubic spline with a pre-defined internal knots. # Here we make the knot vector (k+1)-regular by adding boundary knots: from scipy.interpolate import make_lsq_spline, BSpline t = [-1, 0, 1] k = 3 t = np.r_[(x[0], ) * (k + 1), t, (x[-1], ) * (k + 1)] spl = make_lsq_spline(x, y, t, k) # For comparison, we also construct an interpolating spline for the same # set of data: from scipy.interpolate import make_interp_spline spl_i = make_interp_spline(x, y) # Plot both: import matplotlib.pyplot as plt xs = np.linspace(-3, 3, 100) plt.plot(x, y, 'ro', ms=5) plt.plot(xs, spl(xs), 'g-', lw=3, label='LSQ spline') plt.plot(xs, spl_i(xs), 'b-', lw=3, alpha=0.7, label='interp spline') plt.legend(loc='best') plt.show() # **NaN handling**: If the input arrays contain ``nan`` values, the result is
def test_int_xy(self):
x = np.arange(10).astype(np.int_)
y = np.arange(10).astype(np.int_)
t = _augknt(x, k=1)
# cython chokes on "buffer type mismatch"
make_lsq_spline(x, y, t, k=1)
def apply_fitting_fields(plotgui):
"""
Apply the set fitting parameters from the window.
This routine reads the fitting parameters from the window and then
applies the fitting. Each time it is called a new data set should
be generated, unless the fit order is too large for the number of
points. The fit order must be at most 1 less than the number
of points, otherwise the routine just shows an error message
pop-up and returns.
Parameters
----------
plotgui: by assumption a matplotlib_user_interface object
Returns
-------
None
"""
fit_type = plotgui.set_fitting_fields[0].get()
if 'Cubic Spline' in fit_type:
fit_order = float(plotgui.set_fitting_fields[1].get())
else:
try:
fit_order = int(plotgui.set_fitting_fields[1].get())
except ValueError:
str1 = 'Error: bad fit order (%s). Settng to 4.' % (
plotgui.set_fitting_fields[1].get())
plotgui.fit_text.insert(Tk.END, str1)
plotgui.fit_text.see(Tk.END)
fit_order = 4
set_number = plotgui.set_fitting_list_area.current()
fit_flag = plotgui.fit_option.get()
if fit_flag == 0:
xvalues = numpy.copy(plotgui.xdata[set_number]['values'])
yvalues = numpy.copy(plotgui.ydata[set_number]['values'])
else:
xvalues = numpy.copy(plotgui.ydata[set_number]['values'])
yvalues = numpy.copy(plotgui.xdata[set_number]['values'])
inds = numpy.argsort(xvalues)
xvalues = xvalues[inds]
yvalues = yvalues[inds]
npoints = len(xvalues)
if ('Spline' not in fit_type) and ('Internal' not in fit_type):
if npoints + 1 <= fit_order:
tkinter.messagebox.showinfo(
'Error', 'The number of points is too few for the fit order.' +
' Please check your inputs.')
return
xmin = numpy.min(xvalues)
xmax = numpy.max(xvalues)
delx = xmax - xmin
xstep = 1.2 * delx / 1201.
xout = numpy.arange(xmin - delx / 10., xmax + delx / 10., xstep)
if 'Internal' in fit_type:
if npoints < 2:
tkinter.messagebox.showinfo(
'Error', 'The number of points is too few for a linear fit.' +
' Please check your inputs.')
return
if fit_flag == 0:
yerrors = (plotgui.ydata[set_number]['lowerror'] +
plotgui.ydata[set_number]['higherror']) / 2.
else:
yerrors = (plotgui.xdata[set_number]['lowerror'] +
plotgui.xdata[set_number]['higherror']) / 2.
if (numpy.min(yerrors) == 0.) and (numpy.max(yerrors) == 0.):
yerrors = yerrors + 1.
slope, intercept, slope_error, intercept_error, covariance, \
correlation = general_utilities.slope_calculation(
xvalues, yvalues, yerrors)
if slope is None:
tkinter.messagebox.showinfo('Error',
'Error in the standard slope fit.\n')
return
yfit = intercept + xvalues * slope
yout = intercept + xout * slope
errorterm1 = xout * 0. + intercept_error
errorterm2 = xout * slope_error
youterror = numpy.sqrt(errorterm1 * errorterm1 +
errorterm2 * errorterm2)
labelstring = 'Standard linear fit'
rms = numpy.sqrt(numpy.mean((yvalues - yfit) * (yvalues - yfit)))
str1 = 'Regression calculation results:\n'
str1 = str1 + 'Slope: %g +/- %g\n' % (slope, slope_error)
str1 = str1 + 'Intercept: %g +/- %g\n' % (intercept, intercept_error)
str1 = str1 + 'Covariance: %f\n' % (covariance)
str1 = str1 + 'Correlation: %f\n' % (correlation)
str1 = str1 + 'RMS deviation: %f\n' % (rms)
tkinter.messagebox.showinfo('Information', str1)
outfile = open('fit_values.txt', 'a')
print(str1, file=outfile)
print(' ', file=outfile)
outfile.close()
if fit_type == 'Polynomial':
fitpars = polynomial.polyfit(xvalues, yvalues, fit_order)
yout = polynomial.polyval(xout, fitpars)
yfit = polynomial.polyval(xvalues, fitpars)
labelstring = 'Order %d polynomial fit' % (fit_order)
general_utilities.list_polynomial_fitpars(fit_type, fit_order, fitpars)
if fit_type == 'Legendre':
fitpars = legendre.legfit(xvalues, yvalues, fit_order)
yout = legendre.legval(xout, fitpars)
yfit = legendre.legval(xvalues, fitpars)
labelstring = 'Order %d Legendre polynomial fit' % (fit_order)
general_utilities.list_polynomial_fitpars(fit_type, fit_order, fitpars)
if fit_type == 'Laguerre':
fitpars = laguerre.lagfit(xvalues, yvalues, fit_order)
yout = laguerre.lagval(xout, fitpars)
yfit = laguerre.lagval(xvalues, fitpars)
labelstring = 'Order %d Laguerre polynomial fit' % (fit_order)
general_utilities.list_polynomial_fitpars(fit_type, fit_order, fitpars)
if fit_type == 'Chebyshev':
fitpars = chebyshev.chebfit(xvalues, yvalues, fit_order)
yout = chebyshev.chebval(xout, fitpars)
yfit = chebyshev.chebval(xvalues, fitpars)
labelstring = 'Order %d Chebyshev polynomial fit' % (fit_order)
general_utilities.list_polynomial_fitpars(fit_type, fit_order, fitpars)
if fit_type == 'Least-Squares Spline':
if fit_flag == 0:
yerrors = (plotgui.ydata[set_number]['lowerror'] +
plotgui.ydata[set_number]['higherror']) / 2.
else:
yerrors = (plotgui.xdata[set_number]['lowerror'] +
plotgui.xdata[set_number]['higherror']) / 2.
if (numpy.min(yerrors) == 0.) and (numpy.max(yerrors) == 0.):
yerrors = yerrors + 1.
xmin1 = numpy.min(xvalues)
xmax1 = numpy.max(xvalues)
xrange = xmax1 - xmin1
nknots = int(fit_order)
if (nknots < 3) or (nknots > int(len(xvalues) / 2)):
nknots = 3
xstep = xrange / (nknots - 2)
xknots = numpy.arange(
numpy.min(xvalues) + xstep,
numpy.max(xvalues) * 0.999999999, xstep)
k = 3
# Use cubic splines
knotedges = numpy.r_[(xmin, ) * (k + 1), xknots, (xmax, ) * (k + 1)]
weights = 1. / yerrors
weights[yerrors == 0.] = 0.
fitobject = make_lsq_spline(xvalues, yvalues, knotedges, k, w=weights)
yout = fitobject(xout)
yfit = fitobject(xvalues)
labelstring = 'Least squares spline fit, sections = %d' % (nknots)
if fit_type == 'Spline':
fitpars = UnivariateSpline(xvalues,
yvalues,
k=1,
s=None,
bbox=[xmin - delx, xmax + delx])
labelstring = 'Default spline fit'
yout = fitpars(xout)
yfit = fitpars(xvalues)
if fit_type == 'Cubic Spline':
if fit_order < 0.:
str1 = 'Error: smoothing value %f (< 0) is not allowed.'\
+ ' Settng to 0.0' % (fit_order)
plotgui.fit_text.insert(Tk.END, str1)
plotgui.fit_text.see(Tk.END)
fit_order = 0.0
fitpars = UnivariateSpline(xvalues,
yvalues,
k=3,
bbox=[xmin - delx, xmax + delx],
s=fit_order)
yout = fitpars(xout)
yfit = fitpars(xvalues)
labelstring = 'Cubic spline fit, smoothing = %f' % (fit_order)
rms = fit_statistics(yvalues, yfit)
if 'Internal' in fit_type:
if fit_flag == 0:
xlowerror = youterror * 0.
xhigherror = youterror * 0.
ylowerror = youterror
yhigherror = youterror
else:
xlowerror = youterror
xhigherror = youterror
ylowerror = youterror * 0.
yhigherror = youterror * 0.
else:
xlowerror = xout * 0.
xhigherror = xout * 0.
ylowerror = yout * 0.
yhigherror = yout * 0.
xmin = numpy.min(xout)
xmax = numpy.max(xout)
ymin = numpy.min(yfit)
ymax = numpy.max(yfit)
if rms is not None:
str1 = 'Fit: RMS = %g for %d points\n' % (rms, len(yfit))
plotgui.fit_text.insert(Tk.END, str1)
plotgui.fit_text.see(Tk.END)
if fit_flag == 0:
plotgui.xdata[plotgui.nsets] = {
'values': xout,
'lowerror': xlowerror,
'higherror': xhigherror,
'minimum': xmin,
'maximum': xmax,
'errors': False,
'legend': True
}
plotgui.ydata[plotgui.nsets] = {
'values': yout,
'lowerror': ylowerror,
'higherror': yhigherror,
'minimum': ymin,
'maximum': ymax,
'errors': True,
'legend': True
}
else:
plotgui.xdata[plotgui.nsets] = {
'values': yout,
'lowerror': ylowerror,
'higherror': yhigherror,
'minimum': ymin,
'maximum': ymax,
'errors': False,
'legend': True
}
plotgui.ydata[plotgui.nsets] = {
'values': xout,
'lowerror': xlowerror,
'higherror': xhigherror,
'minimum': xmin,
'maximum': xmax,
'errors': False,
'legend': True
}
m = plotgui.nsets % 10
n = int(math.floor(plotgui.nsets / 10))
plotgui.set_properties[plotgui.nsets]['symbol'] = None
plotgui.set_properties[plotgui.nsets]['linestyle'] = '-'
plotgui.set_properties[plotgui.nsets]['colour'] = plotgui.colourset[m]
plotgui.set_properties[plotgui.nsets]['symbolsize'] = 4.0 + 0.3 * n
plotgui.set_properties[plotgui.nsets]['label'] = labelstring
plotgui.nsets = plotgui.nsets + 1
make_plot.make_plot(plotgui)
"""
===============================
1-Dimensional Least Squares Fit
===============================
"""
import ndsplines
import matplotlib.pyplot as plt
import numpy as np
from scipy import interpolate
x = np.linspace(-3, 3, 50)
y = np.exp(-x**2) + 0.1 * np.random.randn(50)
t = [-1, 0, 1]
k = 3
t = np.r_[(x[0], ) * (k + 1), t, (x[-1], ) * (k + 1)]
ndspl = ndsplines.make_lsq_spline(x[:, None], y[:, None], [t], np.array([k]))
ispl = interpolate.make_lsq_spline(x, y, t, k)
xs = np.linspace(-3, 3, 100)
plt.figure()
plt.plot(x, y, 'o', ms=5)
plt.plot(xs, ndspl(xs).squeeze(), label='LSQ ND spline')
plt.plot(xs, ispl(xs), '--', label='LSQ scipy.interpolate spline')
plt.legend(loc='best')
plt.show()
print("Computed coefficients close?", np.allclose(ndspl.coefficients, ispl.c))
