def test_rw_split():
# I guess there's no good way to test both fork and spawn modes, since split_at()
# is only aware of one type of SharedmemManager at runtime.
array = np.arange(20 * 30).reshape(20, 30)
sh_array = parallel_context.shared_copy(array)
out = np.zeros(20)
sh_out = parallel_context.shared_copy(out)
split_fn = split_at(split_arg=(0, 1), n_jobs=3)(sum_columns_and_write)
# this call should not touch out
split_fn(array, out)
assert np.all(out == 0), 'non-shared output modified?'
# this call should not touch out
split_fn(sh_array, out)
assert np.all(out == 0), 'non-shared output modified?'
# this call should write to shm
split_fn(array, sh_out)
assert np.all(
sh_out == sum_columns(array)), 'shared mem did not get written?'
sh_out.fill(0)
# this call should write to shm
split_fn(sh_array, sh_out)
assert np.all(
sh_out == sum_columns(array)), 'shared mem did not get written?'
def test_kwarg_split():
x = parallel_context.shared_ndarray((50, 100))
x[:] = 1
split_affine = split_at(split_arg=(0, ),
splice_at=(0, 1),
split_kwargs='offset',
n_jobs=3)(affine_computation)
assert (split_affine(x) == 2).all()
y, offset = split_affine(x, offset=np.ones(50))
assert (y == 3).all()
assert (offset == x.mean(axis=1)).all()
def test_ro_split():
# I guess there's no good way to test both fork and spawn modes, since split_at()
# is only aware of one type of SharedmemManager at runtime.
array = np.arange(20 * 30).reshape(20, 30)
sh_array = parallel_context.shared_copy(array)
split_fn = split_at(n_jobs=3)(sum_columns)
# this method should work equally for read-only and shared access
res_1 = split_fn(array)
res_2 = split_fn(sh_array)
res_3 = sum_columns(array)
assert np.all(res_1 == res_3), 'read-only does not match ref'
assert np.all(res_2 == res_3), 'shm does not match ref'
assert np.all(res_1 == res_2), 'para results do not match'
def test_complex_input_output_pattern():
shared_arg = np.random.randn(100)
split_arg = parallel_context.shared_ndarray((50, 20))
split_arg[:] = np.arange(20 * 50).reshape(50, 20)
test_arg = split_arg.copy()
split_fn = split_at(split_arg=(1, ),
splice_at=(1, ),
shared_args=(0, ),
concurrent=True,
n_jobs=3)(very_weird_fun)
size, sum = split_fn(shared_arg, split_arg)
assert size == shared_arg.size, 'single return value wrong'
assert np.all(sum == split_arg.sum(axis=1))
assert np.all(split_arg == test_arg + shared_arg.var())
def __call__(self, array):
if self.f_lo > 0 or self.f_hi > 0:
block_filter = 'parallel' if self.para else 'serial'
fdes = dict(lo=self.f_lo, hi=self.f_hi, Fs=self.sample_rate, ord=3)
farg = dict(filtfilt=self.filtfilt)
y = filter_array(array, inplace=False, block_filter=block_filter, design_kwargs=fdes, filt_kwargs=farg)
tau_samps = self.tau * self.sample_rate
if tau_samps == 0:
return y
if self.square:
y **= 2
else:
np.abs(y, out=y)
if tau_samps > 0:
if self.para:
gauss_para = split_at()(gaussian_filter1d)
y = gauss_para(y, tau_samps, axis=1)
else:
y = gaussian_filter1d(y, tau_samps, axis=-1)
if self.square:
np.sqrt(y, out=y)
return y
y_re = y_flat[..., 0::2]
y_im = y_flat[..., 1::2]
else:
y_re = np.zeros_like(x)
# make dummy array to satisfy Cython signature
y_im = np.empty((x.shape[0], 1), 'd')
for i in range(x.shape[0]):
bandpass_moving_projection(x[i].astype('d'),
dpss,
wf,
wt,
y_re[i],
y_im[i],
f0,
baseband=baseband)
if not baseband:
y = y_re
else:
y = np.zeros(x.shape, 'd')
for i in range(x.shape[0]):
lowpass_moving_projection(x[i].astype('d'), dpss, wf, wt, y[i])
if save_dpss:
return y, (dpss, eigs)
return y
moving_projection = split_at()(moving_projection_serial)
except ImportError:
raise
# moving_projection = split_at()(_moving_projection_preserve)
# def bfilter(b, a, x, out=None, **kwargs):
# if out is None:
# out = x
# return bfilter_para(b, a, x, out, **kwargs)
# this is the wrong wrapper for bfilter -- the default behavior is inplace filtering
bfilter = split_optional_output('out',
runs_inplace=True,
same_shape=True,
split_arg=(2, ))(bfilter_serial)
# WIP
# if hasattr(bfilter_para, 'uses_parallel'):
# bfilter.uses_parallel = bfilter_para.uses_parallel
overlap_add = split_at()(overlap_add_serial)
# multi taper spectral estimation
multi_taper_psd = split_at(splice_at=(1, 2))(multi_taper_psd_serial)
# convolution
convolve1d = split_at(split_arg=0)(convolve1d_serial)
# # Linear filtering wrapper -- converts lfilter to a "void" method rather than returning an array.
# # @split_at(split_arg=(2, 3, 4), splice_at=(0,))
# def lfilter_void(b, a, x, y, zi, **kwargs):
# kwargs['zi'] = zi
# y[:], zi = lfilter_serial(b, a, x, **kwargs)
# return zi
#
# lfilter_para = split_at(split_arg=(2, 3, 4), splice_at=(0,))(lfilter_void)
if envelope:
epochs = signal.hilbert(
epochs, N=ut.nextpow2(epoch_len), axis=-1
)
epochs = np.abs(epochs[..., :epoch_len])**2
avg[:, c - 1, :] = np.sum(epochs, axis=1) / n_avg[:, c - 1][:, None]
x.shape = [x for x in x.shape if x > 1]
if envelope:
np.sqrt(avg, avg)
return avg, n_avg
# create a parallelized version (but only split big jobs)
ep_trigger_avg = array_split.split_at(splice_at=(0, 1), split_over=100)(_ep_trigger_avg)
def iter_epochs(x, pivots, selected=(), pre=0, post=0, fill=np.nan):
"""
Generator that yields epochs pivoted at the specified triggers.
Parameters
----------
x : data (n_chan, n_pt)
pivots : array-like or StimulatedExperiment
A sequence of literal pivot samples, or an experiment wrapper
containing the timestamps.
selected : sequence
Indices into trig_code for a subset of stims. If empty, return *ALL*
epochs (*a potentially very large array*)