def gamma(bijection, dtype, shape=1, scale=1):
"""
Generates random numbers from the gamma distribution
.. math::
P(x) = x^{k-1} \\frac{e^{-x/\\theta}}{\\theta^k \\Gamma(k)},
where :math:`k` is ``shape``, and :math:`\\theta` is ``scale``.
Supported dtypes: ``float(32/64)``.
Returns a :py:class:`~reikna.cbrng.samplers.Sampler` object.
"""
ctype = dtypes.ctype(dtype)
uf = uniform_float(bijection, dtype, low=0, high=1)
nbm = normal_bm(bijection, dtype, mean=0, std=1)
module = Module(TEMPLATE.get_def("gamma"),
render_kwds=dict(dtype=dtype,
ctype=ctype,
bijection=bijection,
shape=shape,
scale=dtypes.c_constant(scale, dtype),
uf=uf,
nbm=nbm))
return Sampler(bijection, module, dtype)
def mod(ff_elem=None, method="c"):
if ff_elem is None:
ff_elem = get_ff_elem()
module = Module(
TEMPLATE.get_def('mod_def'),
render_kwds=dict(method=method, ff=ff_elem, ff_elem=ff_elem.module))
return FiniteFieldMod(ff_elem, module)
def get_drift(state_dtype, U, gamma, dx, wigner=False):
return Drift(
Module.create(
"""
<%
r_dtype = dtypes.real_for(s_dtype)
s_ctype = dtypes.ctype(s_dtype)
r_ctype = dtypes.ctype(r_dtype)
%>
INLINE WITHIN_KERNEL ${s_ctype} ${prefix}0(
const int idx_x,
const ${s_ctype} psi,
${r_ctype} t)
{
return ${mul_cc}(
COMPLEX_CTR(${s_ctype})(
-${gamma},
-(${U} * (${norm}(psi) - ${correction}))),
psi
);
}
""",
render_kwds=dict(
s_dtype=state_dtype,
U=U,
gamma=gamma,
mul_cc=functions.mul(state_dtype, state_dtype),
norm=functions.norm(state_dtype),
correction=1. / dx if wigner else 0
)),
state_dtype, components=1)
def ntt1024(
base_method='c', mul_method='c_from_asm', lsh_method='c', ff_elem=None,
use_constant_memory=False):
if ff_elem is None:
ff_elem = arithmetic.get_ff_elem()
base_kwds = dict(ff_elem=ff_elem, method=base_method)
mul_kwds = dict(ff_elem=ff_elem, method=mul_method, nested_method=base_method)
lsh_kwds = dict(ff_elem=ff_elem, method=lsh_method, nested_method=base_method)
module = Module(
TEMPLATE.get_def('ntt1024'),
render_kwds=dict(
ff=ff_elem,
ff_elem=ff_elem.module,
lsh32=arithmetic.lsh(32, numpy.uint32, **lsh_kwds).module,
lsh64=arithmetic.lsh(64, numpy.uint32, **lsh_kwds).module,
lsh96=arithmetic.lsh(96, numpy.uint32, **lsh_kwds).module,
lsh128=arithmetic.lsh(128, numpy.uint32, **lsh_kwds).module,
lsh160=arithmetic.lsh(160, numpy.uint32, **lsh_kwds).module,
lsh192=arithmetic.lsh(192, numpy.uint32, **lsh_kwds).module,
add=arithmetic.add(**base_kwds).module,
sub=arithmetic.sub(**base_kwds).module,
mul=arithmetic.mul(**mul_kwds).module,
use_constant_memory=use_constant_memory,
))
return NTT1024(ff_elem, module, use_constant_memory)
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 get_nonlinear_wrapper(components, c_dtype, nonlinear_module, dt):
s_dtype = dtypes.real_for(c_dtype)
return Module.create(
"""
%for comp in range(components):
INLINE WITHIN_KERNEL ${c_ctype} ${prefix}${comp}(
%for pcomp in range(components):
${c_ctype} psi${pcomp},
%endfor
${s_ctype} V, ${s_ctype} t)
{
${c_ctype} nonlinear = ${nonlinear}${comp}(
%for pcomp in range(components):
psi${pcomp},
%endfor
V, t);
return ${mul}(
COMPLEX_CTR(${c_ctype})(0, -${dt}),
nonlinear);
}
%endfor
""",
render_kwds=dict(
components=components,
c_ctype=dtypes.ctype(c_dtype),
s_ctype=dtypes.ctype(s_dtype),
mul=functions.mul(c_dtype, c_dtype),
dt=dtypes.c_constant(dt, s_dtype),
nonlinear=nonlinear_module))
def vonmises(bijection, dtype, mu=0, kappa=1):
"""
Generates random numbers from the von Mises distribution
.. math::
P(x) = \\frac{\\exp(\\kappa \\cos(x - \\mu))}{2 \\pi I_0(\\kappa)},
where :math:`\\mu` is the mode, :math:`\\kappa` is the dispersion,
and :math:`I_0` is the modified Bessel function of the first kind.
Supported dtypes: ``float(32/64)``.
Returns a :py:class:`~reikna.cbrng.samplers.Sampler` object.
"""
ctype = dtypes.ctype(dtype)
uf = uniform_float(bijection, dtype, low=0, high=1)
module = Module(TEMPLATE.get_def("vonmises"),
render_kwds=dict(dtype=dtype,
ctype=ctype,
bijection=bijection,
mu=dtypes.c_constant(mu, dtype),
kappa=kappa,
uf=uf))
return Sampler(bijection, module, dtype)
def uniform_float(bijection, dtype, low=0, high=1):
"""
Generates uniformly distributed floating-points numbers in the interval ``[low, high)``.
Supported dtypes: ``float(32/64)``.
A fixed number of counters is used in each thread.
Returns a :py:class:`~reikna.cbrng.samplers.Sampler` object.
"""
assert low < high
ctype = dtypes.ctype(dtype)
bitness = 64 if dtypes.is_double(dtype) else 32
raw_func = 'get_raw_uint' + str(bitness)
raw_max = dtypes.c_constant(2**bitness, dtype)
size = dtypes.c_constant(high - low, dtype)
low = dtypes.c_constant(low, dtype)
module = Module(TEMPLATE.get_def("uniform_float"),
render_kwds=dict(bijection=bijection,
ctype=ctype,
raw_func=raw_func,
raw_max=raw_max,
size=size,
low=low))
return Sampler(bijection, module, dtype, deterministic=True)
def nonlinear_no_potential(dtype, U, nu):
c_dtype = dtype
c_ctype = dtypes.ctype(c_dtype)
s_dtype = dtypes.real_for(dtype)
s_ctype = dtypes.ctype(s_dtype)
return Module.create(
"""
%for comp in (0, 1):
INLINE WITHIN_KERNEL ${c_ctype} ${prefix}${comp}(
${c_ctype} psi0, ${c_ctype} psi1, ${s_ctype} t)
{
return (
${mul}(psi${comp}, (
${dtypes.c_constant(U[comp, 0])} * ${norm}(psi0) +
${dtypes.c_constant(U[comp, 1])} * ${norm}(psi1)
))
- ${mul}(psi${1 - comp}, ${nu})
);
}
%endfor
""",
render_kwds=dict(
mul=functions.mul(c_dtype, s_dtype),
norm=functions.norm(c_dtype),
U=U,
nu=dtypes.c_constant(nu, s_dtype),
c_ctype=c_ctype,
s_ctype=s_ctype))
def module_leaf_macro(output, param):
return Module(_module_leaf_macro,
render_kwds=dict(output=output,
name=param.name,
VALUE_NAME=VALUE_NAME,
lname=leaf_name(param.name),
index_seq=index_cnames_seq(param),
index_expr=flat_index_expr(param)))
def pow(exp_dtype, ff_elem=None):
if ff_elem is None:
ff_elem = get_ff_elem()
module = Module(
TEMPLATE.get_def('pow_def'),
render_kwds=dict(
ff=ff_elem, ff_elem=ff_elem.module, mul=mul(ff_elem).module, exp_dtype=exp_dtype))
return FiniteFieldPow(ff_elem, module, exp_dtype)
def cast(out_dtype, in_dtype):
"""
Returns a :py:class:`~reikna.cluda.Module` with a function of one argument
that casts values of ``in_dtype`` to ``out_dtype``.
"""
return Module(
TEMPLATE.get_def('cast'),
render_kwds=dict(out_dtype=out_dtype, in_dtype=in_dtype))
def prepare_for_mul(ff_elem=None):
if ff_elem is None:
ff_elem = get_ff_elem()
module = Module(
TEMPLATE.get_def('prepare_for_mul_def'),
render_kwds=dict(
ff=ff_elem, ff_elem=ff_elem.module, sub=sub(ff_elem=ff_elem).module,
))
return FiniteFieldPrepareForMul(ff_elem, module)
def module_same_indices(output, param, subtree_parameters, module_idx):
return Module(_module_same_indices,
render_kwds=dict(
output=output,
name=param.name,
VALUE_NAME=VALUE_NAME,
module_idx=module_idx,
nq_indices=index_cnames_seq(param),
nq_params=param_cnames_seq(subtree_parameters)))
def norm(dtype):
"""
Returns a :py:class:`~reikna.cluda.Module` with a function of one argument
that returns the 2-norm of the value of type ``dtype``
(product by the complex conjugate if the value is complex, square otherwise).
"""
return Module(
TEMPLATE.get_def('norm'),
render_kwds=dict(dtype=dtype))
def inv_pow2(exp_dtype, ff_elem=None):
# this can be relaxed when we witch to regular C
assert numpy.dtype(exp_dtype) == numpy.dtype('uint32')
if ff_elem is None:
ff_elem = get_ff_elem()
module = Module(
TEMPLATE.get_def('inv_pow2_def'),
render_kwds=dict(ff=ff_elem, ff_elem=ff_elem.module, exp_dtype=exp_dtype))
return FiniteFieldPow(ff_elem, module, exp_dtype)
def div(in_dtype1, in_dtype2, out_dtype=None):
"""
Returns a :py:class:`~reikna.cluda.Module` with a function of two arguments
that divides values of ``in_dtype1`` and ``in_dtype2``.
If ``out_dtype`` is given, it will be set as a return type for this function.
"""
out_dtype = derive_out_dtype(out_dtype, in_dtype1, in_dtype2)
return Module(
TEMPLATE.get_def('div'),
render_kwds=dict(out_dtype=out_dtype, in_dtype1=in_dtype1, in_dtype2=in_dtype2))
def get_nonlinear_wrapper(state_dtype, grid_dims, drift, diffusion=None):
real_dtype = dtypes.real_for(state_dtype)
if diffusion is not None:
noise_dtype = diffusion.dtype
else:
noise_dtype = real_dtype
return Module.create(
"""
<%
components = drift.components
idx_args = ["idx_" + str(dim) for dim in range(grid_dims)]
psi_args = ["psi_" + str(comp) for comp in range(components)]
if diffusion is not None:
dW_args = ["dW_" + str(ncomp) for ncomp in range(diffusion.noise_sources)]
%>
%for comp in range(components):
INLINE WITHIN_KERNEL ${s_ctype} ${prefix}${comp}(
%for idx in idx_args:
const int ${idx},
%endfor
%for psi in psi_args:
const ${s_ctype} ${psi},
%endfor
%if diffusion is not None:
%for dW in dW_args:
const ${n_ctype} ${dW},
%endfor
%endif
const ${r_ctype} t,
const ${r_ctype} dt)
{
return
${mul_sr}(${drift.module}${comp}(
${", ".join(idx_args)}, ${", ".join(psi_args)}, t), dt)
%if diffusion is not None:
%for ncomp in range(diffusion.noise_sources):
+ ${mul_sn}(${diffusion.module}${comp}_${ncomp}(
${", ".join(idx_args)}, ${", ".join(psi_args)}, t), ${dW_args[ncomp]})
%endfor
%endif
;
}
%endfor
""",
render_kwds=dict(
grid_dims=grid_dims,
s_ctype=dtypes.ctype(state_dtype),
r_ctype=dtypes.ctype(real_dtype),
n_ctype=dtypes.ctype(noise_dtype),
mul_sr=functions.mul(state_dtype, real_dtype),
mul_sn=functions.mul(state_dtype, noise_dtype),
drift=drift,
diffusion=diffusion))
def fft512(use_constant_memory=False):
module = Module(TEMPLATE.get_def('fft512'),
render_kwds=dict(
elem_ctype=dtypes.ctype(numpy.complex128),
temp_ctype=dtypes.ctype(numpy.float64),
cdata_ctype=dtypes.ctype(numpy.complex128),
polar_unit=functions.polar_unit(numpy.float64),
mul=functions.mul(numpy.complex128, numpy.complex128),
use_constant_memory=use_constant_memory,
))
return FFT512(module, use_constant_memory)
def conj(dtype):
"""
Returns a :py:class:`~reikna.cluda.Module` with a function of one argument
that conjugates the value of type ``dtype`` (must be a complex data type).
"""
if not dtypes.is_complex(dtype):
raise NotImplementedError("conj() of " + str(dtype) + " is not supported")
return Module(
TEMPLATE.get_def('conj'),
render_kwds=dict(dtype=dtype))
def mul_prepared(ff_elem=None, method="c", nested_method="c"):
if ff_elem is None:
ff_elem = get_ff_elem()
module = Module(
TEMPLATE.get_def('mul_prepared_def'),
render_kwds=dict(
method=method,
ff=ff_elem, ff_elem=ff_elem.module,
sub=sub(ff_elem=ff_elem, method=nested_method).module,
))
return FiniteFieldMulPrepared(ff_elem, module)
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 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 mul(*in_dtypes, **kwds):
"""mul(*in_dtypes, out_dtype=None)
Returns a :py:class:`~reikna.cluda.Module` with a function of
``len(in_dtypes)`` arguments that multiplies values of types ``in_dtypes``.
If ``out_dtype`` is given, it will be set as a return type for this function.
"""
assert set(kwds.keys()).issubset(['out_dtype'])
out_dtype = derive_out_dtype(kwds.get('out_dtype', None), *in_dtypes)
return Module(
TEMPLATE.get_def('add_or_mul'),
render_kwds=dict(op='mul', out_dtype=out_dtype, in_dtypes=in_dtypes))
def get_nonlinear(dtype, interaction, tunneling):
r"""
Nonlinear module
.. math::
N(\psi_1, ... \psi_C)
= \sum_{n=1}^{C} U_{jn} |\psi_n|^2 \psi_j
- \nu_j psi_{m_j}
``interaction``: a symmetrical ``components x components`` array with interaction strengths.
``tunneling``: a list of (other_comp, coeff) pairs of tunnelling strengths.
"""
c_dtype = dtype
c_ctype = dtypes.ctype(c_dtype)
s_dtype = dtypes.real_for(dtype)
s_ctype = dtypes.ctype(s_dtype)
return Module.create(
"""
%for comp in range(components):
INLINE WITHIN_KERNEL ${c_ctype} ${prefix}${comp}(
%for pcomp in range(components):
${c_ctype} psi${pcomp},
%endfor
${s_ctype} V, ${s_ctype} t)
{
return (
${mul}(psi${comp}, (
%for other_comp in range(components):
+ ${dtypes.c_constant(interaction[comp, other_comp], s_dtype)} *
${norm}(psi${other_comp})
%endfor
+ V
))
- ${mul}(
psi${tunneling[comp][0]},
${dtypes.c_constant(tunneling[comp][1], s_dtype)})
);
}
%endfor
""",
render_kwds=dict(
components=interaction.shape[0],
mul=functions.mul(c_dtype, s_dtype),
norm=functions.norm(c_dtype),
interaction=interaction,
tunneling=tunneling,
s_dtype=s_dtype,
c_ctype=c_ctype,
s_ctype=s_ctype))
def philox(bitness, counter_words, rounds=10):
"""
A CBRNG based on a low number of slow rounds (multiplications).
:param bitness: ``32`` or ``64``, corresponds to the size of generated random integers.
:param counter_words: ``2`` or ``4``, number of integers generated in one go.
:param rounds: ``1`` to ``12``, the more rounds, the better randomness is achieved.
Default values are big enough to qualify as PRNG.
:returns: a :py:class:`Bijection` object.
"""
W_CONSTANTS = {
64: [
numpy.uint64(0x9E3779B97F4A7C15), # golden ratio
numpy.uint64(0xBB67AE8584CAA73B) # sqrt(3)-1
],
32: [
numpy.uint32(0x9E3779B9), # golden ratio
numpy.uint32(0xBB67AE85) # sqrt(3)-1
]
}
M_CONSTANTS = {
(64, 2): [numpy.uint64(0xD2B74407B1CE6E93)],
(64, 4):
[numpy.uint64(0xD2E7470EE14C6C93),
numpy.uint64(0xCA5A826395121157)],
(32, 2): [numpy.uint32(0xD256D193)],
(32, 4): [numpy.uint32(0xD2511F53),
numpy.uint32(0xCD9E8D57)]
}
assert 1 <= rounds <= 12
word_dtype = dtypes.normalize_type(numpy.uint32 if bitness ==
32 else numpy.uint64)
key_words = counter_words // 2
key_dtype, key_ctype, counter_dtype, counter_ctype = create_struct_types(
word_dtype, key_words, counter_words)
module = Module(TEMPLATE.get_def("philox"),
render_kwds=dict(word_dtype=word_dtype,
word_ctype=dtypes.ctype(word_dtype),
key_words=key_words,
counter_words=counter_words,
key_ctype=key_ctype,
counter_ctype=counter_ctype,
rounds=rounds,
w_constants=W_CONSTANTS[bitness],
m_constants=M_CONSTANTS[(bitness,
counter_words)]))
return Bijection(module, word_dtype, key_dtype, counter_dtype)
def mul(ff_elem=None, method="c", nested_method="c"):
if ff_elem is None:
ff_elem = get_ff_elem()
module = Module(
TEMPLATE.get_def('mul_def'),
render_kwds=dict(
method=method,
ff=ff_elem, ff_elem=ff_elem.module,
add=add(ff_elem=ff_elem, method=nested_method).module, # used for "c" method
sub=sub(ff_elem=ff_elem, method=nested_method).module,
mod=mod(ff_elem=ff_elem, method=nested_method).module, # used for "c_from_asm" method
))
return FiniteFieldMul(ff_elem, module)
def uniform_integer(bijection, dtype, low, high=None):
"""
Generates uniformly distributed integer numbers in the interval ``[low, high)``.
If ``high`` is ``None``, the interval is ``[0, low)``.
Supported dtypes: any numpy integers.
If the size of the interval is a power of 2, a fixed number of counters
is used in each thread.
Returns a :py:class:`~reikna.cbrng.samplers.Sampler` object.
"""
if high is None:
low, high = 0, low + 1
else:
assert low < high - 1
dtype = dtypes.normalize_type(dtype)
ctype = dtypes.ctype(dtype)
if dtype.kind == 'i':
assert low >= -2**(dtype.itemsize * 8 - 1)
assert high < 2**(dtype.itemsize * 8 - 1)
else:
assert low >= 0
assert high < 2**(dtype.itemsize * 8)
num = high - low
if num <= 2**32:
raw_dtype = numpy.dtype('uint32')
else:
raw_dtype = numpy.dtype('uint64')
raw_func = bijection.raw_functions[raw_dtype]
max_num = 2**(raw_dtype.itemsize * 8)
raw_ctype = dtypes.ctype(dtypes.normalize_type(raw_dtype))
module = Module(TEMPLATE.get_def("uniform_integer"),
render_kwds=dict(bijection=bijection,
dtype=dtype,
ctype=ctype,
raw_ctype=raw_ctype,
raw_func=raw_func,
max_num=max_num,
num=num,
low=low))
return Sampler(bijection,
module,
dtype,
deterministic=(max_num % num == 0))
def add(*in_dtypes, **kwds):
"""add(*in_dtypes, out_dtype=None)
Returns a :py:class:`~reikna.cluda.Module` with a function of
``len(in_dtypes)`` arguments that adds values of types ``in_dtypes``.
If ``out_dtype`` is given, it will be set as a return type for this function.
This is necessary since on some platforms the ``+`` operator for a complex and a real number
works in an unexpected way (returning ``(a.x + b, a.y + b)`` instead of ``(a.x + b, a.y)``).
"""
assert set(kwds.keys()).issubset(['out_dtype'])
out_dtype = derive_out_dtype(kwds.get('out_dtype', None), *in_dtypes)
return Module(
TEMPLATE.get_def('add_or_mul'),
render_kwds=dict(op='add', out_dtype=out_dtype, in_dtypes=in_dtypes))
def module_transformation(output, param, subtree_parameters, tr_snippet,
tr_args):
return Module(_module_transformation,
render_kwds=dict(
output=output,
name=param.name,
param_cnames_seq=param_cnames_seq,
subtree_parameters=subtree_parameters,
q_indices=index_cnames_seq(param, qualified=True),
VALUE_NAME=VALUE_NAME,
nq_params=param_cnames_seq(subtree_parameters),
nq_indices=index_cnames_seq(param),
connector_ctype=param.annotation.type.ctype,
tr_snippet=tr_snippet,
tr_args=tr_args))
def module_combined(output, param, subtree_parameters, module_idx):
return Module(_module_combined,
render_kwds=dict(
output=output,
VALUE_NAME=VALUE_NAME,
shape=param.annotation.type.shape,
module_idx=module_idx,
disassemble=_snippet_disassemble_combined,
connector_ctype=param.annotation.type.ctype,
nq_indices=index_cnames_seq(param),
q_indices=index_cnames_seq(param, qualified=True),
param_cnames_str=param_cnames_seq,
subtree_parameters=subtree_parameters,
nq_params=param_cnames_seq(subtree_parameters),
indices=index_cnames(param.annotation.type.shape)))
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 lsh(exp_range, exp_dtype, ff_elem=None, method="c", nested_method="c"):
# this can be relaxed when we witch to regular C
assert numpy.dtype(exp_dtype) == numpy.dtype('uint32')
assert exp_range in (32, 64, 96, 128, 160, 192)
if ff_elem is None:
ff_elem = get_ff_elem()
module = Module(
TEMPLATE.get_def('lsh_def'),
render_kwds=dict(
method=method,
sub=sub(ff_elem=ff_elem, method=nested_method).module, # used in "c" method
add=add(ff_elem=ff_elem, method=nested_method).module, # used in "c" method
mod=mod(ff_elem=ff_elem, method=nested_method).module, # used in "c_from_asm" method
exp_range=exp_range, ff=ff_elem, ff_elem=ff_elem.module,
exp_dtype=exp_dtype))
return FiniteFieldLsh(ff_elem, module, exp_dtype)
def get_nonlinear_wrapper(c_dtype, nonlinear_module, dt):
s_dtype = dtypes.real_for(c_dtype)
return Module.create("""
%for comp in (0, 1):
INLINE WITHIN_KERNEL ${c_ctype} ${prefix}${comp}(
${c_ctype} psi0, ${c_ctype} psi1, ${s_ctype} t)
{
${c_ctype} nonlinear = ${nonlinear}${comp}(psi0, psi1, t);
return ${mul}(
COMPLEX_CTR(${c_ctype})(0, -${dt}),
nonlinear);
}
%endfor
""",
render_kwds=dict(c_ctype=dtypes.ctype(c_dtype),
s_ctype=dtypes.ctype(s_dtype),
mul=functions.mul(c_dtype, c_dtype),
dt=dtypes.c_constant(dt, s_dtype),
nonlinear=nonlinear_module))
def get_diffusion(state_dtype, gamma):
return Diffusion(
Module.create(
"""
<%
r_dtype = dtypes.real_for(s_dtype)
s_ctype = dtypes.ctype(s_dtype)
r_ctype = dtypes.ctype(r_dtype)
%>
INLINE WITHIN_KERNEL ${s_ctype} ${prefix}0_0(
const int idx_x,
const ${s_ctype} psi,
${r_ctype} t)
{
return COMPLEX_CTR(${s_ctype})(${numpy.sqrt(gamma)}, 0);
}
""",
render_kwds=dict(
mul_cr=functions.mul(state_dtype, dtypes.real_for(state_dtype)),
s_dtype=state_dtype,
gamma=gamma)),
state_dtype, components=1, noise_sources=1)
def create(cls, bijection, seed=None, reserve_id_space=True):
"""
Creates a generator.
:param bijection: a :py:class:`~reikna.cbrng.bijections.Bijection` object.
:param seed: an integer, or numpy array of 32-bit unsigned integers.
:param reserve_id_space: if ``True``, the last 32 bit of the key will be reserved
for the thread identifier.
As a result, the total size of the key should be 64 bit or more.
If ``False``, the thread identifier will be just added to the key,
which will still result in different keys for different threads,
with the danger that different seeds produce the same sequences.
"""
if reserve_id_space:
if bijection.key_words == 1 and bijection.word_dtype.itemsize == 4:
# It's too hard to compress both global and thread-dependent part
# in a single 32-bit word.
# Let the user handle this himself.
raise ValueError("Cannor reserve ID space in a 32-bit key")
if bijection.word_dtype.itemsize == 4:
key_words32 = bijection.key_words - 1
else:
if bijection.key_words > 1:
key_words32 = (bijection.key_words - 1) * 2
else:
# Philox-2x64 case, the key is a single 64-bit integer.
# We use first 32 bit for the key, and the remaining 32 bit for a thread identifier.
key_words32 = 1
else:
key_words32 = bijection.key_words * (bijection.word_dtype.itemsize // 4)
if isinstance(seed, numpy.ndarray):
# explicit key was provided
assert seed.size == key_words32 and seed.dtype == numpy.uint32
key = seed.copy().flatten()
else:
# use numpy to generate the key from seed
np_rng = numpy.random.RandomState(seed)
# 32-bit Python can only generate random integer up to 2**31-1
key16 = np_rng.randint(0, 2**16, key_words32 * 2)
key = numpy.zeros(key_words32, numpy.uint32)
for i in range(key_words32 * 2):
key[i // 2] += key16[i] << (16 if i % 2 == 0 else 0)
full_key = numpy.zeros(1, bijection.key_dtype)[0]
if bijection.word_dtype.itemsize == 4:
full_key['v'][:key_words32] = key
else:
for i in range(key_words32):
full_key['v'][i // 2] += key[i] << (32 if i % 2 == 0 else 0)
module = Module.create("""
WITHIN_KERNEL ${bijection.module}Key ${prefix}key_from_int(int idx)
{
${bijection.module}Key result;
%for i in range(bijection.key_words):
result.v[${i}] = ${key['v'][i]}
%if i == bijection.key_words - 1:
+ idx
%endif
;
%endfor
return result;
}
""",
render_kwds=dict(
bijection=bijection,
key=full_key))
return cls(module, full_key)