def pow(dtype, power_dtype=None):
"""
Returns a :py:class:`~reikna.cluda.Module` with a function of two arguments
that raises the first argument of type ``dtype`` (must be a real or complex data type)
to the power of the second argument (a corresponding real data type or an integer).
"""
if dtypes.is_complex(power_dtype):
raise NotImplementedError("pow() with a complex power is not supported")
if power_dtype is None:
if dtypes.is_integer(dtype):
raise ValueError("Power dtype must be specified for an integer argument")
elif dtypes.is_real(dtype):
power_dtype = dtype
else:
power_dtype = dtypes.real_for(dtype)
if dtypes.is_complex(dtype):
r_dtype = dtypes.real_for(dtype)
elif dtypes.is_real(dtype):
r_dtype = dtype
elif dtypes.is_real(power_dtype):
r_dtype = power_dtype
else:
r_dtype = numpy.float32
if dtypes.is_integer(dtype) and dtypes.is_real(power_dtype):
dtype = power_dtype
return Module(
TEMPLATE.get_def('pow'),
render_kwds=dict(
dtype=dtype, power_dtype=power_dtype,
mul_=mul(dtype, dtype), div_=div(dtype, dtype),
polar_=polar(r_dtype)))
def __init__(self,
x,
NFFT=256,
noverlap=128,
pad_to=None,
window=hanning_window):
# print("x Data type = %s" % x.dtype)
# print("Is Real = %s" % dtypes.is_real(x.dtype))
# print("dim = %s" % x.ndim)
assert dtypes.is_real(x.dtype)
assert x.ndim == 1
rolling_frame_trf = rolling_frame(x, NFFT, noverlap, pad_to)
complex_dtype = dtypes.complex_for(x.dtype)
fft_arr = Type(complex_dtype, rolling_frame_trf.output.shape)
real_fft_arr = Type(x.dtype, rolling_frame_trf.output.shape)
window_trf = window(real_fft_arr, NFFT)
broadcast_zero_trf = transformations.broadcast_const(real_fft_arr, 0)
to_complex_trf = transformations.combine_complex(fft_arr)
amplitude_trf = transformations.norm_const(fft_arr, 1)
crop_trf = crop_frequencies(amplitude_trf.output)
fft = FFT(fft_arr, axes=(1, ))
fft.parameter.input.connect(to_complex_trf,
to_complex_trf.output,
input_real=to_complex_trf.real,
input_imag=to_complex_trf.imag)
fft.parameter.input_imag.connect(broadcast_zero_trf,
broadcast_zero_trf.output)
fft.parameter.input_real.connect(window_trf,
window_trf.output,
unwindowed_input=window_trf.input)
fft.parameter.unwindowed_input.connect(
rolling_frame_trf,
rolling_frame_trf.output,
flat_input=rolling_frame_trf.input)
fft.parameter.output.connect(amplitude_trf,
amplitude_trf.input,
amplitude=amplitude_trf.output)
fft.parameter.amplitude.connect(crop_trf,
crop_trf.input,
cropped_amplitude=crop_trf.output)
self._fft = fft
self._transpose = Transpose(fft.parameter.cropped_amplitude)
Computation.__init__(self, [
Parameter('output',
Annotation(self._transpose.parameter.output, 'o')),
Parameter('input', Annotation(fft.parameter.flat_input, 'i'))
])
def test_pow(thr, out_code, in_codes):
out_dtype, in_dtypes = generate_dtypes(out_code, in_codes)
if len(in_dtypes) == 1:
func = functions.pow(in_dtypes[0])
if dtypes.is_real(in_dtypes[0]):
in_dtypes.append(in_dtypes[0])
else:
in_dtypes.append(dtypes.real_for(in_dtypes[0]))
else:
func = functions.pow(in_dtypes[0], power_dtype=in_dtypes[1])
check_func(thr, func, numpy.power, out_dtype, in_dtypes)
def test_pow(thr, out_code, in_codes):
out_dtype, in_dtypes = generate_dtypes(out_code, in_codes)
if len(in_dtypes) == 1:
func = functions.pow(in_dtypes[0])
if dtypes.is_real(in_dtypes[0]):
in_dtypes.append(in_dtypes[0])
else:
in_dtypes.append(dtypes.real_for(in_dtypes[0]))
else:
func = functions.pow(in_dtypes[0], power_dtype=in_dtypes[1])
check_func(thr, func, numpy.power, out_dtype, in_dtypes)
def pow(dtype, exponent_dtype=None, output_dtype=None):
"""
Returns a :py:class:`~reikna.cluda.Module` with a function of two arguments
that raises the first argument of type ``dtype``
to the power of the second argument of type ``exponent_dtype``
(an integer or real data type).
If ``exponent_dtype`` or ``output_dtype`` are not given, they default to ``dtype``.
If ``dtype`` is not the same as ``output_dtype``,
the input is cast to ``output_dtype`` *before* exponentiation.
If ``exponent_dtype`` is real, but both ``dtype`` and ``output_dtype`` are integer,
a ``ValueError`` is raised.
"""
if exponent_dtype is None:
exponent_dtype = dtype
if output_dtype is None:
output_dtype = dtype
if dtypes.is_complex(exponent_dtype):
raise NotImplementedError("pow() with a complex exponent is not supported")
if dtypes.is_real(exponent_dtype):
if dtypes.is_complex(output_dtype):
exponent_dtype = dtypes.real_for(output_dtype)
elif dtypes.is_real(output_dtype):
exponent_dtype = output_dtype
else:
raise ValueError("pow(integer, float): integer is not supported")
kwds = dict(
dtype=dtype, exponent_dtype=exponent_dtype, output_dtype=output_dtype,
div_=None, mul_=None, cast_=None, polar_=None)
if output_dtype != dtype:
kwds['cast_'] = cast(output_dtype, dtype)
if dtypes.is_integer(exponent_dtype) and not dtypes.is_real(output_dtype):
kwds['mul_'] = mul(output_dtype, output_dtype)
kwds['div_'] = div(output_dtype, output_dtype)
if dtypes.is_complex(output_dtype):
kwds['polar_'] = polar(dtypes.real_for(output_dtype))
return Module(TEMPLATE.get_def('pow'), render_kwds=kwds)
def polar(dtype):
"""
Returns a :py:class:`~reikna.cluda.Module` with a function of two arguments
that returns the complex-valued ``rho * exp(i * theta)``
for values ``rho, theta`` of type ``dtype`` (must be a real data type).
"""
if not dtypes.is_real(dtype):
raise NotImplementedError("polar() of " + str(dtype) + " is not supported")
return Module(
TEMPLATE.get_def('polar'),
render_kwds=dict(dtype=dtype, polar_unit_=polar_unit(dtype)))
def polar_unit(dtype):
"""
Returns a :py:class:`~reikna.cluda.Module` with a function of one argument
that returns a complex number ``(cos(theta), sin(theta))``
for a value ``theta`` of type ``dtype`` (must be a real data type).
"""
if not dtypes.is_real(dtype):
raise NotImplementedError("polar_unit() of " + str(dtype) + " is not supported")
return Module(
TEMPLATE.get_def('polar_unit'),
render_kwds=dict(dtype=dtype))
def generate_dtypes(out_code, in_codes):
test_dtype = lambda idx: dict(i=numpy.int32, f=numpy.float32, c=numpy.complex64)[idx]
in_dtypes = list(map(test_dtype, in_codes))
out_dtype = dtypes.result_type(*in_dtypes) if out_code == 'auto' else test_dtype(out_code)
if not any(map(dtypes.is_double, in_dtypes)):
# numpy thinks that int32 * float32 == float64,
# but we still need to run this test on older videocards
if dtypes.is_complex(out_dtype):
out_dtype = numpy.complex64
elif dtypes.is_real(out_dtype):
out_dtype = numpy.float32
return out_dtype, in_dtypes
def exp(dtype):
"""
Returns a :py:class:`~reikna.cluda.Module` with a function of one argument
that exponentiates the value of type ``dtype``
(must be a real or complex data type).
"""
if dtypes.is_integer(dtype):
raise NotImplementedError("exp() of " + str(dtype) + " is not supported")
if dtypes.is_real(dtype):
polar_unit_ = None
else:
polar_unit_ = polar_unit(dtypes.real_for(dtype))
return Module(
TEMPLATE.get_def('exp'),
render_kwds=dict(dtype=dtype, polar_unit_=polar_unit_))
def generate_dtypes(out_code, in_codes):
test_dtype = lambda idx: dict(
i=numpy.int32, f=numpy.float32, c=numpy.complex64)[idx]
in_dtypes = list(map(test_dtype, in_codes))
out_dtype = dtypes.result_type(
*in_dtypes) if out_code == 'auto' else test_dtype(out_code)
if not any(map(dtypes.is_double, in_dtypes)):
# numpy thinks that int32 * float32 == float64,
# but we still need to run this test on older videocards
if dtypes.is_complex(out_dtype):
out_dtype = numpy.complex64
elif dtypes.is_real(out_dtype):
out_dtype = numpy.float32
return out_dtype, in_dtypes
def __init__(self, shape, box, drift, trajectories=1, diffusion=None):
if diffusion is not None:
assert diffusion.dtype == drift.dtype
assert diffusion.components == drift.components
if not diffusion.real_noise or dtypes.is_real(drift.dtype):
noise_dtype = drift.dtype
else:
noise_dtype = dtypes.real_for(drift.dtype)
self.noise_type = Type(noise_dtype, (trajectories, diffusion.noise_sources) + shape)
self.noise = True
cell_volume = product(box) / product(shape)
self._noise_normalization = 1. / cell_volume
else:
self.noise_type = None
self.noise = False
def __init__(self, x, NFFT=256, noverlap=128, pad_to=None, window=hanning_window):
assert dtypes.is_real(x.dtype)
assert x.ndim == 1
rolling_frame_trf = rolling_frame(x, NFFT, noverlap, pad_to)
complex_dtype = dtypes.complex_for(x.dtype)
fft_arr = Type(complex_dtype, rolling_frame_trf.output.shape)
real_fft_arr = Type(x.dtype, rolling_frame_trf.output.shape)
window_trf = window(real_fft_arr, NFFT)
broadcast_zero_trf = transformations.broadcast_const(real_fft_arr, 0)
to_complex_trf = transformations.combine_complex(fft_arr)
amplitude_trf = transformations.norm_const(fft_arr, 1)
crop_trf = crop_frequencies(amplitude_trf.output)
fft = FFT(fft_arr, axes=(1,))
fft.parameter.input.connect(
to_complex_trf, to_complex_trf.output,
input_real=to_complex_trf.real, input_imag=to_complex_trf.imag)
fft.parameter.input_imag.connect(
broadcast_zero_trf, broadcast_zero_trf.output)
fft.parameter.input_real.connect(
window_trf, window_trf.output, unwindowed_input=window_trf.input)
fft.parameter.unwindowed_input.connect(
rolling_frame_trf, rolling_frame_trf.output, flat_input=rolling_frame_trf.input)
fft.parameter.output.connect(
amplitude_trf, amplitude_trf.input, amplitude=amplitude_trf.output)
fft.parameter.amplitude.connect(
crop_trf, crop_trf.input, cropped_amplitude=crop_trf.output)
self._fft = fft
self._transpose = Transpose(fft.parameter.cropped_amplitude)
Computation.__init__(self,
[Parameter('output', Annotation(self._transpose.parameter.output, 'o')),
Parameter('input', Annotation(fft.parameter.flat_input, 'i'))])