import numpy as np

a = np.array([10, 20, 30, 40])
a + 5

"""
Explanation: <a name="top"></a>
<div style="width:1000 px">

<div style="float:right; width:98 px; height:98px;">
<img src="https://raw.githubusercontent.com/Unidata/MetPy/master/metpy/plots/_static/unidata_150x150.png" alt="Unidata Logo" style="height: 98px;">
</div>

<h1>NumPy Broadcasting and Vectorization</h1>
<h3>Unidata Python Workshop</h3>

<div style="clear:both"></div>
</div>

<hr style="height:2px;">

<div style="float:right; width:250 px"><img src="http://www.contribute.geeksforgeeks.org/wp-content/uploads/numpy-logo1.jpg" alt="NumPy Logo" style="height: 250px;"></div>

Overview:

Teaching: 20 minutes
Exercises: 15 minutes

Questions

How can we work with arrays of differing shapes without needing to manually loop or copy data?
How can we reframe operations on data to avoid looping in Python?

Objectives

<a href="#broadcast">Use broadcasting to implicitly loop over data</a>
<a href="#vectorizing">Vectorize calculations to avoid explicit loops</a>

<a name="broadcasting"></a>
1. Using broadcasting to implicitly loop over data
Broadcasting is a useful NumPy tool that allows us to perform operations between arrays with different shapes, provided that they are compatible with each other in certain ways. To start, we can create an array below and add 5 to it:
End of explanation
"""


b = np.array([5])
a + b

"""
Explanation: This works even though 5 is not an array; it works like as we would expect, adding 5 to each of the elements in a. This also works if 5 is an array:
End of explanation
"""


b = np.array([5, 5, 10, 10])
a + b

"""
Explanation: This takes the single element in b and adds it to each of the elements in a. This won't work for just any b, though; for instance, the following:
python
b = np.array([5, 6, 7])
a + b
won't work. It does work if a and b are the same shape:
End of explanation
"""


b = np.array([1, 2, 3, 4, 5])

result = np.empty((5, 4), dtype=np.int32)
for row, valb in enumerate(b):
    for col, vala in enumerate(a):
        result[row, col] = vala + valb
result

"""
Explanation: What if what we really want is pairwise addition of a, b? Without broadcasting, we could accomplish this by looping:
End of explanation
"""


aa = np.tile(a, (5, 1))
aa

# Turn b into a column array, then tile it
bb = np.tile(b.reshape(5, 1), (1, 4))
bb

aa + bb

"""
Explanation: We can also do this by manually repeating the arrays to the proper shape for the result, using np.tile. This avoids the need to manually loop:
End of explanation
"""


a.shape

b.shape

"""
Explanation: We can also do this using broadcasting, which is where NumPy implicitly repeats the array without using additional memory. With broadcasting, NumPy takes care of repeating for you, provided dimensions are "compatible". This works as:
1. Check the number of dimensions of the arrays. If they are different, prepend size one dimensions
2. Check if each of the dimensions are compatible: either the same size, or one of them is 1.
End of explanation
"""


bb = b[:, np.newaxis]
bb.shape

a + bb

"""
Explanation: Right now, they have the same number of dimensions, 1, but that dimension is incompatible.  We can solve this by appending a dimension using np.newaxis when indexing:
End of explanation
"""


a + b[:, np.newaxis]

"""
Explanation: This can be written more directly in one line:
End of explanation
"""


x = np.array([1, 2])
y = np.array([3, 4, 5])
z = np.array([6, 7, 8, 9])

d_2d = x[:, np.newaxis]**2 + y**2

d_2d.shape

d_3d = d_2d[..., np.newaxis] + z**2

d_3d.shape

"""
Explanation: This also works 2D and 1D, etc.:
End of explanation
"""


h = x[:, np.newaxis, np.newaxis]**2 + y[np.newaxis, :, np.newaxis]**2 + z**2

"""
Explanation: Or in one line:
End of explanation
"""


h.shape

np.all(h == d_3d)

"""
Explanation: We can see this one-line result has the same shape and same values as the other multi-step calculation.
End of explanation
"""


# Starting data
pressure = np.array([1000, 850, 500, 300])
temps = np.linspace(20, 30, 24).reshape(4, 3, 2)

#
# YOUR CALCULATION HERE
#

# %load solutions/broadcasting.py


"""
Explanation: Broadcasting is often useful when you want to do calculations with coordinate values, which are often given as 1D arrays corresponding to positions along a particular array dimension. For example, taking range and azimiuth values for radar data (1D separable polar coordinates) and converting to x,y pairs relative to the radar location.
Exercise
Given the 3D temperature field and 1-D pressure coordinates below, calculate: $T * exp(P / 1000)$. You will need to use broadcasting to make the arrays compatible.
End of explanation
"""


a = np.linspace(0, 20, 6)
a

"""
Explanation: <a name="vectorizing"></a>
2. Vectorize calculations to avoid explicit loops
When working with arrays of data, loops over the individual array elements is a fact of life. However, for improved runtime performance, it is important to avoid performing these loops in Python as much as possible, and let NumPy handle the looping for you. Avoiding these loops frequently, but not always, results in shorter and clearer code as well.
Look ahead/behind
One common pattern for vectorizing is in converting loops that work over the current point as well as the previous and/or next point. This comes up when doing finite-difference calculations (e.g. approximating derivatives)
End of explanation
"""


d = np.zeros(a.size - 1)
for i in range(len(a) - 1):
    d[i] = a[i + 1] - a[i]
d

"""
Explanation: We can calculate the forward difference for this array with a manual loop as:
End of explanation
"""


a[1:]

"""
Explanation: It would be nice to express this calculation as a loop, if possible. To see how to go about this, let's condsider the values that are involved in calculating d[i], a[i+1] and a[i]. The values over the loop iterations are:
<table>
    <tr> <td>i</td><td>a[i+1]</td><td>a[i]</td> </tr>
    <tr> <td>0</td><td>4</td><td>0</td> </tr>
    <tr> <td>1</td><td>8</td><td>4</td> </tr>
    <tr> <td>2</td><td>12</td><td>8</td> </tr>
    <tr> <td>3</td><td>16</td><td>12</td> </tr>
    <tr> <td>4</td><td>20</td><td>16</td> </tr>
</table>

We can express the series of values for a[i+1] then as:
End of explanation
"""


a[:-1]

"""
Explanation: and a[i] as:
End of explanation
"""


a[1:] - a[:-1]

"""
Explanation: This means that we can express the forward difference as:
End of explanation
"""


# %load solutions/vectorized_diff.py


"""
Explanation: It should be noted that using slices in this way returns only a view on the original array. This means not only can you use the slices to modify the original data (even accidentally), but that this is also a quick operation that does not involve a copy and does not bloat memory usage.
Exercise: 2nd Derivative
A finite difference estimate of the 2nd derivative is given by:
$$f''(x) = 2
f_i - f_{i+1} - f_{i-1}$$
(we're ignoring $\Delta x$ here)

Write vectorized code to calculate this finite difference for a (using slices)

What values should we be expecting to get for the 2nd derivative?
End of explanation
"""


temps = np.round(20 + np.random.randn(10) * 5, 1)
temps

"""
Explanation: Blocking
Another application where vectorization comes into play to make operations more efficient is when operating on blocks of data. Let's start by creating some temperature data (rounding to make it easier to see/recognize the values).
End of explanation
"""


avg = np.zeros_like(temps)
# We're just ignoring the edge effects here
for i in range(1, len(temps) - 1):
    sub = temps[i - 1:i + 2]
    avg[i] = sub.mean()

avg

"""
Explanation: Let's start by writing a loop to take a 3-point running mean of the data. We'll do this by iterating over all the point in the array and average the 3 points centered on that point. We'll simplify the problem by avoiding dealing with the cases at the edges of the array.
End of explanation
"""


# i - 1            i          i + 1
(temps[:-2] + temps[1:-1] + temps[2:]) / 3

"""
Explanation: As with the case of doing finite differences, we can express this using slices of the original array:
End of explanation
"""


temps

"""
Explanation: Another option to solve this is not using slicing but by using a powerful numpy tool: as_strided. This tool can result in some odd behavior, so take care when using--the tradeoff is that this can be used to do some powerful operations. What we're doing here is altering how NumPy is interpreting the values in the memory that underpins the array. So for this array:
End of explanation
"""


block_size = 3
new_shape = (len(temps) - block_size + 1, block_size)
bytes_per_item = temps.dtype.itemsize
temps_strided = np.lib.stride_tricks.as_strided(temps,
                                                shape=new_shape,
                                                strides=(bytes_per_item, bytes_per_item))
temps_strided

"""
Explanation: we can create a view of the array with a new, bigger shape, with rows made up of overlapping values. We do this by specifying a new shape of 8x3, one row for each of the length 3 blocks we can fit in the original 1D array of data. We then use the strides argument to control how numpy walks between items in each dimension. The last item in the strides tuple is just as normal--it says that the number of bytes to walk between items is just the size of an item. (Increasing this would skip items.) The first item says that when we go to a new, in this case row, only advance the size of a single item. This is what gives us overlapping rows.
End of explanation
"""


temps_strided.mean(axis=-1)

"""
Explanation: Now that we have this view of the array with the rows representing overlapping blocks, we can operate across the rows with mean and the axis=-1 argument to get our running average:
End of explanation
"""


temps_strided[0, 2] = 2000
temps_strided

"""
Explanation: It should be noted that there are no copies going on here, so if we change a value at a single indexed location, the change actually shows up in multiple locations:
End of explanation
"""


pressure = np.linspace(1000, 100, 25)
temps = np.random.randn(25, 30, 40) * 3 + np.linspace(25, -100, 25).reshape(-1, 1, 1)

"""
Explanation: Finding the difference between min and max
Another operation that crops up when slicing and dicing data is trying to identify a set of indexes, along a particular axis, within a larger multidimensional array. For instance, say we have a 3D array of temperatures, and want to identify the location of the $-10^oC$ isotherm within each column:
End of explanation
"""


# Using axis=0 to tell it to operate along the pressure dimension
inds = np.argmin(np.abs(temps - -10), axis=0)
inds

inds.shape

"""
Explanation: NumPy has the function argmin() which returns the index of the minium value. We can use this to find the minimum absolute difference between the value and -10:
End of explanation
"""


pressure[inds]

"""
Explanation: Great! We have an array representing the index of the point closest to $-10^oC$ in each column of data. We could use this to look up into our pressure coordinates to find the pressure level for each column:
End of explanation
"""


temps[inds, :, :].shape

"""
Explanation: How about using that to find the actual temperature value that was closest?
End of explanation
"""


output = np.empty(inds.shape, dtype=temps.dtype)
for (i, j), val in np.ndenumerate(inds):
    output[i, j] = temps[val, i, j]
output

"""
Explanation: Unfortunately, this replaced the pressure dimension (size 25) with the shape of our index array (30 x 40), giving us a 30 x 40 x 30 x 40 array (imagine what would have happened with real data!). One solution here would be to loop:
End of explanation
"""


pressure[np.arange(pressure.size)]

np.all(pressure[np.arange(pressure.size)] == pressure)

"""
Explanation: Of course, what we really want to do is avoid the explicit loop. Let's temporarily simplify the problem to a single dimension. If we have a 1D array, we can pass a 1D array of indices (a full) range, and get back the same as the original data array:
End of explanation
"""


y_inds = np.arange(temps.shape[1])[:, np.newaxis]
x_inds = np.arange(temps.shape[2])
temps[inds, y_inds, x_inds]

"""
Explanation: We can use this to select all the indices on the other dimensions of our temperature array. We will also need to use the magic of broadcasting to combine arrays of indices across dimensions:
Now vectorized solution
End of explanation
"""


np.all(output == temps[inds, y_inds, x_inds])

"""
Explanation: Now let's say we want to find the relative humidity at the -10C isotherm
End of explanation
"""


