| 1 | #!/usr/bin/env python |
|---|
| 2 | # -*- coding: utf-8 -*- |
|---|
| 3 | |
|---|
| 4 | import numpy as np |
|---|
| 5 | import math |
|---|
| 6 | import pyopencl as cl |
|---|
| 7 | from weights import GaussianDispersion |
|---|
| 8 | from sasmodel import card |
|---|
| 9 | |
|---|
| 10 | def set_precision(src, qx, qy, dtype): |
|---|
| 11 | qx = np.ascontiguousarray(qx, dtype=dtype) |
|---|
| 12 | qy = np.ascontiguousarray(qy, dtype=dtype) |
|---|
| 13 | if np.dtype(dtype) == np.dtype('float32'): |
|---|
| 14 | header = """\ |
|---|
| 15 | #define real float |
|---|
| 16 | """ |
|---|
| 17 | else: |
|---|
| 18 | header = """\ |
|---|
| 19 | #pragma OPENCL EXTENSION cl_khr_fp64: enable |
|---|
| 20 | #define real double |
|---|
| 21 | """ |
|---|
| 22 | return header+src, qx, qy |
|---|
| 23 | |
|---|
| 24 | class GpuEllipse(object): |
|---|
| 25 | PARS = { |
|---|
| 26 | 'scale':1, 'radius_a':1, 'radius_b':1, 'sldEll':1e-6, 'sldSolv':0, 'background':0, 'axis_theta':0, 'axis_phi':0, |
|---|
| 27 | } |
|---|
| 28 | PD_PARS = ['radius_a', 'radius_b', 'axis_theta', 'axis_phi'] |
|---|
| 29 | |
|---|
| 30 | def __init__(self, qx, qy, dtype='float32'): |
|---|
| 31 | |
|---|
| 32 | ctx,_queue = card() |
|---|
| 33 | src, qx, qy = set_precision(open('Kernel-Ellipse.cpp').read(), qx, qy, dtype=dtype) |
|---|
| 34 | self.prg = cl.Program(ctx, src).build() |
|---|
| 35 | self.qx, self.qy = qx, qy |
|---|
| 36 | |
|---|
| 37 | #buffers |
|---|
| 38 | mf = cl.mem_flags |
|---|
| 39 | self.qx_b = cl.Buffer(ctx, mf.READ_ONLY | mf.COPY_HOST_PTR, hostbuf=self.qx) |
|---|
| 40 | self.qy_b = cl.Buffer(ctx, mf.READ_ONLY | mf.COPY_HOST_PTR, hostbuf=self.qy) |
|---|
| 41 | self.res_b = cl.Buffer(ctx, mf.WRITE_ONLY, qx.nbytes) |
|---|
| 42 | self.res = np.empty_like(self.qx) |
|---|
| 43 | |
|---|
| 44 | def eval(self, pars): |
|---|
| 45 | #b_n = radius_b # want, a_n = radius_a # want, etc |
|---|
| 46 | _ctx,queue = card() |
|---|
| 47 | radius_a, radius_b, axis_theta, axis_phi = \ |
|---|
| 48 | [GaussianDispersion(int(pars[base+'_pd_n']), pars[base+'_pd'], pars[base+'_pd_nsigma']) |
|---|
| 49 | for base in GpuEllipse.PD_PARS] |
|---|
| 50 | |
|---|
| 51 | radius_a.value, radius_a.weight = radius_a.get_weights(pars['radius_a'], 0, 10000, True) |
|---|
| 52 | radius_b.value, radius_b.weight = radius_b.get_weights(pars['radius_b'], 0, 10000, True) |
|---|
| 53 | axis_theta.value, axis_theta.weight = axis_theta.get_weights(pars['axis_theta'], -90, 180, False) |
|---|
| 54 | axis_phi.value, axis_phi.weight = axis_phi.get_weights(pars['axis_phi'], -90, 180, False) |
|---|
| 55 | |
|---|
| 56 | #Perform the computation, with all weight points |
|---|
| 57 | sum, norm, norm_vol, vol = 0.0, 0.0, 0.0, 0.0 |
|---|
| 58 | size = len(axis_theta.weight) |
|---|
| 59 | sub = pars['sldEll'] - pars['sldSolv'] |
|---|
| 60 | real = np.float32 if self.qx.dtype == np.dtype('float32') else np.float64 |
|---|
| 61 | |
|---|
| 62 | #Loop over radius weight points |
|---|
| 63 | for i in xrange(len(radius_a.weight)): |
|---|
| 64 | #Loop over length weight points |
|---|
| 65 | for j in xrange(len(radius_b.weight)): |
|---|
| 66 | #Average over theta distribution |
|---|
| 67 | for k in xrange(len(axis_theta.weight)): |
|---|
| 68 | #Average over phi distribution |
|---|
| 69 | for l in xrange(len(axis_phi.weight)): |
|---|
| 70 | #call the kernel |
|---|
| 71 | self.prg.EllipsoidKernel(queue, self.qx.shape, None, real(radius_a.weight[i]), |
|---|
| 72 | real(radius_b.weight[j]), real(axis_theta.weight[k]), |
|---|
| 73 | real(axis_phi.weight[l]), real(pars['scale']), real(radius_a.value[i]), |
|---|
| 74 | real(radius_b.value[j]), real(sub), real(axis_theta.value[k]), |
|---|
| 75 | real(axis_phi.value[l]), self.qx_b, self.qy_b, self.res_b, |
|---|
| 76 | np.uint32(self.qx.size), np.uint32(len(axis_theta.weight))) |
|---|
| 77 | #copy result back from buffer |
|---|
| 78 | cl.enqueue_copy(queue, self.res, self.res_b) |
|---|
| 79 | sum += self.res |
|---|
| 80 | vol += radius_a.weight[i]*radius_b.weight[j]*pow(radius_b.value[j], 2)*radius_a.value[i] |
|---|
| 81 | norm_vol += radius_a.weight[i]*radius_b.weight[j] |
|---|
| 82 | norm += radius_a.weight[i]*radius_b.weight[j]*axis_theta.weight[k]*axis_phi.weight[l] |
|---|
| 83 | # Averaging in theta needs an extra normalization |
|---|
| 84 | # factor to account for the sin(theta) term in the |
|---|
| 85 | # integration (see documentation). |
|---|
| 86 | |
|---|
| 87 | # if size > 1: |
|---|
| 88 | # norm /= math.asin(1.0) |
|---|
| 89 | if vol != 0.0 and norm_vol != 0.0: |
|---|
| 90 | sum *= norm_vol/vol |
|---|
| 91 | |
|---|
| 92 | return sum/norm+pars['background'] |
|---|
| 93 | |
|---|
| 94 | |
|---|
| 95 | def demo(): |
|---|
| 96 | from time import time |
|---|
| 97 | import matplotlib.pyplot as plt |
|---|
| 98 | |
|---|
| 99 | #create qx and qy evenly spaces |
|---|
| 100 | qx = np.linspace(-.02, .02, 128) |
|---|
| 101 | qy = np.linspace(-.02, .02, 128) |
|---|
| 102 | qx, qy = np.meshgrid(qx, qy) |
|---|
| 103 | |
|---|
| 104 | #saved shape of qx |
|---|
| 105 | r_shape = qx.shape |
|---|
| 106 | #reshape for calculation; resize as float32 |
|---|
| 107 | qx = qx.flatten() |
|---|
| 108 | qy = qy.flatten() |
|---|
| 109 | |
|---|
| 110 | #int main |
|---|
| 111 | pars = EllipsoidParameters(.027, 60, 180, .297e-6, 5.773e-06, 4.9, 0, 90) |
|---|
| 112 | |
|---|
| 113 | t = time() |
|---|
| 114 | result = GpuEllipse(qx, qy) |
|---|
| 115 | result.x = result.ellipsoid_fit(qx, qy, pars, b_n=35, t_n=35, a_n=1, p_n=1, sigma=3, b_w=.1, t_w=.1, a_w=.1, p_w=.1) |
|---|
| 116 | result.x = np.reshape(result.x, r_shape) |
|---|
| 117 | tt = time() |
|---|
| 118 | print("Time taken: %f" % (tt - t)) |
|---|
| 119 | |
|---|
| 120 | plt.pcolormesh(result.x) |
|---|
| 121 | plt.show() |
|---|
| 122 | |
|---|
| 123 | |
|---|
| 124 | if __name__ == "__main__": |
|---|
| 125 | demo() |
|---|
| 126 | |
|---|
| 127 | |
|---|
| 128 | |
|---|
| 129 | |
|---|
