def wrapper(*args):
coeffs = [_get_coeffs(x) for x in args]
out = func(*coeffs)
if all(not isinstance(x, cupy.poly1d) for x in args):
return out
if isinstance(out, cupy.ndarray):
return cupy.poly1d(out)
if isinstance(out, tuple):
return tuple([cupy.poly1d(x) for x in out])
assert False # Never reach
def test_poly1d_set(self, dtype):
arr1 = testing.shaped_arange((10, ), cupy, dtype)
arr2 = numpy.ones(10, dtype=dtype)
a = cupy.poly1d(arr1)
b = numpy.poly1d(arr2, variable='z')
a.set(b)
assert a.variable == b.variable
testing.assert_array_equal(a.coeffs, b.coeffs)
def polyval(p, x):
"""Evaluates a polynomial at specific values.
Args:
p (cupy.ndarray or cupy.poly1d): input polynomial.
x (scalar, cupy.ndarray): values at which the polynomial
is evaluated.
Returns:
cupy.ndarray or cupy.poly1d: polynomial evaluated at x.
.. warning::
This function doesn't currently support poly1d values to evaluate.
.. seealso:: :func:`numpy.polyval`
"""
if isinstance(p, cupy.poly1d):
p = p.coeffs
if not isinstance(p, cupy.ndarray) or p.ndim == 0:
raise TypeError('p must be 1d ndarray or poly1d object')
if p.ndim > 1:
raise ValueError('p must be 1d array')
if isinstance(x, cupy.poly1d):
# TODO(asi1024): Needs performance improvement.
dtype = numpy.result_type(x.coeffs, 1)
res = cupy.poly1d(cupy.array([0], dtype=dtype))
prod = cupy.poly1d(cupy.array([1], dtype=dtype))
for c in p[::-1]:
res = res + prod * c
prod = prod * x
return res
dtype = numpy.result_type(p.dtype.type(0), x)
p = p.astype(dtype, copy=False)
if p.size == 0:
return cupy.zeros(x.shape, dtype)
if dtype == numpy.bool_:
return p.any() * x + p[-1]
if not cupy.isscalar(x):
x = cupy.asarray(x, dtype=dtype)[..., None]
x = x ** cupy.arange(p.size, dtype=dtype)
return (p[::-1] * x).sum(axis=-1, dtype=dtype)
def test_polyval_poly1d_values(self, dtype):
a = testing.shaped_arange((5, ), cupy, dtype)
b = testing.shaped_arange((3, ), cupy, dtype)
b = cupy.poly1d(b)
with pytest.raises(NotImplementedError):
cupy.polyval(a, b)
def test_poly1d_get2(self, dtype):
a1 = testing.shaped_arange((), cupy, dtype)
a2 = testing.shaped_arange((), numpy, dtype)
b1 = cupy.poly1d(a1).get()
b2 = numpy.poly1d(a2)
assert b1 == b2
def test_poly1d_get1(self, dtype):
a1 = testing.shaped_arange((10, ), cupy, dtype)
a2 = testing.shaped_arange((10, ), numpy, dtype)
b1 = cupy.poly1d(a1, variable='z').get()
b2 = numpy.poly1d(a2, variable='z')
assert b1 == b2
def normalize(data, mask=None, poly_fit=0):
""" Apply normalization on GPU
Applies normalisation (data - mean) / stdev
Args:
data (np/cp.array): Data to preprocess
mask (np.cp.array): 1D Channel mask for RFI flagging
return_space ('cpu' or 'gpu'): Returns array in CPU or GPU space
poly_fit (int): Fit polynomial of degree N, 0 = no fit.
Returns: d_gpu (cp.array): Normalized data
"""
# Normalise
t0 = time.time()
d_flag = cp.copy(data)
n_int, n_ifs, n_chan = data.shape
# Setup 1D channel mask -- used for polynomial fitting
if mask is None:
mask = cp.zeros(n_chan, dtype='bool')
# Do polynomial fit and compute stats (with masking)
d_mean_ifs, d_std_ifs = cp.zeros(n_ifs), cp.zeros(n_ifs)
N_masked = mask.sum()
N_flagged = N_masked * n_ifs * n_int
N_tot = np.product(data.shape)
N_unflagged = (N_tot - N_flagged)
t0p = time.time()
for ii in range(n_ifs):
x = cp.arange(n_chan, dtype='float64')
xc = cp.compress(~mask, x)
dfit = cp.compress(~mask, data[:, ii].mean(axis=0))
if poly_fit > 0:
# WAR: int64 dtype causes issues in cupy 10 (19.04.2022)
p = cp.poly1d(cp.polyfit(xc, dfit, poly_fit))
fit = p(x)
dfit -= p(xc)
data[:, ii] = data[:, ii] - fit
# compute mean and stdev
dmean = dfit.mean()
dvar = ((data[:, ii] - dmean)**2).mean(axis=0)
dvar = cp.compress(~mask, dvar).mean()
dstd = cp.sqrt(dvar)
d_mean_ifs[ii] = dmean
d_std_ifs[ii] = dstd
t1p = time.time()
### logger.info(f"Poly fit time: {(t1p-t0p)*1e3:2.2f}ms")
flag_fraction = N_flagged / N_tot
flag_correction = N_tot / (N_tot - N_flagged)
### logger.info(f"Flagged fraction: {flag_fraction:2.4f}")
### if flag_fraction > 0.2:
### logger.warning(f"High flagged fraction: {flag_fraction:2.3f}")
# Apply to original data
for ii in range(n_ifs):
data[:, ii] = ((data[:, ii] - d_mean_ifs[ii]) / d_std_ifs[ii])
t1 = time.time()
### logger.info(f"Normalisation time: {(t1-t0)*1e3:2.2f}ms")
return data
def test_polyval_poly1d_values(self, dtype):
a = testing.shaped_arange((5, ), cupy, dtype)
b = testing.shaped_arange((3, ), cupy, dtype)
b = cupy.poly1d(b)
return cupy.polyval(a, b)
