# coding: utf-8

# # PyOpenCL: Arrays

# ## Setup code

# In[5]:

import pyopencl as cl
import numpy as np
import numpy.linalg as la


# In[6]:

a = np.random.rand(1024, 1024).astype(np.float32)


# In[7]:

ctx = cl.create_some_context()
queue = cl.CommandQueue(ctx)


# ## Creating arrays

# This notebook demonstrates working with PyOpenCL's arrays, which provide a friendlier (and more numpy-like) face on OpenCL's buffers. This is the module where they live:

# In[8]:

import pyopencl.array


# Now transfer to a *device array*.

# In[9]:

a_dev = cl.array.to_device(queue, a)


# Works like a numpy array! (`shape`, `dtype`, `strides`)

# In[10]:

a_dev.shape


# In[11]:

a_dev.dtype


# In[12]:

a_dev.strides


# ## Working with arrays

# **Goal:** Wanted to double all entries.

# In[13]:

twice_a_dev = 2*a_dev


# Easy to turn back into a `numpy` array.

# In[14]:

twice_a = twice_a_dev.get()


# Check!

# In[15]:

print(la.norm(twice_a - 2*a))


# Can just `print` the array, too.

# In[16]:

print(twice_a_dev)


# ----

# Easy to evaluate arbitrary (elementwise) expressions.

# In[13]:

import pyopencl.clmath


# In[14]:

cl.clmath.sin(a_dev)**2 - (1./a_dev) + 5


# ## Low-level Access
# 
# Can still do everything manually though!

# In[15]:

prg = cl.Program(ctx, """
    __kernel void twice(__global float *a)
    {
      int gid0 = get_global_id(0);
      int gid1 = get_global_id(1);
      int i = gid1 * 1024 + gid0;
      a[i] = 2*a[i];
    }
    """).build()
twice = prg.twice


# In[16]:

twice(queue, a_dev.shape, None, a_dev.data)


# In[17]:

print(la.norm(a_dev.get() - 2*a), la.norm(a))


# But the hardcoded 1024 is ... inelegant. So fix that!
# 
# (Also with arg `dtype setting`.)