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 norm_const(arr_t, order): """ Returns a transformation that calculates the ``order``-norm (1 output, 1 input): ``output = abs(input) ** order``. """ if dtypes.is_complex(arr_t.dtype): out_dtype = dtypes.real_for(arr_t.dtype) else: out_dtype = arr_t.dtype return Transformation( [ Parameter('output', Annotation(Type(out_dtype, arr_t.shape), 'o')), Parameter('input', Annotation(arr_t, 'i'))], """ ${input.ctype} val = ${input.load_same}; ${output.ctype} norm = ${norm}(val); %if order != 2: norm = pow(norm, ${dtypes.c_constant(order / 2, output.dtype)}); %endif ${output.store_same}(norm); """, render_kwds=dict( norm=functions.norm(arr_t.dtype), order=order))
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 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 test_norm(thr, out_code, in_codes): out_dtype, in_dtypes = generate_dtypes(out_code, in_codes) check_func(thr, functions.norm(in_dtypes[0]), lambda x: numpy.abs(x)**2, out_dtype, in_dtypes)
def test_norm(thr, out_code, in_codes): out_dtype, in_dtypes = generate_dtypes(out_code, in_codes) check_func( thr, functions.norm(in_dtypes[0]), lambda x: numpy.abs(x) ** 2, out_dtype, in_dtypes)