#!/usr/bin/env python
# coding: utf-8

# ## Faster wedge

# Going to look at faster ways to make wedges.

# In[1]:


get_ipython().run_line_magic('', 'matplotlib inline')
import matplotlib.pyplot as plt

import numpy as np


# Set up the model parameters.

# In[2]:


length = 100.# x range
width = 100. # y range
depth = 240. # z range


# Define the plotting function.

# In[3]:


def plot_slices(data, interpolation='nearest', cmap='Greys', vmin=0., vmax=1.):
    plt.figure(figsize=(16,12))
    
    for j in range(3):
        plt.subplot(3,4,j+1)
        plt.title('y = ' + str(int(j*(length/3))))
        plt.imshow(data[:,:,int(j*(length/3))], interpolation=interpolation, cmap=cmap, vmax=vmax, vmin=vmin)
    plt.subplot(3,4,4)
    plt.title('y = ' + str(int(length-1)))
    plt.imshow(data[:,:,-1], interpolation=interpolation, cmap=cmap, vmax=vmax, vmin=vmin)
    
    for j in range(3):
        plt.subplot(3,4,5+j)
        plt.title('x = ' + str(int(j*(width/3))))
        plt.imshow(data[:,int(j*(width/3)),:], interpolation=interpolation, cmap=cmap, vmax=vmax, vmin=vmin)
    plt.subplot(3,4,8)
    plt.title('x = ' + str(int(width-1)))
    plt.imshow(data[:,-1,:], interpolation=interpolation, cmap=cmap, vmax=vmax, vmin=vmin)
    
    for j in range(3):
        plt.subplot(3,4,9+j)
        plt.title('z = ' + str(1 + int(1 + depth/3 + j*depth/(3*3))))
        plt.imshow(data[1 + int(depth/3 + j*depth/(3*3)),:,:], interpolation=interpolation, cmap=cmap, vmax=vmax, vmin=vmin)
    plt.subplot(3,4,12)
    plt.title('z = ' + str(int(2*depth/3 - 1)))
    plt.imshow(data[int(2*depth/3 - 1),:,:], interpolation=interpolation, cmap=cmap, vmax=vmax, vmin=vmin)
    
    plt.show()


# ## Original formulation

# Let's start with the wedge code from version 1 of <code>Spectral_wedge.ipynb</code>. The main difference is that we will need to use the rock properties directly, rather than building an array of <code>int</code>s first. 

# In[5]:


rocks = [(2800, 2850), (2800, 2850), (2400, 2450), (2400,2450)]

def make_ai(rocks):
    rx = np.array(rocks)
    return rx[:,0] * rx[:,1] / 10e6

ai = make_ai(rocks)
print(ai)


# In[6]:


def make_earth(depth, length, width, ai):
    
    earth = np.ones((depth,length,width))

    for xslice in range(int(length)):
        thickness = (depth/3)*(1 + xslice/length)
        
        value = (ai[1]*xslice/length) + (ai[2]*(length-xslice)/length)
        earth[:,xslice,:] *= value    
        
        # Draw the bottom of the wedge
        value = ai[0]
        for (z,y),v in np.ndenumerate(earth[:,xslice,:]):
            if z > thickness:
                earth[z,xslice,y] = value
    
        # Put the top on
        value = ai[3]
        for (z,y),v in np.ndenumerate(earth[:,xslice,:]):
            if z < depth/3:
                earth[z,xslice,y] = value
        
    return earth


# In[7]:


get_ipython().run_line_magic('timeit', 'earth = make_earth(depth, length, width, ai)')


# In[7]:


earth = make_earth(depth, length, width, ai)


# In[8]:


earth.shape


# ## Index the array with a list - SLOW

# You can index an array with an array-like object — a list or another array. 

# In[8]:


a = np.array([1,2,3,4,5,6,7,8])     # A 1D array
b = [1,3,5,7]
print a[b]

d = np.array([[1,2,3,4],[5,6,7,8]]) # A 2D array...
e = [[0,1],[1,3]]                   # Needs 2 indices
print d[e]


# Now let's try using index lists, building the indices with nested list comprehensions.

# In[9]:


def use_index_lists(depth, length, width, ai):
    
    earth = np.ones((depth,length,width))
    
    # Bottom indices
    indices = [[[x for x in range(int(length))] for y in range(int(width))] for z in range(int(depth/2))]    
    earth[indices] = 2.
    
    return earth


# In[ ]:


# Only getting started, let's test it...
get_ipython().run_line_magic('timeit', 'earth_index_lists = use_index_lists(depth, length, width, ai)')


# In[ ]:


earth_index_lists = use_index_lists(depth, length, width, ai)


# In[ ]:


earth_index_lists.shape


# That was unbelievably slow. Forget it. 

# ## Fewer loops - FAST

# It occurs to me I have too many loops. I think I need one, because the values I want vary with index.
# 
# I just realized that instead of enumerating and deciding if the *z*-index is greater than thickness, say, I can just slice with thickness as the index. Doh! NumPy is awesome...

# In[9]:


def one_loop(depth, length, width, ai):
    
    earth = np.ones((depth,length,width))
    
    d = depth/3
    
    for xslice in range(int(length)):
        xl = xslice/length
        thickness = d*(1 + xl)
        mix = (ai[1]*xl) + (ai[2]*(length-xslice)/length)
        
        # Draw wedge
        earth[:,xslice,:] *= mix    

    # Put the top on
    earth[:d,:,:] = ai[3]
 
    return earth


# In[10]:


get_ipython().run_line_magic('timeit', 'one_loop(depth, length, width, ai)')


# Holy crap, that's a bit better.

# In[11]:


earth = one_loop(depth, length, width, ai)


# ## Check plot

# In[12]:


plot_slices(earth, interpolation='nearest', cmap='jet')


# ## The final frontier

# The last tricky bit is the outer loop. It's only there to loop over the *x*-slices. The tricky part is that the wedge thickness calculation depends on the index of the slice. I think <code>numpy.indices</code> can provide this aspect and avoid the loop...

# In[13]:


zs, xs, ys = np.indices((4,3,2))
print "x values:\n", xs
print "\ny values:\n", ys
print "\nz values:\n", zs


# In[8]:


def no_loop(depth, length, width, ai):
    
    # Make an earth full of ones
    earth = np.ones((depth,length,width))
    
    # Set up the gradient, which we'll then crop, so to speak
    mix = np.array([(ai[1]*x/length) + (ai[2]*(length-x)/length) for x in range(int(length))])
    earth *= mix[:, np.newaxis]
    
    # Now make the wedge
    d = depth/3
    thickness = [d*(1 + x/length) for x in range(int(length))]
    
    # Make a set of index arrays
    zs = np.indices(earth.shape)[0]
    
    # Make a boolean array map of stuff below the wedge
    below = zs > thickness
    
    # I don't know why I have to do this but I do
    below = np.swapaxes(below, axis1=1, axis2=2)
    
    # Set those values to the lowermost rock AI
    earth[below] = ai[0]

    # Put the top layer on
    earth[:d,:,:] = ai[3]
 
    return earth


# In[9]:


get_ipython().run_line_magic('timeit', 'earth = no_loop(depth, length, width, ai)')


# In[ ]:


def no_loop(depth, length, width, ai):
    
    # Make an earth full of ones
    earth = np.ones((depth,length,width))
    
    # Set up the gradient, which we'll then crop, so to speak
    mix = np.array([(ai[1]*x/length) + (ai[2]*(length-x)/length) for x in range(int(length))])
    earth *= mix[:, np.newaxis]
    
    # Now make the wedge
    d = depth/3
    thickness = [d*(1 + x/length) for x in range(int(length))]
    
    # Make a set of index arrays
    zs = np.indices(earth.shape)[0]
    
    # Make a boolean array map of stuff below the wedge
    below = zs > thickness
    
    # I don't know why I have to do this but I do
    below = np.swapaxes(below, axis1=1, axis2=2)
    
    # Set those values to the lowermost rock AI
    earth[below] = ai[0]

    # Put the top layer on
    earth[:d,:,:] = ai[3]
 
    return earth


# In[11]:


depth, length, width, ai


# In[ ]:





# In[16]:


earth = no_loop(depth, length, width, ai)


# In[17]:


plot_slices(earth, interpolation='nearest', cmap='jet')


# Amazing — the <code>for</code> loop is 3 times faster! And I think it's the <code>indices</code> part that is slowing it down, because without that section (the line with <code>np.indices</code> and the three lines after it), it only takes 9.5 ms to execute.

# ## Conclusion

# Assuming I'm not doing something stupid in this code, it seems it's not always faster with indexing.