np.random.seed(0)
xform = np.array(np.random.random_sample((3, 2, 400, 400)), dtype=np.float64)
inputs = np.array(np.random.random_sample((2, 400, 400)), dtype=np.float64)

xform_fp32 = np.array(xform, dtype=np.float32)
inputs_fp32 = np.array(inputs, dtype=np.float32)

input0 = inputs[0,::]
input1 = inputs[1,::]


"""
Explanation: Create test data
This is equivalent to the first-dimension-concatenated array that results from $\nu_e$ and $\nu_\mu$ fluxes binned at $400 \;E \times 400 \cos\theta$ bins.
The resulting input array is then $2\times400\times400$, and the transform that must (effectively) left-multiply that is $3\times 2\times400\times400$.
This hopefully represents a realistic scenario for performing an accurate oscillation calculation.
End of explanation
"""


%timeit np.einsum('ij..., j...', xform, inputs, dtype=np.float64, casting='unsafe');

output_einsum = np.einsum('ij..., j...', xform, inputs)
print output_einsum.shape

"""
Explanation: Numpy using einsum
Float64 math on float64 inputs/transforms
End of explanation
"""


x = xform[:,:,1,10]
x

i = inputs[:,1,10]
i

o = np.dot(x, i)
o

output_einsum[1,10,:]

np.all(o == output_einsum[1,10,:])

"""
Explanation: check that it's doing what I want it to do
End of explanation
"""


%timeit np.einsum('ij..., j...', xform, inputs, dtype=np.float32, casting='unsafe');

"""
Explanation: Float32 math on float64 inputs/transforms
End of explanation
"""


%timeit np.einsum('ij..., j...', xform_fp32, inputs_fp32, dtype=np.float32, casting='no');

"""
Explanation: Float32 math on float32 inputs/transforms
End of explanation
"""


def apply_python(inputs, transform):
    N_k = inputs.shape[1]
    N_l = inputs.shape[2]
    output = np.empty((N_k, N_l, 3), np.float64)
    for k in range(N_k):
        for l in range(N_l):
            output[k,l,0] = (
                transform[0,0,k,l]*inputs[0,k,l] +
                transform[0,1,k,l]*inputs[1,k,l]
            )
            output[k,l,1] = (
                transform[1,0,k,l]*inputs[0,k,l] +
                transform[1,1,k,l]*inputs[1,k,l]
            )
            output[k,l,2] = (
                transform[2,0,k,l]*inputs[0,k,l] +
                transform[2,1,k,l]*inputs[1,k,l]
            )
    return output

"""
Explanation: Python looping
End of explanation
"""


%timeit apply_python(inputs, xform);

output_python = apply_python(inputs, xform)
np.all(output_python == output_einsum)

"""
Explanation: Float64 operands
End of explanation
"""


%timeit apply_python(inputs_fp32, xform_fp32);

"""
Explanation: Float32 operands
End of explanation
"""


@jit("float64[:,:,:](float64[:,:,:], float64[:,:,:,:])", nopython=False, nogil=True, cache=True)
def apply_numba(inputs, transform):
    N_k = inputs.shape[1]
    N_l = inputs.shape[2]
    output = np.empty((N_k, N_l, 3), float64)
    for k in range(N_k):
        for l in range(N_l):
            output[k,l,0] = (
                transform[0,0,k,l]*inputs[0,k,l] +
                transform[0,1,k,l]*inputs[1,k,l]
            )
            output[k,l,1] = (
                transform[1,0,k,l]*inputs[0,k,l] +
                transform[1,1,k,l]*inputs[1,k,l]
            )
            output[k,l,2] = (
                transform[2,0,k,l]*inputs[0,k,l] +
                transform[2,1,k,l]*inputs[1,k,l]
            )
    return output

%timeit apply_numba(inputs, xform)

output_numba = apply_numba(inputs, xform)
np.all(output_numba == output_einsum)

"""
Explanation: Numba
Float64 math on float64 operands
End of explanation
"""


@jit("float32[:,:,:](float32[:,:,:], float32[:,:,:,:])", nopython=True, nogil=True, cache=True)
def apply_numba_fp32(inputs, transform):
    N_k = inputs.shape[1]
    N_l = inputs.shape[2]
    output = np.empty((N_k, N_l, 3), float32)
    for k in range(N_k):
        for l in range(N_l):
            output[k,l,0] = (
                transform[0,0,k,l]*inputs[0,k,l] +
                transform[0,1,k,l]*inputs[1,k,l]
            )
            output[k,l,1] = (
                transform[1,0,k,l]*inputs[0,k,l] +
                transform[1,1,k,l]*inputs[1,k,l]
            )
            output[k,l,2] = (
                transform[2,0,k,l]*inputs[0,k,l] +
                transform[2,1,k,l]*inputs[1,k,l]
            )
    return output

%timeit apply_numba_fp32(inputs_fp32, xform_fp32)

"""
Explanation: Float32 math on float32 operands
End of explanation
"""


@jit("float32[:,:,:](float32[:,:,:], float32[:,:,:,:])", nopython=True, nogil=True, cache=True)
def apply_numba_fp32(inputs, transform):
    N_k = inputs.shape[1]
    N_l = inputs.shape[2]
    output = np.empty((N_k, N_l, 3), float32)
    for k in range(N_k):
        for l in range(N_l):
            output[k,l,0] = (
                transform[0,0,k,l]*inputs[0,k,l] +
                transform[0,1,k,l]*inputs[1,k,l]
            )
            output[k,l,1] = (
                transform[1,0,k,l]*inputs[0,k,l] +
                transform[1,1,k,l]*inputs[1,k,l]
            )
            output[k,l,2] = (
                transform[2,0,k,l]*inputs[0,k,l] +
                transform[2,1,k,l]*inputs[1,k,l]
            )
    return output

"""
Explanation: Parallelize
End of explanation
"""


@jit("float32[:,:,:], float32[:,:,:,:], float32[:,:,:]", target='cuda') #nopython=True, nogil=True, cache=True)
def apply_numba_fp32(inputs, transform, output):
    N_k = inputs.shape[1]
    N_l = inputs.shape[2]
    #output = np.empty((N_k, N_l, 3), float32)
    for k in range(N_k):
        for l in range(N_l):
            output[k,l,0] = (
                transform[0,0,k,l]*inputs[0,k,l] +
                transform[0,1,k,l]*inputs[1,k,l]
            )
            output[k,l,1] = (
                transform[1,0,k,l]*inputs[0,k,l] +
                transform[1,1,k,l]*inputs[1,k,l]
            )
            output[k,l,2] = (
                transform[2,0,k,l]*inputs[0,k,l] +
                transform[2,1,k,l]*inputs[1,k,l]
            )
    return output

@vectorize("float32, float32, float32",
           nopython=True, target='parallel')
def apply_numba_fp32(inputs, transform, output):
    N_k = inputs.shape[1]
    N_l = inputs.shape[2]
    #output = np.empty((N_k, N_l, 3), float32)
    for k in range(N_k):
        for l in range(N_l):
            output[k,l,0] = (
                transform[0,0,k,l]*inputs[0,k,l] +
                transform[0,1,k,l]*inputs[1,k,l]
            )
            output[k,l,1] = (
                transform[1,0,k,l]*inputs[0,k,l] +
                transform[1,1,k,l]*inputs[1,k,l]
            )
            output[k,l,2] = (
                transform[2,0,k,l]*inputs[0,k,l] +
                transform[2,1,k,l]*inputs[1,k,l]
            )
    return output

"""
Explanation: Attempting to compile to GPU
End of explanation
"""


