def da_sub(daa, dab):
"""
subtract 2 DataArrays as cleverly as possible:
* keep the metadata of the first DA in the result
* ensures the result has the right type so that no underflows happen
returns (DataArray): the result of daa - dab
"""
rt = numpy.result_type(daa, dab) # dtype of result of daa-dab
dt = None # default is to let numpy decide
if rt.kind == "f":
# float should always be fine
pass
elif rt.kind in "iub":
# underflow can happen (especially if unsigned)
# find the worse case value (could be improved, but would be longer)
worse_val = int(daa.min()) - int(dab.max())
dt = numpy.result_type(rt, numpy.min_scalar_type(worse_val))
else:
# subtracting such a data is suspicious, but try anyway
logging.warning("Subtraction on data of type %s unsupported", rt.name)
res = numpy.subtract(daa, dab, dtype=dt) # metadata is copied from daa
logging.debug("type = %s, %s", res.dtype.name, daa.dtype.name)
return res
def accum (self, key1, key2, value, weight=1):
index1 = self._key1map[key1]
index2 = self._key2map[key2]
if self._m0 is None:
self._m0 = np.zeros ((self.chunk0size, self.chunk0size), dtype=np.result_type (weight))
self._m1 = np.zeros ((self.chunk0size, self.chunk0size), dtype=np.result_type (value, weight))
self._m2 = np.zeros_like (self._m1)
if index1 >= self._m0.shape[0]:
self._m0 = np.concatenate ((self._m0, np.zeros_like (self._m0)), axis=0)
self._m1 = np.concatenate ((self._m1, np.zeros_like (self._m1)), axis=0)
self._m2 = np.concatenate ((self._m2, np.zeros_like (self._m2)), axis=0)
if index2 >= self._m0.shape[1]:
self._m0 = np.concatenate ((self._m0, np.zeros_like (self._m0)), axis=1)
self._m1 = np.concatenate ((self._m1, np.zeros_like (self._m1)), axis=1)
self._m2 = np.concatenate ((self._m2, np.zeros_like (self._m2)), axis=1)
self._m0[index1,index2] += weight
q = weight * value
self._m1[index1,index2] += q
q *= value
self._m2[index1,index2] += q
return self
def ShiftConvNumbaFFT(h, N, M, xdtype=np.complex_, powerof2=True):
# implements Doppler filter:
# y[n, p] = SUM_k (exp(2*pi*j*n*(k - (L-1))/N) * h[k]) * x[p - k]
# = SUM_k (exp(-2*pi*j*n*k/N) * s*[k]) * x[p - (L-1) + k]
L = len(h)
outlen = M + L - 1
nfft = outlen
if powerof2:
nfft = pow2(nfft)
dopplermat = np.exp(2*np.pi*1j*np.arange(N)[:, np.newaxis]*(np.arange(L) - (L - 1))/N)
dopplermat.astype(np.result_type(h.dtype, np.complex64)) # cast to complex type with precision of h
hbank = h*dopplermat
# speed not critical here, just use numpy fft
hbankpad = zero_pad(hbank, nfft)
H = np.fft.fft(hbankpad) / nfft # divide by nfft b/c FFTW's ifft does not do this
xcdtype = np.result_type(xdtype, np.complex64) # cast to complex type with precision of x
xpad = pyfftw.n_byte_align(np.zeros(nfft, xcdtype), 16)
X = pyfftw.n_byte_align(np.zeros(nfft, xcdtype), 16)
xfft = pyfftw.FFTW(xpad, X, threads=_THREADS)
ydtype = np.result_type(H.dtype, xcdtype)
Y = pyfftw.n_byte_align_empty(H.shape, 16, ydtype)
y = pyfftw.n_byte_align_empty(H.shape, 16, ydtype)
ifft = pyfftw.FFTW(Y, y, direction='FFTW_BACKWARD', threads=_THREADS)
xtype = numba.__getattribute__(str(np.dtype(xdtype)))
#htype = numba.__getattribute__(str(H.dtype))
#xctype = numba.__getattribute__(str(X.dtype))
#ytype = numba.__getattribute__(str(Y.dtype))
#@jit(argtypes=[htype[:, ::1], xctype[::1], ytype[:, ::1], xtype[::1]])
#def fun(H, X, Y, x):
#xpad[:M] = x
#xfft.execute() # input is xpad, output is X
#Y[:, :] = H*X # need expression optimized by numba but that writes into Y
#ifft.execute() # input is Y, output is y
#yc = np.array(y)[:, :outlen] # need a copy, which np.array provides
#return yc
#@dopplerbank_dec(h, N, M, nfft=nfft, H=H)
#def shiftconv_numba_fft(x):
#return fun(H, X, Y, x)
#@jit(argtypes=[xtype[::1]])
@jit
def shiftconv_numba_fft(x):
xpad[:M] = x
xfft.execute() # input is xpad, output is X
Y[:, :] = X*H # need expression optimized by numba but that writes into Y
ifft.execute() # input is Y, output is y
yc = np.array(y[:, :outlen]) # need a copy, which np.array provides
return yc
shiftconv_numba_fft = dopplerbank_dec(h, N, M, nfft=nfft, H=H)(shiftconv_numba_fft)
return shiftconv_numba_fft
def _take_with_fill(self, indices, fill_value=None):
if fill_value is None:
fill_value = self.dtype.na_value
if indices.min() < -1:
raise ValueError("Invalid value in 'indices'. Must be between -1 "
"and the length of the array.")
if indices.max() >= len(self):
raise IndexError("out of bounds value in 'indices'.")
if len(self) == 0:
# Empty... Allow taking only if all empty
if (indices == -1).all():
dtype = np.result_type(self.sp_values, fill_value)
taken = np.empty_like(indices, dtype=dtype)
taken.fill(fill_value)
return taken
else:
raise IndexError('cannot do a non-empty take from an empty '
'axes.')
sp_indexer = self.sp_index.lookup_array(indices)
if self.sp_index.npoints == 0:
# Avoid taking from the empty self.sp_values
taken = np.full(sp_indexer.shape, fill_value=fill_value,
dtype=np.result_type(fill_value))
else:
taken = self.sp_values.take(sp_indexer)
# sp_indexer may be -1 for two reasons
# 1.) we took for an index of -1 (new)
# 2.) we took a value that was self.fill_value (old)
new_fill_indices = indices == -1
old_fill_indices = (sp_indexer == -1) & ~new_fill_indices
# Fill in two steps.
# Old fill values
# New fill values
# potentially coercing to a new dtype at each stage.
m0 = sp_indexer[old_fill_indices] < 0
m1 = sp_indexer[new_fill_indices] < 0
result_type = taken.dtype
if m0.any():
result_type = np.result_type(result_type, self.fill_value)
taken = taken.astype(result_type)
taken[old_fill_indices] = self.fill_value
if m1.any():
result_type = np.result_type(result_type, fill_value)
taken = taken.astype(result_type)
taken[new_fill_indices] = fill_value
return taken
def NumbaFFTW(h, M, xdtype=np.complex_, powerof2=True):
L = len(h)
outlen = M + L - 1
nfft = outlen
if powerof2:
nfft = pow2(nfft)
outdtype = np.result_type(h.dtype, xdtype)
fftdtype = np.result_type(outdtype, np.complex64) # output is always complex, promote using smallest
# speed not critical here, just use numpy fft
# cast to outdtype so we use same type of fft as when transforming x
hpad = zero_pad(h, nfft).astype(outdtype)
if np.iscomplexobj(hpad):
H = np.fft.fft(hpad)
else:
H = np.fft.rfft(hpad)
H = (H / nfft).astype(fftdtype) # divide by nfft b/c FFTW's ifft does not do this
xpad = pyfftw.n_byte_align(np.zeros(nfft, outdtype), 16) # outdtype so same type fft as h->H
X = pyfftw.n_byte_align(np.zeros(len(H), fftdtype), 16) # len(H) b/c rfft may be used
xfft = pyfftw.FFTW(xpad, X, threads=_THREADS)
y = pyfftw.n_byte_align_empty(nfft, 16, outdtype)
ifft = pyfftw.FFTW(X, y, direction='FFTW_BACKWARD', threads=_THREADS)
xtype = numba.__getattribute__(str(np.dtype(xdtype)))
outtype = numba.__getattribute__(str(outdtype))
ffttype = numba.__getattribute__(str(fftdtype))
#@jit(restype=outtype[::1],
#argtypes=[outtype[::1], ffttype[::1], ffttype[::1], outtype[::1], xtype[::1]])
#def filt(xpad, X, H, y, x):
#xpad[:M] = x
#xfft.execute() # input in xpad, result in X
#X[:] = H*X
#ifft.execute() # input in X, result in y
#yc = y[:outlen].copy()
#return yc
#@filter_dec(h, M, nfft=nfft, H=H)
#def numba_fftw(x):
#return filt(xpad, X, H, y, x)
#@jit(argtypes=[xtype[::1]])
@jit
def numba_fftw(x):
xpad[:M] = x
xfft.execute() # input in xpad, result in X
X[:] = H*X # want expression that is optimized by numba but writes into X
ifft.execute() # input in X, result in y
yc = y[:outlen].copy()
return yc
numba_fftw = filter_dec(h, M, nfft=nfft, H=H)(numba_fftw)
return numba_fftw
def test_result_type(self):
self.check_promotion_cases(np.result_type)
f64 = float64(0)
c64 = complex64(0)
## Scalars do not coerce to complex if the value is real
#assert_equal(np.result_type(c64,array([f64])), np.dtype(float64))
# But they do if the value is complex
assert_equal(np.result_type(complex64(3j),array([f64])),
np.dtype(complex128))
# Scalars do coerce to complex even if the value is real
# This is so "a+0j" can be reliably used to make something complex.
assert_equal(np.result_type(c64,array([f64])), np.dtype(complex128))
def _interpolated_dtype(self, dtype):
"""
Determine the minimum base dtype required by the
underlying interpolator.
"""
return np.result_type(_DEFAULT_DTYPE, dtype)
def build_supergather(step, width, bins, dataset): sutype = np.result_type(dataset) cdps = np.unique(dataset['cdp']) dataset['offset'] = np.abs(dataset['offset']) supergather_centres = range(min(cdps)+width, max(cdps)-width, step) supergather_slices = [cdps[a-width:a+width] for a in supergather_centres] for index, inds in enumerate(supergather_slices): for cdpn, cdp in enumerate(dataset['cdp']): if cdp in inds: dataset['ns1'][cdpn] = index dataset = dataset[dataset['ns1'] != 0] output = np.empty(0, dtype=sutype) for ind in np.unique(dataset['ns1']): sg = dataset[dataset['ns1'] == ind] hist = np.digitize(sg['offset'], bins) sg['ep'] = hist vals = np.unique(sg['ep']) holder = np.zeros(len(vals), dtype=sutype) for v in vals: traces = sg[sg['ep'] == v] header = traces[0].copy() header['trace'].fill(0.0) fold = traces.size header['trace'] += np.sum(traces['trace'], axis=-2)/np.float(fold) holder[v-1] = header #~ display(holder) output = np.concatenate([output, holder]) return output
def promote(*operands):
"""
Take an arbitrary number of graph nodes and produce the promoted
dtype by discarding all shape information and just looking at the
measures.
::
5, 5, | int |
2, 5, | int |
1, 3, | float |
>>> promote(IntNode(1), Op(IntNode(2))
int
>>> promote(FloatNode(1), Op(IntNode(2))
float
"""
# Looks something like this...
# (ArrayNode, IntNode...) -> (dshape('2, int'), dshape('int'))
# (dshape('2, int', dshape('int')) -> (dshape('int', dshape('int'))
# (dshape('2, int', dshape('int')) -> (dtype('int', dtype('int'))
types = (op.simple_type() for op in operands if op is not None)
measures = (extract_measure(t) for t in types)
dtypes = (to_numpy(m) for m in measures)
promoted = np.result_type(*dtypes)
datashape = CType.from_dtype(promoted)
return datashape
def _contract_plain(mydf, mos, coulG, phase, max_memory):
cell = mydf.cell
moiT, mojT, mokT, molT = mos
nmoi, nmoj, nmok, nmol = [x.shape[0] for x in mos]
ngrids = moiT.shape[1]
wcoulG = coulG * (cell.vol/ngrids)
dtype = numpy.result_type(phase, *mos)
eri = numpy.empty((nmoi*nmoj,nmok*nmol), dtype=dtype)
blksize = int(min(max(nmoi,nmok), (max_memory*1e6/16 - eri.size)/2/ngrids/max(nmoj,nmol)+1))
assert blksize > 0
buf0 = numpy.empty((blksize,max(nmoj,nmol),ngrids), dtype=dtype)
buf1 = numpy.ndarray((blksize,nmoj,ngrids), dtype=dtype, buffer=buf0)
buf2 = numpy.ndarray((blksize,nmol,ngrids), dtype=dtype, buffer=buf0)
for p0, p1 in lib.prange(0, nmoi, blksize):
mo_pairs = numpy.einsum('ig,jg->ijg', moiT[p0:p1].conj()*phase,
mojT, out=buf1[:p1-p0])
mo_pairs_G = tools.fft(mo_pairs.reshape(-1,ngrids), mydf.mesh)
mo_pairs = None
mo_pairs_G*= wcoulG
v = tools.ifft(mo_pairs_G, mydf.mesh)
mo_pairs_G = None
v *= phase.conj()
if dtype == numpy.double:
v = numpy.asarray(v.real, order='C')
for q0, q1 in lib.prange(0, nmok, blksize):
mo_pairs = numpy.einsum('ig,jg->ijg', mokT[q0:q1].conj(),
molT, out=buf2[:q1-q0])
eri[p0*nmoj:p1*nmoj,q0*nmol:q1*nmol] = lib.dot(v, mo_pairs.reshape(-1,ngrids).T)
v = None
return eri
def test_dtypes(self):
for dtypes in get_dtypes():
dtype_op, dtype_in = dtypes
dtype_out = np.result_type(dtype_op, dtype_in)
A = lo.ZeroOperator(3, 3, dtype=dtype_op)
x = np.array([1, 1, 1]).astype(dtype_in)
assert_((A * x).dtype == dtype_out)
def _convert_list(self, value):
"""Convert a string into a typed numpy array.
If it is not possible it returns a numpy string.
"""
try:
numpy_values = []
values = value.split(" ")
types = set([])
for string_value in values:
v = self._convert_scalar_value(string_value)
numpy_values.append(v)
types.add(v.dtype.type)
result_type = numpy.result_type(*types)
if issubclass(result_type.type, (numpy.string_, six.binary_type)):
# use the raw data to create the result
return numpy.string_(value)
elif issubclass(result_type.type, (numpy.unicode_, six.text_type)):
# use the raw data to create the result
return numpy.unicode_(value)
else:
return numpy.array(numpy_values, dtype=result_type)
except ValueError:
return numpy.string_(value)
def _inequality(self, other, op, op_name, bad_scalar_msg):
# Scalar other.
if isscalarlike(other):
if 0 == other and op_name in ('_le_', '_ge_'):
raise NotImplementedError(" >= and <= don't work with 0.")
elif op(0, other):
warn(bad_scalar_msg, SparseEfficiencyWarning)
other_arr = np.empty(self.shape, dtype=np.result_type(other))
other_arr.fill(other)
other_arr = self.__class__(other_arr)
return self._binopt(other_arr, op_name)
else:
return self._scalar_binopt(other, op)
# Dense other.
elif isdense(other):
return op(self.todense(), other)
# Sparse other.
elif isspmatrix(other):
#TODO sparse broadcasting
if self.shape != other.shape:
raise ValueError("inconsistent shapes")
elif self.format != other.format:
other = other.asformat(self.format)
if op_name not in ('_ge_', '_le_'):
return self._binopt(other, op_name)
warn("Comparing sparse matrices using >= and <= is inefficient, "
"using <, >, or !=, instead.", SparseEfficiencyWarning)
all_true = self.__class__(np.ones(self.shape))
res = self._binopt(other, '_gt_' if op_name == '_le_' else '_lt_')
return all_true - res
else:
raise ValueError("Operands could not be compared.")
def _valid_input(self, value, dtype=None):
if not misc.is_valid_param_value(value):
msg = 'The value must be either a tensorflow variable, an array or a scalar.'
raise ValueError(msg)
cast = not (dtype is None)
is_built = False
shape = None
if hasattr(self, '_value'): # The parameter has not initialized yet.
is_built = self.is_built_coherence() == Build.YES
shape = self.shape
inner_dtype = self.dtype
if dtype is not None and inner_dtype != dtype:
msg = 'Overriding parameter\'s type "{0}" with "{1}" is not possible.'
raise ValueError(msg.format(inner_dtype, dtype))
elif isinstance(value, np.ndarray) and inner_dtype != value.dtype:
msg = 'The value has different data type "{0}". Parameter type is "{1}".'
raise ValueError(msg.format(value.dtype, inner_dtype))
cast = False
dtype = inner_dtype
if misc.is_number(value):
value_type = np.result_type(value).type
num_type = misc.normalize_num_type(value_type)
dtype = num_type if dtype is None else dtype
value = np.array(value, dtype=dtype)
elif misc.is_list(value):
dtype = settings.float_type if dtype is None else dtype
value = np.array(value, dtype=dtype)
elif cast:
value = value.astype(dtype)
if shape is not None and self.fixed_shape and is_built and shape != value.shape:
msg = 'Value has different shape. Parameter shape {0}, value shape {1}.'
raise ValueError(msg.format(shape, value.shape))
return value
def test_upcast():
a0 = csr_matrix([[np.pi, np.pi*1j], [3, 4]], dtype=complex)
b0 = np.array([256+1j, 2**32], dtype=complex)
for a_dtype in supported_dtypes:
for b_dtype in supported_dtypes:
msg = "(%r, %r)" % (a_dtype, b_dtype)
if np.issubdtype(a_dtype, np.complexfloating):
a = a0.copy().astype(a_dtype)
else:
a = a0.real.copy().astype(a_dtype)
if np.issubdtype(b_dtype, np.complexfloating):
b = b0.copy().astype(b_dtype)
else:
b = b0.real.copy().astype(b_dtype)
if not (a_dtype == np.bool_ and b_dtype == np.bool_):
c = np.zeros((2,), dtype=np.bool_)
assert_raises(ValueError, _sparsetools.csr_matvec,
2, 2, a.indptr, a.indices, a.data, b, c)
if ((np.issubdtype(a_dtype, np.complexfloating) and
not np.issubdtype(b_dtype, np.complexfloating)) or
(not np.issubdtype(a_dtype, np.complexfloating) and
np.issubdtype(b_dtype, np.complexfloating))):
c = np.zeros((2,), dtype=np.float64)
assert_raises(ValueError, _sparsetools.csr_matvec,
2, 2, a.indptr, a.indices, a.data, b, c)
c = np.zeros((2,), dtype=np.result_type(a_dtype, b_dtype))
_sparsetools.csr_matvec(2, 2, a.indptr, a.indices, a.data, b, c)
assert_allclose(c, np.dot(a.toarray(), b), err_msg=msg)
def _normalize_vector_type(self, dtype):
"""Normalize the """
if self.__at_least_32bits:
if numpy.issubdtype(dtype, numpy.signedinteger):
dtype = numpy.result_type(dtype, numpy.uint32)
if numpy.issubdtype(dtype, numpy.unsignedinteger):
dtype = numpy.result_type(dtype, numpy.uint32)
elif numpy.issubdtype(dtype, numpy.floating):
dtype = numpy.result_type(dtype, numpy.float32)
elif numpy.issubdtype(dtype, numpy.complexfloating):
dtype = numpy.result_type(dtype, numpy.complex64)
if self.__signed_type:
if numpy.issubdtype(dtype, numpy.unsignedinteger):
signed = numpy.dtype("%s%i" % ('i', dtype.itemsize))
dtype = numpy.result_type(dtype, signed)
return dtype
def promote(lhs, rhs):
"""Promote two scalar dshapes to a possibly larger, but compatible type.
Examples
--------
>>> from datashape import int32, int64, Option
>>> x = Option(int32)
>>> y = int64
>>> promote(x, y)
?int64
>>> promote(int64, int64)
ctype("int64")
Notes
----
This uses ``numpy.result_type`` for type promotion logic. See the numpy
documentation at
http://docs.scipy.org/doc/numpy/reference/generated/numpy.result_type.html
"""
if lhs == rhs:
return lhs
else:
left, right = getattr(lhs, "ty", lhs), getattr(rhs, "ty", rhs)
dtype = np.result_type(datashape.to_numpy_dtype(left), datashape.to_numpy_dtype(right))
return optionify(lhs, rhs, datashape.CType.from_numpy_dtype(dtype))
def rfft(a): n = a.shape[-1] b = a.reshape(np.prod(a.shape[:-1]),n) fb = np.empty((b.shape[0],b.shape[1]/2+1),dtype=np.result_type(a,0j)) for i in range(b.shape[0]): fb[i] = myfft.rfft(b[i])*n**-0.5 return np.reshape(fb, list(a.shape[:-1]) + [fb.shape[-1]])
def blk_matvec(x, blks):
nargins = [[blk.shape[-1] for blk in blkrow] for blkrow in blks]
nargouts = [[blk.shape[0] for blk in blkrow] for blkrow in blks]
nargin = sum(nargins[0])
nargout = sum([out[0] for out in nargouts])
nx = len(x)
self.logger.debug('Multiplying with a vector of size %d' % nx)
self.logger.debug('nargin=%d, nargout=%d' % (nargin, nargout))
if nx != nargin:
raise ShapeError('Multiplying with vector of wrong shape.')
result_type = np.result_type(self.dtype, x.dtype)
y = np.zeros(nargout, dtype=result_type)
nblk_row = len(blks)
nblk_col = len(blks[0])
row_start = col_start = 0
for row in range(nblk_row):
row_end = row_start + nargouts[row][0]
yout = y[row_start:row_end]
for col in range(nblk_col):
col_end = col_start + nargins[0][col]
xin = x[col_start:col_end]
B = blks[row][col]
yout[:] += B * xin
col_start = col_end
row_start = row_end
col_start = 0
return y
def blk_matvec(x, blks):
nx = len(x)
nargins = [blk.shape[-1] for blk in blocks]
nargin = sum(nargins)
nargouts = [blk.shape[0] for blk in blocks]
nargout = sum(nargouts)
self.logger.debug('Multiplying with a vector of size %d' % nx)
self.logger.debug('nargin=%d, nargout=%d' % (nargin, nargout))
if nx != nargin:
raise ShapeError('Multiplying with vector of wrong shape.')
result_type = np.result_type(self.dtype, x.dtype)
y = np.empty(nargout, dtype=result_type)
nblks = len(blks)
row_start = col_start = 0
for blk in range(nblks):
row_end = row_start + nargouts[blk]
yout = y[row_start:row_end]
col_end = col_start + nargins[blk]
xin = x[col_start:col_end]
B = blks[blk]
yout[:] = B * xin
col_start = col_end
row_start = row_end
return y
def call(self, args, axis=0, out=None, chunksize=1024 * 1024, **kwargs):
""" axis is the axis to chop it off.
if self.altreduce is set, the results will
be reduced with altreduce and returned
otherwise will be saved to out, then return out.
"""
if self.altreduce is not None:
ret = [None]
else:
if out is None :
if self.outdtype is not None:
dtype = self.outdtype
else:
try:
dtype = numpy.result_type(*[args[i] for i in self.ins] * 2)
except:
dtype = None
out = sharedmem.empty(
numpy.broadcast(*[args[i] for i in self.ins] * 2).shape,
dtype=dtype)
if axis != 0:
for i in self.ins:
args[i] = numpy.rollaxis(args[i], axis)
out = numpy.rollaxis(out, axis)
size = numpy.max([len(args[i]) for i in self.ins])
with sharedmem.MapReduce() as pool:
def work(i):
sl = slice(i, i+chunksize)
myargs = args[:]
for j in self.ins:
try:
tmp = myargs[j][sl]
a, b, c = sl.indices(len(args[j]))
myargs[j] = tmp
except Exception as e:
print tmp
print j, e
pass
if b == a: return None
rt = self.ufunc(*myargs, **kwargs)
if self.altreduce is not None:
return rt
else:
out[sl] = rt
def reduce(rt):
if self.altreduce is None:
return
if ret[0] is None:
ret[0] = rt
elif rt is not None:
ret[0] = self.altreduce(ret[0], rt)
pool.map(work, range(0, size, chunksize), reduce=reduce)
if self.altreduce is None:
if axis != 0:
out = numpy.rollaxis(out, 0, axis + 1)
return out
else:
return ret[0]
def __new__(cls, *arrays):
"""Create combined array with views to separate arrays.
The provided arrays are flattened and concatenated in the order
given to make the combined array. The array views accessible with the
`parts` attribute and through iteration provide views into the
combined array that correspond to the location of the original arrays.
Parameters
----------
arrays : iterable
Individual arrays to combine.
"""
dtype = np.result_type(*arrays)
sizes = [arr.size for arr in arrays]
size = np.sum(sizes)
self = np.ndarray.__new__(cls, size, dtype=dtype)
idxs = [0] + list(np.cumsum(sizes))
self._slices = tuple(slice(idxs[k], idxs[k + 1])
for k in xrange(len(idxs) - 1))
self._shapes = tuple(arr.shape for arr in arrays)
# copy original arrays into corresponding views of the combined array
for view, arr in zip(iter(self), arrays):
view[...] = arr
return self
def covariance(m, ddof=1.5):
"""
A simplified version of numpy's 'cov' which allows ddof=1.5 and assumes
variables to be in columns (rowvar=0).
"""
# Handles complex arrays too
m = np.asarray(m)
dtype = np.result_type(m, np.float64)
X = np.array(m, ndmin=2, dtype=dtype)
if X.shape[0] != 1:
X = X.T
if X.shape[0] == 0:
return np.array([]).reshape(0, 0)
avg = np.mean(X, axis=1)
# Determine the normalization
fact = float(X.shape[1] - ddof)
if fact <= 0:
warn("Degrees of freedom <= 0 for slice", RuntimeWarning)
fact = 0.0
X -= avg[:, None]
X_T = X.T
return (np.dot(X, X_T.conj())/fact).squeeze()
def cast(arrays, dtype=None, order='c'):
"""
Cast a list of arrays into a same data type.
Parameters
----------
arrays : sequence of array-like or None
The list of arrays to be cast.
dtype : numpy.dtype
If specified, all arrays will be cast to this data type. Otherwise,
the data types is inferred from the arrays.
Example
-------
>>> cast([[1., 2.], None, np.array(2j)])
(array([ 1.+0.j, 2.+0.j]), None, array(2j))
"""
arrays = tuple(arrays)
if dtype is None:
arrays_ = [np.array(a, copy=False) for a in arrays if a is not None]
dtype = np.result_type(*arrays_)
result = (np.array(a, dtype=dtype, order=order, copy=False)
if a is not None else None for a in arrays)
return tuple(result)
def use_hmm(observations, state_count, symbol_count, maxit=1000, accuracy=-1, retries=10, dtype=numpy.float32):
curr_A, curr_B, curr_pi = None, None, None
curr_eps = None
# try:
# importlib.import_module('spscicomp.hmm.kernel.opencl')
# kernel = spscicomp.hmm.kernel.opencl
# LOG.debug('OpenCL-Kernel used')
# except:
# LOG.debug('OpenCL-Kernel not available')
try:
importlib.import_module('spscicomp.hmm.kernel.c')
kernel = spscicomp.hmm.kernel.c
if numpy.result_type(observations) != numpy.int16:
LOG.debug('Observations data type was not int16, thus casting it for c-extension.')
observations = numpy.array(observations, dtype=numpy.int16)
LOG.debug('C-Kernel used')
except:
LOG.debug('C-Kernel not available')
kernel = spscicomp.hmm.kernel.python
LOG.debug('Python-Kernel used')
for _ in range(0, retries):
A = spscicomp.hmm.utility.generate_random_matrice(state_count, state_count)
B = spscicomp.hmm.utility.generate_random_matrice(state_count, symbol_count)
pi = spscicomp.hmm.utility.generate_random_array(state_count)
A, B, pi, eps, it = spscicomp.hmm.algorithms.baum_welch_multiple(obs=observations, A=A, B=B, pi=pi,
kernel=kernel,
dtype=dtype, maxit=maxit, accuracy=accuracy)
if curr_eps is None or curr_eps < eps:
curr_A, curr_B, curr_pi, curr_eps = A, B, pi, eps
return curr_A, curr_B, curr_pi
def _coo_matmul_gradsp(a, b, c_row, c_col, c_shape, transa, transb, transc,
dtype):
if dtype is None:
dtype = numpy.result_type(a.dtype, b.dtype)
if transa:
A = a.swapaxes(-1, -2)
else:
A = a
if transb:
B = b.swapaxes(-1, -2)
else:
B = b
if transc:
C_row = c_col
C_col = c_row
else:
C_row = c_row
C_col = c_col
xp = cuda.get_array_module(A, B)
if xp is numpy:
return _coo_matmul_gradsp_cpu(A, B, C_row, C_col, dtype)
else:
return _coo_matmul_gradsp_gpu(A, B, C_row, C_col, dtype)
def matvec_transp(x):
if x.shape != (nargout,):
msg = 'Input has shape ' + str(x.shape)
msg += ' instead of (%d,)' % self.nargout
raise ValueError(msg)
result_type = np.result_type(self.dtype, x.dtype)
return np.zeros(nargin, dtype=result_type)
def test_dtypes(self):
for dtype_op in allowed_types:
for dtype_in in allowed_types:
dtype_out = np.result_type(dtype_op, dtype_in)
A = lo.IdentityOperator(3, dtype=dtype_op)
x = np.array([1, 1, 1]).astype(dtype_in)
assert_((A * x).dtype == dtype_out)
def _coo_matmul(sp_data, sp_row, sp_col, sp_shape, dn, transa, transb, transc,
dtype=None):
if dtype is None:
dtype = numpy.result_type(sp_data.dtype, dn.dtype)
A_data = sp_data
if transa:
A_row = sp_col
A_col = sp_row
A_shape = (sp_shape[1], sp_shape[0])
else:
A_row = sp_row
A_col = sp_col
A_shape = sp_shape
if transb:
B = dn.swapaxes(-1, -2)
else:
B = dn
xp = cuda.get_array_module(A_data, B)
if xp is numpy:
C = _coo_matmul_cpu(A_data, A_row, A_col, A_shape, B, dtype)
else:
C = _coo_matmul_gpu(A_data, A_row, A_col, A_shape, B, dtype)
if transc:
C = C.swapaxes(-1, -2)
return C
def trisolve(dl, d, du, b, inplace=False):
"""
The tridiagonal matrix (Thomas) algorithm for solving tridiagonal systems
of equations:
a_{i}x_{i-1} + b_{i}x_{i} + c_{i}x_{i+1} = y_{i}
in matrix form:
Mx = b
TDMA is O(n), whereas standard Gaussian elimination is O(n^3).
Arguments:
-----------
dl: (n - 1,) vector
the lower diagonal of M
d: (n,) vector
the main diagonal of M
du: (n - 1,) vector
the upper diagonal of M
b: (n,) vector
the result of Mx
inplace:
if True, and if d and b are both float64 vectors, they will be
modified in place (may be faster)
Returns:
-----------
x: (n,) vector
the solution to Mx = b
References:
-----------
http://en.wikipedia.org/wiki/Tridiagonal_matrix_algorithm
http://www.netlib.org/lapack/explore-html/d1/db3/dgtsv_8f.html
"""
if (dl.shape[0] != du.shape[0] or (d.shape[0] != dl.shape[0] + 1)
or d.shape[0] != b.shape[0]):
raise ValueError('Invalid diagonal shapes')
bshape_in = b.shape
rtype = np.result_type(dl, d, du, b)
if not inplace:
# force a copy
dl = np.array(dl, dtype=rtype, copy=True, order='F')
d = np.array(d, dtype=rtype, copy=True, order='F')
du = np.array(du, dtype=rtype, copy=True, order='F')
b = np.array(b, dtype=rtype, copy=True, order='F')
# this may also force copies if arrays have inconsistent types / incorrect
# order
dl, d, du, b = (np.array(v, dtype=rtype, copy=False, order='F')
for v in (dl, d, du, b))
# use the LAPACK implementation
_lapack_trisolve(dl, d, du, b, rtype)
return b.reshape(bshape_in)
def get_k_e1_kpts(mydf,
dm_kpts,
kpts=np.zeros((1, 3)),
kpts_band=None,
exxdiv=None):
cell = mydf.cell
mesh = mydf.mesh
coords = cell.gen_uniform_grids(mesh)
ngrids = coords.shape[0]
if getattr(dm_kpts, 'mo_coeff', None) is not None:
mo_coeff = dm_kpts.mo_coeff
mo_occ = dm_kpts.mo_occ
else:
mo_coeff = None
kpts = np.asarray(kpts)
dm_kpts = lib.asarray(dm_kpts, order='C')
dms = _format_dms(dm_kpts, kpts)
nset, nkpts, nao = dms.shape[:3]
weight = 1. / nkpts * (cell.vol / ngrids)
kpts_band, input_band = _format_kpts_band(kpts_band, kpts), kpts_band
nband = len(kpts_band)
if gamma_point(kpts_band) and gamma_point(kpts):
vk_kpts = np.zeros((3, nset, nband, nao, nao), dtype=dms.dtype)
else:
vk_kpts = np.zeros((3, nset, nband, nao, nao), dtype=np.complex128)
coords = mydf.grids.coords
if input_band is None:
ao2_kpts = [
np.asarray(ao.transpose(0, 2, 1), order='C')
for ao in mydf._numint.eval_ao(cell, coords, kpts=kpts, deriv=1)
]
ao1_kpts = ao2_kpts
ao2_kpts = [ao2_kpt[0] for ao2_kpt in ao2_kpts]
else:
ao2_kpts = [
np.asarray(ao.T, order='C')
for ao in mydf._numint.eval_ao(cell, coords, kpts=kpts)
]
ao1_kpts = [
np.asarray(ao.transpose(0, 2, 1), order='C') for ao in
mydf._numint.eval_ao(cell, coords, kpts=kpts_band, deriv=1)
]
if mo_coeff is not None and nset == 1:
mo_coeff = [
mo_coeff[k][:, occ > 0] * np.sqrt(occ[occ > 0])
for k, occ in enumerate(mo_occ)
]
ao2_kpts = [np.dot(mo_coeff[k].T, ao) for k, ao in enumerate(ao2_kpts)]
mem_now = lib.current_memory()[0]
max_memory = mydf.max_memory - mem_now
blksize = int(
min(nao,
max(1, (max_memory - mem_now) * 1e6 / 16 / 4 / 3 / ngrids / nao)))
lib.logger.debug1(mydf, 'fft_jk: get_k_kpts max_memory %s blksize %d',
max_memory, blksize)
ao1_dtype = np.result_type(*ao1_kpts)
ao2_dtype = np.result_type(*ao2_kpts)
vR_dm = np.empty((3, nset, nao, ngrids), dtype=vk_kpts.dtype)
t1 = (time.clock(), time.time())
for k2, ao2T in enumerate(ao2_kpts):
if ao2T.size == 0:
continue
kpt2 = kpts[k2]
naoj = ao2T.shape[0]
if mo_coeff is None or nset > 1:
ao_dms = [lib.dot(dms[i, k2], ao2T.conj()) for i in range(nset)]
else:
ao_dms = [ao2T.conj()]
for k1, ao1T in enumerate(ao1_kpts):
kpt1 = kpts_band[k1]
# If we have an ewald exxdiv, we add the G=0 correction near the
# end of the function to bypass any discretization errors
# that arise from the FFT.
mydf.exxdiv = exxdiv
if exxdiv == 'ewald' or exxdiv is None:
coulG = tools.get_coulG(cell, kpt2 - kpt1, False, mydf, mesh)
else:
coulG = tools.get_coulG(cell, kpt2 - kpt1, True, mydf, mesh)
if is_zero(kpt1 - kpt2):
expmikr = np.array(1.)
else:
expmikr = np.exp(-1j * np.dot(coords, kpt2 - kpt1))
for p0, p1 in lib.prange(0, nao, blksize):
rho1 = np.einsum('aig,jg->aijg',
ao1T[1:, p0:p1].conj() * expmikr, ao2T)
vG = tools.fft(rho1.reshape(-1, ngrids), mesh)
rho1 = None
vG *= coulG
vR = tools.ifft(vG, mesh).reshape(3, p1 - p0, naoj, ngrids)
vG = None
if vR_dm.dtype == np.double:
vR = vR.real
for i in range(nset):
np.einsum('aijg,jg->aig',
vR,
ao_dms[i],
out=vR_dm[:, i, p0:p1])
vR = None
vR_dm *= expmikr.conj()
for i in range(nset):
vk_kpts[:, i, k1] -= weight * np.einsum(
'aig,jg->aij', vR_dm[:, i], ao1T[0])
t1 = lib.logger.timer_debug1(mydf, 'get_k_kpts: make_kpt (%d,*)' % k2,
*t1)
# Ewald correction has no contribution to nuclear gradient unless range separted Coulomb is used
# The gradient correction part is not added in the vk matrix
if exxdiv == 'ewald' and cell.omega != 0:
raise NotImplementedError("Range Separated Coulomb")
# when cell.omega !=0: madelung constant will have a non-zero derivative
vk_kpts = np.asarray(
[_format_jks(vk, dm_kpts, input_band, kpts) for vk in vk_kpts])
return vk_kpts
def _dtype(x):
return canonicalize_dtype(onp.result_type(x))
def _tile_dataframe(cls, op):
from ...tensor.merge.stack import TensorStack
df = op.outputs[0]
if df.ndim == 1:
if op.axis == 0:
ts = []
for name in df.index_value.to_pandas():
a = tensor_from_series(op.input[name])
t = tensor_quantile(a,
op.q,
interpolation=op.interpolation,
handle_non_numeric=not op.numeric_only)
ts.append(t)
try:
dtype = np.result_type(*[it.dtype for it in ts])
except TypeError:
dtype = np.dtype(object)
stack_op = TensorStack(axis=0, dtype=dtype)
tr = stack_op(ts)
r = series_from_tensor(tr,
index=df.index_value.to_pandas(),
name=np.asscalar(ts[0].op.q))
else:
assert op.axis == 1
empty_df = build_empty_df(op.input.dtypes)
fields = empty_df._get_numeric_data().columns.tolist()
t = tensor_from_dataframe(op.input[fields])
tr = tensor_quantile(t,
op.q,
axis=1,
interpolation=op.interpolation,
handle_non_numeric=not op.numeric_only)
r = series_from_tensor(tr, name=np.asscalar(tr.op.q))
r._index_value = op.input.index_value
else:
assert df.ndim == 2
if op.axis == 0:
d = OrderedDict()
for name in df.dtypes.index:
a = tensor_from_series(op.input[name])
t = tensor_quantile(a,
op.q,
interpolation=op.interpolation,
handle_non_numeric=not op.numeric_only)
d[name] = t
r = create_df(d, index=op.q)
else:
assert op.axis == 1
empty_df = build_empty_df(op.input.dtypes)
fields = empty_df._get_numeric_data().columns.tolist()
t = tensor_from_dataframe(op.input[fields])
tr = tensor_quantile(t,
op.q,
axis=1,
interpolation=op.interpolation,
handle_non_numeric=not op.numeric_only)
if not op.input.index_value.has_value():
raise NotImplementedError
# TODO(xuye.qin): use index=op.input.index when we support DataFrame.index
r = dataframe_from_tensor(
tr, index=op.q, columns=op.input.index_value.to_pandas())
return (yield from recursive_tile(r))
def average(a, axis=None, weights=None, returned=False):
"""Returns the weighted average along an axis.
Args:
a (cupy.ndarray): Array to compute average.
axis (int): Along which axis to compute average. The flattened array
is used by default.
weights (cupy.ndarray): Array of weights where each element
corresponds to the value in ``a``. If ``None``, all the values
in ``a`` have a weight equal to one.
returned (bool): If ``True``, a tuple of the average and the sum
of weights is returned, otherwise only the average is returned.
Returns:
cupy.ndarray or tuple of cupy.ndarray: The average of the input array
along the axis and the sum of weights.
.. seealso:: :func:`numpy.average`
"""
a = cupy.asarray(a)
if weights is None:
avg = a.mean(axis)
scl = avg.dtype.type(a.size / avg.size)
else:
wgt = cupy.asarray(weights)
if issubclass(a.dtype.type, (numpy.integer, numpy.bool_)):
result_dtype = numpy.result_type(a.dtype, wgt.dtype, 'f8')
else:
result_dtype = numpy.result_type(a.dtype, wgt.dtype)
# Sanity checks
if a.shape != wgt.shape:
if axis is None:
raise TypeError(
'Axis must be specified when shapes of a and weights '
'differ.')
if wgt.ndim != 1:
raise TypeError(
'1D weights expected when shapes of a and weights differ.')
if wgt.shape[0] != a.shape[axis]:
raise ValueError(
'Length of weights not compatible with specified axis.')
# setup wgt to broadcast along axis
wgt = cupy.broadcast_to(wgt, (a.ndim - 1) * (1,) + wgt.shape)
wgt = wgt.swapaxes(-1, axis)
scl = wgt.sum(axis=axis, dtype=result_dtype)
if cupy.any(scl == 0.0):
raise ZeroDivisionError(
'Weights sum to zero, can\'t be normalized')
avg = cupy.multiply(a, wgt, dtype=result_dtype).sum(axis) / scl
if returned:
if scl.shape != avg.shape:
scl = cupy.broadcast_to(cupy.array(scl), avg.shape).copy()
return avg, scl
else:
return avg
def pv_like(x):
aval = ShapedArray(onp.shape(x), onp.result_type(x))
return pe.PartialVal((aval, unit))
def make_shaped_array(x): dtype = xla_bridge.canonicalize_dtype(onp.result_type(x)) return ShapedArray(onp.shape(x), dtype)
def __init__(self, val):
self.val = val
self.shape = onp.shape(val)
# canonicalized self.dtype doesn't necessarily match self.val
self.dtype = onp.dtype(xla_bridge.canonicalize_dtype(onp.result_type(val)))
assert self.dtype != onp.dtype('O')
def _multiply_no_nan(x, y, name=None): # pylint: disable=unused-argument
dtype = np.result_type(x, y)
# TODO(b/146385087): The gradient should be
# `lambda dz: [multiply_no_nan(dz, y), multiply_no_nan(x, dz)]`.
return np.where(np.equal(y, 0.), np.zeros((), dtype=dtype),
np.multiply(x, y))
def _divide_no_nan(x, y, name=None): # pylint: disable=unused-argument
dtype = np.result_type(x, y)
y_is_zero = np.equal(y, 0.)
div = np.divide(x, np.where(y_is_zero, np.ones((), dtype=dtype), y))
return np.where(y_is_zero, np.zeros((), dtype=dtype), div)
def __init__(
self,
arr_shape,
order="F",
arr_dtype=np.float32,
use_fft_shifts=True,
sample_mask=None,
ortho=False,
coil_sensitivities=None,
force_real_image=False,
debug=False,
preplan_pyfftw=True,
pyfftw_threads=None,
fft_axes=None,
fftshift_axes=None,
planner_effort="FFTW_ESTIMATE",
loop_over_coils=False,
preserve_memory=False,
disable_warnings=False,
im_mask=None,
pixel_basis="dirac",
rel_fov=None,
**kwargs,
):
"""Cartesian MRI Operator (with partial FFT and coil maps).
Parameters
----------
arr_shape : int
shape of the array
order : {'C','F'}, optional
array ordering that will be assumed if inputs/outputs need to be
reshaped
arr_dtype : numpy.dtype, optional
dtype for the array
sample_mask : array_like, optional
boolean mask of which FFT coefficients to keep
coil_sensitivities : array, optional
Array of coil sensitivities.
ortho : bool, optional
if True, change the normalizeation to the orthogonal case
preplan_pyfftw : bool, optional
if True, precompute the pyFFTW plan upon object creation
pyfftw_threads : int, optional
number of threads to be used by pyFFTW. defaults to
multiprocessing.cpu_count() // 2.
use_fft_shifts : bool, optional
If False, do not apply any FFT shifts
fft_axes : tuple or None, optional
Specify a subset of the axes to transform. The default is to
transform all axes.
fftshift_axes : tuple or None, optional
Specify a subset of the axes to fftshift. The default is to
shift all axes.
im_mask : ndarray or None, optional
Image domain mask
force_real_image : bool, optional
loop_over_coils : bool, optional
If True, memory required is lower, but speed will be slower.
preserve_memory : bool, optional
If False, conjugate copy of coils won't be precomputed
debug : bool, optional
Additional Parameters
---------------------
nd_input : bool, optional
nd_output : bool, optional
"""
if isinstance(arr_shape, (np.ndarray, list)):
# retrieve shape from array
arr_shape = tuple(arr_shape)
if not isinstance(arr_shape, tuple):
raise ValueError("expected array_shape to be a tuple or list")
self.arr_shape = arr_shape
self.ndim = len(arr_shape)
self.order = order
self.use_fft_shifts = use_fft_shifts
self.disable_warnings = disable_warnings
if "loc_in" in kwargs and kwargs["loc_in"] == "gpu":
xp = cupy
on_gpu = True
else:
xp = np
on_gpu = False
self._on_gpu = on_gpu
if sample_mask is not None:
# masking faster if continguity of mask matches
if self.order == "F":
sample_mask = xp.asfortranarray(sample_mask)
elif self.order == "C":
sample_mask = xp.ascontiguousarray(sample_mask)
else:
raise ValueError("order must be C or F")
self.sample_mask = sample_mask
self.force_real_image = force_real_image
self.debug = debug
if self.sample_mask is not None:
if sample_mask.shape != arr_shape:
raise ValueError("sample mask shape must match arr_shape")
# make sure it is boolean
self.sample_mask = self.sample_mask > 0
# # prestore raveled mask indices to save time later during masking
# self.sample_mask_idx = xp.where(
# self.sample_mask.ravel(order=self.order)
# )
self.preserve_memory = preserve_memory
self.loop_over_coils = loop_over_coils
# can specify a subset of the axes to perform the FFT/FFTshifts over
self.fft_axes = fft_axes
if self.fft_axes is None:
if self.order == "C":
# last ndim axes
self.fft_axes = tuple([-ax for ax in range(self.ndim, 0, -1)])
else:
# first ndim axes
self.fft_axes = tuple(range(self.ndim))
if fftshift_axes is None:
self.fftshift_axes = self.fft_axes
else:
self.fftshift_axes = fftshift_axes
# configure scaling (e.g. unitary operator or not)
self.ortho = ortho
if self.fft_axes is None:
Ntrans = prod(self.arr_shape)
else:
Ntrans = prod(
np.asarray(self.arr_shape)[np.asarray(self.fft_axes)])
if self.ortho:
# sqrt of product of shape along axes where FFT is performed
self.scale_ortho = sqrt(Ntrans)
self.gpu_scale_inverse = 1 # self.scale_ortho / Ntrans
# self.gpu_scale_forward = self.scale_ortho
else:
self.scale_ortho = None
self.gpu_scale_inverse = 1 / Ntrans
if "mask_out" in kwargs:
raise ValueError("This operator specifies `mask_out` via the "
"parameter `sample_mask")
if ("mask_in" in kwargs) or ("mask_out" in kwargs):
raise ValueError("This operator specifies `mask_in` via the "
"parameter `im_mask")
if coil_sensitivities is not None:
if coil_sensitivities.ndim == self.ndim + 1:
# self.arr_shape + (Ncoils, )
if self.order == "C":
Ncoils = coil_sensitivities.shape[0]
else:
Ncoils = coil_sensitivities.shape[-1]
Nmaps = 1
elif coil_sensitivities.ndim == self.ndim + 2:
# case with multiple maps (e.g. ESPIRIT soft-SENSE)
# self.arr_shape + (Ncoils, Nmaps)
if self.order == "C":
Ncoils = coil_sensitivities.shape[1]
Nmaps = coil_sensitivities.shape[0]
else:
Ncoils = coil_sensitivities.shape[-2]
Nmaps = coil_sensitivities.shape[-1]
else:
# determine based on size
Ncoils = coil_sensitivities.size / prod(self.arr_shape)
if (Ncoils % 1) != 0:
raise ValueError("sensitivity map size mismatch")
self.Ncoils = int(Ncoils)
self.Nmaps = Nmaps
if self.order == "C":
cmap_shape = (self.Nmaps, self.Ncoils) + self.arr_shape
if not coil_sensitivities.flags.c_contiguous:
msg = (
"Converting coil_sensitivities to be C contiguous"
" (requires a copy). To avoid the copy, convert to "
"Fortran contiguous order prior to calling "
"MRI_Cartesian (see np.ascontiguousarray)")
if not self.disable_warnings:
warnings.warn(msg)
coil_sensitivities = xp.ascontiguousarray(
coil_sensitivities)
else:
cmap_shape = self.arr_shape + (self.Ncoils, self.Nmaps)
if not coil_sensitivities.flags.f_contiguous:
msg = (
"Converting coil_sensitivities to be Fortan contiguous"
" (requires a copy). To avoid the copy, convert to "
"Fortran contiguous order prior to calling "
"MRI_Cartesian (see np.asfortranarray)")
if not self.disable_warnings:
warnings.warn(msg)
coil_sensitivities = xp.asfortranarray(coil_sensitivities)
if tuple(coil_sensitivities.shape) != tuple(cmap_shape):
coil_sensitivities = coil_sensitivities.reshape(
cmap_shape, order=self.order)
if not self.preserve_memory:
self.coil_sensitivities_conj = xp.conj(coil_sensitivities)
else:
self.Ncoils = 1
self.Nmaps = 1
if self.Ncoils == 1:
# TODO: currently has a shape bug if Ncoils == 1
# and loop_over_coils = False.
self.loop_over_coils = True
self.coil_sensitivities = coil_sensitivities
if im_mask is not None:
if im_mask.shape != arr_shape:
raise ValueError("im_mask shape mismatch")
if order != "F":
raise ValueError("only order='F' supported for im_mask case")
nargin = xp.count_nonzero(im_mask)
self.im_mask = im_mask
else:
nargin = prod(arr_shape)
self.im_mask = None
nargin *= self.Nmaps
# nargout = # of k-space samples
if sample_mask is not None:
nargout = xp.count_nonzero(sample_mask) * self.Ncoils
else:
nargout = nargin // self.Nmaps * self.Ncoils
nargout = int(nargout)
self.idx_orig = None
self.idx_conj = None
self.sample_mask_conj = None
# output of FFTs will be complex, regardless of input type
self.result_dtype = np.result_type(arr_dtype, np.complex64)
matvec_allows_repetitions = kwargs.pop("matvec_allows_repetitions",
True)
squeeze_reps = kwargs.pop("squeeze_reps", True)
nd_input = kwargs.pop("nd_input", False)
nd_output = kwargs.pop("nd_output", False)
if (self.sample_mask is not None) and nd_output:
raise ValueError("cannot have both nd_output and sample_mask")
if nd_output:
if self.Ncoils == 1:
shape_out = self.arr_shape
else:
if Nmaps == 1:
if self.order == "C":
shape_out = (self.Ncoils, ) + self.arr_shape
else:
shape_out = self.arr_shape + (self.Ncoils, )
else:
if self.order == "C":
shape_out = (self.Nmaps, self.Ncoils) + self.arr_shape
else:
shape_out = self.arr_shape + (self.Ncoils, self.Nmaps)
else:
shape_out = (nargout, 1)
if self.Nmaps == 1:
shape_in = self.arr_shape
else:
if self.order == "C":
shape_in = (self.Nmaps, ) + self.arr_shape
else:
shape_in = self.arr_shape + (self.Nmaps, )
if self.order == "C":
self.shape_inM = (self.Nmaps, ) + self.arr_shape
else:
self.shape_inM = self.arr_shape + (self.Nmaps, )
self.shape_in1 = self.arr_shape
self.have_pyfftw = config.have_pyfftw
if self.on_gpu:
self.preplan_pyfftw = False
else:
self.preplan_pyfftw = preplan_pyfftw if self.have_pyfftw else False
if self.preplan_pyfftw:
self._preplan_fft(pyfftw_threads, planner_effort)
# raise ValueError("Implementation Incomplete")
if self.on_gpu:
self.fftn = partial(fftn, xp=cupy)
self.ifftn = partial(ifftn, xp=cupy)
else:
if self.preplan_pyfftw:
self._preplan_fft(pyfftw_threads, planner_effort)
else:
self.fftn = partial(fftn, xp=np)
self.ifftn = partial(ifftn, xp=np)
self.fftshift = xp.fft.fftshift
self.ifftshift = xp.fft.ifftshift
self.rel_fov = rel_fov
self.mask = None # TODO: implement or remove (expected by CUDA code)
matvec = self.forward
matvec_adj = self.adjoint
self.norm_available = False
self.norm = self._norm
super(MRI_Cartesian, self).__init__(
nargin=nargin,
nargout=nargout,
matvec=matvec,
matvec_transp=matvec_adj,
matvec_adj=matvec_adj,
nd_input=nd_input or (im_mask is not None),
nd_output=nd_output,
shape_in=shape_in,
shape_out=shape_out,
order=self.order,
matvec_allows_repetitions=matvec_allows_repetitions,
squeeze_reps=squeeze_reps,
mask_in=im_mask,
mask_out=None, # mask_out,
symmetric=False, # TODO: set properly
hermetian=False, # TODO: set properly
dtype=self.result_dtype,
**kwargs,
)
self._init_pixel_basis(pixel_basis=pixel_basis)
def einsum(*operands, **kwargs):
casting = kwargs.pop('casting', 'safe')
dtype = kwargs.pop('dtype', None)
optimize = kwargs.pop('optimize', False)
order = kwargs.pop('order', 'K')
split_every = kwargs.pop('split_every', None)
if kwargs:
raise TypeError("einsum() got unexpected keyword "
"argument(s) %s" % ",".join(kwargs))
einsum_dtype = dtype
inputs, outputs, ops = parse_einsum_input(operands)
subscripts = '->'.join((inputs, outputs))
# Infer the output dtype from operands
if dtype is None:
dtype = np.result_type(*[o.dtype for o in ops])
if einsum_can_optimize:
if optimize is not False:
# Avoid computation of dask arrays within np.einsum_path
# by passing in small numpy arrays broadcasted
# up to the right shape
fake_ops = [
np.broadcast_to(o.dtype.type(0), shape=o.shape) for o in ops
]
optimize, _ = np.einsum_path(subscripts,
*fake_ops,
optimize=optimize)
kwargs = {'optimize': optimize}
else:
kwargs = {}
inputs = [tuple(i) for i in inputs.split(",")]
# Set of all indices
all_inds = set(a for i in inputs for a in i)
# Which indices are contracted?
contract_inds = all_inds - set(outputs)
ncontract_inds = len(contract_inds)
# Introduce the contracted indices into the blockwise product
# so that we get numpy arrays, not lists
result = blockwise(
chunk.einsum,
tuple(outputs) + tuple(contract_inds),
*(a for ap in zip(ops, inputs) for a in ap),
# blockwise parameters
adjust_chunks={ind: 1
for ind in contract_inds},
dtype=dtype,
# np.einsum parameters
subscripts=subscripts,
kernel_dtype=einsum_dtype,
ncontract_inds=ncontract_inds,
order=order,
casting=casting,
**kwargs)
# Now reduce over any extra contraction dimensions
if ncontract_inds > 0:
size = len(outputs)
return result.sum(axis=list(range(size, size + ncontract_inds)),
split_every=split_every)
return result
def check_output(output, unit, inputs, function=None):
"""Check that function output can be stored in the output array given.
Parameters
----------
output : array or `~astropy.units.Quantity` or tuple
Array that should hold the function output (or tuple of such arrays).
unit : `~astropy.units.Unit` or None, or tuple
Unit that the output will have, or `None` for pure numbers (should be
tuple of same if output is a tuple of outputs).
inputs : tuple
Any input arguments. These should be castable to the output.
function : callable
The function that will be producing the output. If given, used to
give a more informative error message.
Returns
-------
arrays : `~numpy.ndarray` view of ``output`` (or tuple of such views).
Raises
------
UnitTypeError : If ``unit`` is inconsistent with the class of ``output``
TypeError : If the ``inputs`` cannot be cast safely to ``output``.
"""
if isinstance(output, tuple):
return tuple(
check_output(output_, unit_, inputs, function)
for output_, unit_ in zip(output, unit))
# ``None`` indicates no actual array is needed. This can happen, e.g.,
# with np.modf(a, out=(None, b)).
if output is None:
return None
if hasattr(output, '__quantity_subclass__'):
# Check that we're not trying to store a plain Numpy array or a
# Quantity with an inconsistent unit (e.g., not angular for Angle).
if unit is None:
raise TypeError("Cannot store non-quantity output{0} in {1} "
"instance".format(
(" from {0} function".format(function.__name__)
if function is not None else ""),
type(output)))
if output.__quantity_subclass__(unit)[0] is not type(output):
raise UnitTypeError(
"Cannot store output with unit '{0}'{1} "
"in {2} instance. Use {3} instance instead.".format(
unit, (" from {0} function".format(function.__name__)
if function is not None else ""), type(output),
output.__quantity_subclass__(unit)[0]))
# Turn into ndarray, so we do not loop into array_wrap/array_ufunc
# if the output is used to store results of a function.
output = output.view(np.ndarray)
else:
# output is not a Quantity, so cannot obtain a unit.
if not (unit is None or unit is dimensionless_unscaled):
raise UnitTypeError(
"Cannot store quantity with dimension "
"{0}in a non-Quantity instance.".format(
"" if function is None else "resulting from {0} function ".
format(function.__name__)))
# check we can handle the dtype (e.g., that we are not int
# when float is required).
if not np.can_cast(
np.result_type(*inputs), output.dtype, casting='same_kind'):
raise TypeError("Arguments cannot be cast safely to inplace "
"output with dtype={0}".format(output.dtype))
return output
def numeric_normalize_types(*args):
"""Cast all args to a common type using numpy promotion logic
"""
dtype = np.result_type(*[a.dtype for a in args])
return [a.astype(dtype) for a in args]
def get_k_kpts(mydf,
dm_kpts,
hermi=1,
kpts=np.zeros((1, 3)),
kpts_band=None,
exxdiv=None):
'''Get the Coulomb (J) and exchange (K) AO matrices at sampled k-points.
Args:
dm_kpts : (nkpts, nao, nao) ndarray
Density matrix at each k-point
kpts : (nkpts, 3) ndarray
Kwargs:
hermi : int
Whether K matrix is hermitian
| 0 : not hermitian and not symmetric
| 1 : hermitian
kpts_band : (3,) ndarray or (*,3) ndarray
A list of arbitrary "band" k-points at which to evalute the matrix.
Returns:
vj : (nkpts, nao, nao) ndarray
vk : (nkpts, nao, nao) ndarray
or list of vj and vk if the input dm_kpts is a list of DMs
'''
cell = mydf.cell
mesh = mydf.mesh
coords = cell.gen_uniform_grids(mesh)
ngrids = coords.shape[0]
if getattr(dm_kpts, 'mo_coeff', None) is not None:
mo_coeff = dm_kpts.mo_coeff
mo_occ = dm_kpts.mo_occ
else:
mo_coeff = None
kpts = np.asarray(kpts)
dm_kpts = lib.asarray(dm_kpts, order='C')
dms = _format_dms(dm_kpts, kpts)
nset, nkpts, nao = dms.shape[:3]
weight = 1. / nkpts * (cell.vol / ngrids)
kpts_band, input_band = _format_kpts_band(kpts_band, kpts), kpts_band
nband = len(kpts_band)
if gamma_point(kpts_band) and gamma_point(kpts):
vk_kpts = np.zeros((nset, nband, nao, nao), dtype=dms.dtype)
else:
vk_kpts = np.zeros((nset, nband, nao, nao), dtype=np.complex128)
coords = mydf.grids.coords
ao2_kpts = [
np.asarray(ao.T, order='C')
for ao in mydf._numint.eval_ao(cell, coords, kpts=kpts)
]
if input_band is None:
ao1_kpts = ao2_kpts
else:
ao1_kpts = [
np.asarray(ao.T, order='C')
for ao in mydf._numint.eval_ao(cell, coords, kpts=kpts_band)
]
if mo_coeff is not None and nset == 1:
mo_coeff = [
mo_coeff[k][:, occ > 0] * np.sqrt(occ[occ > 0])
for k, occ in enumerate(mo_occ)
]
ao2_kpts = [np.dot(mo_coeff[k].T, ao) for k, ao in enumerate(ao2_kpts)]
mem_now = lib.current_memory()[0]
max_memory = mydf.max_memory - mem_now
blksize = int(
min(nao, max(1, (max_memory - mem_now) * 1e6 / 16 / 4 / ngrids / nao)))
lib.logger.debug1(mydf, 'fft_jk: get_k_kpts max_memory %s blksize %d',
max_memory, blksize)
ao1_dtype = np.result_type(*ao1_kpts)
ao2_dtype = np.result_type(*ao2_kpts)
vR_dm = np.empty((nset, nao, ngrids), dtype=vk_kpts.dtype)
t1 = (time.clock(), time.time())
for k2, ao2T in enumerate(ao2_kpts):
if ao2T.size == 0:
continue
kpt2 = kpts[k2]
naoj = ao2T.shape[0]
if mo_coeff is None or nset > 1:
ao_dms = [lib.dot(dms[i, k2], ao2T.conj()) for i in range(nset)]
else:
ao_dms = [ao2T.conj()]
for k1, ao1T in enumerate(ao1_kpts):
kpt1 = kpts_band[k1]
# If we have an ewald exxdiv, we add the G=0 correction near the
# end of the function to bypass any discretization errors
# that arise from the FFT.
mydf.exxdiv = exxdiv
if exxdiv == 'ewald' or exxdiv is None:
coulG = tools.get_coulG(cell, kpt2 - kpt1, False, mydf, mesh)
else:
coulG = tools.get_coulG(cell, kpt2 - kpt1, True, mydf, mesh)
if is_zero(kpt1 - kpt2):
expmikr = np.array(1.)
else:
expmikr = np.exp(-1j * np.dot(coords, kpt2 - kpt1))
for p0, p1 in lib.prange(0, nao, blksize):
rho1 = np.einsum('ig,jg->ijg', ao1T[p0:p1].conj() * expmikr,
ao2T)
vG = tools.fft(rho1.reshape(-1, ngrids), mesh)
rho1 = None
vG *= coulG
vR = tools.ifft(vG, mesh).reshape(p1 - p0, naoj, ngrids)
vG = None
if vR_dm.dtype == np.double:
vR = vR.real
for i in range(nset):
np.einsum('ijg,jg->ig', vR, ao_dms[i], out=vR_dm[i, p0:p1])
vR = None
vR_dm *= expmikr.conj()
for i in range(nset):
vk_kpts[i, k1] += weight * lib.dot(vR_dm[i], ao1T.T)
t1 = lib.logger.timer_debug1(mydf, 'get_k_kpts: make_kpt (%d,*)' % k2,
*t1)
# Function _ewald_exxdiv_for_G0 to add back in the G=0 component to vk_kpts
# Note in the _ewald_exxdiv_for_G0 implementation, the G=0 treatments are
# different for 1D/2D and 3D systems. The special treatments for 1D and 2D
# can only be used with AFTDF/GDF/MDF method. In the FFTDF method, 1D, 2D
# and 3D should use the ewald probe charge correction.
if exxdiv == 'ewald':
_ewald_exxdiv_for_G0(cell, kpts, dms, vk_kpts, kpts_band=kpts_band)
return _format_jks(vk_kpts, dm_kpts, input_band, kpts)
def asfcarray(a):
a = np.asarray(a)
return np.asarray(a, np.result_type(a, 0.0))
def _result_type(*arrays_and_dtypes):
dtype = np.result_type(*arrays_and_dtypes)
return canonicalize_dtype(dtype)
def where(condition, x=None, y=None):
"""
Return elements, either from `x` or `y`, depending on `condition`.
If only `condition` is given, return ``condition.nonzero()``.
Parameters
----------
condition : array_like, bool
When True, yield `x`, otherwise yield `y`.
x, y : array_like, optional
Values from which to choose. `x`, `y` and `condition` need to be
broadcastable to some shape.
Returns
-------
out : Tensor or tuple of Tensors
If both `x` and `y` are specified, the output tensor contains
elements of `x` where `condition` is True, and elements from
`y` elsewhere.
If only `condition` is given, return the tuple
``condition.nonzero()``, the indices where `condition` is True.
See Also
--------
nonzero, choose
Notes
-----
If `x` and `y` are given and input arrays are 1-D, `where` is
equivalent to::
[xv if c else yv for (c,xv,yv) in zip(condition,x,y)]
Examples
--------
>>> import mars.tensor as mt
>>> mt.where([[True, False], [True, True]],
... [[1, 2], [3, 4]],
... [[9, 8], [7, 6]]).execute()
array([[1, 8],
[3, 4]])
>>> mt.where([[0, 1], [1, 0]]).execute()
(array([0, 1]), array([1, 0]))
>>> x = mt.arange(9.).reshape(3, 3)
>>> mt.where( x > 5 ).execute()
(array([2, 2, 2]), array([0, 1, 2]))
>>> mt.where(x < 5, x, -1).execute() # Note: broadcasting.
array([[ 0., 1., 2.],
[ 3., 4., -1.],
[-1., -1., -1.]])
Find the indices of elements of `x` that are in `goodvalues`.
>>> goodvalues = [3, 4, 7]
>>> ix = mt.isin(x, goodvalues)
>>> ix.execute()
array([[False, False, False],
[ True, True, False],
[False, True, False]])
>>> mt.where(ix).execute()
(array([1, 1, 2]), array([0, 1, 1]))
"""
if (x is None) != (y is None):
raise ValueError('either both or neither of x and y should be given')
if x is None and y is None:
return astensor(condition).nonzero()
x, y = astensor(x), astensor(y)
dtype = np.result_type(x.dtype, y.dtype)
shape = broadcast_shape(x.shape, y.shape)
if np.isscalar(condition):
return broadcast_to(x if condition else y, shape).astype(dtype)
else:
condition = astensor(condition)
op = TensorWhere(dtype=dtype)
return op(condition, x, y, shape=shape)
def linspace(start,
stop,
num=50,
endpoint=True,
retstep=False,
dtype=None,
axis=0):
"""Implements a simplified linspace function as of numpy verion >= 1.16.
As of numpy 1.16, the arguments start and stop can be array-like and
there is an optional argument `axis`.
For simplicity, we only allow 1d array-like to be passed to start and stop.
See: https://github.com/numpy/numpy/pull/12388 and numpy 1.16 release
notes about start and stop arrays for linspace logspace and geomspace.
Returns
-------
out : ndarray of shape (num, n_start) or (num,)
The output array with `n_start=start.shape[0]` columns.
"""
if np_version < parse_version('1.16'):
start = np.asanyarray(start) * 1.0
stop = np.asanyarray(stop) * 1.0
dt = np.result_type(start, stop, float(num))
if dtype is None:
dtype = dt
if start.ndim == 0 == stop.ndim:
return np.linspace(start=start,
stop=stop,
num=num,
endpoint=endpoint,
retstep=retstep,
dtype=dtype)
if start.ndim != 1 or stop.ndim != 1 or start.shape != stop.shape:
raise ValueError("start and stop must be 1d array-like of same"
" shape.")
n_start = start.shape[0]
out = np.empty((num, n_start), dtype=dtype)
step = np.empty(n_start, dtype=np.float)
for i in range(n_start):
out[:, i], step[i] = np.linspace(start=start[i],
stop=stop[i],
num=num,
endpoint=endpoint,
retstep=True,
dtype=dtype)
if axis != 0:
out = np.moveaxis(out, 0, axis)
if retstep:
return out, step
else:
return out
else:
return np.linspace(start=start,
stop=stop,
num=num,
endpoint=endpoint,
retstep=retstep,
dtype=dtype,
axis=axis)
def check_array(
array,
accept_sparse=False,
accept_large_sparse=True,
dtype="numeric",
order=None,
copy=False,
force_all_finite=True,
ensure_2d=True,
allow_nd=False,
ensure_min_samples=1,
ensure_min_features=1,
warn_on_dtype=None,
estimator=None,
):
"""Input validation on an array, list, sparse matrix or similar.
By default, the input is checked to be a non-empty 2D array containing
only finite values. If the dtype of the array is object, attempt
converting to float, raising on failure.
Parameters
----------
array : object
Input object to check / convert.
accept_sparse : string, boolean or list/tuple of strings (default=False)
String[s] representing allowed sparse matrix formats, such as 'csc',
'csr', etc. If the input is sparse but not in the allowed format,
it will be converted to the first listed format. True allows the input
to be any format. False means that a sparse matrix input will
raise an error.
accept_large_sparse : bool (default=True)
If a CSR, CSC, COO or BSR sparse matrix is supplied and accepted by
accept_sparse, accept_large_sparse=False will cause it to be accepted
only if its indices are stored with a 32-bit dtype.
.. versionadded:: 0.20
dtype : string, type, list of types or None (default="numeric")
Data type of result. If None, the dtype of the input is preserved.
If "numeric", dtype is preserved unless array.dtype is object.
If dtype is a list of types, conversion on the first type is only
performed if the dtype of the input is not in the list.
order : 'F', 'C' or None (default=None)
Whether an array will be forced to be fortran or c-style.
When order is None (default), then if copy=False, nothing is ensured
about the memory layout of the output array; otherwise (copy=True)
the memory layout of the returned array is kept as close as possible
to the original array.
copy : boolean (default=False)
Whether a forced copy will be triggered. If copy=False, a copy might
be triggered by a conversion.
force_all_finite : boolean or 'allow-nan', (default=True)
Whether to raise an error on np.inf and np.nan in array. The
possibilities are:
- True: Force all values of array to be finite.
- False: accept both np.inf and np.nan in array.
- 'allow-nan': accept only np.nan values in array. Values cannot
be infinite.
For object dtyped data, only np.nan is checked and not np.inf.
.. versionadded:: 0.20
``force_all_finite`` accepts the string ``'allow-nan'``.
ensure_2d : boolean (default=True)
Whether to raise a value error if array is not 2D.
allow_nd : boolean (default=False)
Whether to allow array.ndim > 2.
ensure_min_samples : int (default=1)
Make sure that the array has a minimum number of samples in its first
axis (rows for a 2D array). Setting to 0 disables this check.
ensure_min_features : int (default=1)
Make sure that the 2D array has some minimum number of features
(columns). The default value of 1 rejects empty datasets.
This check is only enforced when the input data has effectively 2
dimensions or is originally 1D and ``ensure_2d`` is True. Setting to 0
disables this check.
warn_on_dtype : boolean or None, optional (default=None)
Raise DataConversionWarning if the dtype of the input data structure
does not match the requested dtype, causing a memory copy.
.. deprecated:: 0.21
``warn_on_dtype`` is deprecated in version 0.21 and will be
removed in 0.23.
estimator : str or estimator instance (default=None)
If passed, include the name of the estimator in warning messages.
Returns
-------
array_converted : object
The converted and validated array.
"""
# warn_on_dtype deprecation
if warn_on_dtype is not None:
warnings.warn(
"'warn_on_dtype' is deprecated in version 0.21 and will be "
"removed in 0.23. Don't set `warn_on_dtype` to remove this "
"warning.",
FutureWarning,
stacklevel=2,
)
# store reference to original array to check if copy is needed when
# function returns
array_orig = array
# store whether originally we wanted numeric dtype
dtype_numeric = isinstance(dtype, str) and dtype == "numeric"
dtype_orig = getattr(array, "dtype", None)
if not hasattr(dtype_orig, "kind"):
# not a data type (e.g. a column named dtype in a pandas DataFrame)
dtype_orig = None
# check if the object contains several dtypes (typically a pandas
# DataFrame), and store them. If not, store None.
dtypes_orig = None
if hasattr(array, "dtypes") and hasattr(array.dtypes, "__array__"):
dtypes_orig = np.array(array.dtypes)
if all(isinstance(dtype, np.dtype) for dtype in dtypes_orig):
dtype_orig = np.result_type(*array.dtypes)
if dtype_numeric:
if dtype_orig is not None and dtype_orig.kind == "O":
# if input is object, convert to float.
dtype = np.float64
else:
dtype = None
if isinstance(dtype, (list, tuple)):
if dtype_orig is not None and dtype_orig in dtype:
# no dtype conversion required
dtype = None
else:
# dtype conversion required. Let's select the first element of the
# list of accepted types.
dtype = dtype[0]
if force_all_finite not in (True, False, "allow-nan"):
raise ValueError(
'force_all_finite should be a bool or "allow-nan"'
". Got {!r} instead".format(force_all_finite)
)
if estimator is not None:
if isinstance(estimator, str):
estimator_name = estimator
else:
estimator_name = estimator.__class__.__name__
else:
estimator_name = "Estimator"
context = " by %s" % estimator_name if estimator is not None else ""
if sp.issparse(array):
_ensure_no_complex_data(array)
array = _ensure_sparse_format(
array,
accept_sparse=accept_sparse,
dtype=dtype,
copy=copy,
force_all_finite=force_all_finite,
accept_large_sparse=accept_large_sparse,
)
else:
# If np.array(..) gives ComplexWarning, then we convert the warning
# to an error. This is needed because specifying a non complex
# dtype to the function converts complex to real dtype,
# thereby passing the test made in the lines following the scope
# of warnings context manager.
with warnings.catch_warnings():
try:
warnings.simplefilter("error", ComplexWarning)
if dtype is not None and np.dtype(dtype).kind in "iu":
# Conversion float -> int should not contain NaN or
# inf (numpy#14412). We cannot use casting='safe' because
# then conversion float -> int would be disallowed.
array = np.asarray(array, order=order)
if array.dtype.kind == "f":
_assert_all_finite(array, allow_nan=False, msg_dtype=dtype)
array = array.astype(dtype, casting="unsafe", copy=False)
else:
array = np.asarray(array, order=order, dtype=dtype)
except ComplexWarning:
raise ValueError("Complex data not supported\n" "{}\n".format(array))
# It is possible that the np.array(..) gave no warning. This happens
# when no dtype conversion happened, for example dtype = None. The
# result is that np.array(..) produces an array of complex dtype
# and we need to catch and raise exception for such cases.
_ensure_no_complex_data(array)
if ensure_2d:
# If input is scalar raise error
if array.ndim == 0:
raise ValueError(
"Expected 2D array, got scalar array instead:\narray={}.\n"
"Reshape your data either using array.reshape(-1, 1) if "
"your data has a single feature or array.reshape(1, -1) "
"if it contains a single sample.".format(array)
)
# If input is 1D raise error
if array.ndim == 1:
raise ValueError(
"Expected 2D array, got 1D array instead:\narray={}.\n"
"Reshape your data either using array.reshape(-1, 1) if "
"your data has a single feature or array.reshape(1, -1) "
"if it contains a single sample.".format(array)
)
# in the future np.flexible dtypes will be handled like object dtypes
if dtype_numeric and np.issubdtype(array.dtype, np.flexible):
warnings.warn(
"Beginning in version 0.22, arrays of bytes/strings will be "
"converted to decimal numbers if dtype='numeric'. "
"It is recommended that you convert the array to "
"a float dtype before using it in scikit-learn, "
"for example by using "
"your_array = your_array.astype(np.float64).",
FutureWarning,
stacklevel=2,
)
# make sure we actually converted to numeric:
if dtype_numeric and array.dtype.kind == "O":
array = array.astype(np.float64)
if not allow_nd and array.ndim >= 3:
raise ValueError(
"Found array with dim %d. %s expected <= 2."
% (array.ndim, estimator_name)
)
if force_all_finite:
_assert_all_finite(array, allow_nan=force_all_finite == "allow-nan")
if ensure_min_samples > 0:
n_samples = _num_samples(array)
if n_samples < ensure_min_samples:
raise ValueError(
"Found array with %d sample(s) (shape=%s) while a"
" minimum of %d is required%s."
% (n_samples, array.shape, ensure_min_samples, context)
)
if ensure_min_features > 0 and array.ndim == 2:
n_features = array.shape[1]
if n_features < ensure_min_features:
raise ValueError(
"Found array with %d feature(s) (shape=%s) while"
" a minimum of %d is required%s."
% (n_features, array.shape, ensure_min_features, context)
)
if warn_on_dtype and dtype_orig is not None and array.dtype != dtype_orig:
msg = "Data with input dtype %s was converted to %s%s." % (
dtype_orig,
array.dtype,
context,
)
warnings.warn(msg, DataConversionWarning, stacklevel=2)
if copy and np.may_share_memory(array, array_orig):
array = np.array(array, dtype=dtype, order=order)
if warn_on_dtype and dtypes_orig is not None and {array.dtype} != set(dtypes_orig):
# if there was at the beginning some other types than the final one
# (for instance in a DataFrame that can contain several dtypes) then
# some data must have been converted
msg = "Data with input dtype %s were all converted to %s%s." % (
", ".join(map(str, sorted(set(dtypes_orig)))),
array.dtype,
context,
)
warnings.warn(msg, DataConversionWarning, stacklevel=3)
return array
def canonicalize_ndarray_dtype(x):
return onp.asarray(x, xb.canonicalize_dtype(onp.result_type(x)))
def _expm_multiply_interval(A,
B,
start=None,
stop=None,
num=None,
endpoint=None,
balance=False,
status_only=False):
"""
Compute the action of the matrix exponential at multiple time points.
Parameters
----------
A : transposable linear operator
The operator whose exponential is of interest.
B : ndarray
The matrix to be multiplied by the matrix exponential of A.
start : scalar, optional
The starting time point of the sequence.
stop : scalar, optional
The end time point of the sequence, unless `endpoint` is set to False.
In that case, the sequence consists of all but the last of ``num + 1``
evenly spaced time points, so that `stop` is excluded.
Note that the step size changes when `endpoint` is False.
num : int, optional
Number of time points to use.
endpoint : bool, optional
If True, `stop` is the last time point. Otherwise, it is not included.
balance : bool
Indicates whether or not to apply balancing.
status_only : bool
A flag that is set to True for some debugging and testing operations.
Returns
-------
F : ndarray
:math:`e^{t_k A} B`
status : int
An integer status for testing and debugging.
Notes
-----
This is algorithm (5.2) in Al-Mohy and Higham (2011).
There seems to be a typo, where line 15 of the algorithm should be
moved to line 6.5 (between lines 6 and 7).
"""
if balance:
raise NotImplementedError
if len(A.shape) != 2 or A.shape[0] != A.shape[1]:
raise ValueError('expected A to be like a square matrix')
if A.shape[1] != B.shape[0]:
raise ValueError('the matrices A and B have incompatible shapes')
ident = _ident_like(A)
n = A.shape[0]
if len(B.shape) == 1:
n0 = 1
elif len(B.shape) == 2:
n0 = B.shape[1]
else:
raise ValueError('expected B to be like a matrix or a vector')
u_d = 2**-53
tol = u_d
mu = _trace(A) / float(n)
# Get the linspace samples, attempting to preserve the linspace defaults.
linspace_kwargs = {'retstep': True}
if num is not None:
linspace_kwargs['num'] = num
if endpoint is not None:
linspace_kwargs['endpoint'] = endpoint
samples, step = np.linspace(start, stop, **linspace_kwargs)
# Convert the linspace output to the notation used by the publication.
nsamples = len(samples)
if nsamples < 2:
raise ValueError('at least two time points are required')
q = nsamples - 1
h = step
t_0 = samples[0]
t_q = samples[q]
# Define the output ndarray.
# Use an ndim=3 shape, such that the last two indices
# are the ones that may be involved in level 3 BLAS operations.
X_shape = (nsamples, ) + B.shape
X = np.empty(X_shape, dtype=np.result_type(A.dtype, B.dtype, float))
t = t_q - t_0
A = A - mu * ident
A_1_norm = _exact_1_norm(A)
if t * A_1_norm == 0:
m_star, s = 0, 1
else:
ell = 2
norm_info = LazyOperatorNormInfo(t * A, A_1_norm=t * A_1_norm, ell=ell)
m_star, s = _fragment_3_1(norm_info, n0, tol, ell=ell)
# Compute the expm action up to the initial time point.
X[0] = _expm_multiply_simple_core(A, B, t_0, mu, m_star, s)
# Compute the expm action at the rest of the time points.
if q <= s:
if status_only:
return 0
else:
return _expm_multiply_interval_core_0(A, X, h, mu, m_star, s, q)
elif q > s and not (q % s):
if status_only:
return 1
else:
return _expm_multiply_interval_core_1(A, X, h, mu, m_star, s, q,
tol)
elif q > s and (q % s):
if status_only:
return 2
else:
return _expm_multiply_interval_core_2(A, X, h, mu, m_star, s, q,
tol)
else:
raise Exception('internal error')
def stack(tensors, axis=0):
"""
Join a sequence of tensors along a new axis.
The `axis` parameter specifies the index of the new axis in the dimensions
of the result. For example, if ``axis=0`` it will be the first dimension
and if ``axis=-1`` it will be the last dimension.
Parameters
----------
tensors : sequence of array_like
Each tensor must have the same shape.
axis : int, optional
The axis in the result tensor along which the input tensors are stacked.
out : Tensor, optional
If provided, the destination to place the result. The shape must be
correct, matching that of what stack would have returned if no
out argument were specified.
Returns
-------
stacked : Tensor
The stacked tensor has one more dimension than the input tensors.
See Also
--------
concatenate : Join a sequence of tensors along an existing axis.
split : Split tensor into a list of multiple sub-tensors of equal size.
block : Assemble tensors from blocks.
Examples
--------
>>> import mars.tensor as mt
>>> arrays = [mt.random.randn(3, 4) for _ in range(10)]
>>> mt.stack(arrays, axis=0).shape
(10, 3, 4)
>>> mt.stack(arrays, axis=1).shape
(3, 10, 4)
>>> mt.stack(arrays, axis=2).shape
(3, 4, 10)
>>> a = mt.array([1, 2, 3])
>>> b = mt.array([2, 3, 4])
>>> mt.stack((a, b)).execute()
array([[1, 2, 3],
[2, 3, 4]])
>>> mt.stack((a, b), axis=-1).execute()
array([[1, 2],
[2, 3],
[3, 4]])
"""
tensors = [astensor(t) for t in tensors]
if len(set(t.shape for t in tensors)) != 1:
raise ValueError('all input tensors must have the same shape')
ndim = len(tensors[0].shape)
raw_axis = axis
if axis < 0:
axis = ndim + axis + 1
if axis > ndim or axis < 0:
raise np.AxisError('axis {0} is out of bounds for tensor '
'of dimension {1}'.format(raw_axis, ndim))
dtype = np.result_type(*[t.dtype for t in tensors])
sparse = all(t.issparse() for t in tensors)
op = TensorStack(axis=axis, dtype=dtype, sparse=sparse)
return op(tensors)
def rand_map(shape,
wcs,
ps_lensinput,
lmax=None,
maplmax=None,
dtype=np.float64,
seed=None,
oversample=2.0,
spin=2,
output="l",
geodesic=True,
verbose=False,
delta_theta=None,
phi_seed=None,
separate_phi_from_cmb=False):
import curvedsky, sharp
ctype = np.result_type(dtype, 0j)
# Restrict to target number of components
oshape = shape[-3:]
if len(oshape) == 2: shape = (1, ) + tuple(shape)
#van Engelen added this:
# First draw a random lensing field, and use it to compute the undeflected positions
if not separate_phi_from_cmb:
#AVE - this was the default option.
if verbose: print("Generating alms")
alm, ainfo = curvedsky.rand_alm(ps_lensinput,
lmax=lmax,
seed=seed,
dtype=ctype,
return_ainfo=True)
phi_alm, cmb_alm = alm[0], alm[1:1 + shape[-3]]
del alm
else:
if verbose: print("Generating alms, separating phi from cmb")
phi_alm, phi_ainfo = curvedsky.rand_alm(ps_lensinput[0, 0, :],
lmax=lmax,
seed=phi_seed,
dtype=ctype,
return_ainfo=True)
cmb_alm, cmb_ainfo = curvedsky.rand_alm(ps_lensinput[1:, 1:, :],
lmax=lmax,
seed=seed,
dtype=ctype,
return_ainfo=True)
# Truncate alm if we want a smoother map. In taylens, it was necessary to truncate
# to a lower lmax for the map than for phi, to avoid aliasing. The appropriate lmax
# for the cmb was the one that fits the resolution. FIXME: Can't slice alm this way.
#if maplmax: cmb_alm = cmb_alm[:,:maplmax]
if delta_theta is None: bsize = shape[-2]
else:
bsize = utils.nint(abs(delta_theta / utils.degree / wcs.wcs.cdelt[1]))
# Adjust bsize so we don't get any tiny blocks at the end
nblock = shape[-2] // bsize
bsize = int(shape[-2] / (nblock + 0.5))
# Allocate output maps
if "p" in output: phi_map = enmap.empty(shape[-2:], wcs, dtype=dtype)
if "k" in output:
kappa_map = enmap.empty(shape[-2:], wcs, dtype=dtype)
l = np.arange(ainfo.lmax + 1.0)
kappa_alm = ainfo.lmul(phi_alm, l * (l + 1) / 2)
for i1 in range(0, shape[-2], bsize):
curvedsky.alm2map(kappa_alm, kappa_map[..., i1:i1 + bize, :])
del kappa_alm
if "a" in output:
grad_map = enmap.empty((2, ) + shape[-2:], wcs, dtype=dtype)
if "u" in output: cmb_raw = enmap.empty(shape, wcs, dtype=dtype)
if "l" in output: cmb_obs = enmap.empty(shape, wcs, dtype=dtype)
# Then loop over dec bands
for i1 in range(0, shape[-2], bsize):
i2 = min(i1 + bsize, shape[-2])
lshape, lwcs = enmap.slice_geometry(shape, wcs,
(slice(i1, i2), slice(None)))
if "p" in output:
if verbose: print("Computing phi map")
curvedsky.alm2map(phi_alm, phi_map[..., i1:i2, :])
if verbose: print("Computing grad map")
if "a" in output: grad = grad_map[..., i1:i2, :]
else: grad = enmap.zeros((2, ) + lshape[-2:], lwcs, dtype=dtype)
curvedsky.alm2map(phi_alm, grad, deriv=True)
if "l" not in output: continue
if verbose: print("Computing observed coordinates")
obs_pos = enmap.posmap(lshape, lwcs)
if verbose: print("Computing alpha map")
raw_pos = enmap.samewcs(
offset_by_grad(obs_pos, grad, pol=shape[-3] > 1,
geodesic=geodesic), obs_pos)
del obs_pos, grad
if "u" in output:
if verbose: print("Computing unlensed map")
curvedsky.alm2map(cmb_alm, cmb_raw[..., i1:i2, :], spin=spin)
if verbose: print("Computing lensed map")
cmb_obs[..., i1:i2, :] = curvedsky.alm2map_pos(cmb_alm,
raw_pos[:2],
oversample=oversample,
spin=spin)
if raw_pos.shape[0] > 2 and np.any(raw_pos[2]):
if verbose: print("Rotating polarization")
cmb_obs[..., i1:i2, :] = enmap.rotate_pol(cmb_obs[..., i1:i2, :],
raw_pos[2])
del raw_pos
del cmb_alm, phi_alm
# Output in same order as specified in output argument
res = []
for c in output:
if c == "l": res.append(cmb_obs.reshape(oshape))
elif c == "u": res.append(cmb_raw.reshape(oshape))
elif c == "p": res.append(phi_map)
elif c == "k": res.append(kappa_map)
elif c == "a": res.append(grad_map)
return tuple(res)
def _get_bin_edges(op, a, bins, range, weights):
# parse the overloaded bins argument
n_equal_bins = None
bin_edges = None
first_edge = None
last_edge = None
if isinstance(bins, str):
# when `bins` is str, x.min() and x.max()
# will be calculated in advance
bin_name = bins
if a.size > 0:
assert range is not None
raw_range = range
first_edge, last_edge = _get_outer_edges(a, range)
if a.size == 0:
n_equal_bins = 1
else:
# Do not call selectors on empty arrays
selector = _hist_bin_selectors[bin_name](op, a,
(first_edge, last_edge),
raw_range)
selector.check()
width = selector.get_result()
if width:
n_equal_bins = int(
np.ceil(_unsigned_subtract(last_edge, first_edge) / width))
else:
# Width can be zero for some estimators, e.g. FD when
# the IQR of the data is zero.
n_equal_bins = 1
elif mt.ndim(bins) == 0:
first_edge, last_edge = _get_outer_edges(a, range)
n_equal_bins = bins
else:
# cannot be Tensor, must be calculated first
assert mt.ndim(bins) == 1 and not isinstance(bins, TENSOR_TYPE)
bin_edges = np.asarray(bins)
if not is_asc_sorted(bin_edges):
raise ValueError(
'`bins` must increase monotonically, when an array')
if n_equal_bins is not None:
# numpy gh-10322 means that type resolution rules are dependent on array
# shapes. To avoid this causing problems, we pick a type now and stick
# with it throughout.
bin_type = np.result_type(first_edge, last_edge, a)
if np.issubdtype(bin_type, np.integer):
bin_type = np.result_type(bin_type, float)
# bin edges must be computed
bin_edges = mt.linspace(first_edge,
last_edge,
n_equal_bins + 1,
endpoint=True,
dtype=bin_type,
gpu=op.gpu)
return bin_edges, (first_edge, last_edge, n_equal_bins)
else:
return mt.tensor(bin_edges), None
def einsum(idx_str, *tensors, **kwargs):
'''Perform a more efficient einsum via reshaping to a matrix multiply.
Current differences compared to numpy.einsum:
This assumes that each repeated index is actually summed (i.e. no 'i,i->i')
and appears only twice (i.e. no 'ij,ik,il->jkl'). The output indices must
be explicitly specified (i.e. 'ij,j->i' and not 'ij,j').
'''
DEBUG = kwargs.get('DEBUG', False)
idx_str = idx_str.replace(' ', '')
indices = "".join(re.split(',|->', idx_str))
if '->' not in idx_str or any(
indices.count(x) > 2 for x in set(indices)):
return numpy.einsum(idx_str, *tensors)
if idx_str.count(',') > 1:
indices = re.split(',|->', idx_str)
indices_in = indices[:-1]
idx_final = indices[-1]
n_shared_max = 0
for i in range(len(indices_in)):
for j in range(i):
tmp = list(set(indices_in[i]).intersection(indices_in[j]))
n_shared_indices = len(tmp)
if n_shared_indices > n_shared_max:
n_shared_max = n_shared_indices
shared_indices = tmp
[a, b] = [i, j]
tensors = list(tensors)
A, B = tensors[a], tensors[b]
idxA, idxB = indices[a], indices[b]
idx_out = list(idxA + idxB)
idx_out = "".join([x for x in idx_out if x not in shared_indices])
C = einsum(idxA + "," + idxB + "->" + idx_out, A, B)
indices_in.pop(a)
indices_in.pop(b)
indices_in.append(idx_out)
tensors.pop(a)
tensors.pop(b)
tensors.append(C)
return einsum(",".join(indices_in) + "->" + idx_final, *tensors)
A, B = tensors
# Call numpy.asarray because A or B may be HDF5 Datasets
A = numpy.asarray(A, order='A')
B = numpy.asarray(B, order='A')
if A.size < 2000 or B.size < 2000:
return numpy.einsum(idx_str, *tensors)
# Split the strings into a list of idx char's
idxA, idxBC = idx_str.split(',')
idxB, idxC = idxBC.split('->')
idxA, idxB, idxC = [list(x) for x in [idxA, idxB, idxC]]
assert (len(idxA) == A.ndim)
assert (len(idxB) == B.ndim)
if DEBUG:
print("*** Einsum for", idx_str)
print(" idxA =", idxA)
print(" idxB =", idxB)
print(" idxC =", idxC)
# Get the range for each index and put it in a dictionary
rangeA = dict()
rangeB = dict()
#rangeC = dict()
for idx, rnge in zip(idxA, A.shape):
rangeA[idx] = rnge
for idx, rnge in zip(idxB, B.shape):
rangeB[idx] = rnge
#for idx,rnge in zip(idxC,C.shape):
# rangeC[idx] = rnge
if DEBUG:
print("rangeA =", rangeA)
print("rangeB =", rangeB)
# Find the shared indices being summed over
shared_idxAB = list(set(idxA).intersection(idxB))
#if len(shared_idxAB) == 0:
# return np.einsum(idx_str,A,B)
idxAt = list(idxA)
idxBt = list(idxB)
inner_shape = 1
insert_B_loc = 0
for n in shared_idxAB:
if rangeA[n] != rangeB[n]:
err = ('ERROR: In index string %s, the range of index %s is '
'different in A (%d) and B (%d)' %
(idx_str, n, rangeA[n], rangeB[n]))
raise RuntimeError(err)
# Bring idx all the way to the right for A
# and to the left (but preserve order) for B
idxA_n = idxAt.index(n)
idxAt.insert(len(idxAt) - 1, idxAt.pop(idxA_n))
idxB_n = idxBt.index(n)
idxBt.insert(insert_B_loc, idxBt.pop(idxB_n))
insert_B_loc += 1
inner_shape *= rangeA[n]
if DEBUG:
print("shared_idxAB =", shared_idxAB)
print("inner_shape =", inner_shape)
# Transpose the tensors into the proper order and reshape into matrices
new_orderA = [idxA.index(idx) for idx in idxAt]
new_orderB = [idxB.index(idx) for idx in idxBt]
if DEBUG:
print("Transposing A as", new_orderA)
print("Transposing B as", new_orderB)
print("Reshaping A as (-1,", inner_shape, ")")
print("Reshaping B as (", inner_shape, ",-1)")
shapeCt = list()
idxCt = list()
for idx in idxAt:
if idx in shared_idxAB:
break
shapeCt.append(rangeA[idx])
idxCt.append(idx)
for idx in idxBt:
if idx in shared_idxAB:
continue
shapeCt.append(rangeB[idx])
idxCt.append(idx)
new_orderCt = [idxCt.index(idx) for idx in idxC]
if A.size == 0 or B.size == 0:
shapeCt = [shapeCt[i] for i in new_orderCt]
return numpy.zeros(shapeCt, dtype=numpy.result_type(A, B))
At = A.transpose(new_orderA)
Bt = B.transpose(new_orderB)
if At.flags.f_contiguous:
At = numpy.asarray(At.reshape(-1, inner_shape), order='F')
else:
At = numpy.asarray(At.reshape(-1, inner_shape), order='C')
if Bt.flags.f_contiguous:
Bt = numpy.asarray(Bt.reshape(inner_shape, -1), order='F')
else:
Bt = numpy.asarray(Bt.reshape(inner_shape, -1), order='C')
return dot(At, Bt).reshape(shapeCt, order='A').transpose(new_orderCt)
def get_dtypes(*args):
return [canonicalize_dtype(onp.result_type(arg)) for arg in args]
def result_type(*args):
args = [a if is_scalar_for_elemwise(a) else a.dtype for a in args]
return np.result_type(*args)
def dot(self, v, work_array=None, overwrite_v=False, tol=None):
"""Calculates the action of :math:`\\mathrm{e}^{aA}` on a vector :math:`v`.
Examples
--------
.. literalinclude:: ../../doc_examples/expm_multiply_parallel-example.py
:linenos:
:language: python
:lines: 37-
Parameters
-----------
v : contiguous numpy.ndarray, 1d or 2d array
array to apply :math:`\\mathrm{e}^{aA}` on.
work_array : contiguous numpy.ndarray, optional
array can be any shape but must contain 2*v.size contiguous elements.
This array is used as temporary memory space for the underlying c-code. This saves extra memory allocation for function operations.
overwrite_v : bool, optoinal
if set to `True`, the data in `v` is overwritten by the function. This saves extra memory allocation for the results.
tol: float, optoinal
tolerance value used to truncate Taylor expansion of matrix exponential.
Returns
--------
numpy.ndarray
result of :math:`\\mathrm{e}^{aA}v`.
If `overwrite_v = True` the dunction returns `v` with the data overwritten, otherwise the result is stored in a new array.
"""
v = _np.asarray(v)
if v.ndim > 2:
raise ValueError("array must have ndim <= 2.")
if v.shape[0] != self._A.shape[1]:
raise ValueError("dimension mismatch {}, {}".format(
self._A.shape, v.shape))
v_dtype = _np.result_type(self._dtype, v.dtype)
if overwrite_v:
if v_dtype != v.dtype:
raise ValueError(
"if overwrite_v is True, the input array must match correct output dtype for matrix multiplication."
)
if not v.flags["CARRAY"]:
raise TypeError("input array must a contiguous and writable.")
else:
v = v.astype(v_dtype, order="C", copy=True)
if work_array is None:
if v.ndim == 1:
work_array = _np.zeros((2 * self._A.shape[0], ), dtype=v.dtype)
else:
work_array = _np.zeros((2 * self._A.shape[0], v.shape[1]),
dtype=v.dtype)
else:
work_array = work_array.ravel(order="A")
if work_array.size != 2 * v.size:
raise ValueError(
"work_array must have twice the number of elements as in v."
)
if work_array.dtype != v_dtype:
raise ValueError(
"work_array must be array of dtype which matches the result of the matrix-vector multiplication."
)
a = _np.array(self._a, dtype=v_dtype)
mu = _np.array(self._mu, dtype=v_dtype)
if tol is not None:
tol = _np.array(tol, dtype=mu.real.dtype)
else:
tol = _np.array(self._tol, dtype=mu.real.dtype)
if v.ndim == 1:
_wrapper_expm_multiply(self._A.indptr, self._A.indices,
self._A.data, self._s, self._m_star, a, tol,
mu, v, work_array.ravel())
else:
work_array = work_array.reshape((-1, v.shape[1]))
_wrapper_expm_multiply_batch(self._A.indptr, self._A.indices,
self._A.data, self._s, self._m_star,
a, tol, mu, v, work_array)
return v
def zeros_like_array(x): dtype = xla_bridge.canonicalize_dtype(onp.result_type(x)) return onp.broadcast_to(onp.array(0, dtype), onp.shape(x))
def _spectral_helper(x, y, fs=1.0, window='hann', nperseg=None, noverlap=None,
nfft=None, detrend='constant', return_onesided=True,
scaling='spectrum', axis=-1, mode='psd', boundary=None,
padded=False):
""" Calculate various forms of windowed FFTs for PSD, CSD, etc. """
if mode not in ['psd', 'stft']:
raise ValueError("Unknown value for mode %s, must be one of: "
"{'psd', 'stft'}" % mode)
boundary_funcs = {'even': even_ext,
'odd': odd_ext,
'constant': const_ext,
'zeros': zero_ext,
None: None}
if boundary not in boundary_funcs:
raise ValueError("Unknown boundary option '{0}', must be one of: {1}"
.format(boundary, list(boundary_funcs.keys())))
# If x and y are the same object we can save ourselves some computation.
same_data = y is x
if not same_data and mode != 'psd':
raise ValueError("x and y must be equal if mode is 'stft'")
axis = int(axis)
# Ensure we have np.arrays, get outdtype
x = np.asarray(x)
if not same_data:
y = np.asarray(y)
outdtype = np.result_type(x, y, np.complex64)
else:
outdtype = np.result_type(x, np.complex64)
if not same_data:
# Check if we can broadcast the outer axes together
xouter = list(x.shape)
youter = list(y.shape)
xouter.pop(axis)
youter.pop(axis)
try:
outershape = np.broadcast(np.empty(xouter), np.empty(youter)).shape
except ValueError:
raise ValueError('x and y cannot be broadcast together.')
if same_data:
if x.size == 0:
return np.empty(x.shape), np.empty(x.shape), np.empty(x.shape)
else:
if x.size == 0 or y.size == 0:
outshape = outershape + (min([x.shape[axis], y.shape[axis]]),)
emptyout = np.rollaxis(np.empty(outshape), -1, axis)
return emptyout, emptyout, emptyout
if x.ndim > 1:
if axis != -1:
x = np.rollaxis(x, axis, len(x.shape))
if not same_data and y.ndim > 1:
y = np.rollaxis(y, axis, len(y.shape))
# Check if x and y are the same length, zero-pad if necessary
if not same_data:
if x.shape[-1] != y.shape[-1]:
if x.shape[-1] < y.shape[-1]:
pad_shape = list(x.shape)
pad_shape[-1] = y.shape[-1] - x.shape[-1]
x = np.concatenate((x, np.zeros(pad_shape)), -1)
else:
pad_shape = list(y.shape)
pad_shape[-1] = x.shape[-1] - y.shape[-1]
y = np.concatenate((y, np.zeros(pad_shape)), -1)
if nperseg is not None: # if specified by user
nperseg = int(nperseg)
if nperseg < 1:
raise ValueError('nperseg must be a positive integer')
# parse window; if array like, then set nperseg = win.shape
win, nperseg = _triage_segments(window, nperseg,input_length=x.shape[-1])
if nfft is None:
nfft = nperseg
elif nfft < nperseg:
raise ValueError('nfft must be greater than or equal to nperseg.')
else:
nfft = int(nfft)
if noverlap is None:
noverlap = nperseg//2
else:
noverlap = int(noverlap)
if noverlap >= nperseg:
raise ValueError('noverlap must be less than nperseg.')
nstep = nperseg - noverlap
# Padding occurs after boundary extension, so that the extended signal ends
# in zeros, instead of introducing an impulse at the end.
# I.e. if x = [..., 3, 2]
# extend then pad -> [..., 3, 2, 2, 3, 0, 0, 0]
# pad then extend -> [..., 3, 2, 0, 0, 0, 2, 3]
if boundary is not None:
ext_func = boundary_funcs[boundary]
x = ext_func(x, nperseg//2, axis=-1)
if not same_data:
y = ext_func(y, nperseg//2, axis=-1)
if padded:
# Pad to integer number of windowed segments
# I.e make x.shape[-1] = nperseg + (nseg-1)*nstep, with integer nseg
nadd = (-(x.shape[-1]-nperseg) % nstep) % nperseg
zeros_shape = list(x.shape[:-1]) + [nadd]
x = np.concatenate((x, np.zeros(zeros_shape)), axis=-1)
if not same_data:
zeros_shape = list(y.shape[:-1]) + [nadd]
y = np.concatenate((y, np.zeros(zeros_shape)), axis=-1)
# Handle detrending and window functions
if not detrend:
def detrend_func(d):
return d
elif not hasattr(detrend, '__call__'):
def detrend_func(d):
return signaltools_detrend(d, type=detrend, axis=-1)
elif axis != -1:
# Wrap this function so that it receives a shape that it could
# reasonably expect to receive.
def detrend_func(d):
d = np.rollaxis(d, -1, axis)
d = detrend(d)
return np.rollaxis(d, axis, len(d.shape))
else:
detrend_func = detrend
if np.result_type(win,np.complex64) != outdtype:
win = win.astype(outdtype)
if scaling == 'density':
scale = 1.0 / (fs * (win*win).sum())
elif scaling == 'spectrum':
scale = 1.0 / win.sum()**2
else:
raise ValueError('Unknown scaling: %r' % scaling)
if mode == 'stft':
scale = np.sqrt(scale)
if return_onesided:
if np.iscomplexobj(x):
sides = 'twosided'
#warnings.warn('Input data is complex, switching to ' 'return_onesided=False')
else:
sides = 'onesided'
if not same_data:
if np.iscomplexobj(y):
sides = 'twosided'
#warnings.warn('Input data is complex, switching to return_onesided=False')
else:
sides = 'twosided'
if sides == 'twosided':
raise Exception('NOT IMPLEMENTED')
#freqs = fftpack.fftfreq(nfft, 1/fs)
elif sides == 'onesided':
freqs = np.fft.rfftfreq(nfft, 1/fs)
# Perform the windowed FFTs
result = _fft_helper(x, win, detrend_func, nperseg, noverlap, nfft, sides)
if not same_data:
# All the same operations on the y data
result_y = _fft_helper(y, win, detrend_func, nperseg, noverlap, nfft,
sides)
result = np.conjugate(result) * result_y
elif mode == 'psd':
result = np.conjugate(result) * result
result *= scale
if sides == 'onesided' and mode == 'psd':
if nfft % 2:
result[..., 1:] *= 2
else:
# Last point is unpaired Nyquist freq point, don't double
result[..., 1:-1] *= 2
time = np.arange(nperseg/2, x.shape[-1] - nperseg/2 + 1,
nperseg - noverlap)/float(fs)
if boundary is not None:
time -= (nperseg/2) / fs
result = result.astype(outdtype)
# All imaginary parts are zero anyways
if same_data and mode != 'stft':
result = result.real
# Output is going to have new last axis for time/window index, so a
# negative axis index shifts down one
if axis < 0:
axis -= 1
# Roll frequency axis back to axis where the data came from
result = np.rollaxis(result, -1, axis)
# TODO
class InfoClass():
pass
Info = InfoClass();
Info.df=freqs[1]-freqs[0]
Info.fMax=freqs[-1]
Info.LFreq=len(freqs)
Info.LSeg=nperseg
Info.LWin=len(win)
Info.LOvlp=noverlap
Info.nFFT=nfft
Info.nseg=-1
#print('df:{:.3f} - fm:{:.2f} - nseg:{} - Lf:{:5d} - Lseg:{:5d} - Lwin:{:5d} - Lovlp:{:5d} - Nfft:{:5d} - Lsig:{}'.format(freqs[1]-freqs[0],freqs[-1],-1,len(freqs),nperseg,len(win),noverlap,nfft,x.shape[-1]))
return freqs, time, result, Info
