# Licensed under a 3-clause BSD style license - see PYDL_LICENSE.rst
# -*- coding: utf-8 -*-
# Also cite https://doi.org/10.5281/zenodo.1095150 when referencing PYDL
"""
Implements pure python support methods for
:class:`pypeit.bspline.bspline.bspline`.
.. include:: ../include/links.rst
"""
import warnings
from IPython import embed
import numpy as np
[docs]def bspline_model(x, action, lower, upper, coeff, n, nord, npoly):
"""
Calculate the bspline model.
Args:
x (`numpy.ndarray`_):
The independent variable in the fit.
action (`numpy.ndarray`_):
Action matrix. See
:func:`pypeit.bspline.bspline.bspline.action`. The shape
of the array is expected to be ``nd`` by ``npoly*nord``.
lower (`numpy.ndarray`_):
Vector with the starting indices along the second axis of
action used to construct the model.
upper (`numpy.ndarray`_):
Vector with the (inclusive) ending indices along the
second axis of action used to construct the model.
coeff (`numpy.ndarray`_):
The model coefficients used for each action.
n (:obj:`int`):
Number of unmasked measurements included in the fit.
nord (:obj:`int`):
Fit order.
npoly (:obj:`int`):
Polynomial per fit order.
Returns:
`numpy.ndarray`_: The best fitting bspline model at all
provided :math:`x`.
"""
# TODO: Can we save some of these objects to self so that we
# don't have to recreate them?
# TODO: x is always 1D right?
# TODO: Used for testing bspline
# np.savez_compressed('bspline_model.npz', x=x, action=action, lower=lower, upper=upper,
# coeff=coeff, n=n, nord=nord, npoly=npoly)
# raise ValueError('Entered bspline_model')
yfit = np.zeros(x.shape, dtype=x.dtype)
spot = np.arange(npoly * nord, dtype=int)
nowidth = np.invert(upper+1 > lower)
for i in range(n-nord+1):
if nowidth[i]:
continue
yfit[lower[i]:upper[i]+1] = np.dot(action[lower[i]:upper[i]+1,:],
coeff.flatten('F')[i*npoly+spot])
return yfit
[docs]def intrv(nord, breakpoints, x):
"""
Find the segment between breakpoints which contain each value in
the array x.
The minimum breakpoint is nbkptord -1, and the maximum
is nbkpt - nbkptord - 1.
Parameters
----------
nord : :obj:`int`
Order of the fit.
breakpoints : `numpy.ndarray`_
Locations of good breakpoints
x : `numpy.ndarray`_
Data values, assumed to be monotonically increasing.
Returns
-------
`numpy.ndarray`_
Position of array elements with respect to breakpoints.
"""
# TODO: Used for testing bspline
# np.savez_compressed('intrv.npz', nord=nord, breakpoints=breakpoints, x=x)
# raise ValueError('Entered solution_arrays')
n = breakpoints.size - nord
indx = np.zeros(x.size, dtype=int)
ileft = nord - 1
for i in range(x.size):
while x[i] > breakpoints[ileft+1] and ileft < n - 1:
ileft += 1
indx[i] = ileft
return indx
[docs]def solution_arrays(nn, npoly, nord, ydata, action, ivar, upper, lower):
"""
Support function that builds the arrays for Cholesky
decomposition.
Args:
nn (:obj:`int`):
Number of good break points.
npoly (:obj:`int`):
Polynomial per fit order.
nord (:obj:`int`):
Fit order.
ydata (`numpy.ndarray`_):
Data to fit.
action (`numpy.ndarray`_):
Action matrix. See
:func:`pypeit.bspline.bspline.bspline.action`. The shape
of the array is expected to be ``nd`` by ``npoly*nord``.
ivar (`numpy.ndarray`_):
Inverse variance in the data to fit.
upper (`numpy.ndarray`_):
Vector with the (inclusive) ending indices along the
second axis of action used to construct the model.
lower (`numpy.ndarray`_):
Vector with the starting indices along the second axis of
action used to construct the model.
Returns:
:obj:`tuple`: Returns (1) matrix :math:`A` and (2) vector
:math:`b` prepared for Cholesky decomposition and used in the
solution to the equation :math:`Ax=b`.
"""
# TODO: Used for testing bspline
# np.savez_compressed('solution_arrays.npz', nn=nn, npoly=npoly, nord=nord, ydata=ydata,
# action=action, ivar=ivar, upper=upper, lower=lower)
# raise ValueError('Entered solution_arrays')
nfull = nn * npoly
bw = npoly * nord
a2 = action * np.sqrt(ivar)[:,None]
alpha = np.zeros((bw, nfull+bw), dtype=float)
beta = np.zeros((nfull+bw,), dtype=float)
bi = np.concatenate([np.arange(i)+(bw-i)*(bw+1) for i in range(bw,0,-1)])
bo = np.concatenate([np.arange(i)+(bw-i)*bw for i in range(bw,0,-1)])
upper += 1
nowidth = np.invert(upper > lower)
for k in range(nn-nord+1):
if nowidth[k]:
continue
itop = k*npoly
alpha.T.flat[bo+itop*bw] \
+= np.dot(a2[lower[k]:upper[k],:].T, a2[lower[k]:upper[k],:]).flat[bi]
beta[itop:min(itop,nfull)+bw] \
+= np.dot(ydata[lower[k]:upper[k]] * np.sqrt(ivar[lower[k]:upper[k]]),
a2[lower[k]:upper[k],:])
upper -= 1
return alpha, beta
[docs]def cholesky_band(l, mininf=0.0, verbose=False):
"""
Compute Cholesky decomposition of banded matrix.
This function is pure python.
Parameters
----------
l : `numpy.ndarray`_
A matrix on which to perform the Cholesky decomposition.
mininf : :class:`float`, optional
Entries in the `l` matrix are considered negative if they are less
than this value (default 0.0).
Returns
-------
:obj:`tuple`
If problems were detected, the first item will be the index or
indexes where the problem was detected, and the second item will simply
be the input matrix. If no problems were detected, the first item
will be -1, and the second item will be the Cholesky decomposition.
"""
# # TODO: Used for testing bspline
# np.savez_compressed('cholesky_band_l.npz', l=l, mininf=mininf)
# print(l.shape)
# raise ValueError('Entered band')
bw, nn = l.shape
n = nn - bw
negative = (l[0,:n] <= mininf) | np.invert(np.isfinite(l[0,:n]))
# JFH changed this below to make it more consistent with IDL version. Not sure
# why the np.all(np.isfinite(lower)) was added. The code could return an empty
# list for negative.nonzero() and crash if all elements in lower are NaN.
# KBW: Added the "or not finite" flags to negative.
if negative.any():
nz = negative.nonzero()[0]
if verbose:
warnings.warn('Found {0} bad entries: {1}'.format(nz.size, nz))
return nz, l
lower = l.copy()
kn = bw - 1
spot = np.arange(kn, dtype=int) + 1
bi = np.concatenate([np.arange(i)+(kn-i)*(kn+1) for i in range(kn,0,-1)])
here = bi[:,None] + (np.arange(n)[None,:] + 1)*bw
for j in range(n):
lower[0,j] = np.sqrt(lower[0,j])
lower[spot,j] /= lower[0,j]
if not np.all(np.isfinite(lower[spot,j])):
warnings.warn('NaN found in cholesky_band.')
return j, l
hmm = lower[spot,j,None] * lower[None,spot,j]
lower.T.flat[here[:,j]] -= hmm.flat[bi]
return -1, lower
[docs]def cholesky_solve(a, bb):
"""
Solve the equation Ax=b where A is a Cholesky-banded matrix.
This function is pure python.
Parameters
----------
a : `numpy.ndarray`_
:math:`A` in :math:`A x = b`.
bb : `numpy.ndarray`_
:math:`b` in :math:`A x = b`.
Returns
-------
:obj:`tuple`
A tuple containing the status and the result of the solution. The
status is always -1.
"""
# # TODO: Used for testing bspline
# np.savez_compressed('cholesky_solve_abb.npz', a=a, bb=bb)
# print(a.shape)
# print(bb.shape)
# raise ValueError('Entered solve')
b = bb.copy()
n = b.shape[0] - a.shape[0]
kn = a.shape[0] - 1
spot = np.arange(kn, dtype=int) + 1
for j in range(n):
b[j] /= a[0,j]
b[j+spot] -= b[j]*a[spot,j]
for j in range(n-1, -1, -1):
b[j] = (b[j] - np.sum(a[spot,j] * b[j+spot]))/a[0,j]
return -1, b