gistfile1.py

import numpy

from reikna.cluda import ocl_api
from reikna.core import Annotation, Type, Parameter
from reikna.algorithms import PureParallel
import reikna.cluda.functions as functions


api = ocl_api()
thr = api.Thread.create()

size = (100, 100)
dtype = numpy.complex64
ftype = numpy.float64
ls_alpha = 0.5


lplfilter = PureParallel(
    [
        Parameter('output', Annotation(Type(dtype, shape=size), 'o')),
        Parameter('otfi', Annotation(Type(dtype, shape=size), 'i')),
        Parameter('fft_lpli', Annotation(Type(dtype, shape=size), 'i'))],
    """
    ${otfi.ctype} otf = ${otfi.load_same};
    ${fft_lpli.ctype} fft_lpl = ${fft_lpli.load_same};
    float divf = ${mul}(${norm}(fft_lpl) + ${norm}(otf), ${lsa});
    ${output.store_same}(${div}(${conj}(otf), divf));
    """,
    render_kwds=dict(
        mul=functions.mul(ftype, ftype, out_dtype=ftype),
        div=functions.div(dtype, ftype, out_dtype=dtype),
        norm=functions.norm(dtype),
        conj=functions.conj(dtype),
        lsa=ls_alpha))

lplfilter_c = lplfilter.compile(thr)

otfi = (numpy.random.normal(size=size) + 1j * numpy.random.normal(size=size)).astype(dtype)
fft_lpli = (numpy.random.normal(size=size) + 1j *numpy.random.normal(size=size)).astype(dtype)

otfi_dev = thr.to_device(otfi)
fft_lpli_dev = thr.to_device(fft_lpli)
output_dev = thr.array(size, dtype)

lplfilter_c(output_dev, otfi_dev, fft_lpli_dev)

output = output_dev.get()
output_ref = otfi.conj() / (ls_alpha * (numpy.abs(fft_lpli) ** 2 + numpy.abs(otfi) ** 2))

print(numpy.allclose(output, output_ref))